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

Fix the using of num_targets[] in run_group_hstu_attention

Browse Source
This commit is contained in:
Qianfeng Zhang
2026-06-10 08:57:41 +00:00
parent 73965593f9
commit 5204dc7580
1 changed files with 14 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -847,11 +847,23 @@ bool run_group_hstu(const ck_tile::ArgParser& arg_parser, int num_group)
for(int i_grp = 0; i_grp < num_group; i_grp++)
{
int max_num_target = 0;
if(!num_targets.empty())
{
for(int i_batch = 0; i_batch < num_batch_per_group; i_batch++)
{
int i_global_batch = i_grp * num_batch_per_group * i_batch;
max_num_target = max(max_num_target, num_targets[i_global_batch]);
};
};
group_max_seqlens_q[i_grp] =
group_max_uih_seqlens_q[i_grp] + group_contextual_seqlens[i_grp] + num_targets[i_grp];
group_max_uih_seqlens_q[i_grp] + group_contextual_seqlens[i_grp] + max_num_target;
max_max_seqlen_q = max(max_max_seqlen_q, group_max_seqlens_q[i_grp]);
group_max_seqlens_kv[i_grp] =
group_max_uih_seqlens_kv[i_grp] + group_contextual_seqlens[i_grp] + num_targets[i_grp];
group_max_uih_seqlens_kv[i_grp] + group_contextual_seqlens[i_grp] + max_num_target;
max_max_seqlen_kv = max(max_max_seqlen_kv, group_max_seqlens_kv[i_grp]);
};