ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-05-16 02:54:21 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add static_assert and comments in the with_softmax pipelines

Browse Source
This commit is contained in:
Qianfeng Zhang
2025-11-28 14:49:33 +00:00
parent f952d3571c
commit d99493606e
2 changed files with 7 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -161,6 +161,9 @@ struct HstuAttentionWithSoftmaxFwdPipelineQRKSVS
constexpr index_t k1_loops = kN0 / kK1;
static_assert(k1_loops >= 2,
"k1_loops >= 2 required due to pre-storing two v_tiles to Lds");
constexpr auto NumKVLdsBuffers = Policy::template GetNumKVLdsBuffers<Problem>();
// Block GEMM
@@ -585,6 +588,7 @@ struct HstuAttentionWithSoftmaxFwdPipelineQRKSVS
auto p = cast_tile<PDataType>(tile_elementwise_in(p_compute_element_func, pcomp_tile));
// k1_loops >= 2 required
shuffle_tile(v_shuffled_tile, v_tiles[number<1>{}]);
store_tile(

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

@@ -160,6 +160,9 @@ struct HstuAttentionWithSoftmaxFwdPipelineQRKSVSTrLoad
constexpr index_t k1_loops = kN0 / kK1;
static_assert(k1_loops >= 2,
"k1_loops >= 2 required due to pre-storing two v_tiles to Lds");
constexpr auto NumKVLdsBuffers = Policy::template GetNumKVLdsBuffers<Problem>();
// Block GEMM