From b99c50a5d5d2156a8ff092f08651f23a22b76d93 Mon Sep 17 00:00:00 2001 From: "Ding, Yi" Date: Wed, 28 May 2025 03:35:33 +0000 Subject: [PATCH] pad ascale --- .../67_gemm_microscaling/gemm_mx_common.hpp | 85 ++++++++----------- example/67_gemm_microscaling/gemm_mx_fp4.cpp | 4 +- .../device_gemm_mx_xdl_f4_f4_f16_mk_nk_mn.hpp | 4 +- .../include/profiler/profile_gemm_mx_impl.hpp | 85 ++++++++----------- 4 files changed, 78 insertions(+), 100 deletions(-) diff --git a/example/67_gemm_microscaling/gemm_mx_common.hpp b/example/67_gemm_microscaling/gemm_mx_common.hpp index 16bd685fde..29d1cafc6a 100644 --- a/example/67_gemm_microscaling/gemm_mx_common.hpp +++ b/example/67_gemm_microscaling/gemm_mx_common.hpp @@ -110,7 +110,7 @@ bool parse_cmd_args(int argc, #if 1 template -void preShuffleScaleBuffer(ck::e8m0_bexp_t* src, ck::e8m0_bexp_t* dst, int MN, int MN_dst, int K) +void preShuffleScaleBuffer(ck::e8m0_bexp_t* src, ck::e8m0_bexp_t* dst, int MN, int K) { int MNXdlPack = 2; int KXdlPack = 2; @@ -128,40 +128,32 @@ void preShuffleScaleBuffer(ck::e8m0_bexp_t* src, ck::e8m0_bexp_t* dst, int MN, i // To XdlKThread-> XdlMNThread -> KXdlPack -> MNXdlPack // Then, MNRepeat->KRepeat - const auto fn = [&](auto IS_PADDED, int n, int k) { - constexpr auto IS_PADDED_V = decltype(IS_PADDED)::value; - - int n0 = n / (XdlMNThread * MNXdlPack); // i MNRepeat - int tempn = n % (XdlMNThread * MNXdlPack); - int n1 = tempn % XdlMNThread; // i XdlMNThread - int 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 - - int 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, n2 + - // k2 * MNXdlPack))); - if constexpr(IS_PADDED_V) - dst[outputIndex] = 0; - else if constexpr(KLast) - dst[outputIndex] = src[n * K + k]; - else - dst[outputIndex] = src[k * MN + n]; - }; - - int n = 0; - for(; n < MN; ++n) + for(int n = 0; n < MN; ++n) + { for(int k = 0; k < K; ++k) - fn(ck::false_type{}, n, k); - for(; n < MN_dst; ++n) - for(int k = 0; k < K; ++k) - fn(ck::true_type{}, n, 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 + + int k0 = k / (XdlKThread * KXdlPack); // i KRepeat + int tempk = k % (XdlKThread * KXdlPack); + int k1 = tempk % XdlKThread; // i XdlKThread + int 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; + // src[n * K + k] = ck::type_convert(static_cast(powf(2.0f, n2 + + // k2 * MNXdlPack))); + if constexpr(KLast) + dst[outputIndex] = src[n * K + k]; + else + dst[outputIndex] = src[k * MN + n]; + } + } } void preShuffleBuffer(const ck::f4x2_pk_t* src, ck::f4x2_pk_t* dst, int N, int K, int NXdl) @@ -258,11 +250,10 @@ bool run_mx_gemm(const ProblemSizeSplitK& problem_size, const ExecutionConfig& c using AScaleLayout = Row; using BScaleLayout = Col; - auto Scale_Stride_AM = f_get_default_stride(M, K / ScaleBlockSize, -1, AScaleLayout{}); - auto Scale_Stride_BN = f_get_default_stride(K / ScaleBlockSize, N, -1, BScaleLayout{}); - auto Scale_Padded_M = (M + 32 - 1) / 32 * 32; - auto Scale_Padded_Stride_AM = + auto Scale_Padded_M = (M + 32 - 1) / 32 * 32; + auto Scale_Stride_AM = f_get_default_stride(Scale_Padded_M, K / ScaleBlockSize, -1, AScaleLayout{}); + auto Scale_Stride_BN = f_get_default_stride(K / ScaleBlockSize, N, -1, BScaleLayout{}); Tensor a_m_k(f_host_tensor_descriptor(M, K, StrideA, ALayout{})); auto b_k_n = @@ -273,14 +264,14 @@ bool run_mx_gemm(const ProblemSizeSplitK& problem_size, const ExecutionConfig& c f_host_tensor_descriptor(K, N, StrideB, BLayout{})); // use layout only for size // scales for A and B - Tensor a_m_k_scale( - f_host_tensor_descriptor(M, K / ScaleBlockSize, Scale_Stride_AM, AScaleLayout{})); + Tensor a_m_k_scale(f_host_tensor_descriptor( + Scale_Padded_M, K / ScaleBlockSize, Scale_Stride_AM, AScaleLayout{})); Tensor b_k_n_scale( f_host_tensor_descriptor(K / ScaleBlockSize, N, Scale_Stride_BN, BScaleLayout{})); // shuffled scales for A and B Tensor a_shuffled_scale(f_host_tensor_descriptor( - Scale_Padded_M, K / ScaleBlockSize, Scale_Padded_Stride_AM, AScaleLayout{})); + Scale_Padded_M, K / ScaleBlockSize, Scale_Stride_AM, AScaleLayout{})); Tensor b_shuffled_scale( f_host_tensor_descriptor(K / ScaleBlockSize, N, Scale_Stride_BN, BScaleLayout{})); @@ -355,14 +346,12 @@ bool run_mx_gemm(const ProblemSizeSplitK& problem_size, const ExecutionConfig& c } #if 1 - preShuffleScaleBuffer>( // - a_m_k_scale.mData.data(), - a_shuffled_scale.mData.data(), - M, - Scale_Padded_M, - K / ScaleBlockSize); + preShuffleScaleBuffer>(a_m_k_scale.mData.data(), + a_shuffled_scale.mData.data(), + Scale_Padded_M, + K / ScaleBlockSize); preShuffleScaleBuffer>( - b_k_n_scale.mData.data(), b_shuffled_scale.mData.data(), N, N, K / ScaleBlockSize); + b_k_n_scale.mData.data(), b_shuffled_scale.mData.data(), N, K / ScaleBlockSize); if constexpr(BPreShuffle) { int NPerXdl = 16; // Fixed 16 diff --git a/example/67_gemm_microscaling/gemm_mx_fp4.cpp b/example/67_gemm_microscaling/gemm_mx_fp4.cpp index cff5148fa7..eeb459b5be 100644 --- a/example/67_gemm_microscaling/gemm_mx_fp4.cpp +++ b/example/67_gemm_microscaling/gemm_mx_fp4.cpp @@ -50,14 +50,14 @@ using DeviceOpInstance = ck::tensor_operation::device::DeviceGemmMX_Xdl_CShuffle GemmSpec, // GemmSpec ScaleBlockSize, // ScaleBlockSize: Scaling block size 256, // BlockSize: Thread block size - 256, // MPerBlock + 128, // MPerBlock 256, // NPerBlock KPerBlock, // KPerBlock 16, // AK1 16, // BK1 16, // MPerXDL 16, // NPerXDL - 8, // MXdlPerWave + 4, // MXdlPerWave 8, // NXdlPerWave S<8, 32, 1>, // ABlockTransferThreadClusterLengths_AK0_M_AK1 S<1, 0, 2>, // ABlockTransferThreadClusterArrangeOrder diff --git a/library/src/tensor_operation_instance/gpu/gemm_mx/device_gemm_mx_xdl_f4_f4_f16/device_gemm_mx_xdl_f4_f4_f16_mk_nk_mn.hpp b/library/src/tensor_operation_instance/gpu/gemm_mx/device_gemm_mx_xdl_f4_f4_f16/device_gemm_mx_xdl_f4_f4_f16_mk_nk_mn.hpp index 2b4c18787a..4c2d768ce9 100644 --- a/library/src/tensor_operation_instance/gpu/gemm_mx/device_gemm_mx_xdl_f4_f4_f16/device_gemm_mx_xdl_f4_f4_f16_mk_nk_mn.hpp +++ b/library/src/tensor_operation_instance/gpu/gemm_mx/device_gemm_mx_xdl_f4_f4_f16/device_gemm_mx_xdl_f4_f4_f16_mk_nk_mn.hpp @@ -49,8 +49,8 @@ using device_gemm_mx_xdl_f4_f4_f16_mk_nk_mn_instances = std::tuple< DeviceGemmMX_Xdl_CShuffleV3< Row, Col, Row, F4, E8M0PK, F4, E8M0PK, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, ScaleBlockSize, 256, 32, 256, 128, 16, 16, 16, 16, 2, 4, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, true, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, true, 2, 2, S<1, 32, 1, 8>, 8, BlkGemmPipeSched, BlockGemmPipelineVersion::v1>, DeviceGemmMX_Xdl_CShuffleV3< Row, Col, Row, F4, E8M0PK, F4, E8M0PK, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, ScaleBlockSize, 256, 64, 128, 128, 16, 16, 16, 16, 4, 2, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, true, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, true, 2, 2, S<1, 32, 1, 8>, 8, BlkGemmPipeSched, BlockGemmPipelineVersion::v1>, DeviceGemmMX_Xdl_CShuffleV3< Row, Col, Row, F4, E8M0PK, F4, E8M0PK, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, ScaleBlockSize, 256, 64, 256, 128, 16, 16, 16, 16, 4, 4, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, true, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, true, 2, 2, S<1, 32, 1, 8>, 8, BlkGemmPipeSched, BlockGemmPipelineVersion::v1>, - DeviceGemmMX_Xdl_CShuffleV3< Row, Col, Row, F4, E8M0PK, F4, E8M0PK, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, ScaleBlockSize, 256, 96, 128, 128, 16, 16, 16, 16, 6, 2, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, true, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, true, 2, 2, S<1, 32, 1, 8>, 8, BlkGemmPipeSched, BlockGemmPipelineVersion::v1>, - DeviceGemmMX_Xdl_CShuffleV3< Row, Col, Row, F4, E8M0PK, F4, E8M0PK, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, ScaleBlockSize, 256, 96, 256, 128, 16, 16, 16, 16, 6, 4, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, true, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, true, 2, 2, S<1, 32, 1, 8>, 8, BlkGemmPipeSched, BlockGemmPipelineVersion::v1>, + // DeviceGemmMX_Xdl_CShuffleV3< Row, Col, Row, F4, E8M0PK, F4, E8M0PK, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, ScaleBlockSize, 256, 96, 128, 128, 16, 16, 16, 16, 6, 2, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, true, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, true, 2, 2, S<1, 32, 1, 8>, 8, BlkGemmPipeSched, BlockGemmPipelineVersion::v1>, + // DeviceGemmMX_Xdl_CShuffleV3< Row, Col, Row, F4, E8M0PK, F4, E8M0PK, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, ScaleBlockSize, 256, 96, 256, 128, 16, 16, 16, 16, 6, 4, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, true, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, true, 2, 2, S<1, 32, 1, 8>, 8, BlkGemmPipeSched, BlockGemmPipelineVersion::v1>, DeviceGemmMX_Xdl_CShuffleV3< Row, Col, Row, F4, E8M0PK, F4, E8M0PK, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, ScaleBlockSize, 256, 256, 256, 128, 16, 16, 16, 16, 8, 8, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, true, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, true, 2, 2, S<1, 32, 1, 8>, 8, BlkGemmPipeSched, BlockGemmPipelineVersion::v3>, DeviceGemmMX_Xdl_CShuffleV3< Row, Col, Row, F4, E8M0PK, F4, E8M0PK, F16, F32, F16, PassThrough, PassThrough, PassThrough, GemmSpec, ScaleBlockSize, 256, 128, 256, 128, 16, 16, 16, 16, 4, 8, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, true, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, true, 2, 2, S<1, 32, 1, 8>, 8, BlkGemmPipeSched, BlockGemmPipelineVersion::v3>, diff --git a/profiler/include/profiler/profile_gemm_mx_impl.hpp b/profiler/include/profiler/profile_gemm_mx_impl.hpp index 081b7fc597..2c6f778697 100644 --- a/profiler/include/profiler/profile_gemm_mx_impl.hpp +++ b/profiler/include/profiler/profile_gemm_mx_impl.hpp @@ -26,7 +26,7 @@ namespace profiler { #if 1 template -void preShuffleScaleBuffer(ck::e8m0_bexp_t* src, ck::e8m0_bexp_t* dst, int MN, int MN_dst, int K) +void preShuffleScaleBuffer(ck::e8m0_bexp_t* src, ck::e8m0_bexp_t* dst, int MN, int K) { int MNXdlPack = 2; int KXdlPack = 2; @@ -44,40 +44,32 @@ void preShuffleScaleBuffer(ck::e8m0_bexp_t* src, ck::e8m0_bexp_t* dst, int MN, i // To XdlKThread-> XdlMNThread -> KXdlPack -> MNXdlPack // Then, MNRepeat->KRepeat - const auto fn = [&](auto IS_PADDED, int n, int k) { - constexpr auto IS_PADDED_V = decltype(IS_PADDED)::value; - - int n0 = n / (XdlMNThread * MNXdlPack); // i MNRepeat - int tempn = n % (XdlMNThread * MNXdlPack); - int n1 = tempn % XdlMNThread; // i XdlMNThread - int 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 - - int 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, n2 + - // k2 * MNXdlPack))); - if constexpr(IS_PADDED_V) - dst[outputIndex] = 0; - else if constexpr(KLast) - dst[outputIndex] = src[n * K + k]; - else - dst[outputIndex] = src[k * MN + n]; - }; - - int n = 0; - for(; n < MN; ++n) + for(int n = 0; n < MN; ++n) + { for(int k = 0; k < K; ++k) - fn(ck::false_type{}, n, k); - for(; n < MN_dst; ++n) - for(int k = 0; k < K; ++k) - fn(ck::true_type{}, n, 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 + + int k0 = k / (XdlKThread * KXdlPack); // i KRepeat + int tempk = k % (XdlKThread * KXdlPack); + int k1 = tempk % XdlKThread; // i XdlKThread + int 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; + // src[n * K + k] = ck::type_convert(static_cast(powf(2.0f, n2 + + // k2 * MNXdlPack))); + if constexpr(KLast) + dst[outputIndex] = src[n * K + k]; + else + dst[outputIndex] = src[k * MN + n]; + } + } } #endif @@ -140,24 +132,23 @@ bool profile_gemm_mx_impl(int do_verification, return static_cast(stride); }; - auto Scale_Stride_AM = f_get_default_stride(M, K / ScaleBlockSize, -1, AScaleLayout{}); - auto Scale_Stride_BN = f_get_default_stride(K / ScaleBlockSize, N, -1, BScaleLayout{}); - auto Scale_Padded_M = (M + 32 - 1) / 32 * 32; - auto Scale_Padded_Stride_AM = + auto Scale_Padded_M = (M + 32 - 1) / 32 * 32; + auto Scale_Stride_AM = f_get_default_stride(Scale_Padded_M, K / ScaleBlockSize, -1, AScaleLayout{}); + auto Scale_Stride_BN = f_get_default_stride(K / ScaleBlockSize, N, -1, BScaleLayout{}); Tensor a_m_k(f_host_tensor_descriptor(M, K, StrideA, ALayout{})); Tensor b_k_n(f_host_tensor_descriptor(K, N, StrideB, BLayout{})); // scales for A and B - Tensor a_m_k_scale( - f_host_tensor_descriptor(M, K / ScaleBlockSize, Scale_Stride_AM, AScaleLayout{})); + Tensor a_m_k_scale(f_host_tensor_descriptor( + Scale_Padded_M, K / ScaleBlockSize, Scale_Stride_AM, AScaleLayout{})); Tensor b_k_n_scale( f_host_tensor_descriptor(K / ScaleBlockSize, N, Scale_Stride_BN, BScaleLayout{})); // shuffled scales for A and B Tensor a_shuffled_scale(f_host_tensor_descriptor( - Scale_Padded_M, K / ScaleBlockSize, Scale_Padded_Stride_AM, AScaleLayout{})); + Scale_Padded_M, K / ScaleBlockSize, Scale_Stride_AM, AScaleLayout{})); Tensor b_shuffled_scale( f_host_tensor_descriptor(K / ScaleBlockSize, N, Scale_Stride_BN, BScaleLayout{})); @@ -232,14 +223,12 @@ bool profile_gemm_mx_impl(int do_verification, } #if 1 - preShuffleScaleBuffer>( // - a_m_k_scale.mData.data(), - a_shuffled_scale.mData.data(), - M, - Scale_Padded_M, - K / ScaleBlockSize); + preShuffleScaleBuffer>(a_m_k_scale.mData.data(), + a_shuffled_scale.mData.data(), + Scale_Padded_M, + K / ScaleBlockSize); preShuffleScaleBuffer>( - b_k_n_scale.mData.data(), b_shuffled_scale.mData.data(), N, N, K / ScaleBlockSize); + b_k_n_scale.mData.data(), b_shuffled_scale.mData.data(), N, K / ScaleBlockSize); #endif using AElementOp = ck::tensor_operation::element_wise::PassThrough;