From 5204dc7580eea8e7214f9e91d7cdf4ce76a1bc0d Mon Sep 17 00:00:00 2001 From: Qianfeng Zhang Date: Wed, 10 Jun 2026 08:57:41 +0000 Subject: [PATCH] Fix the using of num_targets[] in run_group_hstu_attention --- .../example_hstu_attention_fwd.cpp | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 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 7e1cd9e643..acff2ecfd6 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 @@ -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]); };