ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-05-16 10:59:55 +00:00
Code Issues Packages Projects Releases Wiki Activity

Convert P to fp16/bf16 before doing second gemm in reference hstu implementation

Browse Source
This commit is contained in:
Qianfeng Zhang
2025-05-29 01:04:19 +00:00
parent 36a0f2020c
commit 781cba355a
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)
{
GemmAccDataType preg =
ck_tile::type_convert<GemmAccDataType>(locals[sk]);
InOutDataType preg = ck_tile::type_convert<InOutDataType>(locals[sk]);
InOutDataType vreg =
v_batch_seq_nhead_hdim(0, seq_offsets[i_batch] + sk, i_head, k);
dot_prod += preg * ck_tile::type_convert<GemmAccDataType>(vreg);
dot_prod += ck_tile::type_convert<GemmAccDataType>(preg) *
ck_tile::type_convert<GemmAccDataType>(vreg);
}
else
{
GemmAccDataType preg =
ck_tile::type_convert<GemmAccDataType>(locals[sk]);
InOutDataType preg = ck_tile::type_convert<InOutDataType>(locals[sk]);
InOutDataType vreg = v_batch_seq_nhead_hdim(i_batch, sk, i_head, k);
dot_prod += preg * ck_tile::type_convert<GemmAccDataType>(vreg);
dot_prod += ck_tile::type_convert<GemmAccDataType>(preg) *
ck_tile::type_convert<GemmAccDataType>(vreg);
};
};