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

added new int4->fp8 conversion, performance has been improved by 6%.

Browse Source
This commit is contained in:
mtgu0705
2025-03-24 18:00:58 +08:00
parent f6d6b1c214
commit facd5560e7
8 changed files with 256 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
example/01_gemm/CMakeLists.txt
View File

@@ -42,6 +42,9 @@ foreach(gpu IN LISTS GPU_TARGETS)
if(gpu IN_LIST gpu_list AND target EQUAL 0)
add_example_executable(example_gemm_xdl_fp8_pk_i4_bpreshuffle_v3 gemm_xdl_fp8_pk_i4_bpreshuffle_v3.cpp)
add_example_executable(example_gemm_xdl_fp8_pk_i4_v3 gemm_xdl_fp8_pk_i4_v3.cpp)
set(EXAMPLE_COMPILE_OPTIONS)
list(APPEND EXAMPLE_COMPILE_OPTIONS -v --save-temps -Wno-gnu-line-marker)
target_compile_options(example_gemm_xdl_fp8_pk_i4_v3 PRIVATE ${EXAMPLE_COMPILE_OPTIONS})
set(target 1)
endif()
endforeach()

52
example/01_gemm/gemm_xdl_fp8_pk_i4_bpreshuffle_v3.cpp
View File

@@ -206,7 +206,9 @@ bool run_gemm(const ProblemType& problem_size, const ExecutionConfig& config)
}
}
#if CK_USE_PK4_LAYOUT_SHUFFLE
// vector pk_i4x4 permute
#if !CK_USE_PK4_LAYOUT_SHUFFLE_V2
for(int i = 0; i < N; i++)
{
for(int j = 0; j < K; j += 8)
@@ -254,6 +256,56 @@ bool run_gemm(const ProblemType& problem_size, const ExecutionConfig& config)
}
}
}
#else
for(int i = 0; i < N; i++)
{
for(int j = 0; j < K; j += 8)
{
int input[8];
for(int k = 0; k < 4; k++)
{
int i4x2 = b_k_n_preshuffled(j + k * 2, i).data;
input[k * 2 + 0] = (i4x2 >> 4) & 0xf;
input[k * 2 + 1] = (i4x2 >> 0) & 0xf;
}
// permute 01234567->20643175
{
int hi = input[4];
int lo = input[0];
int i4x2 = (hi << 4) | lo;
b_k_n_preshuffled(j + 0, i) = i4x2;
}
{
int hi = input[5];
int lo = input[1];
int i4x2 = (hi << 4) | lo;
b_k_n_preshuffled(j + 2, i) = i4x2;
}
{
int hi = input[6];
int lo = input[2];
int i4x2 = (hi << 4) | lo;
b_k_n_preshuffled(j + 4, i) = i4x2;
}
{
int hi = input[7];
int lo = input[3];
int i4x2 = (hi << 4) | lo;
b_k_n_preshuffled(j + 6, i) = i4x2;
}
}
}
#endif
#endif
a_m_k_device_buf.ToDevice(a_m_k.mData.data());
b_k_n_device_buf.ToDevice(b_k_n_preshuffled.mData.data());

56
example/01_gemm/gemm_xdl_fp8_pk_i4_v3.cpp
View File

@@ -46,7 +46,7 @@ using CElementOp = MultiplyConst;
static constexpr auto GemmDefault = ck::tensor_operation::device::GemmSpecialization::Default;
static constexpr bool PermuteA = false;
static constexpr bool PermuteB = true;
static constexpr bool PermuteB = false;
static constexpr ck::index_t KPerBlock = 128;
// clang-format off
@@ -65,7 +65,7 @@ using DeviceGemmV2Instance =
S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>,
2, 32, 32, 0,
1, 1, S<1, 32, 1, 8>, 8,
ck::BlockGemmPipelineScheduler::Interwave, ck::BlockGemmPipelineVersion::v2, ADataType, ADataType, PermuteA, PermuteB>;
ck::BlockGemmPipelineScheduler::Interwave, ck::BlockGemmPipelineVersion::v1, ADataType, ADataType, PermuteA, PermuteB>;
// clang-format on
@@ -183,7 +183,9 @@ bool run_gemm(const ProblemType& problem_size, const ExecutionConfig& config)
}
}
#if CK_USE_PK4_LAYOUT_SHUFFLE
// vector pk_i4x4 permute
#if !CK_USE_PK4_LAYOUT_SHUFFLE_V2
for(int i = 0; i < N; i++)
{
for(int j = 0; j < K; j += 8)
@@ -231,6 +233,56 @@ bool run_gemm(const ProblemType& problem_size, const ExecutionConfig& config)
}
}
}
#else
for(int i = 0; i < N; i++)
{
for(int j = 0; j < K; j += 8)
{
int input[8];
for(int k = 0; k < 4; k++)
{
int i4x2 = b_k_n_permute(j + k * 2, i).data;
input[k * 2 + 0] = (i4x2 >> 4) & 0xf;
input[k * 2 + 1] = (i4x2 >> 0) & 0xf;
}
// permute 01234567->04152637
{
int hi = input[4];
int lo = input[0];
int i4x2 = (hi << 4) | lo;
b_k_n_permute(j + 0, i) = i4x2;
}
{
int hi = input[5];
int lo = input[1];
int i4x2 = (hi << 4) | lo;
b_k_n_permute(j + 2, i) = i4x2;
}
{
int hi = input[6];
int lo = input[2];
int i4x2 = (hi << 4) | lo;
b_k_n_permute(j + 4, i) = i4x2;
}
{
int hi = input[7];
int lo = input[3];
int i4x2 = (hi << 4) | lo;
b_k_n_permute(j + 6, i) = i4x2;
}
}
}
#endif
#endif
a_m_k_device_buf.ToDevice(a_m_k.mData.data());
b_k_n_device_buf.ToDevice(b_k_n_permute.mData.data());

