From 521471c956ae597d336247f690a9675ca9461bcf Mon Sep 17 00:00:00 2001 From: "Ding, Yi" Date: Tue, 13 May 2025 10:13:18 +0000 Subject: [PATCH] Fix fp8/bf8 B-row --- example/67_gemm_microscaling/CMakeLists.txt | 40 +------- example/67_gemm_microscaling/gemm_mx_bf8.cpp | 98 +++++++++++++++++++ .../67_gemm_microscaling/gemm_mx_common.hpp | 23 +++-- .../{gemm_mx_f8.cpp => gemm_mx_fp8.cpp} | 31 ++---- .../67_gemm_microscaling/gemm_mx_fp8_bf8.cpp | 97 ++++++++++++++++++ 5 files changed, 217 insertions(+), 72 deletions(-) create mode 100644 example/67_gemm_microscaling/gemm_mx_bf8.cpp rename example/67_gemm_microscaling/{gemm_mx_f8.cpp => gemm_mx_fp8.cpp} (85%) create mode 100644 example/67_gemm_microscaling/gemm_mx_fp8_bf8.cpp diff --git a/example/67_gemm_microscaling/CMakeLists.txt b/example/67_gemm_microscaling/CMakeLists.txt index eb2c0ec820..643dd1cb6b 100644 --- a/example/67_gemm_microscaling/CMakeLists.txt +++ b/example/67_gemm_microscaling/CMakeLists.txt @@ -1,47 +1,13 @@ add_custom_target(example_gemm_mx) -add_example_executable(example_gemm_mx_fp8 gemm_mx_f8.cpp) +add_example_executable(example_gemm_mx_fp8 gemm_mx_fp8.cpp) add_example_dependencies(example_gemm_mx example_gemm_mx_fp8) -if (TARGET example_gemm_mx_fp8) -target_compile_definitions(example_gemm_mx_fp8 PRIVATE - EXAMPLE_A_DATA_TYPE=f8_t - EXAMPLE_B_DATA_TYPE=f8_t - EXAMPLE_C_DATA_TYPE=half_t - EXAMPLE_A_LAYOUT=Row - EXAMPLE_B_LAYOUT=Col - EXAMPLE_C_LAYOUT=Row -) -endif() - -add_example_executable(example_gemm_mx_bf8 gemm_mx_f8.cpp) +add_example_executable(example_gemm_mx_bf8 gemm_mx_bf8.cpp) add_example_dependencies(example_gemm_mx example_gemm_mx_bf8) -if (TARGET example_gemm_mx_bf8) -# TODO: Fix B_LAYOUT=ROW -target_compile_definitions(example_gemm_mx_bf8 PRIVATE - EXAMPLE_A_DATA_TYPE=bf8_t - EXAMPLE_B_DATA_TYPE=bf8_t - EXAMPLE_C_DATA_TYPE=bhalf_t - EXAMPLE_A_LAYOUT=Row - EXAMPLE_B_LAYOUT=Col - EXAMPLE_C_LAYOUT=Row -) -endif() - -add_example_executable(example_gemm_mx_fp8_bf8 gemm_mx_f8.cpp) +add_example_executable(example_gemm_mx_fp8_bf8 gemm_mx_fp8_bf8.cpp) add_example_dependencies(example_gemm_mx example_gemm_mx_fp8_bf8) -if (TARGET example_gemm_mx_fp8_bf8) -# TODO: Fix B_LAYOUT=ROW -target_compile_definitions(example_gemm_mx_fp8_bf8 PRIVATE - EXAMPLE_A_DATA_TYPE=f8_t - EXAMPLE_B_DATA_TYPE=bf8_t - EXAMPLE_C_DATA_TYPE=bhalf_t - EXAMPLE_A_LAYOUT=Row - EXAMPLE_B_LAYOUT=Col - EXAMPLE_C_LAYOUT=Row -) -endif() add_example_executable(example_gemm_mx_fp4 gemm_mx_fp4.cpp) add_example_dependencies(example_gemm_mx example_gemm_mx_fp4) diff --git a/example/67_gemm_microscaling/gemm_mx_bf8.cpp b/example/67_gemm_microscaling/gemm_mx_bf8.cpp new file mode 100644 index 0000000000..495b0bbcb7 --- /dev/null +++ b/example/67_gemm_microscaling/gemm_mx_bf8.cpp @@ -0,0 +1,98 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. + +#include "gemm_mx_common.hpp" + +using ADataType = ck::bf8_t; +using BDataType = ck::bf8_t; + +using XDataType = ck::e8m0_bexp_t; + +using CDataType = ck::bhalf_t; +using AccDataType = float; +using CShuffleDataType = CDataType; + +using ALayout = Row; +using BLayout = Col; +using CLayout = Row; + +using AElementOp = PassThrough; // elementwise transformation for A matrix +using BElementOp = PassThrough; // elementwise transformation for B matrix +using CElementOp = PassThrough; // elementwise transformation for C matrix + +constexpr ck::index_t ScaleBlockSize = 32; // scaling block size +constexpr ck::index_t KPerBlock = 256; + +constexpr auto GemmSpec = ck::tensor_operation::device::GemmSpecialization::Default; +constexpr auto BlkGemmPSched = ck::BlockGemmPipelineScheduler::Intrawave; +constexpr auto BlkGemmPVer = ck::BlockGemmPipelineVersion::v3; + +using DeviceOpInstance = ck::tensor_operation::device::DeviceGemmMX_Xdl_CShuffleV3< + ALayout, // ALayout + BLayout, // BLayout + CLayout, // CLayout + ADataType, // ADataType + XDataType, // AScaleDataType + BDataType, // BDataType + XDataType, // BScaleDataType + CDataType, // CDataType + AccDataType, // GemmAccDataType + CShuffleDataType, // CShuffleDataType + AElementOp, // AElementwiseOperation + BElementOp, // BElementwiseOperation + CElementOp, // CElementwiseOperation + GemmSpec, // GemmSpec + ScaleBlockSize, // ScaleBlockSize: Scaling block size + 128, // BlockSize: Thread block size + 128, // MPerBlock + 32, // NPerBlock + KPerBlock, // KPerBlock + 16, // AK1 + 16, // BK1 + 16, // MPerXDL + 16, // NPerXDL + 4, // MXdlPerWave + 2, // NXdlPerWave + S<8, 16, 1>, // ABlockTransferThreadClusterLengths_AK0_M_AK1 + S<1, 0, 2>, // ABlockTransferThreadClusterArrangeOrder + S<1, 0, 2>, // ABlockTransferSrcAccessOrder + 2, // ABlockTransferSrcVectorDim + 16, // ABlockTransferSrcScalarPerVector + 16, // ABlockTransferDstScalarPerVector_AK1 + false, // ABlockLdsExtraM + S<8, 16, 1>, // BBlockTransferThreadClusterLengths_BK0_N_BK1 + S<1, 0, 2>, // BBlockTransferThreadClusterArrangeOrder + S<1, 0, 2>, // BBlockTransferSrcAccessOrder + 2, // BBlockTransferSrcVectorDim + 16, // BBlockTransferSrcScalarPerVector + 16, // BBlockTransferDstScalarPerVector_BK1 + false, // BBlockLdsExtraN + 2, // CShuffleMXdlPerWavePerShuffle + 2, // CShuffleNXdlPerWavePerShuffle + S<1, 16, 1, 8>, // CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock + 4, // CShuffleBlockTransferScalarPerVector_NPerBlock + BlkGemmPSched, // BlkGemmPipeSched + BlkGemmPVer, // BlkGemmPipelineVer + ADataType, // ComputeTypeA + BDataType // ComputeTypeB + >; + +int main(int argc, char* argv[]) +{ + return run_mx_gemm_example(argc, argv) + ? 0 + : -1; +} diff --git a/example/67_gemm_microscaling/gemm_mx_common.hpp b/example/67_gemm_microscaling/gemm_mx_common.hpp index 2b11ac86f5..737e2d16bc 100644 --- a/example/67_gemm_microscaling/gemm_mx_common.hpp +++ b/example/67_gemm_microscaling/gemm_mx_common.hpp @@ -104,6 +104,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 K) { int MNXdlPack = 2; @@ -142,8 +143,10 @@ void preShuffleScaleBuffer(ck::e8m0_bexp_t* src, ck::e8m0_bexp_t* dst, int MN, i k2 * MNXdlPack + n2; // src[n * K + k] = ck::type_convert(static_cast(powf(2.0f, n2 + // k2 * MNXdlPack))); - dst[outputIndex] = src[n * K + k]; - // printf("Src: %0d, Dst: %d\n", n * K + k, outputIndex);; + if constexpr(KLast) + dst[outputIndex] = src[n * K + k]; + else + dst[outputIndex] = src[k * MN + n]; } } } @@ -287,12 +290,12 @@ bool run_mx_gemm(const ProblemSizeSplitK& problem_size, const ExecutionConfig& c if constexpr(ck::is_same_v) { - a_m_k_scale.GenerateTensorValue( - GeneratorTensor_2{120, 129}); // scales: {0.25, 0.5, 1, 2} - b_k_n_scale.GenerateTensorValue( - GeneratorTensor_2{125, 129}); // scales: {0.25, 0.5, 1, 2} - // ck::utils::FillConstant{ck::type_convert(1.0f)}(a_m_k_scale); - // ck::utils::FillConstant{ck::type_convert(1.0f)}(b_k_n_scale); + // a_m_k_scale.GenerateTensorValue( + // GeneratorTensor_2{120, 129}); // scales: {0.25, 0.5, 1, 2} + // b_k_n_scale.GenerateTensorValue( + // GeneratorTensor_2{125, 129}); // scales: {0.25, 0.5, 1, 2} + ck::utils::FillConstant{ck::type_convert(1.0f)}(a_m_k_scale); + ck::utils::FillConstant{ck::type_convert(1.0f)}(b_k_n_scale); } else { @@ -355,9 +358,9 @@ bool run_mx_gemm(const ProblemSizeSplitK& problem_size, const ExecutionConfig& c } #if 1 - preShuffleScaleBuffer( + preShuffleScaleBuffer>( a_m_k_scale.mData.data(), a_shuffled_scale.mData.data(), M, K / ScaleBlockSize); - preShuffleScaleBuffer( + preShuffleScaleBuffer>( b_k_n_scale.mData.data(), b_shuffled_scale.mData.data(), N, K / ScaleBlockSize); #endif // printf("a_scale:\n"); diff --git a/example/67_gemm_microscaling/gemm_mx_f8.cpp b/example/67_gemm_microscaling/gemm_mx_fp8.cpp similarity index 85% rename from example/67_gemm_microscaling/gemm_mx_f8.cpp rename to example/67_gemm_microscaling/gemm_mx_fp8.cpp index 6dfaaca8ca..405b72a6c4 100644 --- a/example/67_gemm_microscaling/gemm_mx_f8.cpp +++ b/example/67_gemm_microscaling/gemm_mx_fp8.cpp @@ -3,37 +3,18 @@ #include "gemm_mx_common.hpp" -#ifndef EXAMPLE_A_DATA_TYPE -#define EXAMPLE_A_DATA_TYPE f8_t -#endif -#ifndef EXAMPLE_B_DATA_TYPE -#define EXAMPLE_B_DATA_TYPE f8_t -#endif -#ifndef EXAMPLE_C_DATA_TYPE -#define EXAMPLE_C_DATA_TYPE half_t -#endif -#ifndef EXAMPLE_A_LAYOUT -#define EXAMPLE_A_LAYOUT Row -#endif -#ifndef EXAMPLE_B_LAYOUT -#define EXAMPLE_B_LAYOUT Col -#endif -#ifndef EXAMPLE_C_LAYOUT -#define EXAMPLE_C_LAYOUT Row -#endif - -using ADataType = ck::EXAMPLE_A_DATA_TYPE; -using BDataType = ck::EXAMPLE_B_DATA_TYPE; +using ADataType = ck::f8_t; +using BDataType = ck::f8_t; using XDataType = ck::e8m0_bexp_t; -using CDataType = ck::EXAMPLE_C_DATA_TYPE; +using CDataType = ck::half_t; using AccDataType = float; using CShuffleDataType = CDataType; -using ALayout = EXAMPLE_A_LAYOUT; -using BLayout = EXAMPLE_B_LAYOUT; -using CLayout = EXAMPLE_C_LAYOUT; +using ALayout = Row; +using BLayout = Col; +using CLayout = Row; using AElementOp = PassThrough; // elementwise transformation for A matrix using BElementOp = PassThrough; // elementwise transformation for B matrix diff --git a/example/67_gemm_microscaling/gemm_mx_fp8_bf8.cpp b/example/67_gemm_microscaling/gemm_mx_fp8_bf8.cpp new file mode 100644 index 0000000000..2684c65420 --- /dev/null +++ b/example/67_gemm_microscaling/gemm_mx_fp8_bf8.cpp @@ -0,0 +1,97 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. + +#include "gemm_mx_common.hpp" + +using ADataType = ck::f8_t; +using BDataType = ck::bf8_t; + +using XDataType = ck::e8m0_bexp_t; + +using CDataType = ck::bhalf_t; +using AccDataType = float; +using CShuffleDataType = CDataType; + +using ALayout = Row; +using BLayout = Row; +using CLayout = Row; + +using AElementOp = PassThrough; // elementwise transformation for A matrix +using BElementOp = PassThrough; // elementwise transformation for B matrix +using CElementOp = PassThrough; // elementwise transformation for C matrix + +constexpr ck::index_t ScaleBlockSize = 32; // scaling block size + +constexpr auto GemmSpec = ck::tensor_operation::device::GemmSpecialization::Default; +constexpr auto BlkGemmPSched = ck::BlockGemmPipelineScheduler::Intrawave; +constexpr auto BlkGemmPVer = ck::BlockGemmPipelineVersion::v3; + +using DeviceOpInstance = ck::tensor_operation::device::DeviceGemmMX_Xdl_CShuffleV3< + ALayout, // ALayout + BLayout, // BLayout + CLayout, // CLayout + ADataType, // ADataType + XDataType, // AScaleDataType + BDataType, // BDataType + XDataType, // BScaleDataType + CDataType, // CDataType + AccDataType, // GemmAccDataType + CShuffleDataType, // CShuffleDataType + AElementOp, // AElementwiseOperation + BElementOp, // BElementwiseOperation + CElementOp, // CElementwiseOperation + GemmSpec, // GemmSpec + ScaleBlockSize, // ScaleBlockSize: Scaling block size + 256, // BlockSize: Thread block size + 128, // MPerBlock + 128, // NPerBlock + 256, // KPerBlock + 16, // AK1 + 8, // BK1 + 16, // MPerXDL + 16, // NPerXDL + 4, // MXdlPerWave + 4, // NXdlPerWave + S<16, 16, 1>, // ABlockTransferThreadClusterLengths_AK0_M_AK1 + S<1, 0, 2>, // ABlockTransferThreadClusterArrangeOrder + S<1, 0, 2>, // ABlockTransferSrcAccessOrder + 2, // ABlockTransferSrcVectorDim + 16, // ABlockTransferSrcScalarPerVector + 16, // ABlockTransferDstScalarPerVector_AK1 + false, // ABlockLdsExtraM + S<32, 8, 1>, // BBlockTransferThreadClusterLengths_BK0_N_BK1 + S<0, 2, 1>, // BBlockTransferThreadClusterArrangeOrder + S<0, 2, 1>, // BBlockTransferSrcAccessOrder + 1, // BBlockTransferSrcVectorDim + 16, // BBlockTransferSrcScalarPerVector + 8, // BBlockTransferDstScalarPerVector_BK1 + false, // BBlockLdsExtraN + 2, // CShuffleMXdlPerWavePerShuffle + 2, // CShuffleNXdlPerWavePerShuffle + S<1, 32, 1, 8>, // CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock + 8, // CShuffleBlockTransferScalarPerVector_NPerBlock + BlkGemmPSched, // BlkGemmPipeSched + BlkGemmPVer, // BlkGemmPipelineVer + ADataType, // ComputeTypeA + BDataType // ComputeTypeB + >; + +int main(int argc, char* argv[]) +{ + return run_mx_gemm_example(argc, argv) + ? 0 + : -1; +}