ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-13 02:27:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

Enable splitk for mxfp4; clang format;

Browse Source
This commit is contained in:
aska-0096
2025-06-02 12:23:01 +00:00
parent 5696e3c9f5
commit dd24786f78
12 changed files with 191 additions and 2486 deletions
Download Patch File Download Diff File Expand all files Collapse all files

57
example/67_gemm_microscaling/gemm_mx_common.hpp
View File

@@ -108,7 +108,6 @@ bool parse_cmd_args(int argc,
return true;
}
#if 1
template <bool KLast>
void preShuffleScaleBuffer(ck::e8m0_bexp_t* src, ck::e8m0_bexp_t* dst, int MN, int K)
{
@@ -146,8 +145,9 @@ void preShuffleScaleBuffer(ck::e8m0_bexp_t* src, ck::e8m0_bexp_t* dst, int MN, i
k0 * MNXdlPack * KXdlPack * XdlMNThread * XdlKThread +
k1 * MNXdlPack * KXdlPack * XdlMNThread + n1 * MNXdlPack * KXdlPack +
k2 * MNXdlPack + n2;
// src[n * K + k] = ck::type_convert<ck::e8m0_bexp_t>(static_cast<float>(powf(2.0f, n2 +
// k2 * MNXdlPack)));
// src[n * K + k] = ck::type_convert<ck::e8m0_bexp_t>(static_cast<float>(powf(2.0f,
// 2-k)));
if constexpr(KLast)
dst[outputIndex] = src[n * K + k];
else
@@ -186,7 +186,6 @@ void preShuffleBuffer(const ck::f4x2_pk_t* src, ck::f4x2_pk_t* dst, int N, int K
}
}
}
#endif
template <typename DeviceOpInstance,
typename ADataType,
@@ -346,7 +345,6 @@ bool run_mx_gemm(const ProblemSizeSplitK& problem_size, const ExecutionConfig& c
}
}
#if 1
preShuffleScaleBuffer<ck::is_same_v<ALayout, Row>>(a_m_k_scale.mData.data(),
a_shuffled_scale.mData.data(),
Scale_Padded_M,
@@ -358,48 +356,6 @@ bool run_mx_gemm(const ProblemSizeSplitK& problem_size, const ExecutionConfig& c
int NPerXdl = 16; // Fixed 16
preShuffleBuffer(b_k_n->mData.data(), b_input->mData.data(), N, K, NPerXdl);
}
#endif
// printf("a_scale:\n");
// for(ck::index_t i = 0; i < M; i++)
// {
// for(ck::index_t j = 0; j < K / ScaleBlockSize; j++)
// {
// // a_m_k_scale(i, j) =
// // ck::type_convert<XDataType>(static_cast<float>(powf(2.0f, (j / 4) % 4)));
// // a_m_k_scale(i, j) =ck::type_convert<XDataType>(static_cast<float>(1.0f));
// // a_shuffled_scale(i, j) =ck::type_convert<XDataType>(static_cast<float>(1.0f));
// printf("%02x ", *reinterpret_cast<uint8_t*>(&a_m_k_scale(i, j)));
// }
// printf("\n");
// }
// printf("b_scale:\n");
// for(ck::index_t i = 0; i < N; i++)
// {
// for(ck::index_t j = 0; j < K / ScaleBlockSize; j++)
// {
// // // b_k_n_scale(j, i) =
// // // ck::type_convert<XDataType>(static_cast<float>(powf(2.0f, (j / 4) % 4)));
// // b_k_n_scale(j, i) =ck::type_convert<XDataType>(static_cast<float>(1.0f));
// // b_shuffled_scale(j, i) =ck::type_convert<XDataType>(static_cast<float>(1.0f));
// printf("%02x ", *reinterpret_cast<uint8_t*>(&b_k_n_scale(j, i)));
// }
// printf("\n");
// }
// printf("a_shuffled_scale:\n");
// for(ck::index_t i = 0; i < M * K / ScaleBlockSize; i++)
// {
// printf("%02x ", *reinterpret_cast<uint8_t*>(&(a_shuffled_scale.mData.data()[i])));
// if(i % 64 == 63)
// printf("\n");
// }
// printf("b_shuffled_scale:\n");
// for(ck::index_t i = 0; i < N * K / ScaleBlockSize; i++)
// {
// printf("%02x ", *reinterpret_cast<uint8_t*>(&(b_shuffled_scale.mData.data()[i])));
// if(i % 64 == 63)
// printf("\n");
// }
if(config.verbosity > 0)
std::cout << "Device memory allocation..." << std::endl;
@@ -524,9 +480,10 @@ bool run_mx_gemm(const ProblemSizeSplitK& problem_size, const ExecutionConfig& c
// << std::endl;
// }
res_verified = res_verified && ck::utils::check_err(c_m_n_device_result,
c_m_n_host_result,
"Error: Incorrect results!");
res_verified =
res_verified &&
ck::utils::check_err(
c_m_n_device_result, c_m_n_host_result, "Error: Incorrect results!", 5e-1, 5e-1);
if(config.verbosity > 0 && res_verified)
std::cout << "Verification Successful!" << std::endl;