2
example/65_gemm_multiply_multiply/CMakeLists.txt
View File

@@ -14,7 +14,7 @@ foreach(gpu IN LISTS GPU_TARGETS)
add_example_executable(example_moe_gemm2_xdl_pk_i4 moe_gemm2_xdl_pk_i4.cpp)
if(CK_hip_VERSION VERSION_LESS_EQUAL 6.3.42132)
set(EXAMPLE_COMPILE_OPTIONS)
list(APPEND EXAMPLE_COMPILE_OPTIONS -mllvm --amdgpu-enable-max-ilp-scheduling-strategy=1)
list(APPEND EXAMPLE_COMPILE_OPTIONS -v --save-temps -Wno-gnu-line-marker -mllvm --amdgpu-enable-max-ilp-scheduling-strategy=1)
target_compile_options(example_moe_gemm1_xdl_pk_i4 PRIVATE ${EXAMPLE_COMPILE_OPTIONS})
target_compile_options(example_moe_gemm2_xdl_pk_i4 PRIVATE ${EXAMPLE_COMPILE_OPTIONS})
endif()

59
example/65_gemm_multiply_multiply/moe_gemm1_xdl_pk_i4.cpp
View File

@@ -191,14 +191,14 @@ int main(int argc, char* argv[])
// experts = 8
// per expert:
// GEMM shape
ck::index_t N = 14336 * 2;
ck::index_t K = 4096;
ck::index_t N = 4096 * 2;
ck::index_t K = 6144;
ck::index_t experts = 8;
ck::index_t sorted_tile_num = 16;
ck::index_t valid_tile_num = 13;
ck::index_t sorted_size = sorted_tile_num * MPerBlock;
ck::index_t valid_size = valid_tile_num * MPerBlock;
ck::index_t tokens = 64;
ck::index_t tokens = 832;
ck::index_t topk = 2;
if(argc == 1)
@@ -362,6 +362,7 @@ int main(int argc, char* argv[])
#if CK_USE_PK4_LAYOUT_SHUFFLE
// vector pk_i4x4 permute
#if !CK_USE_PK4_LAYOUT_SHUFFLE_V2
for(int e = 0; e < experts; e++)
{
for(int i = 0; i < N; i++)
@@ -412,6 +413,58 @@ int main(int argc, char* argv[])
}
}
}
#else
for(int e = 0; e < experts; e++)
{
for(int i = 0; i < N; i++)
{
for(int j = 0; j < K; j += 8)
{
int input[8];
for(int k = 0; k < 4; k++)
{
int i4x2 = b0_preshuffled(e, j + k * 2, i).data;
input[k * 2 + 0] = (i4x2 >> 4) & 0xf;
input[k * 2 + 1] = (i4x2 >> 0) & 0xf;
}
// permute 01234567->04152637
{
int hi = input[4];
int lo = input[0];
int i4x2 = (hi << 4) | lo;
b0_preshuffled(e, j + 0, i) = i4x2;
}
{
int hi = input[5];
int lo = input[1];
int i4x2 = (hi << 4) | lo;
b0_preshuffled(e, j + 2, i) = i4x2;
}
{
int hi = input[6];
int lo = input[2];
int i4x2 = (hi << 4) | lo;
b0_preshuffled(e, j + 4, i) = i4x2;
}
{
int hi = input[7];
int lo = input[3];
int i4x2 = (hi << 4) | lo;
b0_preshuffled(e, j + 6, i) = i4x2;
}
}
}
}
#endif
#endif
b0_device_buf.ToDevice(b0_preshuffled.mData.data());

54
example/65_gemm_multiply_multiply/moe_gemm2_xdl_pk_i4.cpp
View File

