From 53e69c1404972e0c2230ff8c2292cfd303360fda Mon Sep 17 00:00:00 2001 From: Qianfeng Zhang Date: Tue, 30 Jun 2026 14:05:01 +0000 Subject: [PATCH] Use run-time checking (almost_invariant_seqlen) to replace pre-compiling setting (HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM) for switching grid-schedule --- .../ck_tile/18_hstu_attention/CMakeLists.txt | 4 - .../example_hstu_attention_fwd.cpp | 25 +++ ...stu_attention_batched_forward_dispatch.hpp | 10 +- ...ntion_batched_forward_splitkv_dispatch.hpp | 6 +- .../hstu_attention_fwd_kernel.hpp | 151 +++++++++++------- ...u_attention_fwd_splitkv_combine_kernel.hpp | 66 ++++---- .../hstu_attention_fwd_splitkv_kernel.hpp | 143 ++++++++++------- .../hstu_attention_group_forward_dispatch.hpp | 10 +- ...tention_group_forward_splitkv_dispatch.hpp | 21 ++- .../hstu_attention_host_util.hpp | 52 ++++++ ...hstu_attention_jagged_forward_dispatch.hpp | 5 + ...ention_jagged_forward_splitkv_dispatch.hpp | 14 +- .../hstu_attention_params.hpp | 4 + 13 files changed, 341 insertions(+), 170 deletions(-) diff --git a/example/ck_tile/18_hstu_attention/CMakeLists.txt b/example/ck_tile/18_hstu_attention/CMakeLists.txt index 5b53c79201..46f4da5bb8 100644 --- a/example/ck_tile/18_hstu_attention/CMakeLists.txt +++ b/example/ck_tile/18_hstu_attention/CMakeLists.txt @@ -12,10 +12,6 @@ set(EXAMPLE_HSTU_ATTENTION_COMPILE_OPTIONS) list(APPEND EXAMPLE_HSTU_ATTENTION_COMPILE_OPTIONS -Wno-undefined-func-template -Wno-float-equal -DCK_TILE_FLOAT_TO_BFLOAT16_DEFAULT=3) -if (DEFINED ENV{ASSUME_HIGHLY_VARIED_SEQLEN}) - list(APPEND EXAMPLE_HSTU_ATTENTION_COMPILE_OPTIONS -DHSTU_SCHED_BATCH_AS_FIRST_GRID_DIM=0) -endif() - if(GPU_TARGETS MATCHES "gfx95" AND NOT GPU_TARGETS MATCHES "gfx94" AND NOT GPU_TARGETS MATCHES "gfx90") ## disable slp-vectorize improve pipelines performance on gfx950 list(APPEND EXAMPLE_HSTU_ATTENTION_COMPILE_OPTIONS -DBUILD_HSTU_FOR_GFX95 -fno-slp-vectorize) diff --git a/example/ck_tile/18_hstu_attention/example_hstu_attention_fwd.cpp b/example/ck_tile/18_hstu_attention/example_hstu_attention_fwd.cpp index fd69e86795..3d4e597a47 100644 --- a/example/ck_tile/18_hstu_attention/example_hstu_attention_fwd.cpp +++ b/example/ck_tile/18_hstu_attention/example_hstu_attention_fwd.cpp @@ -277,6 +277,18 @@ bool run_no_group_hstu_forward(const ck_tile::ArgParser& arg_parser, bool is_jag phy_seqlen_kv = max_seqlen_kv; }; + int min_seqlen_q = std::numeric_limits::max(); + int min_seqlen_kv = std::numeric_limits::max(); + + if(is_jagged) + { + for(int i = 0; i < num_batch; i++) + { + min_seqlen_q = min(min_seqlen_q, seq_offsets_q[i + 1] - seq_offsets_q[i]); + min_seqlen_kv = min(min_seqlen_kv, seq_offsets_kv[i + 1] - seq_offsets_kv[i]); + }; + }; + long total_flops = 0; // estimate the total flops occurred, ignoring the scaling and SiLu @@ -379,6 +391,8 @@ bool run_no_group_hstu_forward(const ck_tile::ArgParser& arg_parser, bool is_jag params.seq_kv_offsets_ptr = seq_offsets_kv_dev.GetDeviceBuffer(); params.max_seqlen_q = max_seqlen_q; params.max_seqlen_kv = max_seqlen_kv; + params.min_seqlen_q = min_seqlen_q; + params.min_seqlen_kv = min_seqlen_kv; params.q_ptr = q_dev.GetDeviceBuffer(); params.k_ptr = k_dev.GetDeviceBuffer(); params.v_ptr = v_dev.GetDeviceBuffer(); @@ -798,6 +812,15 @@ bool run_group_hstu_forward(const ck_tile::ArgParser& arg_parser, int num_group) } }; + int min_seqlen_q = std::numeric_limits::max(); + int min_seqlen_kv = std::numeric_limits::max(); + + for(int i = 0; i < num_batch; i++) + { + min_seqlen_q = min(min_seqlen_q, seq_offsets_q[i + 1] - seq_offsets_q[i]); + min_seqlen_kv = min(min_seqlen_kv, seq_offsets_kv[i + 1] - seq_offsets_kv[i]); + }; + long total_flops = 0; // estimate the total flops occurred, ignoring the scaling and SILu @@ -900,6 +923,8 @@ bool run_group_hstu_forward(const ck_tile::ArgParser& arg_parser, int num_group) params.seq_kv_offsets_ptr = seq_offsets_kv_dev.GetDeviceBuffer(); params.max_seqlen_q = max_max_seqlen_q; params.max_seqlen_kv = max_max_seqlen_kv; + params.min_seqlen_q = min_seqlen_q; + params.min_seqlen_kv = min_seqlen_kv; params.q_ptr = q_dev.GetDeviceBuffer(); params.k_ptr = k_dev.GetDeviceBuffer(); params.v_ptr = v_dev.GetDeviceBuffer(); diff --git a/example/ck_tile/18_hstu_attention/hstu_attention_batched_forward_dispatch.hpp b/example/ck_tile/18_hstu_attention/hstu_attention_batched_forward_dispatch.hpp index 5adfe17811..2041f3e7eb 100644 --- a/example/ck_tile/18_hstu_attention/hstu_attention_batched_forward_dispatch.hpp +++ b/example/ck_tile/18_hstu_attention/hstu_attention_batched_forward_dispatch.hpp @@ -176,9 +176,13 @@ struct batched_forward_dispatch param.philox_offset); }(); - bool has_minfull_attn_seqlen = (param.min_full_attn_seqlen > 0); - dim3 kGridSize = HstuKernel::GridSize( - param.num_batch, param.num_head, param.seqlen_q, param.hdim_v, has_minfull_attn_seqlen); + bool has_minfull_attn_seqlen = (param.min_full_attn_seqlen > 0); + dim3 kGridSize = HstuKernel::GridSize(param.num_batch, + param.num_head, + param.seqlen_q, + param.hdim_v, + true /* almost_invariant_seqlen */, + has_minfull_attn_seqlen); dim3 kBlockSize = HstuKernel::BlockSize(); constexpr ck_tile::index_t kBlockPerCu = HstuKernel::kBlockPerCu; diff --git a/example/ck_tile/18_hstu_attention/hstu_attention_batched_forward_splitkv_dispatch.hpp b/example/ck_tile/18_hstu_attention/hstu_attention_batched_forward_splitkv_dispatch.hpp index b9f532e85f..62f11cc00c 100644 --- a/example/ck_tile/18_hstu_attention/hstu_attention_batched_forward_splitkv_dispatch.hpp +++ b/example/ck_tile/18_hstu_attention/hstu_attention_batched_forward_splitkv_dispatch.hpp @@ -316,6 +316,7 @@ struct batched_forward_splitkv_dispatch param.seqlen_q, param.hdim_v, ws.num_splits, + true, // almost_invariant_seqlen has_minfull_attn_seqlen); dim3 kBlockSize = HstuKernel::BlockSize(); constexpr ck_tile::index_t kBlockPerCu = HstuKernel::kBlockPerCu; @@ -347,8 +348,9 @@ struct batched_forward_splitkv_dispatch param.hdim_v); }(); - dim3 kGridSize = HstuKernel::GridSize(param.num_batch, param.num_head, param.seqlen_q); - dim3 kBlockSize = HstuKernel::BlockSize(); + dim3 kGridSize = HstuKernel::GridSize( + param.num_batch, param.num_head, param.seqlen_q, true /* almost_invariant_seqlen */); + dim3 kBlockSize = HstuKernel::BlockSize(); constexpr ck_tile::index_t kBlockPerCu = HstuKernel::kBlockPerCu; (void)ck_tile::launch_kernel( 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 dac1e3a888..fd7e285704 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 @@ -14,10 +14,6 @@ #include "hstu_block_masking.hpp" #include "hstu_attention_kernel_util.hpp" -#ifndef HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM -#define HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM 1 -#endif - // S[seqlen_q, seqlen_k] = Q[seqlen_q, hdim_q] @ K[seqlen_k, hdim_q] // S'[seqlen_q, seqlen_k] = S[seqlen_q, seqlen_k] * Scale[1] // S''[seqlen_q, seqlen_k] = S'[seqlen_q, seqlen_k] + Bias[seqlen_q, seqlen_k] @@ -101,6 +97,8 @@ struct HstuAttentionFwdKernel float scale_s; // scaling value exerted on the immediate Q@K result float scale_p; // scaling value exerted on the SiLU result + bool almost_invariant_seqlen; // should always be true for batched mode + ck_tile::index_t contextual_seqlen; ck_tile::index_t window_size; ck_tile::index_t min_full_attn_seqlen; @@ -138,6 +136,8 @@ struct HstuAttentionFwdKernel float scale_s; // scaling value exerted on the immediate Q@K result float scale_p; // scaling value exerted on the SiLU result + bool almost_invariant_seqlen; + ck_tile::index_t contextual_seqlen; ck_tile::index_t window_size; ck_tile::index_t min_full_attn_seqlen; @@ -177,6 +177,8 @@ struct HstuAttentionFwdKernel float scale_s; // scaling value exerted on the immediate Q@K result float scale_p; // scaling value exerted on the SiLU result + bool almost_invariant_seqlen; + int32_t contextual_seqlen; // to be set by the per-group contextual_seqlen int32_t window_size; // to be set by the per-group window_size int32_t min_full_attn_seqlen; // to be set by the per-group min_full_attn_seqlen @@ -356,6 +358,7 @@ struct HstuAttentionFwdKernel num_head, scale_s, attn_scale ? attn_scale : 1.0f / static_cast(seqlen_q), + true, // almost_invariant_seqlen contextual_seqlen, window_size, min_full_attn_seqlen}, // args for common karg @@ -402,6 +405,7 @@ struct HstuAttentionFwdKernel ck_tile::index_t num_head, float scale_s, float attn_scale, + bool almost_invariant_seqlen, ck_tile::index_t seq_stride_q, ck_tile::index_t seq_stride_k, ck_tile::index_t seq_stride_v, @@ -445,6 +449,7 @@ struct HstuAttentionFwdKernel num_head, scale_s, attn_scale ? attn_scale : 1.0f / static_cast(max_seqlen_q), + almost_invariant_seqlen, contextual_seqlen, window_size, min_full_attn_seqlen}, // args for common karg @@ -493,6 +498,7 @@ struct HstuAttentionFwdKernel ck_tile::index_t hdim_v, ck_tile::index_t num_head, float scale_s, + bool almost_invariant_seqlen, ck_tile::index_t seq_stride_q, ck_tile::index_t seq_stride_k, ck_tile::index_t seq_stride_v, @@ -534,9 +540,10 @@ struct HstuAttentionFwdKernel num_head, scale_s, 1.0f, // to be set according to the per-group attn_scale and max_seqlen - 0, // to be set by the per-group contextual_seqlen - 0, // to be set by the per-group window_size - 0, // to be set by the per-group min_full_attn_seqlen + almost_invariant_seqlen, + 0, // to be set by the per-group contextual_seqlen + 0, // to be set by the per-group window_size + 0, // to be set by the per-group min_full_attn_seqlen reinterpret_cast(group_max_seqlen_q_ptr), reinterpret_cast(group_contextual_seqlen_ptr), reinterpret_cast(group_window_size_ptr), @@ -571,6 +578,7 @@ struct HstuAttentionFwdKernel ck_tile::index_t nhead_, ck_tile::index_t seqlen_, ck_tile::index_t hdim_v_, + bool almost_invariant_seqlen, bool has_minfull_attn_seqlen = false) { // The Q sequence [0, seqlen) will be split to two parts for allocating workgroups: @@ -589,27 +597,33 @@ struct HstuAttentionFwdKernel num_tile_in_seqlen += 1; }; - if constexpr(HstuAttentionPipeline::kN1 < HstuAttentionPipeline::kQKHeaddim) + if(almost_invariant_seqlen) { -#if HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM - return dim3(batch_size_, - nhead_, - num_tile_in_seqlen * - ck_tile::integer_divide_ceil(hdim_v_, HstuAttentionPipeline::kN1)); -#else - return dim3(num_tile_in_seqlen * - ck_tile::integer_divide_ceil(hdim_v_, HstuAttentionPipeline::kN1), - nhead_, - batch_size_); -#endif + if constexpr(HstuAttentionPipeline::kN1 < HstuAttentionPipeline::kQKHeaddim) + { + return dim3(batch_size_, + nhead_, + num_tile_in_seqlen * + ck_tile::integer_divide_ceil(hdim_v_, HstuAttentionPipeline::kN1)); + } + else + { + return dim3(batch_size_, nhead_, num_tile_in_seqlen); + } } else { -#if HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM - return dim3(batch_size_, nhead_, num_tile_in_seqlen); -#else - return dim3(num_tile_in_seqlen, nhead_, batch_size_); -#endif + if constexpr(HstuAttentionPipeline::kN1 < HstuAttentionPipeline::kQKHeaddim) + { + return dim3(num_tile_in_seqlen * + ck_tile::integer_divide_ceil(hdim_v_, HstuAttentionPipeline::kN1), + nhead_, + batch_size_); + } + else + { + return dim3(num_tile_in_seqlen, nhead_, batch_size_); + } } } @@ -620,52 +634,65 @@ struct HstuAttentionFwdKernel const index_t num_tile_n1 = ck_tile::integer_divide_ceil(kargs.hdim_v, HstuAttentionPipeline::kN1); -#if HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM - const index_t i_batch = blockIdx.x; - const index_t i_nhead = blockIdx.y; - const index_t i_block = blockIdx.z; -#else - const index_t i_block = blockIdx.x; - const index_t i_nhead = blockIdx.y; - const index_t i_batch = blockIdx.z; -#endif + if(kargs.almost_invariant_seqlen) + { + const index_t i_batch = blockIdx.x; + const index_t i_nhead = blockIdx.y; + const index_t i_block = blockIdx.z; - const auto f = [](index_t dividend, index_t divisor) { - index_t quotient = dividend / divisor; - index_t modulus = dividend - quotient * divisor; - return ck_tile::make_tuple(quotient, modulus); - }; + const auto f = [](index_t dividend, index_t divisor) { + index_t quotient = dividend / divisor; + index_t modulus = dividend - quotient * divisor; + return ck_tile::make_tuple(quotient, modulus); + }; -#if HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM - auto [i_tile_m, i_tile_n] = f(i_block, num_tile_n1); - i_tile_m = gridDim.z / num_tile_n1 - 1 - i_tile_m; -#else - const auto [i_tile_m, i_tile_n] = f(i_block, num_tile_n1); -#endif + auto [i_tile_m, i_tile_n] = f(i_block, num_tile_n1); + i_tile_m = gridDim.z / num_tile_n1 - 1 - i_tile_m; - return ck_tile::make_tuple(i_tile_m, i_tile_n, i_nhead, i_batch); + return ck_tile::make_tuple(i_tile_m, i_tile_n, i_nhead, i_batch); + } + else + { + const index_t i_block = blockIdx.x; + const index_t i_nhead = blockIdx.y; + const index_t i_batch = blockIdx.z; + + const auto f = [](index_t dividend, index_t divisor) { + index_t quotient = dividend / divisor; + index_t modulus = dividend - quotient * divisor; + return ck_tile::make_tuple(quotient, modulus); + }; + + const auto [i_tile_m, i_tile_n] = f(i_block, num_tile_n1); + + return ck_tile::make_tuple(i_tile_m, i_tile_n, i_nhead, i_batch); + } } else { -#if HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM - const index_t i_batch = blockIdx.x; - const index_t i_nhead = blockIdx.y; - const index_t i_block = blockIdx.z; -#else - const index_t i_block = blockIdx.x; - const index_t i_nhead = blockIdx.y; - const index_t i_batch = blockIdx.z; -#endif + if(kargs.almost_invariant_seqlen) + { + const index_t i_batch = blockIdx.x; + const index_t i_nhead = blockIdx.y; + const index_t i_block = blockIdx.z; -#if HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM - index_t i_tile_m = i_block; - i_tile_m = gridDim.z - 1 - i_tile_m; -#else - const index_t i_tile_m = i_block; -#endif - const index_t i_tile_n = 0; + index_t i_tile_m = i_block; + i_tile_m = gridDim.z - 1 - i_tile_m; + const index_t i_tile_n = 0; - return ck_tile::make_tuple(i_tile_m, i_tile_n, i_nhead, i_batch); + return ck_tile::make_tuple(i_tile_m, i_tile_n, i_nhead, i_batch); + } + else + { + const index_t i_block = blockIdx.x; + const index_t i_nhead = blockIdx.y; + const index_t i_batch = blockIdx.z; + + const index_t i_tile_m = i_block; + const index_t i_tile_n = 0; + + return ck_tile::make_tuple(i_tile_m, i_tile_n, i_nhead, i_batch); + } } } diff --git a/example/ck_tile/18_hstu_attention/hstu_attention_fwd_splitkv_combine_kernel.hpp b/example/ck_tile/18_hstu_attention/hstu_attention_fwd_splitkv_combine_kernel.hpp index c93102f1f0..b558774612 100644 --- a/example/ck_tile/18_hstu_attention/hstu_attention_fwd_splitkv_combine_kernel.hpp +++ b/example/ck_tile/18_hstu_attention/hstu_attention_fwd_splitkv_combine_kernel.hpp @@ -14,10 +14,6 @@ #include "hstu_block_masking.hpp" -#ifndef HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM -#define HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM 1 -#endif - // S[seqlen_q, seqlen_k] = Q[seqlen_q, hdim_q] @ K[seqlen_k, hdim_q] // S'[seqlen_q, seqlen_k] = S[seqlen_q, seqlen_k] * Scale[1] // S''[seqlen_q, seqlen_k] = S'[seqlen_q, seqlen_k] + Bias[seqlen_q, seqlen_k] @@ -71,6 +67,8 @@ struct HstuAttentionFwdSplitKVCombineKernel ck_tile::index_t num_head; ck_tile::index_t num_splits; ck_tile::index_t hdim_v; + + bool almost_invariant_seqlen; // should always be true for batched mode }; struct HstuAttentionJaggedCombineBaseKargs @@ -87,6 +85,8 @@ struct HstuAttentionFwdSplitKVCombineKernel ck_tile::index_t hdim_v; ck_tile::index_t seqlen_q; + + bool almost_invariant_seqlen; }; struct HstuAttentionCombineSoftmaxKargs @@ -160,9 +160,10 @@ struct HstuAttentionFwdSplitKVCombineKernel seqlen_q, num_head, num_splits, - hdim_v}, - {}, // place holder for softmax - {}, // place holder for LSE + hdim_v, + true}, // almost_invariant_seqlen + {}, // place holder for softmax + {}, // place holder for LSE }; if constexpr(kUseSoftmax) @@ -193,7 +194,8 @@ struct HstuAttentionFwdSplitKVCombineKernel const void* seq_q_offsets_ptr, ck_tile::index_t num_head, ck_tile::index_t num_splits, // number of splitted seqlen_kv - ck_tile::index_t hdim_v) + ck_tile::index_t hdim_v, + bool almost_invariant_seqlen) { Kargs kargs{ {o_acc_ptr, @@ -204,7 +206,8 @@ struct HstuAttentionFwdSplitKVCombineKernel num_head, num_splits, hdim_v, - 0 /* seqlen_q will be updated later*/}, + 0, /* seqlen_q will be updated later*/ + almost_invariant_seqlen}, {}, // place holder for softmax {}, // place holder for LSE }; @@ -223,33 +226,40 @@ struct HstuAttentionFwdSplitKVCombineKernel return kargs; } - CK_TILE_HOST static constexpr auto - GridSize(ck_tile::index_t batch_size_, ck_tile::index_t nhead_, ck_tile::index_t seqlen_) + CK_TILE_HOST static constexpr auto GridSize(ck_tile::index_t batch_size_, + ck_tile::index_t nhead_, + ck_tile::index_t seqlen_, + bool almost_invariant_seqlen) { ck_tile::index_t num_tile_in_seqlen = ck_tile::integer_divide_ceil(seqlen_, HstuAttentionPipeline::kM); -#if HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM - return dim3(batch_size_, nhead_, num_tile_in_seqlen); -#else - return dim3(num_tile_in_seqlen, nhead_, batch_size_); -#endif + if(almost_invariant_seqlen) + { + return dim3(batch_size_, nhead_, num_tile_in_seqlen); + } + else + { + return dim3(num_tile_in_seqlen, nhead_, batch_size_); + } } CK_TILE_DEVICE static constexpr auto GetTileIndex(const Kargs& kargs) { - ignore = kargs; - -#if HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM - const index_t i_batch = blockIdx.x; - const index_t i_nhead = blockIdx.y; - const index_t i_tile_m = blockIdx.z; -#else - const index_t i_tile_m = blockIdx.x; - const index_t i_nhead = blockIdx.y; - const index_t i_batch = blockIdx.z; -#endif - return ck_tile::make_tuple(i_tile_m, i_nhead, i_batch); + if(kargs.almost_invariant_seqlen) + { + const index_t i_batch = blockIdx.x; + const index_t i_nhead = blockIdx.y; + const index_t i_tile_m = blockIdx.z; + return ck_tile::make_tuple(i_tile_m, i_nhead, i_batch); + } + else + { + const index_t i_tile_m = blockIdx.x; + const index_t i_nhead = blockIdx.y; + const index_t i_batch = blockIdx.z; + return ck_tile::make_tuple(i_tile_m, i_nhead, i_batch); + } } CK_TILE_HOST static constexpr auto BlockSize() 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 add49574de..551d3e6d46 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 @@ -14,10 +14,6 @@ #include "hstu_block_masking.hpp" #include "hstu_attention_kernel_util.hpp" -#ifndef HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM -#define HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM 1 -#endif - // S[seqlen_q, seqlen_k] = Q[seqlen_q, hdim_q] @ K[seqlen_k, hdim_q] // S'[seqlen_q, seqlen_k] = S[seqlen_q, seqlen_k] * Scale[1] // S''[seqlen_q, seqlen_k] = S'[seqlen_q, seqlen_k] + Bias[seqlen_q, seqlen_k] @@ -99,6 +95,8 @@ struct HstuAttentionFwdSplitKVKernel float scale_s; // scaling value exerted on the immediate Q@K result float scale_p; // scaling value exerted on the SiLU result + bool almost_invariant_seqlen; // should always be true for batched mode + ck_tile::index_t contextual_seqlen; ck_tile::index_t window_size; ck_tile::index_t min_full_attn_seqlen; @@ -135,6 +133,8 @@ struct HstuAttentionFwdSplitKVKernel float scale_s; // scaling value exerted on the immediate Q@K result float scale_p; // scaling value exerted on the SiLU result + bool almost_invariant_seqlen; + ck_tile::index_t contextual_seqlen; ck_tile::index_t window_size; ck_tile::index_t min_full_attn_seqlen; @@ -173,6 +173,8 @@ struct HstuAttentionFwdSplitKVKernel float scale_s; // scaling value exerted on the immediate Q@K result float scale_p; // scaling value exerted on the SiLU result + bool almost_invariant_seqlen; + int32_t contextual_seqlen; // to be set by the per-group contextual_seqlen int32_t window_size; // to be set by the per-group window_size int32_t min_full_attn_seqlen; // to be set by the per-group min_full_attn_seqlen @@ -334,6 +336,7 @@ struct HstuAttentionFwdSplitKVKernel num_head, scale_s, attn_scale ? attn_scale : 1.0f / static_cast(seqlen_q), + true, // almost_invariant_seqlen contextual_seqlen, window_size, min_full_attn_seqlen}, // args for common karg @@ -377,6 +380,7 @@ struct HstuAttentionFwdSplitKVKernel ck_tile::index_t num_head, float scale_s, float attn_scale, + bool almost_invariant_seqlen, ck_tile::index_t seq_stride_q, ck_tile::index_t seq_stride_k, ck_tile::index_t seq_stride_v, @@ -415,6 +419,7 @@ struct HstuAttentionFwdSplitKVKernel num_head, scale_s, attn_scale ? attn_scale : 1.0f / static_cast(max_seqlen_q), + almost_invariant_seqlen, contextual_seqlen, window_size, min_full_attn_seqlen}, // args for common karg @@ -461,6 +466,7 @@ struct HstuAttentionFwdSplitKVKernel ck_tile::index_t hdim_v, ck_tile::index_t num_head, float scale_s, + bool almost_invariant_seqlen, ck_tile::index_t seq_stride_q, ck_tile::index_t seq_stride_k, ck_tile::index_t seq_stride_v, @@ -497,9 +503,10 @@ struct HstuAttentionFwdSplitKVKernel num_head, scale_s, 1.0f, // to be set according to the per-group attn_scale and max_seqlen - 0, // to be set by the per-group contextual_seqlen - 0, // to be set by the per-group window_size - 0, // to be set by the per-group min_full_attn_seqlen + almost_invariant_seqlen, + 0, // to be set by the per-group contextual_seqlen + 0, // to be set by the per-group window_size + 0, // to be set by the per-group min_full_attn_seqlen reinterpret_cast(group_max_seqlen_q_ptr), reinterpret_cast(group_contextual_seqlen_ptr), reinterpret_cast(group_window_size_ptr), @@ -532,6 +539,7 @@ struct HstuAttentionFwdSplitKVKernel ck_tile::index_t seqlen_, ck_tile::index_t hdim_v_, ck_tile::index_t num_splits, + bool almost_invariant_seqlen, bool has_minfull_attn_seqlen = false) { // The Q sequence [0, seqlen) will be split to two parts for allocating workgroups: @@ -550,27 +558,33 @@ struct HstuAttentionFwdSplitKVKernel num_tile_in_seqlen += 1; }; - if constexpr(HstuAttentionPipeline::kN1 < HstuAttentionPipeline::kQKHeaddim) + if(almost_invariant_seqlen) { -#if HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM - return dim3(batch_size_, - nhead_, - num_tile_in_seqlen * num_splits * - ck_tile::integer_divide_ceil(hdim_v_, HstuAttentionPipeline::kN1)); -#else - return dim3(num_tile_in_seqlen * num_splits * - ck_tile::integer_divide_ceil(hdim_v_, HstuAttentionPipeline::kN1), - nhead_, - batch_size_); -#endif + if constexpr(HstuAttentionPipeline::kN1 < HstuAttentionPipeline::kQKHeaddim) + { + return dim3(batch_size_, + nhead_, + num_tile_in_seqlen * num_splits * + ck_tile::integer_divide_ceil(hdim_v_, HstuAttentionPipeline::kN1)); + } + else + { + return dim3(batch_size_, nhead_, num_tile_in_seqlen * num_splits); + } } else { -#if HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM - return dim3(batch_size_, nhead_, num_tile_in_seqlen * num_splits); -#else - return dim3(num_tile_in_seqlen * num_splits, nhead_, batch_size_); -#endif + if constexpr(HstuAttentionPipeline::kN1 < HstuAttentionPipeline::kQKHeaddim) + { + return dim3(num_tile_in_seqlen * num_splits * + ck_tile::integer_divide_ceil(hdim_v_, HstuAttentionPipeline::kN1), + nhead_, + batch_size_); + } + else + { + return dim3(num_tile_in_seqlen * num_splits, nhead_, batch_size_); + } } } @@ -587,51 +601,58 @@ struct HstuAttentionFwdSplitKVKernel const index_t num_tile_n1 = ck_tile::integer_divide_ceil(kargs.hdim_v, HstuAttentionPipeline::kN1); -#if HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM - const index_t i_batch = blockIdx.x; - const index_t i_nhead = blockIdx.y; - const index_t i_block = blockIdx.z; -#else - const index_t i_block = blockIdx.x; - const index_t i_nhead = blockIdx.y; - const index_t i_batch = blockIdx.z; -#endif + if(kargs.almost_invariant_seqlen) + { + const index_t i_batch = blockIdx.x; + const index_t i_nhead = blockIdx.y; + const index_t i_block = blockIdx.z; -#if HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM - auto [i_tile_m_i_split, i_tile_n] = f(i_block, num_tile_n1); - auto [i_tile_m, i_split] = f(i_tile_m_i_split, kargs.num_splits); + auto [i_tile_m_i_split, i_tile_n] = f(i_block, num_tile_n1); + auto [i_tile_m, i_split] = f(i_tile_m_i_split, kargs.num_splits); - i_tile_m = gridDim.z / num_tile_n1 / kargs.num_splits - 1 - i_tile_m; -#else - auto [i_tile_m_i_split, i_tile_n] = f(i_block, num_tile_n1); - auto [i_tile_m, i_split] = f(i_tile_m_i_split, kargs.num_splits); -#endif + i_tile_m = gridDim.z / num_tile_n1 / kargs.num_splits - 1 - i_tile_m; - return ck_tile::make_tuple(i_tile_m, i_tile_n, i_nhead, i_batch, i_split); + return ck_tile::make_tuple(i_tile_m, i_tile_n, i_nhead, i_batch, i_split); + } + else + { + const index_t i_block = blockIdx.x; + const index_t i_nhead = blockIdx.y; + const index_t i_batch = blockIdx.z; + + auto [i_tile_m_i_split, i_tile_n] = f(i_block, num_tile_n1); + auto [i_tile_m, i_split] = f(i_tile_m_i_split, kargs.num_splits); + + return ck_tile::make_tuple(i_tile_m, i_tile_n, i_nhead, i_batch, i_split); + } } else { -#if HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM - const index_t i_batch = blockIdx.x; - const index_t i_nhead = blockIdx.y; - const index_t i_block = blockIdx.z; -#else - const index_t i_block = blockIdx.x; - const index_t i_nhead = blockIdx.y; - const index_t i_batch = blockIdx.z; -#endif + if(kargs.almost_invariant_seqlen) + { + const index_t i_batch = blockIdx.x; + const index_t i_nhead = blockIdx.y; + const index_t i_block = blockIdx.z; -#if HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM - index_t i_tile_m_i_split = i_block; - auto [i_tile_m, i_split] = f(i_tile_m_i_split, kargs.num_splits); - i_tile_m = gridDim.z / kargs.num_splits - 1 - i_tile_m; -#else - index_t i_tile_m_i_split = i_block; - auto [i_tile_m, i_split] = f(i_tile_m_i_split, kargs.num_splits); -#endif - const index_t i_tile_n = 0; + index_t i_tile_m_i_split = i_block; + auto [i_tile_m, i_split] = f(i_tile_m_i_split, kargs.num_splits); + i_tile_m = gridDim.z / kargs.num_splits - 1 - i_tile_m; + const index_t i_tile_n = 0; - return ck_tile::make_tuple(i_tile_m, i_tile_n, i_nhead, i_batch, i_split); + return ck_tile::make_tuple(i_tile_m, i_tile_n, i_nhead, i_batch, i_split); + } + else + { + const index_t i_block = blockIdx.x; + const index_t i_nhead = blockIdx.y; + const index_t i_batch = blockIdx.z; + + index_t i_tile_m_i_split = i_block; + auto [i_tile_m, i_split] = f(i_tile_m_i_split, kargs.num_splits); + const index_t i_tile_n = 0; + + return ck_tile::make_tuple(i_tile_m, i_tile_n, i_nhead, i_batch, i_split); + } } } diff --git a/example/ck_tile/18_hstu_attention/hstu_attention_group_forward_dispatch.hpp b/example/ck_tile/18_hstu_attention/hstu_attention_group_forward_dispatch.hpp index caec9ab7e1..0320f55ff9 100644 --- a/example/ck_tile/18_hstu_attention/hstu_attention_group_forward_dispatch.hpp +++ b/example/ck_tile/18_hstu_attention/hstu_attention_group_forward_dispatch.hpp @@ -124,6 +124,8 @@ struct group_forward_dispatch template static void RunWithKernel(HstuAttentionGroupFwdParams& param, hipStream_t stream) { + bool almost_invariant_seqlen = is_almost_invariant_seqlen(param); + const auto kargs = [&] { return HstuKernel::MakeKargs(param.q_ptr, param.k_ptr, @@ -144,6 +146,7 @@ struct group_forward_dispatch param.hdim_v, param.num_head, param.scale_s, + almost_invariant_seqlen, param.seq_stride_q, param.seq_stride_k, param.seq_stride_v, @@ -162,8 +165,11 @@ struct group_forward_dispatch param.philox_offset); }(); - dim3 kGridSize = - HstuKernel::GridSize(param.num_batch, param.num_head, param.max_seqlen_q, param.hdim_v); + dim3 kGridSize = HstuKernel::GridSize(param.num_batch, + param.num_head, + param.max_seqlen_q, + param.hdim_v, + almost_invariant_seqlen); dim3 kBlockSize = HstuKernel::BlockSize(); constexpr ck_tile::index_t kBlockPerCu = HstuKernel::kBlockPerCu; diff --git a/example/ck_tile/18_hstu_attention/hstu_attention_group_forward_splitkv_dispatch.hpp b/example/ck_tile/18_hstu_attention/hstu_attention_group_forward_splitkv_dispatch.hpp index 36f20ccc07..35fdb95cb1 100644 --- a/example/ck_tile/18_hstu_attention/hstu_attention_group_forward_splitkv_dispatch.hpp +++ b/example/ck_tile/18_hstu_attention/hstu_attention_group_forward_splitkv_dispatch.hpp @@ -265,6 +265,8 @@ struct group_forward_splitkv_dispatch HIP_CHECK_ERROR(hipMallocAsync(&ws.lse_acc_ptr, workspace_bytes, stream)); } + bool almost_invariant_seqlen = is_almost_invariant_seqlen(param); + const auto kargs = [&] { return HstuKernel::MakeKargs(param.q_ptr, param.k_ptr, @@ -286,6 +288,7 @@ struct group_forward_splitkv_dispatch param.hdim_v, param.num_head, param.scale_s, + almost_invariant_seqlen, param.seq_stride_q, param.seq_stride_k, param.seq_stride_v, @@ -300,8 +303,12 @@ struct group_forward_splitkv_dispatch param.philox_offset); }(); - dim3 kGridSize = HstuKernel::GridSize( - param.num_batch, param.num_head, param.max_seqlen_q, param.hdim_v, ws.num_splits); + dim3 kGridSize = HstuKernel::GridSize(param.num_batch, + param.num_head, + param.max_seqlen_q, + param.hdim_v, + ws.num_splits, + almost_invariant_seqlen); dim3 kBlockSize = HstuKernel::BlockSize(); constexpr ck_tile::index_t kBlockPerCu = HstuKernel::kBlockPerCu; @@ -315,6 +322,8 @@ struct group_forward_splitkv_dispatch SplitkvWorkspace& ws, hipStream_t stream) { + bool almost_invariant_seqlen = is_almost_invariant_seqlen_q(param); + const auto kargs = [&] { return HstuKernel::MakeKargs(ws.o_acc_ptr, ws.lse_acc_ptr, @@ -327,11 +336,13 @@ struct group_forward_splitkv_dispatch param.seq_q_offsets_ptr, param.num_head, ws.num_splits, - param.hdim_v); + param.hdim_v, + almost_invariant_seqlen); }(); - dim3 kGridSize = HstuKernel::GridSize(param.num_batch, param.num_head, param.max_seqlen_q); - dim3 kBlockSize = HstuKernel::BlockSize(); + dim3 kGridSize = HstuKernel::GridSize( + param.num_batch, param.num_head, param.max_seqlen_q, almost_invariant_seqlen); + dim3 kBlockSize = HstuKernel::BlockSize(); constexpr ck_tile::index_t kBlockPerCu = HstuKernel::kBlockPerCu; (void)ck_tile::launch_kernel( diff --git a/example/ck_tile/18_hstu_attention/hstu_attention_host_util.hpp b/example/ck_tile/18_hstu_attention/hstu_attention_host_util.hpp index 26d266d7f0..6965e25d96 100644 --- a/example/ck_tile/18_hstu_attention/hstu_attention_host_util.hpp +++ b/example/ck_tile/18_hstu_attention/hstu_attention_host_util.hpp @@ -9,6 +9,7 @@ #include #include +#include "hstu_attention_params.hpp" #define HSTU_CHECK(COND, ERR) \ if(!(COND)) \ @@ -30,3 +31,54 @@ static inline int get_number_of_cu() return props.multiProcessorCount; } + +static inline bool is_almost_invariant_seqlen(HstuAttentionNoGroupFwdParams& param) +{ + float threshold = 0.7f; + + if(param.is_jagged) + { + bool res = (static_cast(param.min_seqlen_q) / param.max_seqlen_q) > threshold; + if(param.is_cross_attention) + res = res && + ((static_cast(param.min_seqlen_kv) / param.max_seqlen_kv) > threshold); + + return res; + } + else + return true; +}; + +static inline bool is_almost_invariant_seqlen(HstuAttentionGroupFwdParams& param) +{ + float threshold = 0.7f; + + bool res = (static_cast(param.min_seqlen_q) / param.max_seqlen_q) > threshold; + if(param.is_cross_attention) + res = res && ((static_cast(param.min_seqlen_kv) / param.max_seqlen_kv) > threshold); + + return res; +}; + +static inline bool is_almost_invariant_seqlen_q(HstuAttentionNoGroupFwdParams& param) +{ + float threshold = 0.7f; + + if(param.is_jagged) + { + bool res = (static_cast(param.min_seqlen_q) / param.max_seqlen_q) > threshold; + + return res; + } + else + return true; +}; + +static inline bool is_almost_invariant_seqlen_q(HstuAttentionGroupFwdParams& param) +{ + float threshold = 0.7f; + + bool res = (static_cast(param.min_seqlen_q) / param.max_seqlen_q) > threshold; + + return res; +}; diff --git a/example/ck_tile/18_hstu_attention/hstu_attention_jagged_forward_dispatch.hpp b/example/ck_tile/18_hstu_attention/hstu_attention_jagged_forward_dispatch.hpp index dc7c749fcf..98871c9ba1 100644 --- a/example/ck_tile/18_hstu_attention/hstu_attention_jagged_forward_dispatch.hpp +++ b/example/ck_tile/18_hstu_attention/hstu_attention_jagged_forward_dispatch.hpp @@ -23,6 +23,7 @@ #include "hstu_attention_epilogue.hpp" #include "hstu_attention_jagged_forward_splitkv_dispatch.hpp" +#include "hstu_attention_host_util.hpp" template static void RunWithKernel(HstuAttentionNoGroupFwdParams& param, hipStream_t stream) { + bool almost_invariant_seqlen = is_almost_invariant_seqlen(param); + const auto kargs = [&] { return HstuKernel::MakeKargs(param.q_ptr, param.k_ptr, @@ -140,6 +143,7 @@ struct jagged_forward_dispatch param.num_head, param.scale_s, param.attn_scale, + almost_invariant_seqlen, param.seq_stride_q, param.seq_stride_k, param.seq_stride_v, @@ -166,6 +170,7 @@ struct jagged_forward_dispatch param.num_head, param.max_seqlen_q, param.hdim_v, + almost_invariant_seqlen, has_minfull_attn_seqlen); dim3 kBlockSize = HstuKernel::BlockSize(); constexpr ck_tile::index_t kBlockPerCu = HstuKernel::kBlockPerCu; diff --git a/example/ck_tile/18_hstu_attention/hstu_attention_jagged_forward_splitkv_dispatch.hpp b/example/ck_tile/18_hstu_attention/hstu_attention_jagged_forward_splitkv_dispatch.hpp index bcaa3ca76c..708e16047c 100644 --- a/example/ck_tile/18_hstu_attention/hstu_attention_jagged_forward_splitkv_dispatch.hpp +++ b/example/ck_tile/18_hstu_attention/hstu_attention_jagged_forward_splitkv_dispatch.hpp @@ -264,6 +264,8 @@ struct jagged_forward_splitkv_dispatch HIP_CHECK_ERROR(hipMallocAsync(&ws.lse_acc_ptr, workspace_bytes, stream)); } + bool almost_invariant_seqlen = is_almost_invariant_seqlen(param); + const auto kargs = [&] { return HstuKernel::MakeKargs(param.q_ptr, param.k_ptr, @@ -281,6 +283,7 @@ struct jagged_forward_splitkv_dispatch param.num_head, param.scale_s, param.attn_scale, + almost_invariant_seqlen, param.seq_stride_q, param.seq_stride_k, param.seq_stride_v, @@ -304,6 +307,7 @@ struct jagged_forward_splitkv_dispatch param.max_seqlen_q, param.hdim_v, ws.num_splits, + almost_invariant_seqlen, has_minfull_attn_seqlen); dim3 kBlockSize = HstuKernel::BlockSize(); constexpr ck_tile::index_t kBlockPerCu = HstuKernel::kBlockPerCu; @@ -318,6 +322,8 @@ struct jagged_forward_splitkv_dispatch SplitkvWorkspace& ws, hipStream_t stream) { + bool almost_invariant_seqlen = is_almost_invariant_seqlen_q(param); + const auto kargs = [&] { return HstuKernel::MakeKargs(ws.o_acc_ptr, ws.lse_acc_ptr, @@ -330,11 +336,13 @@ struct jagged_forward_splitkv_dispatch param.seq_q_offsets_ptr, param.num_head, ws.num_splits, - param.hdim_v); + param.hdim_v, + almost_invariant_seqlen); }(); - dim3 kGridSize = HstuKernel::GridSize(param.num_batch, param.num_head, param.max_seqlen_q); - dim3 kBlockSize = HstuKernel::BlockSize(); + dim3 kGridSize = HstuKernel::GridSize( + param.num_batch, param.num_head, param.max_seqlen_q, almost_invariant_seqlen); + dim3 kBlockSize = HstuKernel::BlockSize(); constexpr ck_tile::index_t kBlockPerCu = HstuKernel::kBlockPerCu; (void)ck_tile::launch_kernel( diff --git a/example/ck_tile/18_hstu_attention/hstu_attention_params.hpp b/example/ck_tile/18_hstu_attention/hstu_attention_params.hpp index 32d43abf81..0b86053cfe 100644 --- a/example/ck_tile/18_hstu_attention/hstu_attention_params.hpp +++ b/example/ck_tile/18_hstu_attention/hstu_attention_params.hpp @@ -25,6 +25,8 @@ struct HstuAttentionNoGroupFwdParams const void* seq_kv_offsets_ptr; // jagged mode only ck_tile::index_t max_seqlen_q; // jagged mode only ck_tile::index_t max_seqlen_kv; // jagged mode only + ck_tile::index_t min_seqlen_q; // jagged mode only + ck_tile::index_t min_seqlen_kv; // jagged mode only const void* q_ptr; const void* k_ptr; @@ -90,6 +92,8 @@ struct HstuAttentionGroupFwdParams const void* seq_kv_offsets_ptr; ck_tile::index_t max_seqlen_q; // the maximum of all the groups' max_seqlen_q ck_tile::index_t max_seqlen_kv; // the maximum of all the groups' max_seqlen_kv + ck_tile::index_t min_seqlen_q; // jagged mode only + ck_tile::index_t min_seqlen_kv; // jagged mode only const void* q_ptr; const void* k_ptr;