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

Tiny fix in trload with_softmax/no_softmax pipeline

Browse Source
This commit is contained in:
Qianfeng Zhang
2025-11-05 14:18:37 +00:00
parent 99993acca4
commit d190af2ef5
2 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
example/ck_tile/18_hstu_attention/hstu_attention_no_softmax_fwd_trload_pipeline.hpp
View File

@@ -380,7 +380,7 @@ struct HstuAttentionNoSoftmaxFwdPipelineQRKSVSTrLoad
{
// STAGE 1, Gemm_0 ( S = Q@K )
static_for<0, k1_loops, 1>{}([&](auto i_k1) {
store_tile(k_lds_write_windows[i_k1],
store_tile(k_lds_write_windows[number<i_k1 % NumKVLdsBuffers>{}],
tile_elementwise_in(k_element_func, k_tiles[i_k1]));
__builtin_amdgcn_sched_barrier(0x00000001);
@@ -493,7 +493,7 @@ struct HstuAttentionNoSoftmaxFwdPipelineQRKSVSTrLoad
gemm_1(
o_acc,
get_slice_tile(p, sequence<0, i_k1 * kK1>{}, sequence<kM0, (i_k1 + 1) * kK1>{}),
v_lds_windows[number<i_k1 + 2>{}]);
v_lds_windows[number<(i_k1 + 2) % NumKVLdsBuffers>{}]);
});
} while(seqlen_k_curr < seqlen_k_end);

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

@@ -391,7 +391,7 @@ struct HstuAttentionWithSoftmaxFwdPipelineQRKSVSTrLoad
{
// STAGE 1, Gemm_0 ( S = Q@K )
static_for<0, k1_loops, 1>{}([&](auto i_k1) {
store_tile(k_lds_write_windows[i_k1],
store_tile(k_lds_write_windows[number<i_k1 % NumKVLdsBuffers>{}],
tile_elementwise_in(k_element_func, k_tiles[i_k1]));
__builtin_amdgcn_sched_barrier(0x00000001);
@@ -571,7 +571,7 @@ struct HstuAttentionWithSoftmaxFwdPipelineQRKSVSTrLoad
gemm_1(
o_acc,
get_slice_tile(p, sequence<0, i_k1 * kK1>{}, sequence<kM0, (i_k1 + 1) * kK1>{}),
v_lds_windows[number<i_k1 + 2>{}]);
v_lds_windows[number<(i_k1 + 2) % NumKVLdsBuffers>{}]);
});
} while(seqlen_k_curr < seqlen_k_end);