ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-06-06 15:54:31 +00:00
Code Issues Packages Projects Releases Wiki Activity

Update reference mx gemm

Browse Source
This commit is contained in:
Rostyslav Geyyer
2025-03-10 15:46:51 +00:00
parent ac0224532e
commit 2576e64bf6
1 changed files with 42 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

48
library/include/ck/library/reference_tensor_operation/cpu/reference_mx_gemm.hpp
View File

@@ -89,9 +89,27 @@ struct ReferenceMXGemm : public device::BaseOperator
{
for(size_t k = 0; k < K; k++)
{
a_m_k_scaled(m, k) =
type_convert<ComputeTypeA>(arg.a_m_k_(m, k)) *
type_convert<ComputeTypeA>(arg.a_m_kblock_scales_(m, k / SCALE_BLOCK));
if constexpr(is_same_v<ADataType, f4x2_pk_t>)
{
if(k % 2 == 1)
a_m_k_scaled(m, k) =
type_convert<ComputeTypeA>(
arg.a_m_k_(m, k).template unpack<>(Number<1>{})) *
type_convert<ComputeTypeA>(
arg.a_m_kblock_scales_(m, k / SCALE_BLOCK));
else
a_m_k_scaled(m, k) =
type_convert<ComputeTypeA>(
arg.a_m_k_(m, k).template unpack<>(Number<0>{})) *
type_convert<ComputeTypeA>(
arg.a_m_kblock_scales_(m, k / SCALE_BLOCK));
}
else
{
a_m_k_scaled(m, k) =
type_convert<ComputeTypeA>(arg.a_m_k_(m, k)) *
type_convert<ComputeTypeA>(arg.a_m_kblock_scales_(m, k / SCALE_BLOCK));
}
}
}
@@ -99,9 +117,27 @@ struct ReferenceMXGemm : public device::BaseOperator
{
for(size_t k = 0; k < K; k++)
{
b_k_n_scaled(k, n) =
type_convert<ComputeTypeB>(arg.b_k_n_(k, n)) *
type_convert<ComputeTypeB>(arg.b_kblock_n_scales_(k / SCALE_BLOCK, n));
if constexpr(is_same_v<BDataType, f4x2_pk_t>)
{
if(k % 2 == 1)
b_k_n_scaled(k, n) =
type_convert<ComputeTypeB>(
arg.b_k_n_(k, n).template unpack<>(Number<1>{})) *
type_convert<ComputeTypeB>(
arg.b_kblock_n_scales_(k / SCALE_BLOCK, n));
else
b_k_n_scaled(k, n) =
type_convert<ComputeTypeB>(
arg.b_k_n_(k, n).template unpack<>(Number<0>{})) *
type_convert<ComputeTypeB>(
arg.b_kblock_n_scales_(k / SCALE_BLOCK, n));
}
else
{
b_k_n_scaled(k, n) =
type_convert<ComputeTypeB>(arg.b_k_n_(k, n)) *
type_convert<ComputeTypeB>(arg.b_kblock_n_scales_(k / SCALE_BLOCK, n));
}
}
}