From 26400def378478a8d84c7903244ebc00817a0926 Mon Sep 17 00:00:00 2001 From: Qianfeng Zhang Date: Mon, 29 Jun 2026 10:21:47 +0000 Subject: [PATCH] Add max_seqlen_kv as host API parameter and adjust the rules for using splitkv for cross-attention --- .../example_hstu_attention_fwd.cpp | 2 + ...stu_attention_batched_forward_dispatch.hpp | 65 +++++++++++++----- ...ntion_batched_forward_splitkv_dispatch.hpp | 5 +- .../hstu_attention_fwd_setting.hpp | 19 ------ .../hstu_attention_group_forward_dispatch.hpp | 66 ++++++++++++++----- ...tention_group_forward_splitkv_dispatch.hpp | 6 +- ...hstu_attention_jagged_forward_dispatch.hpp | 66 ++++++++++++++----- ...ention_jagged_forward_splitkv_dispatch.hpp | 6 +- .../hstu_attention_params.hpp | 4 +- .../hstu_attention_splitkv_helper.hpp | 48 +++++++++++--- 10 files changed, 194 insertions(+), 93 deletions(-) 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 821ef87c49..fd69e86795 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 @@ -378,6 +378,7 @@ bool run_no_group_hstu_forward(const ck_tile::ArgParser& arg_parser, bool is_jag params.seq_q_offsets_ptr = seq_offsets_q_dev.GetDeviceBuffer(); 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.q_ptr = q_dev.GetDeviceBuffer(); params.k_ptr = k_dev.GetDeviceBuffer(); params.v_ptr = v_dev.GetDeviceBuffer(); @@ -898,6 +899,7 @@ bool run_group_hstu_forward(const ck_tile::ArgParser& arg_parser, int num_group) params.seq_q_offsets_ptr = seq_offsets_q_dev.GetDeviceBuffer(); 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.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 60318d884f..5adfe17811 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 @@ -197,7 +197,10 @@ template void run_batched_forward_dispatch(HstuAttentionNoGroupFwdParams& param, hipStream_t stream) { - if(get_hstu_attention_fwd_mtile(param.num_batch, param.num_head, param.seqlen_q) == 128) + int mtile_size = get_hstu_attention_fwd_mtile( + param.num_batch, param.num_head, param.seqlen_q, param.seqlen_kv); + + if(!param.is_cross_attention && mtile_size == 128) batched_forward_dispatch::Run(param, stream); else { + // for cross-attention, we should give more opportunity to use split-kv since the seqlen_kv + // is usually much bigger than seqlen_q, so the main-loop along the seqlen_kv have enough + // iterations to counter-act the cost brought by splitting + const bool disable_fwd_splitkv = []() { const char* env_p = std::getenv("HSTU_DISABLE_SPLITKV"); if(env_p == nullptr) @@ -216,25 +223,47 @@ void run_batched_forward_dispatch(HstuAttentionNoGroupFwdParams& param, hipStrea }(); if(!disable_fwd_splitkv && - shall_use_splitkv(param.num_batch, param.num_head, param.seqlen_q)) + shall_use_splitkv(param.num_batch, param.num_head, param.seqlen_q, param.seqlen_kv)) { - batched_forward_splitkv_dispatch::Run(param, stream); + if(mtile_size == 128) + batched_forward_splitkv_dispatch::Run(param, stream); + else + batched_forward_splitkv_dispatch::Run(param, stream); } else - batched_forward_dispatch::Run(param, stream); + { + if(mtile_size == 128) + batched_forward_dispatch::Run(param, stream); + else + batched_forward_dispatch::Run(param, stream); + } }; }; 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 5c1afb4523..b9f532e85f 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 @@ -38,8 +38,6 @@ template struct batched_forward_splitkv_dispatch { - static_assert(MTile == 64, "MTile must be 64 to get to fwd splitkv path!"); - using HstuAttentionFwdTileSetting = typename std::conditional_t, @@ -254,7 +252,8 @@ struct batched_forward_splitkv_dispatch SplitkvWorkspace& ws, hipStream_t stream) { - ws.num_splits = get_suggested_num_splits(param.num_batch, param.num_head, param.seqlen_q); + ws.num_splits = get_suggested_num_splits( + param.num_batch, param.num_head, param.seqlen_q, param.seqlen_kv); // assume the workspace for o_acc is in compact shape of [num_batch, seqlen_q, num_head, // num_splits, hdim] diff --git a/example/ck_tile/18_hstu_attention/hstu_attention_fwd_setting.hpp b/example/ck_tile/18_hstu_attention/hstu_attention_fwd_setting.hpp index a362b282a0..071586ac65 100644 --- a/example/ck_tile/18_hstu_attention/hstu_attention_fwd_setting.hpp +++ b/example/ck_tile/18_hstu_attention/hstu_attention_fwd_setting.hpp @@ -465,22 +465,3 @@ template struct HstuAttentionWithSoftmaxFwdTileSetting<256, 64>; template struct HstuAttentionWithSoftmaxFwdTileSetting<256, 128>; #endif - -static int get_hstu_attention_fwd_mtile(int num_batches, int num_heads, int max_seqlen_q) -{ - int num_CUs = get_number_of_cu(); - auto ceildiv = [](int a, int b) { return (a + b - 1) / b; }; - - if(max_seqlen_q <= 64) - return 64; - - int nbatch_nhead_mblocks = num_batches * num_heads * ceildiv(max_seqlen_q, 128); - - // assuming each CU is assigned two work-groups - if(nbatch_nhead_mblocks >= static_cast(0.85f * num_CUs * 2.0f)) - return 128; - - // currently, only hdim-128 actually uses mtile-64, for other hdim, the settings for - // mtile-64 can be added through tuning/verification - return 64; -}; 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 50e7622364..caec9ab7e1 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 @@ -182,7 +182,10 @@ template void run_group_forward_dispatch(HstuAttentionGroupFwdParams& param, hipStream_t stream) { - if(get_hstu_attention_fwd_mtile(param.num_batch, param.num_head, param.max_seqlen_q) == 128) + int mtile_size = get_hstu_attention_fwd_mtile( + param.num_batch, param.num_head, param.max_seqlen_q, param.max_seqlen_kv); + + if(!param.is_cross_attention && mtile_size == 128) group_forward_dispatch::Run(param, stream); else { + // for cross-attention, we should give more opportunity to use split-kv since the seqlen_kv + // is usually much bigger than seqlen_q, so the main-loop along the seqlen_kv have enough + // iterations to counter-act the cost brought by splitting + const bool disable_fwd_splitkv = []() { const char* env_p = std::getenv("HSTU_DISABLE_SPLITKV"); if(env_p == nullptr) @@ -201,25 +208,48 @@ void run_group_forward_dispatch(HstuAttentionGroupFwdParams& param, hipStream_t }(); if(!disable_fwd_splitkv && - shall_use_splitkv(param.num_batch, param.num_head, param.max_seqlen_q)) + shall_use_splitkv( + param.num_batch, param.num_head, param.max_seqlen_q, param.max_seqlen_kv)) { - group_forward_splitkv_dispatch::Run(param, stream); + if(mtile_size == 128) + group_forward_splitkv_dispatch::Run(param, stream); + else + group_forward_splitkv_dispatch::Run(param, stream); } else - group_forward_dispatch::Run(param, stream); + { + if(mtile_size == 128) + group_forward_dispatch::Run(param, stream); + else + group_forward_dispatch::Run(param, stream); + } }; }; 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 57b7aa71d4..36f20ccc07 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 @@ -38,8 +38,6 @@ template struct group_forward_splitkv_dispatch { - static_assert(MTile == 64, "MTile must be 64 to get to fwd splitkv path!"); - using HstuAttentionFwdTileSetting = typename std::conditional_t, @@ -246,8 +244,8 @@ struct group_forward_splitkv_dispatch SplitkvWorkspace& ws, hipStream_t stream) { - ws.num_splits = - get_suggested_num_splits(param.num_batch, param.num_head, param.max_seqlen_q); + ws.num_splits = get_suggested_num_splits( + param.num_batch, param.num_head, param.max_seqlen_q, param.max_seqlen_kv); // assume the workspace for o_acc is in compact shape of [num_batch, max_seqlen, num_head, // num_splits, hdim] 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 1092cf047c..dc7c749fcf 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 @@ -185,7 +185,10 @@ template void run_jagged_forward_dispatch(HstuAttentionNoGroupFwdParams& param, hipStream_t stream) { - if(get_hstu_attention_fwd_mtile(param.num_batch, param.num_head, param.max_seqlen_q) == 128) + int mtile_size = get_hstu_attention_fwd_mtile( + param.num_batch, param.num_head, param.max_seqlen_q, param.max_seqlen_kv); + + if(!param.is_cross_attention && mtile_size == 128) jagged_forward_dispatch::Run(param, stream); else { + // for cross-attention, we should give more opportunity to use split-kv since the seqlen_kv + // is usually much bigger than seqlen_q, so the main-loop along the seqlen_kv have enough + // iterations to counter-act the cost brought by splitting + const bool disable_fwd_splitkv = []() { const char* env_p = std::getenv("HSTU_DISABLE_SPLITKV"); if(env_p == nullptr) @@ -204,25 +211,48 @@ void run_jagged_forward_dispatch(HstuAttentionNoGroupFwdParams& param, hipStream }(); if(!disable_fwd_splitkv && - shall_use_splitkv(param.num_batch, param.num_head, param.max_seqlen_q)) + shall_use_splitkv( + param.num_batch, param.num_head, param.max_seqlen_q, param.max_seqlen_kv)) { - jagged_forward_splitkv_dispatch::Run(param, stream); + if(mtile_size == 128) + jagged_forward_splitkv_dispatch::Run(param, stream); + else + jagged_forward_splitkv_dispatch::Run(param, stream); } else - jagged_forward_dispatch::Run(param, stream); + { + if(mtile_size == 128) + jagged_forward_dispatch::Run(param, stream); + else + jagged_forward_dispatch::Run(param, stream); + } }; }; 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 1806f4b186..bcaa3ca76c 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 @@ -38,8 +38,6 @@ template struct jagged_forward_splitkv_dispatch { - static_assert(MTile == 64, "MTile must be 64 to get to fwd splitkv path!"); - using HstuAttentionFwdTileSetting = typename std::conditional_t, @@ -245,8 +243,8 @@ struct jagged_forward_splitkv_dispatch SplitkvWorkspace& ws, hipStream_t stream) { - ws.num_splits = - get_suggested_num_splits(param.num_batch, param.num_head, param.max_seqlen_q); + ws.num_splits = get_suggested_num_splits( + param.num_batch, param.num_head, param.max_seqlen_q, param.max_seqlen_kv); // assume the workspace for o_acc is in compact shape of [num_batch, max_seqlen, num_head, // num_splits, hdim] 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 f89c53ce37..32d43abf81 100644 --- a/example/ck_tile/18_hstu_attention/hstu_attention_params.hpp +++ b/example/ck_tile/18_hstu_attention/hstu_attention_params.hpp @@ -24,6 +24,7 @@ struct HstuAttentionNoGroupFwdParams const void* seq_q_offsets_ptr; // jagged mode only 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 const void* q_ptr; const void* k_ptr; @@ -87,7 +88,8 @@ struct HstuAttentionGroupFwdParams ck_tile::index_t num_batch; const void* seq_q_offsets_ptr; 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_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 const void* q_ptr; const void* k_ptr; diff --git a/example/ck_tile/18_hstu_attention/hstu_attention_splitkv_helper.hpp b/example/ck_tile/18_hstu_attention/hstu_attention_splitkv_helper.hpp index 8fec82f739..9a4d8b58c4 100644 --- a/example/ck_tile/18_hstu_attention/hstu_attention_splitkv_helper.hpp +++ b/example/ck_tile/18_hstu_attention/hstu_attention_splitkv_helper.hpp @@ -7,34 +7,66 @@ #include "hstu_attention_host_util.hpp" -static float get_estimated_cu_coverage_ratio(int num_batches, int num_heads, int max_seqlen_q) +static int +get_hstu_attention_fwd_mtile(int num_batches, int num_heads, int max_seqlen_q, int max_seqlen_kv) { int num_CUs = get_number_of_cu(); auto ceildiv = [](int a, int b) { return (a + b - 1) / b; }; - int nbatch_nhead_mblocks = num_batches * num_heads * ceildiv(max_seqlen_q, 64); + if(max_seqlen_q <= 64) + return 64; + + // for cross-attention where max_seqlen_kv is much bigger than max_seqlen_q, we always use + // mtile_size 128, not to worry about the CU coverage, since split-kv can help us to solve + if(max_seqlen_q >= 128 && static_cast(max_seqlen_kv) / max_seqlen_q >= 5.0) + return 128; + + int nbatch_nhead_mblocks = num_batches * num_heads * ceildiv(max_seqlen_q, 128); + + // assuming each CU is assigned two work-groups + if(nbatch_nhead_mblocks >= static_cast(0.85f * num_CUs * 2.0f)) + return 128; + + // currently, only hdim-128 actually uses mtile-64, for other hdim, the settings for + // mtile-64 can be added through tuning/verification + return 64; +}; + +static float +get_estimated_cu_coverage_ratio(int num_batches, int num_heads, int max_seqlen_q, int max_seqlen_kv) +{ + int num_CUs = get_number_of_cu(); + auto ceildiv = [](int a, int b) { return (a + b - 1) / b; }; + + int m_tile_size = + get_hstu_attention_fwd_mtile(num_batches, num_heads, max_seqlen_q, max_seqlen_kv); + + int nbatch_nhead_mblocks = num_batches * num_heads * ceildiv(max_seqlen_q, m_tile_size); // assume each CU can run two work-groups, common cases for hdim128 return static_cast(nbatch_nhead_mblocks) / (2.0f * num_CUs); }; -static bool shall_use_splitkv(int num_batches, int num_heads, int max_seqlen_q) +static bool shall_use_splitkv(int num_batches, int num_heads, int max_seqlen_q, int max_seqlen_kv) { // Please tune the threshold here - const float threshold = 0.8f; + const float threshold = (max_seqlen_kv >= 2048) ? 1.5f : 0.8f; - if(get_estimated_cu_coverage_ratio(num_batches, num_heads, max_seqlen_q) < threshold) + if(get_estimated_cu_coverage_ratio(num_batches, num_heads, max_seqlen_q, max_seqlen_kv) < + threshold) return true; return false; }; -static int get_suggested_num_splits(int num_batches, int num_heads, int max_seqlen_q) +static int +get_suggested_num_splits(int num_batches, int num_heads, int max_seqlen_q, int max_seqlen_kv) { int i = 2; // Please tune the threshold here - const float threshold = 1.5f; - while(get_estimated_cu_coverage_ratio(num_batches, num_heads, max_seqlen_q) * i < threshold) + const float threshold = 3.0f; + while(get_estimated_cu_coverage_ratio(num_batches, num_heads, max_seqlen_q, max_seqlen_kv) * i < + threshold) i++; // the num_splits shall not be bigger than 64