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

Fix crash on small M

Browse Source
This commit is contained in:
Ding, Yi
2025-10-15 05:11:05 +00:00
parent 0b6b9dbd1b
commit 69e95239f4
3 changed files with 44 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files

57
example/ck_tile/18_flatmm/mxgemm/mx_flatmm.cpp
View File

@@ -297,7 +297,7 @@ float invoke_mx_flatmm(ck_tile::DeviceMem& a_dev_buf,
float tflops = static_cast<float>(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 <class FlatmmConfig, bool KLast, class IterSrc, class IterDst>
void preShuffleScale(const IterSrc src, IterDst dst, int MN, int K)
template <class FlatmmConfig, bool KLast, typename Src>
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<dtype> 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<ck::e8m0_bexp_t>(static_cast<float>(powf(2.0f,
// 2-k)));
if constexpr(KLast)
dst[outputIndex] = src[n * K + k];
shuffled(outputIndex) = n < MN ? src(n, k) : static_cast<dtype>(0);
else
dst[outputIndex] = src[k * MN + n];
shuffled(outputIndex) = n < MN ? src(k, n) : static_cast<dtype>(0);
}
}
return shuffled;
}
#else
template <class FlatmmConfig, class T>

13
example/ck_tile/18_flatmm/mxgemm/run_mx_flatmm.inc
View File

@@ -74,11 +74,6 @@ int run_mx_flatmm_with_layouts(int argc,
ck_tile::HostTensor<ScaleDataType> scale_b(ck_tile::host_tensor_descriptor(
K / ScaleGranularityK, N / ScaleGranularityN, scale_stride_B, is_row_major(b_layout)));
ck_tile::HostTensor<ScaleDataType> scale_a_shuffled(ck_tile::host_tensor_descriptor(
M / ScaleGranularityM, K / ScaleGranularityK, scale_stride_A, is_row_major(a_layout)));
ck_tile::HostTensor<ScaleDataType> 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<ADataType>{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<FlatmmConfig>(b_origin_host.begin(), b_shuffled_host.begin(), N, K);
preShuffleScale<FlatmmConfig, is_row_major(a_layout)>(
scale_a.begin(), scale_a_shuffled.begin(), M, K / ScaleGranularityK);
preShuffleScale<FlatmmConfig, !is_row_major(b_layout)>(
scale_b.begin(), scale_b_shuffled.begin(), N, K / ScaleGranularityK);
// ck_tile::HostTensor<ScaleDataType> scale_a_shuffled = preShuffleScale<FlatmmConfig>(scale_a);
// ck_tile::HostTensor<ScaleDataType> scale_b_shuffled = preShuffleScale<FlatmmConfig>(scale_b);
const auto scale_a_shuffled = preShuffleScale<FlatmmConfig, true>(scale_a);
const auto scale_b_shuffled = preShuffleScale<FlatmmConfig, false>(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());

27
include/ck_tile/ops/flatmm/kernel/mx_flatmm_kernel.hpp
View File

@@ -202,22 +202,19 @@ struct MXFlatmmKernel : FlatmmKernel<TilePartitioner_, MXFlatmmPipeline_, Epilog
auto scale_b = kargs.scale_n_ptr;
static constexpr int BlockScaleSize = 32; // decltype(scale_n)::GranularityK;
const auto&& scale_packs_m = integer_divide_ceil(kargs.M, (MXdlPack * MThreadPerXdl));
const auto&& scale_packs_n = integer_divide_ceil(kargs.N, (NXdlPack * NThreadPerXdl));
const auto&& scale_packs_k = kargs.K / BlockScaleSize / (KXdlPack * KThreadPerXdl);
// A scale tensor view
const auto& scale_a_tensor_view = [&]() {
// Pack 2x2 e8m0 over M/K dimension into 1 int32_t to trigger dword width load
const auto scale_a_naive_desc = make_naive_tensor_descriptor_packed(
make_tuple(kargs.M / (MXdlPack * MThreadPerXdl),
kargs.K / BlockScaleSize / (KXdlPack * KThreadPerXdl),
KThreadPerXdl,
MThreadPerXdl));
make_tuple(scale_packs_m, scale_packs_k, KThreadPerXdl, MThreadPerXdl));
const auto scale_a_desc = transform_tensor_descriptor(
scale_a_naive_desc,
make_tuple(
make_merge_transform(
make_tuple(kargs.M / (MXdlPack * MThreadPerXdl), MThreadPerXdl)),
make_merge_transform(make_tuple(
kargs.K / BlockScaleSize / (KXdlPack * KThreadPerXdl), KThreadPerXdl))),
make_tuple(make_merge_transform(make_tuple(scale_packs_m, MThreadPerXdl)),
make_merge_transform(make_tuple(scale_packs_k, KThreadPerXdl))),
make_tuple(sequence<0, 3>{}, sequence<1, 2>{}),
make_tuple(sequence<0>{}, sequence<1>{}));
@@ -228,17 +225,11 @@ struct MXFlatmmKernel : FlatmmKernel<TilePartitioner_, MXFlatmmPipeline_, Epilog
// B scale tensor view
const auto& scale_b_tensor_view = [&]() {
const auto scale_b_navie_desc = make_naive_tensor_descriptor_packed(
make_tuple(kargs.N / (NXdlPack * NThreadPerXdl),
kargs.K / BlockScaleSize / (KXdlPack * KThreadPerXdl),
KThreadPerXdl,
NThreadPerXdl));
make_tuple(scale_packs_n, scale_packs_k, KThreadPerXdl, NThreadPerXdl));
const auto scale_b_desc = transform_tensor_descriptor(
scale_b_navie_desc,
make_tuple(
make_merge_transform(
make_tuple(kargs.N / (NXdlPack * NThreadPerXdl), NThreadPerXdl)),
make_merge_transform(make_tuple(
kargs.K / BlockScaleSize / (KXdlPack * KThreadPerXdl), KThreadPerXdl))),
make_tuple(make_merge_transform(make_tuple(scale_packs_n, NThreadPerXdl)),
make_merge_transform(make_tuple(scale_packs_k, KThreadPerXdl))),
make_tuple(sequence<0, 3>{}, sequence<1, 2>{}),
make_tuple(sequence<0>{}, sequence<1>{}));