From 69e95239f43c8c1946716692066848dff5eeb378 Mon Sep 17 00:00:00 2001 From: "Ding, Yi" Date: Wed, 15 Oct 2025 05:11:05 +0000 Subject: [PATCH] Fix crash on small M --- .../ck_tile/18_flatmm/mxgemm/mx_flatmm.cpp | 57 +++++++++++-------- .../18_flatmm/mxgemm/run_mx_flatmm.inc | 13 +---- .../ops/flatmm/kernel/mx_flatmm_kernel.hpp | 27 +++------ 3 files changed, 44 insertions(+), 53 deletions(-) diff --git a/example/ck_tile/18_flatmm/mxgemm/mx_flatmm.cpp b/example/ck_tile/18_flatmm/mxgemm/mx_flatmm.cpp index 6577d8ce6c..a3680164f4 100644 --- a/example/ck_tile/18_flatmm/mxgemm/mx_flatmm.cpp +++ b/example/ck_tile/18_flatmm/mxgemm/mx_flatmm.cpp @@ -297,7 +297,7 @@ float invoke_mx_flatmm(ck_tile::DeviceMem& a_dev_buf, float tflops = static_cast(flop) / 1.E9 / ave_time; float gb_per_sec = num_byte / 1.E6 / ave_time; - std::cout << "Run MXFP4_Flatmm kernel " + std::cout << "Run MXFP4_Flatmm kernel " // << " M =" << M << " N =" << N << " K =" << K << " StrideA =" << stride_A << " StrideB =" << stride_B << " StrideC =" << stride_C << " : " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s, " << std::endl; @@ -366,16 +366,24 @@ void preShuffleWeight(const IterSrc src, IterDst dst, int N, int K) } #if 1 -template -void preShuffleScale(const IterSrc src, IterDst dst, int MN, int K) +template +auto preShuffleScale(Src& src) { - int MNXdlPack = 2; - int KXdlPack = 2; + using dtype = typename Src::Data::value_type; + auto src_lengths = src.get_lengths(); + const auto MN = KLast ? src_lengths[0] : src_lengths[1]; + const auto K = KLast ? src_lengths[1] : src_lengths[0]; - int XdlMNThread = FlatmmConfig::N_Warp_Tile; // 16 - int XdlKThread = 64 / XdlMNThread; + size_t MNXdlPack = 2; + size_t KXdlPack = 2; + size_t XdlMNThread = FlatmmConfig::N_Warp_Tile; // 16 + size_t XdlKThread = 64 / XdlMNThread; - int K0 = K / KXdlPack / XdlKThread; // KRepeat + const auto MN_Paded = ck_tile::integer_least_multiple(MN, XdlMNThread * MNXdlPack); + + ck_tile::HostTensor shuffled(ck_tile::HostTensorDescriptor({MN_Paded * K}, {1})); + + size_t K0 = K / KXdlPack / XdlKThread; // KRepeat // The 4 16x128 building blocks will be packed into 1 32x256 for F4 // The 8 16x16x128 mfma will be packed into 1 32x32x256 for F4 @@ -385,33 +393,34 @@ void preShuffleScale(const IterSrc src, IterDst dst, int MN, int K) // To XdlKThread-> XdlMNThread -> KXdlPack -> MNXdlPack // Then, MNRepeat->KRepeat - for(int n = 0; n < MN; ++n) + for(size_t n = 0; n < MN_Paded; ++n) { - for(int k = 0; k < K; ++k) + for(size_t k = 0; k < K; ++k) { - int n0 = n / (XdlMNThread * MNXdlPack); // i MNRepeat - int tempn = n % (XdlMNThread * MNXdlPack); - int n1 = tempn % XdlMNThread; // i XdlMNThread - int n2 = tempn / XdlMNThread; // i MNXdlPack + auto n0 = n / (XdlMNThread * MNXdlPack); // i MNRepeat + auto tempn = n % (XdlMNThread * MNXdlPack); + auto n1 = tempn % XdlMNThread; // i XdlMNThread + auto n2 = tempn / XdlMNThread; // i MNXdlPack - int k0 = k / (XdlKThread * KXdlPack); // i KRepeat - int tempk = k % (XdlKThread * KXdlPack); - int k1 = tempk % XdlKThread; // i XdlKThread - int k2 = tempk / XdlKThread; // i KXdlPack + auto k0 = k / (XdlKThread * KXdlPack); // i KRepeat + auto tempk = k % (XdlKThread * KXdlPack); + auto k1 = tempk % XdlKThread; // i XdlKThread + auto k2 = tempk / XdlKThread; // i KXdlPack - int outputIndex = n0 * MNXdlPack * KXdlPack * XdlMNThread * XdlKThread * K0 + - k0 * MNXdlPack * KXdlPack * XdlMNThread * XdlKThread + - k1 * MNXdlPack * KXdlPack * XdlMNThread + n1 * MNXdlPack * KXdlPack + - k2 * MNXdlPack + n2; + auto outputIndex = n0 * MNXdlPack * KXdlPack * XdlMNThread * XdlKThread * K0 + + k0 * MNXdlPack * KXdlPack * XdlMNThread * XdlKThread + + k1 * MNXdlPack * KXdlPack * XdlMNThread + n1 * MNXdlPack * KXdlPack + + k2 * MNXdlPack + n2; // src[n * K + k] = ck::type_convert(static_cast(powf(2.0f, // 2-k))); if constexpr(KLast) - dst[outputIndex] = src[n * K + k]; + shuffled(outputIndex) = n < MN ? src(n, k) : static_cast(0); else - dst[outputIndex] = src[k * MN + n]; + shuffled(outputIndex) = n < MN ? src(k, n) : static_cast(0); } } + return shuffled; } #else template diff --git a/example/ck_tile/18_flatmm/mxgemm/run_mx_flatmm.inc b/example/ck_tile/18_flatmm/mxgemm/run_mx_flatmm.inc index f449cd631c..e95ed0c238 100644 --- a/example/ck_tile/18_flatmm/mxgemm/run_mx_flatmm.inc +++ b/example/ck_tile/18_flatmm/mxgemm/run_mx_flatmm.inc @@ -74,11 +74,6 @@ int run_mx_flatmm_with_layouts(int argc, ck_tile::HostTensor scale_b(ck_tile::host_tensor_descriptor( K / ScaleGranularityK, N / ScaleGranularityN, scale_stride_B, is_row_major(b_layout))); - ck_tile::HostTensor scale_a_shuffled(ck_tile::host_tensor_descriptor( - M / ScaleGranularityM, K / ScaleGranularityK, scale_stride_A, is_row_major(a_layout))); - ck_tile::HostTensor scale_b_shuffled(ck_tile::host_tensor_descriptor( - K / ScaleGranularityK, N / ScaleGranularityN, scale_stride_B, is_row_major(b_layout))); - if(init_method == 0) { ck_tile::FillUniformDistribution{0.0f, 1.0f}(a_host); @@ -243,12 +238,8 @@ int run_mx_flatmm_with_layouts(int argc, ck_tile::host_tensor_descriptor(K, N, stride_B, is_row_major(b_layout))); preShuffleWeight(b_origin_host.begin(), b_shuffled_host.begin(), N, K); - preShuffleScale( - scale_a.begin(), scale_a_shuffled.begin(), M, K / ScaleGranularityK); - preShuffleScale( - scale_b.begin(), scale_b_shuffled.begin(), N, K / ScaleGranularityK); - // ck_tile::HostTensor scale_a_shuffled = preShuffleScale(scale_a); - // ck_tile::HostTensor scale_b_shuffled = preShuffleScale(scale_b); + const auto scale_a_shuffled = preShuffleScale(scale_a); + const auto scale_b_shuffled = preShuffleScale(scale_b); ck_tile::DeviceMem a_dev_buf(a_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem b_shuffled_dev_buf(b_shuffled_host.get_element_space_size_in_bytes()); diff --git a/include/ck_tile/ops/flatmm/kernel/mx_flatmm_kernel.hpp b/include/ck_tile/ops/flatmm/kernel/mx_flatmm_kernel.hpp index 763a2a9d6e..d898870ecf 100644 --- a/include/ck_tile/ops/flatmm/kernel/mx_flatmm_kernel.hpp +++ b/include/ck_tile/ops/flatmm/kernel/mx_flatmm_kernel.hpp @@ -202,22 +202,19 @@ struct MXFlatmmKernel : FlatmmKernel{}, sequence<1, 2>{}), make_tuple(sequence<0>{}, sequence<1>{})); @@ -228,17 +225,11 @@ struct MXFlatmmKernel : FlatmmKernel{}, sequence<1, 2>{}), make_tuple(sequence<0>{}, sequence<1>{}));