ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-12 18:17:58 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix fp8/bf8 B-row

Browse Source
This commit is contained in:
Ding, Yi
2025-05-13 10:13:18 +00:00
parent 178e361101
commit 521471c956
5 changed files with 217 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files

40
example/67_gemm_microscaling/CMakeLists.txt
View File

@@ -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)

98
example/67_gemm_microscaling/gemm_mx_bf8.cpp Normal file
View File

@@ -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<DeviceOpInstance,
ADataType,
BDataType,
XDataType,
CDataType,
ALayout,
BLayout,
CLayout,
AElementOp,
BElementOp,
CElementOp,
AccDataType,
CShuffleDataType,
ScaleBlockSize>(argc, argv)
? 0
: -1;
}

23
example/67_gemm_microscaling/gemm_mx_common.hpp
View File

@@ -104,6 +104,7 @@ bool parse_cmd_args(int argc,
}
#if 1
template <bool KLast>
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<ck::e8m0_bexp_t>(static_cast<float>(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<XDataType, ck::e8m0_bexp_t>)
{
a_m_k_scale.GenerateTensorValue(
GeneratorTensor_2<XDataType>{120, 129}); // scales: {0.25, 0.5, 1, 2}
b_k_n_scale.GenerateTensorValue(
GeneratorTensor_2<XDataType>{125, 129}); // scales: {0.25, 0.5, 1, 2}
// ck::utils::FillConstant<XDataType>{ck::type_convert<XDataType>(1.0f)}(a_m_k_scale);
// ck::utils::FillConstant<XDataType>{ck::type_convert<XDataType>(1.0f)}(b_k_n_scale);
// a_m_k_scale.GenerateTensorValue(
// GeneratorTensor_2<XDataType>{120, 129}); // scales: {0.25, 0.5, 1, 2}
// b_k_n_scale.GenerateTensorValue(
// GeneratorTensor_2<XDataType>{125, 129}); // scales: {0.25, 0.5, 1, 2}
ck::utils::FillConstant<XDataType>{ck::type_convert<XDataType>(1.0f)}(a_m_k_scale);
ck::utils::FillConstant<XDataType>{ck::type_convert<XDataType>(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<ck::is_same_v<ALayout, Row>>(
a_m_k_scale.mData.data(), a_shuffled_scale.mData.data(), M, K / ScaleBlockSize);
preShuffleScaleBuffer(
preShuffleScaleBuffer<ck::is_same_v<BLayout, Col>>(
b_k_n_scale.mData.data(), b_shuffled_scale.mData.data(), N, K / ScaleBlockSize);
#endif
// printf("a_scale:\n");

31
example/67_gemm_microscaling/gemm_mx_f8.cpp → example/67_gemm_microscaling/gemm_mx_fp8.cpp
View File

@@ -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

97
example/67_gemm_microscaling/gemm_mx_fp8_bf8.cpp Normal file
View File

@@ -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<DeviceOpInstance,
ADataType,
BDataType,
XDataType,
CDataType,
ALayout,
BLayout,
CLayout,
AElementOp,
BElementOp,
CElementOp,
AccDataType,
CShuffleDataType,
ScaleBlockSize>(argc, argv)
? 0
: -1;
}