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

Smally re-arrangement in nwarp_sshuffle pipeline codes

Browse Source
This commit is contained in:
Qianfeng Zhang
2025-09-08 14:54:22 +00:00
parent fd9de331c1
commit 0d2d2bb946
1 changed files with 24 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

33
include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_splitkv_pipeline_nwarp_sshuffle_qr_ks_vs.hpp
View File

@@ -398,19 +398,20 @@ struct BlockFmhaFwdSplitKVPipelineNWarpSShuffleQRKSVS
[&, &i_page_block_v_ = i_page_block_v, &v_dram_window_ = v_dram_window](auto i_k0) {
store_tile(k_lds_window,
tile_elementwise_in(k_element_func, k_tiles[number<i_k0 % 2>{}]));
__builtin_amdgcn_sched_barrier(0);
if constexpr(i_k0 < k0_loops - 2)
{
k_tiles[number<(i_k0 + 2) % 2>{}] = load_tile(k_dram_window);
move_tile_window(k_dram_window, {0, kK0});
}
else
{
v_tiles[number<i_k0 - (k0_loops - 2)>{}] = load_tile(v_dram_window_);
i_page_block_v_ = v_page_block_navigator.move_tile_window(
i_page_block_v_, v_dram_window_, {0, kK1});
}
__builtin_amdgcn_sched_barrier(0);
block_sync_lds();
gemm_0(s_acc,
get_slice_tile(q_tile,
@@ -418,6 +419,18 @@ struct BlockFmhaFwdSplitKVPipelineNWarpSShuffleQRKSVS
sequence<kM0, (i_k0 + 1) * kK0>{}),
k_lds_window);
if constexpr(i_k0 < k0_loops - 2)
{
move_tile_window(k_dram_window, {0, kK0});
}
else
{
i_page_block_v_ = v_page_block_navigator.move_tile_window(
i_page_block_v_, v_dram_window_, {0, kK1});
}
__builtin_amdgcn_sched_barrier(0);
if constexpr(i_k0 < k0_loops - 1)
block_sync_lds();
});
@@ -539,7 +552,13 @@ struct BlockFmhaFwdSplitKVPipelineNWarpSShuffleQRKSVS
__builtin_amdgcn_sched_barrier(0);
// load the first tile for next iteration
const auto s = cast_tile<SMPLComputeDataType>(s_acc); // S{j}
// shuffle through LDS so that the tile layout is consistent with required by Gemm1
store_tile(s_write_lds_window, s);
__builtin_amdgcn_sched_barrier(0);
if(i_total_loops < num_total_loop - 1)
{
// move K tile windows
@@ -553,10 +572,6 @@ struct BlockFmhaFwdSplitKVPipelineNWarpSShuffleQRKSVS
__builtin_amdgcn_sched_barrier(0);
const auto s = cast_tile<SMPLComputeDataType>(s_acc); // S{j}
// shuffle through LDS so that the tile layout is consistent with required by Gemm1
store_tile(s_write_lds_window, s);
// at this point, all LDS operations are done including ds_read for k and ds_write for s
block_sync_lds();
auto s_new = load_tile(s_read_lds_window);