ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-05-17 19:40:04 +00:00
Code Issues Packages Projects Releases Wiki Activity

Remove un-needed constexpr checking for loading v_tiles in Gemm0 loop

Browse Source
This commit is contained in:
Qianfeng Zhang
2025-12-26 15:13:28 +00:00
parent df902c6a06
commit d2dadc22a7
2 changed files with 10 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
example/ck_tile/18_hstu_attention/hstu_attention_with_softmax_fwd_pipeline.hpp
View File

@@ -342,14 +342,11 @@ struct HstuAttentionWithSoftmaxFwdPipelineQRKSVS
}
else
{
// We assume NumPrefetchV >= NumPrefetchK
if constexpr(i_n0 - (n0_loops - NumPrefetchK) < NumPrefetchK)
{
// load v_tiles used in current iteration
v_tiles[number<i_n0 - (n0_loops - NumPrefetchK)>{}] =
load_tile(v_dram_window);
move_tile_window(v_dram_window, {0, kK1});
}
// Since NumPrefetchV >= NumPrefetchK, we are able to have NumPrefetchK
// prefetchings of v_tile arranged in n0_loops
v_tiles[number<i_n0 - (n0_loops - NumPrefetchK)>{}] = load_tile(v_dram_window);
move_tile_window(v_dram_window, {0, kK1});
};
__builtin_amdgcn_sched_barrier(0x00000001);
@@ -449,7 +446,6 @@ struct HstuAttentionWithSoftmaxFwdPipelineQRKSVS
__builtin_amdgcn_sched_barrier(0x00000001);
static_for<NumPrefetchK, NumPrefetchV, 1>{}([&](auto i_k1) {
// load v_tiles used in current iteration
v_tiles[i_k1] = load_tile(v_dram_window);
move_tile_window(v_dram_window, {0, kK1});
});

13
example/ck_tile/18_hstu_attention/hstu_attention_with_softmax_fwd_trload_pipeline.hpp
View File

@@ -342,14 +342,11 @@ struct HstuAttentionWithSoftmaxFwdPipelineQRKSVSTrLoad
}
else
{
// We assume NumPrefetchV >= NumPrefetchK
if constexpr(i_n0 - (n0_loops - NumPrefetchK) < NumPrefetchK)
{
// load v_tiles used in current iteration
v_tiles[number<i_n0 - (n0_loops - NumPrefetchK)>{}] =
load_tile(v_dram_window);
move_tile_window(v_dram_window, {kK1, 0});
}
// Since NumPrefetchV >= NumPrefetchK, we are able to have NumPrefetchK
// prefetchings of v_tile arranged in n0_loops
v_tiles[number<i_n0 - (n0_loops - NumPrefetchK)>{}] = load_tile(v_dram_window);
move_tile_window(v_dram_window, {kK1, 0});
};
__builtin_amdgcn_sched_barrier(0x00000001);