From facd5560e7616b242ea40930f0af4bb4c595ec10 Mon Sep 17 00:00:00 2001 From: mtgu0705 Date: Mon, 24 Mar 2025 18:00:58 +0800 Subject: [PATCH] added new int4->fp8 conversion, performance has been improved by 6%. --- example/01_gemm/CMakeLists.txt | 3 + .../gemm_xdl_fp8_pk_i4_bpreshuffle_v3.cpp | 52 ++++++++++++++++ example/01_gemm/gemm_xdl_fp8_pk_i4_v3.cpp | 56 +++++++++++++++++- .../65_gemm_multiply_multiply/CMakeLists.txt | 2 +- .../moe_gemm1_xdl_pk_i4.cpp | 59 ++++++++++++++++++- .../moe_gemm2_xdl_pk_i4.cpp | 54 +++++++++++++++++ include/ck/ck.hpp | 1 + .../element/unary_element_wise_operation.hpp | 35 +++++++++++ 8 files changed, 256 insertions(+), 6 deletions(-) diff --git a/example/01_gemm/CMakeLists.txt b/example/01_gemm/CMakeLists.txt index 09288cd4ad..65f4b7e923 100755 --- a/example/01_gemm/CMakeLists.txt +++ b/example/01_gemm/CMakeLists.txt @@ -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() diff --git a/example/01_gemm/gemm_xdl_fp8_pk_i4_bpreshuffle_v3.cpp b/example/01_gemm/gemm_xdl_fp8_pk_i4_bpreshuffle_v3.cpp index 43512ff4b3..50071dfe4e 100644 --- a/example/01_gemm/gemm_xdl_fp8_pk_i4_bpreshuffle_v3.cpp +++ b/example/01_gemm/gemm_xdl_fp8_pk_i4_bpreshuffle_v3.cpp @@ -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()); diff --git a/example/01_gemm/gemm_xdl_fp8_pk_i4_v3.cpp b/example/01_gemm/gemm_xdl_fp8_pk_i4_v3.cpp index b52f9d0114..cce2b164a6 100644 --- a/example/01_gemm/gemm_xdl_fp8_pk_i4_v3.cpp +++ b/example/01_gemm/gemm_xdl_fp8_pk_i4_v3.cpp @@ -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()); diff --git a/example/65_gemm_multiply_multiply/CMakeLists.txt b/example/65_gemm_multiply_multiply/CMakeLists.txt index fdc961f716..4d4957b6ba 100755 --- a/example/65_gemm_multiply_multiply/CMakeLists.txt +++ b/example/65_gemm_multiply_multiply/CMakeLists.txt @@ -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() diff --git a/example/65_gemm_multiply_multiply/moe_gemm1_xdl_pk_i4.cpp b/example/65_gemm_multiply_multiply/moe_gemm1_xdl_pk_i4.cpp index 5b694460f2..aba3da8b52 100755 --- a/example/65_gemm_multiply_multiply/moe_gemm1_xdl_pk_i4.cpp +++ b/example/65_gemm_multiply_multiply/moe_gemm1_xdl_pk_i4.cpp @@ -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()); diff --git a/example/65_gemm_multiply_multiply/moe_gemm2_xdl_pk_i4.cpp b/example/65_gemm_multiply_multiply/moe_gemm2_xdl_pk_i4.cpp index e975c9bbf6..b05b8c0001 100644 --- a/example/65_gemm_multiply_multiply/moe_gemm2_xdl_pk_i4.cpp +++ b/example/65_gemm_multiply_multiply/moe_gemm2_xdl_pk_i4.cpp @@ -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()); diff --git a/include/ck/ck.hpp b/include/ck/ck.hpp index 6f510e735e..be2dcee0ec 100644 --- a/include/ck/ck.hpp +++ b/include/ck/ck.hpp @@ -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 diff --git a/include/ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp b/include/ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp index f602e36e73..d53d54bff1 100644 --- a/include/ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp +++ b/include/ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp @@ -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 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()(Number<0>{}) = bit_cast(tmp_res); + + return res.template AsType()[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(x)); +#else + vector_type result; + + result.template AsType()(Number<0>{}) = i4_to_f8x4(bit_cast(x)); + result.template AsType()(Number<1>{}) = i4_to_f8x4(bit_cast(x) >> 4); + + y = result.template AsType()[Number<0>{}]; +#endif #else // Added pk_i4_t to f8x2_fnuz_t conversion