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

Update the reference hstu to not do fp32 to fp16/bf16 conversion before P@V gemm

Browse Source
This commit is contained in:
Qianfeng Zhang
2025-05-20 07:50:56 +00:00
parent 0a8ea6bd02
commit a1346aaf3e
1 changed files with 6 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
example/ck_tile/18_hstu_attention/reference_hstu_attention.hpp
View File

@@ -182,20 +182,20 @@ struct reference_hstu_attention
{
if constexpr(kIsJagged)
{
InOutDataType preg = ck_tile::type_convert<InOutDataType>(locals[sk]);
GemmAccDataType preg =
ck_tile::type_convert<GemmAccDataType>(locals[sk]);
InOutDataType vreg =
v_batch_seq_nhead_hdim(0, seq_offsets[i_batch] + sk, i_head, k);
dot_prod += ck_tile::type_convert<GemmAccDataType>(preg) *
ck_tile::type_convert<GemmAccDataType>(vreg);
dot_prod += preg * ck_tile::type_convert<GemmAccDataType>(vreg);
}
else
{
InOutDataType preg = ck_tile::type_convert<InOutDataType>(locals[sk]);
GemmAccDataType preg =
ck_tile::type_convert<GemmAccDataType>(locals[sk]);
InOutDataType vreg = v_batch_seq_nhead_hdim(i_batch, sk, i_head, k);
dot_prod += ck_tile::type_convert<GemmAccDataType>(preg) *
ck_tile::type_convert<GemmAccDataType>(vreg);
dot_prod += preg * ck_tile::type_convert<GemmAccDataType>(vreg);
};
};