@@ -347,6 +347,8 @@ int main(int argc, char* argv[])
#if CK_USE_PK4_LAYOUT_SHUFFLE
// vector pk_i4x4 permute
#if !CK_USE_PK4_LAYOUT_SHUFFLE_V2
for(int e = 0; e < experts; e++)
{
for(int i = 0; i < N; i++)
@@ -397,6 +399,58 @@ int main(int argc, char* argv[])
}
}
}
#else
for(int e = 0; e < experts; e++)
{
for(int i = 0; i < N; i++)
{
for(int j = 0; j < K; j += 8)
{
int input[8];
for(int k = 0; k < 4; k++)
{
int i4x2 = b0_preshuffled(e, j + k * 2, i).data;
input[k * 2 + 0] = (i4x2 >> 4) & 0xf;
input[k * 2 + 1] = (i4x2 >> 0) & 0xf;
}
// permute 01234567->20643175
{
int hi = input[2];
int lo = input[0];
int i4x2 = (hi << 4) | lo;
b0_preshuffled(e, j + 0, i) = i4x2;
}
{
int hi = input[6];
int lo = input[4];
int i4x2 = (hi << 4) | lo;
b0_preshuffled(e, j + 2, i) = i4x2;
}
{
int hi = input[3];
int lo = input[1];
int i4x2 = (hi << 4) | lo;
b0_preshuffled(e, j + 4, i) = i4x2;
}
{
int hi = input[7];
int lo = input[5];
int i4x2 = (hi << 4) | lo;
b0_preshuffled(e, j + 6, i) = i4x2;
}
}
}
}
#endif
#endif
b0_device_buf.ToDevice(b0_preshuffled.mData.data());

1
include/ck/ck.hpp
View File

@@ -172,6 +172,7 @@ CK_DECLARE_ENV_VAR_BOOL(CK_LOGGING)
// shuffle pk_i4 values during conversion to optimize number of binary
// operations
#define CK_USE_PK4_LAYOUT_SHUFFLE 1
#define CK_USE_PK4_LAYOUT_SHUFFLE_V2 0
// block synchronization only s_wait lgkmcnt(0), not vmcnt(0)
#define CK_EXPERIMENTAL_BLOCK_SYNC_LDS_WITHOUT_SYNC_VMEM 1

35
include/ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp
View File

@@ -81,6 +81,7 @@ __device__ inline half4_t i4_to_half4_scale(int q, const ck::half2_t& scale)
__device__ inline f8x4_t i4_to_f8x4(int q)
{
#if 0
const int LO = 0x000f000f;
const int HI = 0x00f000f0;
@@ -93,6 +94,31 @@ __device__ inline f8x4_t i4_to_f8x4(int q)
float f32_3 = amd_assemble_cvt_f32_i4(hi >> 16);
return amd_assembly_cvt_f8_to_f32(f32_0, f32_1, f32_2, f32_3);
#else
// [0, 1, 2, 3] encoded as FP8
static constexpr uint32_t POS_E4M3s_REG1 = 0x2C282000;
// [4, 5, 6, 7] encoded as FP8
static constexpr uint32_t POS_E4M3s_REG2 = 0x36343230;
// [-8, -7, -6, -5] encoded as FP8
static constexpr uint32_t NEG_E4M3s_REG1 = 0xB2B4B6B8;
// [-4, -3, -2, -1] encoded as FP8
static constexpr uint32_t NEG_E4M3s_REG2 = 0xA0A8ACB0;
uint32_t tmp_pos, tmp_neg, tmp_res;
uint32_t sign= q & 0x08080808;
uint32_t dict_sel = q & 0x07070707;
uint32_t final_sel = 0x03020100 | (sign >> 1);
vector_type<f8_t, 4> res;
tmp_pos = __builtin_amdgcn_perm(POS_E4M3s_REG2, POS_E4M3s_REG1, dict_sel);
tmp_neg = __builtin_amdgcn_perm(NEG_E4M3s_REG2, NEG_E4M3s_REG1, dict_sel);
tmp_res = __builtin_amdgcn_perm(tmp_neg, tmp_pos, final_sel);
res.template AsType<f8x4_t>()(Number<0>{}) = bit_cast<f8x4_t>(tmp_res);
return res.template AsType<f8x4_t>()[Number<0>{}];
#endif
}
__device__ inline f8x8_t i4_to_fp8x8(int q) { return amd_assembly_i4_to_fp8x8(q); }
@@ -163,7 +189,16 @@ struct PassThroughPack8
__host__ __device__ constexpr void operator()(ck::f8x8_t& y, const ck::pk_i4x4_t& x) const
{
#if CK_USE_PK4_LAYOUT_SHUFFLE
#if !CK_USE_PK4_LAYOUT_SHUFFLE_V2
y = i4_to_fp8x8(bit_cast<int>(x));
#else
vector_type<f8_t, 8> result;
result.template AsType<f8x4_t>()(Number<0>{}) = i4_to_f8x4(bit_cast<int>(x));
result.template AsType<f8x4_t>()(Number<1>{}) = i4_to_f8x4(bit_cast<int>(x) >> 4);
y = result.template AsType<f8x8_t>()[Number<0>{}];
#endif
#else
// Added pk_i4_t to f8x2_fnuz_t conversion