ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-05-05 06:01:23 +00:00
Code Issues Packages Projects Releases Wiki Activity

[CK-Tile] Add the API to load SGPR (#2878)

Browse Source 
* Have a workable version for SGPR

* have a workable version for atomic add

* Revert "have a workable version for atomic add"

This reverts commit 792377a590c26cfff9c8f545d9a9e8484a7422eb.

* substitute with the new sgpr read api

* update the CHANGELOG

* have a workable version for atomic add

* Revert "have a workable version for atomic add"

This reverts commit 792377a590c26cfff9c8f545d9a9e8484a7422eb.

* change to static for logic

* have a workable version for atomic add

* Revert "have a workable version for atomic add"

This reverts commit 792377a590c26cfff9c8f545d9a9e8484a7422eb.
This commit is contained in:
Thomas Ning
2025-09-23 01:23:56 -07:00
committed by GitHub
parent b6e8994386
commit 2cbbf5dcb3
40 changed files with 273 additions and 167 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_pagedkv_pipeline_qr_ks_vs.hpp
View File

@@ -320,9 +320,9 @@ struct BlockFmhaFwdPagedKVPipelineQRKSVS
k_block_tile = load_tile(k_dram_window);
}
auto physical_next_block_id_k =
__builtin_amdgcn_readfirstlane(k_page_block_navigator.prefetch_table_id(
amd_wave_read_first_lane(k_page_block_navigator.prefetch_table_id(
i_page_block_k, k_dram_block_window, {kN0, 0}));
auto physical_next_block_id_v = __builtin_amdgcn_readfirstlane(
auto physical_next_block_id_v = amd_wave_read_first_lane(
v_page_block_navigator.prefetch_table_id(i_page_block_v, v_dram_window, {0, kK1}));
if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS)

6
include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_qr_ks_vs.hpp
View File

@@ -321,9 +321,9 @@ struct BlockFmhaFwdSplitKVPipelineQRKSVS
k_block_tile = load_tile(k_dram_window);
}
auto physical_next_block_id_k =
__builtin_amdgcn_readfirstlane(k_page_block_navigator.prefetch_table_id(
amd_wave_read_first_lane(k_page_block_navigator.prefetch_table_id(
i_page_block_k, k_dram_block_window, {kN0, 0}));
auto physical_next_block_id_v = __builtin_amdgcn_readfirstlane(
auto physical_next_block_id_v = amd_wave_read_first_lane(
v_page_block_navigator.prefetch_table_id(i_page_block_v, v_dram_window, {0, kK1}));
if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS)
@@ -618,7 +618,7 @@ struct BlockFmhaFwdSplitKVPipelineQRKSVS
&i_page_block_v_ = i_page_block_v,
&v_dram_window_ = v_dram_window](auto i_k1) {
auto physical_next_block_id_v_ =
__builtin_amdgcn_readfirstlane(v_page_block_navigator.prefetch_table_id(
amd_wave_read_first_lane(v_page_block_navigator.prefetch_table_id(
i_page_block_v_, v_dram_window_, {0, kK1}));
const auto v = load_tile(v_dram_window_); // load next v
block_sync_lds();