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

Remove the using of kSubQKHeaddim

Browse Source
This commit is contained in:
Qianfeng Zhang
2026-06-19 05:16:40 +00:00
parent a625506698
commit 2694adbd55
9 changed files with 40 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -591,7 +591,7 @@ struct HstuAttentionFwdKernel
num_tile_in_seqlen += 1;
};
if constexpr(HstuAttentionPipeline::kN1 < HstuAttentionPipeline::kSubQKHeaddim)
if constexpr(HstuAttentionPipeline::kN1 < HstuAttentionPipeline::kQKHeaddim)
{
#if HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM
return dim3(batch_size_,
@@ -619,7 +619,7 @@ struct HstuAttentionFwdKernel
CK_TILE_DEVICE static constexpr auto GetTileIndex(const Kargs& kargs)
{
if constexpr(HstuAttentionPipeline::kN1 < HstuAttentionPipeline::kSubQKHeaddim)
if constexpr(HstuAttentionPipeline::kN1 < HstuAttentionPipeline::kQKHeaddim)
{
const index_t num_tile_n1 =
ck_tile::integer_divide_ceil(kargs.hdim_v, HstuAttentionPipeline::kN1);

8
example/ck_tile/18_hstu_attention/hstu_attention_fwd_pipeline_policy.hpp
View File

@@ -15,6 +15,8 @@
#include "block_gemm_areg_bsmem_creg_v2_hack_1.hpp"
#include "block_gemm_areg_bsmem_trload_creg_v2_hack_1.hpp"
#include "hstu_attention_kernel_util.hpp"
namespace ck_tile {
struct HstuAttentionFwdPipelineQRKSVSPolicy
@@ -153,7 +155,7 @@ struct HstuAttentionFwdPipelineQRKSVSPolicy
constexpr index_t kKVector = GetAlignmentK<Problem>();
// for hdim96 and hdim160
if constexpr(kKPerBlock < Problem::HstuAttentionTileSetting::kSubQKHeaddim)
if constexpr(!detail::IsPerfectHeaddimSize(kKPerBlock))
{
return kKPerBlock * kNPerBlock;
}
@@ -211,7 +213,7 @@ struct HstuAttentionFwdPipelineQRKSVSPolicy
GetSingleSmemElementSpaceSize<Problem, kPipelineUseTrLoad>();
// for hdim96 and hdim160, use simplest layout
if constexpr(kKPerBlock < Problem::HstuAttentionTileSetting::kSubQKHeaddim)
if constexpr(!detail::IsPerfectHeaddimSize(kKPerBlock))
{
constexpr index_t KSingleSmemElementSpaceSize = kNPerBlock * kKPerBlock;
@@ -347,7 +349,7 @@ struct HstuAttentionFwdPipelineQRKSVSPolicy
constexpr index_t kKVector = GetAlignmentK<Problem>();
constexpr index_t OtherK = kKPerBlock / kKVector;
if constexpr(kKPerBlock == Problem::HstuAttentionTileSetting::kSubQKHeaddim)
if constexpr(detail::IsPerfectHeaddimSize(kKPerBlock))
// for kKPerBlock=32,64,128,256
{
static_assert((OtherK & (OtherK - 1)) == 0, "Check failed!");

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

@@ -552,7 +552,7 @@ struct HstuAttentionFwdSplitKVKernel
num_tile_in_seqlen += 1;
};
if constexpr(HstuAttentionPipeline::kN1 < HstuAttentionPipeline::kSubQKHeaddim)
if constexpr(HstuAttentionPipeline::kN1 < HstuAttentionPipeline::kQKHeaddim)
{
#if HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM
return dim3(batch_size_,
@@ -586,7 +586,7 @@ struct HstuAttentionFwdSplitKVKernel
return ck_tile::make_tuple(quotient, modulus);
};
if constexpr(HstuAttentionPipeline::kN1 < HstuAttentionPipeline::kSubQKHeaddim)
if constexpr(HstuAttentionPipeline::kN1 < HstuAttentionPipeline::kQKHeaddim)
{
const index_t num_tile_n1 =
ck_tile::integer_divide_ceil(kargs.hdim_v, HstuAttentionPipeline::kN1);

5
example/ck_tile/18_hstu_attention/hstu_attention_kernel_util.hpp
View File

@@ -107,6 +107,11 @@ CK_TILE_DEVICE static void scale_tile_in_pack(InOutDstrTensor& in_out_dstr_tenso
};
};
CK_TILE_DEVICE static constexpr bool IsPerfectHeaddimSize(index_t kHeaddim)
{
return kHeaddim > 0 && ((kHeaddim & (kHeaddim - 1)) == 0);
};
} // namespace detail
} // namespace ck_tile

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

@@ -29,13 +29,12 @@ struct HstuAttentionNoSoftmaxFwdPipelineQRKSVS
static constexpr index_t kBlockSize = Problem::kBlockSize;
static constexpr index_t kM0 = HstuAttentionTileSetting::kM0;
static constexpr index_t kN0 = HstuAttentionTileSetting::kN0;
static constexpr index_t kN0Sub = HstuAttentionTileSetting::kN0Sub;
static constexpr index_t kN1 = HstuAttentionTileSetting::kN1;
static constexpr index_t kK1 = HstuAttentionTileSetting::kK1;
static constexpr index_t kQKHeaddim = HstuAttentionTileSetting::kQKHeaddim;
static constexpr index_t kSubQKHeaddim = HstuAttentionTileSetting::kSubQKHeaddim;
static constexpr index_t kM0 = HstuAttentionTileSetting::kM0;
static constexpr index_t kN0 = HstuAttentionTileSetting::kN0;
static constexpr index_t kN0Sub = HstuAttentionTileSetting::kN0Sub;
static constexpr index_t kN1 = HstuAttentionTileSetting::kN1;
static constexpr index_t kK1 = HstuAttentionTileSetting::kK1;
static constexpr index_t kQKHeaddim = HstuAttentionTileSetting::kQKHeaddim;
static_assert(kQKHeaddim <= 256, "hdim bigger than 256 is not suitable for this pipeline!");

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

@@ -30,15 +30,14 @@ struct HstuAttentionNoSoftmaxFwdPipelineQRKSVSTrLoad
static constexpr index_t kBlockSize = Problem::kBlockSize;
static constexpr index_t kM0 = HstuAttentionTileSetting::kM0;
static constexpr index_t kN0 = HstuAttentionTileSetting::kN0;
static constexpr index_t kN0Sub = HstuAttentionTileSetting::kN0Sub;
static constexpr index_t kN1 = HstuAttentionTileSetting::kN1;
static constexpr index_t kK1 = HstuAttentionTileSetting::kK1;
static constexpr index_t kQKHeaddim = HstuAttentionTileSetting::kQKHeaddim;
static constexpr index_t kSubQKHeaddim = HstuAttentionTileSetting::kSubQKHeaddim;
static constexpr index_t kM0 = HstuAttentionTileSetting::kM0;
static constexpr index_t kN0 = HstuAttentionTileSetting::kN0;
static constexpr index_t kN0Sub = HstuAttentionTileSetting::kN0Sub;
static constexpr index_t kN1 = HstuAttentionTileSetting::kN1;
static constexpr index_t kK1 = HstuAttentionTileSetting::kK1;
static constexpr index_t kQKHeaddim = HstuAttentionTileSetting::kQKHeaddim;
static_assert(kSubQKHeaddim <= 256, "hdim bigger than 256 is not suitable for this pipeline!");
static_assert(kQKHeaddim <= 256, "hdim bigger than 256 is not suitable for this pipeline!");
static_assert(Problem::kUseSoftmax == false, "This pipeline only works with not-using softmax");

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

@@ -59,10 +59,6 @@ struct HstuAttentionFwdTileSettingClass
BlockTile::at(number<5>{}); // total length of K0, used for pipeline that need load Q at
// once (or repeately load Q as a whole tile)
static_assert((kN0 % kN0Sub == 0) && (kN0 % kK1 == 0), "Check failed!");
static constexpr index_t kSubQKHeaddim = ceil_to_qualified_tile_length(kQKHeaddim);
static_assert(kSubQKHeaddim % kN1 == 0, "Check failed!");
};
template <index_t kM_, // tile size in seqlen_q dimension

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

@@ -29,13 +29,12 @@ struct HstuAttentionWithSoftmaxFwdPipelineQRKSVS
static constexpr index_t kBlockSize = Problem::kBlockSize;
static constexpr index_t kM0 = HstuAttentionTileSetting::kM0;
static constexpr index_t kN0 = HstuAttentionTileSetting::kN0;
static constexpr index_t kN0Sub = HstuAttentionTileSetting::kN0Sub;
static constexpr index_t kN1 = HstuAttentionTileSetting::kN1;
static constexpr index_t kK1 = HstuAttentionTileSetting::kK1;
static constexpr index_t kQKHeaddim = HstuAttentionTileSetting::kQKHeaddim;
static constexpr index_t kSubQKHeaddim = HstuAttentionTileSetting::kSubQKHeaddim;
static constexpr index_t kM0 = HstuAttentionTileSetting::kM0;
static constexpr index_t kN0 = HstuAttentionTileSetting::kN0;
static constexpr index_t kN0Sub = HstuAttentionTileSetting::kN0Sub;
static constexpr index_t kN1 = HstuAttentionTileSetting::kN1;
static constexpr index_t kK1 = HstuAttentionTileSetting::kK1;
static constexpr index_t kQKHeaddim = HstuAttentionTileSetting::kQKHeaddim;
static_assert(kQKHeaddim <= 256, "hdim bigger than 256 is not suitable for this pipeline!");

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

@@ -29,15 +29,14 @@ struct HstuAttentionWithSoftmaxFwdPipelineQRKSVSTrLoad
static constexpr index_t kBlockSize = Problem::kBlockSize;
static constexpr index_t kM0 = HstuAttentionTileSetting::kM0;
static constexpr index_t kN0 = HstuAttentionTileSetting::kN0;
static constexpr index_t kN0Sub = HstuAttentionTileSetting::kN0Sub;
static constexpr index_t kN1 = HstuAttentionTileSetting::kN1;
static constexpr index_t kK1 = HstuAttentionTileSetting::kK1;
static constexpr index_t kQKHeaddim = HstuAttentionTileSetting::kQKHeaddim;
static constexpr index_t kSubQKHeaddim = HstuAttentionTileSetting::kSubQKHeaddim;
static constexpr index_t kM0 = HstuAttentionTileSetting::kM0;
static constexpr index_t kN0 = HstuAttentionTileSetting::kN0;
static constexpr index_t kN0Sub = HstuAttentionTileSetting::kN0Sub;
static constexpr index_t kN1 = HstuAttentionTileSetting::kN1;
static constexpr index_t kK1 = HstuAttentionTileSetting::kK1;
static constexpr index_t kQKHeaddim = HstuAttentionTileSetting::kQKHeaddim;
static_assert(kSubQKHeaddim <= 256, "hdim bigger than 256 is not suitable for this pipeline!");
static_assert(kQKHeaddim <= 256, "hdim bigger than 256 is not suitable for this pipeline!");
static_assert(Problem::kUseSoftmax == true, "This pipeline only works with using softmax");