diff --git a/example/ck_tile/18_hstu_attention/hstu_attention_fwd_kernel.hpp b/example/ck_tile/18_hstu_attention/hstu_attention_fwd_kernel.hpp index 1ea3309163..ea63042735 100644 --- a/example/ck_tile/18_hstu_attention/hstu_attention_fwd_kernel.hpp +++ b/example/ck_tile/18_hstu_attention/hstu_attention_fwd_kernel.hpp @@ -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); diff --git a/example/ck_tile/18_hstu_attention/hstu_attention_fwd_pipeline_policy.hpp b/example/ck_tile/18_hstu_attention/hstu_attention_fwd_pipeline_policy.hpp index 849d244463..9e78f3429d 100644 --- a/example/ck_tile/18_hstu_attention/hstu_attention_fwd_pipeline_policy.hpp +++ b/example/ck_tile/18_hstu_attention/hstu_attention_fwd_pipeline_policy.hpp @@ -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(); // 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(); // 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(); 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!"); diff --git a/example/ck_tile/18_hstu_attention/hstu_attention_fwd_splitkv_kernel.hpp b/example/ck_tile/18_hstu_attention/hstu_attention_fwd_splitkv_kernel.hpp index a91b667619..3b38b3e07c 100644 --- a/example/ck_tile/18_hstu_attention/hstu_attention_fwd_splitkv_kernel.hpp +++ b/example/ck_tile/18_hstu_attention/hstu_attention_fwd_splitkv_kernel.hpp @@ -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); diff --git a/example/ck_tile/18_hstu_attention/hstu_attention_kernel_util.hpp b/example/ck_tile/18_hstu_attention/hstu_attention_kernel_util.hpp index 45da8261a3..5c4d2327c2 100644 --- a/example/ck_tile/18_hstu_attention/hstu_attention_kernel_util.hpp +++ b/example/ck_tile/18_hstu_attention/hstu_attention_kernel_util.hpp @@ -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 diff --git a/example/ck_tile/18_hstu_attention/hstu_attention_no_softmax_fwd_pipeline.hpp b/example/ck_tile/18_hstu_attention/hstu_attention_no_softmax_fwd_pipeline.hpp index 529b4079a4..90440d998a 100644 --- a/example/ck_tile/18_hstu_attention/hstu_attention_no_softmax_fwd_pipeline.hpp +++ b/example/ck_tile/18_hstu_attention/hstu_attention_no_softmax_fwd_pipeline.hpp @@ -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!"); diff --git a/example/ck_tile/18_hstu_attention/hstu_attention_no_softmax_fwd_trload_pipeline.hpp b/example/ck_tile/18_hstu_attention/hstu_attention_no_softmax_fwd_trload_pipeline.hpp index 289b410abf..ec1bcdad2d 100644 --- a/example/ck_tile/18_hstu_attention/hstu_attention_no_softmax_fwd_trload_pipeline.hpp +++ b/example/ck_tile/18_hstu_attention/hstu_attention_no_softmax_fwd_trload_pipeline.hpp @@ -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"); diff --git a/example/ck_tile/18_hstu_attention/hstu_attention_tile_setting_define.hpp b/example/ck_tile/18_hstu_attention/hstu_attention_tile_setting_define.hpp index e306c4fd6b..74a3f30ec6 100644 --- a/example/ck_tile/18_hstu_attention/hstu_attention_tile_setting_define.hpp +++ b/example/ck_tile/18_hstu_attention/hstu_attention_tile_setting_define.hpp @@ -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