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

Add max_seqlen_kv as host API parameter and adjust the rules for using splitkv for cross-attention

Browse Source
This commit is contained in:
Qianfeng Zhang
2026-06-29 10:21:47 +00:00
parent 7c1c61b710
commit 26400def37
10 changed files with 194 additions and 93 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
example/ck_tile/18_hstu_attention/example_hstu_attention_fwd.cpp
View File

@@ -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();

65
example/ck_tile/18_hstu_attention/hstu_attention_batched_forward_dispatch.hpp
View File

@@ -197,7 +197,10 @@ template <typename InOutDataType,
ck_tile::index_t MaxK>
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<InOutDataType,
kUseCausal,
kUseSoftmax,
@@ -208,6 +211,10 @@ void run_batched_forward_dispatch(HstuAttentionNoGroupFwdParams& param, hipStrea
128>::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<InOutDataType,
kUseCausal,
kUseSoftmax,
kStoreLSE,
kHasBias,
kHasDropout,
MaxK,
64>::Run(param, stream);
if(mtile_size == 128)
batched_forward_splitkv_dispatch<InOutDataType,
kUseCausal,
kUseSoftmax,
kStoreLSE,
kHasBias,
kHasDropout,
MaxK,
128>::Run(param, stream);
else
batched_forward_splitkv_dispatch<InOutDataType,
kUseCausal,
kUseSoftmax,
kStoreLSE,
kHasBias,
kHasDropout,
MaxK,
64>::Run(param, stream);
}
else
batched_forward_dispatch<InOutDataType,
kUseCausal,
kUseSoftmax,
kStoreLSE,
kHasBias,
kHasDropout,
MaxK,
64>::Run(param, stream);
{
if(mtile_size == 128)
batched_forward_dispatch<InOutDataType,
kUseCausal,
kUseSoftmax,
kStoreLSE,
kHasBias,
kHasDropout,
MaxK,
128>::Run(param, stream);
else
batched_forward_dispatch<InOutDataType,
kUseCausal,
kUseSoftmax,
kStoreLSE,
kHasBias,
kHasDropout,
MaxK,
64>::Run(param, stream);
}
};
};

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

@@ -38,8 +38,6 @@ template <typename InOutDataType,
ck_tile::index_t MTile>
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<kUseSoftmax,
HstuAttentionWithSoftmaxFwdTileSetting<MaxK, MTile>,
@@ -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]

19
example/ck_tile/18_hstu_attention/hstu_attention_fwd_setting.hpp
View File

@@ -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<int>(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;
};

66
example/ck_tile/18_hstu_attention/hstu_attention_group_forward_dispatch.hpp
View File

@@ -182,7 +182,10 @@ template <typename InOutDataType,
ck_tile::index_t MaxK>
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<InOutDataType,
kUseCausal,
kUseSoftmax,
@@ -193,6 +196,10 @@ void run_group_forward_dispatch(HstuAttentionGroupFwdParams& param, hipStream_t
128>::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<InOutDataType,
kUseCausal,
kUseSoftmax,
kStoreLSE,
kHasBias,
kHasDropout,
MaxK,
64>::Run(param, stream);
if(mtile_size == 128)
group_forward_splitkv_dispatch<InOutDataType,
kUseCausal,
kUseSoftmax,
kStoreLSE,
kHasBias,
kHasDropout,
MaxK,
128>::Run(param, stream);
else
group_forward_splitkv_dispatch<InOutDataType,
kUseCausal,
kUseSoftmax,
kStoreLSE,
kHasBias,
kHasDropout,
MaxK,
64>::Run(param, stream);
}
else
group_forward_dispatch<InOutDataType,
kUseCausal,
kUseSoftmax,
kStoreLSE,
kHasBias,
kHasDropout,
MaxK,
64>::Run(param, stream);
{
if(mtile_size == 128)
group_forward_dispatch<InOutDataType,
kUseCausal,
kUseSoftmax,
kStoreLSE,
kHasBias,
kHasDropout,
MaxK,
128>::Run(param, stream);
else
group_forward_dispatch<InOutDataType,
kUseCausal,
kUseSoftmax,
kStoreLSE,
kHasBias,
kHasDropout,
MaxK,
64>::Run(param, stream);
}
};
};

6
example/ck_tile/18_hstu_attention/hstu_attention_group_forward_splitkv_dispatch.hpp
View File

@@ -38,8 +38,6 @@ template <typename InOutDataType,
ck_tile::index_t MTile>
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<kUseSoftmax,
HstuAttentionWithSoftmaxFwdTileSetting<MaxK, MTile>,
@@ -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]

66
example/ck_tile/18_hstu_attention/hstu_attention_jagged_forward_dispatch.hpp
View File

@@ -185,7 +185,10 @@ template <typename InOutDataType,
ck_tile::index_t MaxK>
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<InOutDataType,
kUseCausal,
kUseSoftmax,
@@ -196,6 +199,10 @@ void run_jagged_forward_dispatch(HstuAttentionNoGroupFwdParams& param, hipStream
128>::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<InOutDataType,
kUseCausal,
kUseSoftmax,
kStoreLSE,
kHasBias,
kHasDropout,
MaxK,
64>::Run(param, stream);
if(mtile_size == 128)
jagged_forward_splitkv_dispatch<InOutDataType,
kUseCausal,
kUseSoftmax,
kStoreLSE,
kHasBias,
kHasDropout,
MaxK,
128>::Run(param, stream);
else
jagged_forward_splitkv_dispatch<InOutDataType,
kUseCausal,
kUseSoftmax,
kStoreLSE,
kHasBias,
kHasDropout,
MaxK,
64>::Run(param, stream);
}
else
jagged_forward_dispatch<InOutDataType,
kUseCausal,
kUseSoftmax,
kStoreLSE,
kHasBias,
kHasDropout,
MaxK,
64>::Run(param, stream);
{
if(mtile_size == 128)
jagged_forward_dispatch<InOutDataType,
kUseCausal,
kUseSoftmax,
kStoreLSE,
kHasBias,
kHasDropout,
MaxK,
128>::Run(param, stream);
else
jagged_forward_dispatch<InOutDataType,
kUseCausal,
kUseSoftmax,
kStoreLSE,
kHasBias,
kHasDropout,
MaxK,
64>::Run(param, stream);
}
};
};

6
example/ck_tile/18_hstu_attention/hstu_attention_jagged_forward_splitkv_dispatch.hpp
View File

@@ -38,8 +38,6 @@ template <typename InOutDataType,
ck_tile::index_t MTile>
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<kUseSoftmax,
HstuAttentionWithSoftmaxFwdTileSetting<MaxK, MTile>,
@@ -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]

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

@@ -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;

48
example/ck_tile/18_hstu_attention/hstu_attention_splitkv_helper.hpp
View File

@@ -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<float>(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<int>(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<float>(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