diff --git a/example/01_gemm/CMakeLists.txt b/example/01_gemm/CMakeLists.txt index 96678d275a..bc5b7619ec 100755 --- a/example/01_gemm/CMakeLists.txt +++ b/example/01_gemm/CMakeLists.txt @@ -48,6 +48,9 @@ foreach(gpu IN LISTS GPU_TARGETS) add_example_executable(example_gemm_xdl_bf16_pk_i4_v3 gemm_xdl_bf16_pk_i4_v3.cpp) 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 80f7e95d30..225ca1864f 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 @@ -30,7 +30,7 @@ static constexpr bool PermuteA = false; static constexpr bool PermuteB = false; // clang-format off -#if 0 +#if 1 using DeviceGemmV2Instance = ck::tensor_operation::device::DeviceGemm_Xdl_CShuffleV3_BPreshuffle< ALayout, BLayout, CLayout, @@ -38,14 +38,14 @@ using DeviceGemmV2Instance = AElementOp, BElementOp, CElementOp, GemmDefault, 256, 128, 128, - 256, 16, 32, + 128, 16, 32, 32, 32, 4, 1, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 32, 32, 0, - 1, 1, S<1, 32, 1, 8>, 4, + 1, 1, S<1, 32, 1, 8>, 8, ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v1, F8, F8, PermuteA, PermuteB>; #else @@ -187,6 +187,7 @@ bool run_gemm(const ProblemType& problem_size, const ExecutionConfig& config) } } +#if CK_USE_PK4_LAYOUT_SHUFFLE // vector pk_i4x4 permute for(int i = 0; i < N; i++) { @@ -235,6 +236,8 @@ bool run_gemm(const ProblemType& problem_size, const ExecutionConfig& config) } } } + +#endif a_m_k_device_buf.ToDevice(a_m_k.mData.data()); b_k_n_device_buf.ToDevice(b_k_n_preshuffled.mData.data()); @@ -285,8 +288,12 @@ bool run_gemm(const ProblemType& problem_size, const ExecutionConfig& config) i4 = (i4x2.data >> 0) & 0xf; else i4 = (i4x2.data >> 4) & 0xf; +#if CK_USE_PK4_LAYOUT_SHUFFLE + float v_b = i4_to_f32_gfx9(i4) * 16; +#else + float v_b = i4 - 8; +#endif - float v_b = i4_to_f32_gfx9(i4); b_k_n_f32(k, n) = v_b; } } 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 7b72461dd9..695e1af554 100644 --- a/example/01_gemm/gemm_xdl_fp8_pk_i4_v3.cpp +++ b/example/01_gemm/gemm_xdl_fp8_pk_i4_v3.cpp @@ -27,7 +27,7 @@ using CElementOp = PassThrough; 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 @@ -40,13 +40,13 @@ using DeviceGemmV2Instance = 128, 128, KPerBlock, 16, 32, 32, 32, - 2, 2, + 4, 1, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, 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 @@ -164,6 +164,7 @@ bool run_gemm(const ProblemType& problem_size, const ExecutionConfig& config) } } +#if CK_USE_PK4_LAYOUT_SHUFFLE // vector pk_i4x4 permute for(int i = 0; i < N; i++) { @@ -178,7 +179,7 @@ bool run_gemm(const ProblemType& problem_size, const ExecutionConfig& config) input[k * 2 + 1] = (i4x2 >> 0) & 0xf; } - // permute 01234567->20643175 + // permute 01234567->02461357 { int hi = input[2]; int lo = input[0]; @@ -212,6 +213,7 @@ bool run_gemm(const ProblemType& problem_size, const ExecutionConfig& config) } } } +#endif a_m_k_device_buf.ToDevice(a_m_k.mData.data()); b_k_n_device_buf.ToDevice(b_k_n_permute.mData.data()); @@ -265,7 +267,11 @@ bool run_gemm(const ProblemType& problem_size, const ExecutionConfig& config) else i4 = (i4x2.data >> 4) & 0xf; - float v_b = i4_to_f32_gfx9(i4); +#if CK_USE_PK4_LAYOUT_SHUFFLE + float v_b = i4_to_f32_gfx9(i4) * 16; +#else + float v_b = i4 - 8; +#endif b_k_n_f32(k, n) = v_b; } } diff --git a/example/65_gemm_multiply_multiply/CMakeLists.txt b/example/65_gemm_multiply_multiply/CMakeLists.txt old mode 100644 new mode 100755 index 95fd8bace8..4d4957b6ba --- a/example/65_gemm_multiply_multiply/CMakeLists.txt +++ b/example/65_gemm_multiply_multiply/CMakeLists.txt @@ -12,6 +12,12 @@ foreach(gpu IN LISTS GPU_TARGETS) if(gpu IN_LIST gpu_list AND target EQUAL 0) add_example_executable(example_moe_gemm1_xdl_pk_i4 moe_gemm1_xdl_pk_i4.cpp) 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 -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() set(target 1) endif() endforeach() 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 old mode 100644 new mode 100755 index 1102ce1054..9bd454452b --- a/example/65_gemm_multiply_multiply/moe_gemm1_xdl_pk_i4.cpp +++ b/example/65_gemm_multiply_multiply/moe_gemm1_xdl_pk_i4.cpp @@ -59,11 +59,7 @@ struct MulABScale __host__ __device__ constexpr void operator()( EDataType& e, const float& c, const float& d0, const float& d1) const { -#if CK_USE_PK4_LAYOUT_SHUFFLE - e = ck::type_convert(c * d1 * d0 * 16); -#else e = ck::type_convert(c * d1 * d0); -#endif } }; @@ -82,11 +78,8 @@ struct MulABScaleSilu { // act float x0 = 0; -#if CK_USE_PK4_LAYOUT_SHUFFLE - ck::tensor_operation::element_wise::Silu{}(x0, c * d1 * d0 * 16); -#else + ck::tensor_operation::element_wise::Silu{}(x0, c * d1 * d0); -#endif e = ck::type_convert(x0); } }; @@ -172,7 +165,7 @@ using DeviceOpInstance = ck::tensor_operation::device::DeviceMoeGemm< 256, MPerBlock, 128, 128, 16, 32, 32, 32, - 4, 1, + 4, 1, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 32, 32, 0, 1, 1, S<1, 32, 1, 8>, S<8, 1, 1>, @@ -198,7 +191,8 @@ int main(int argc, char* argv[]) 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 = 644; + ck::index_t tokens = 832; + ck::index_t topk = 2; if(argc == 1) @@ -232,7 +226,7 @@ int main(int argc, char* argv[]) ck::index_t StrideB = K; ck::index_t StrideE = N; constexpr ck::index_t NumDTensor = DsDataType::Size(); - constexpr auto StrideDs = std::array{0, 0}; + constexpr auto StrideDs = std::array{1, 0}; ck::index_t KBatch = 1; @@ -245,12 +239,12 @@ int main(int argc, char* argv[]) { expert_ids.mData[i] = eids[i]; } - int token_per_tile = tokens * topk / valid_tile_num; + int token_per_tile = (tokens * topk + valid_tile_num - 1) / valid_tile_num; int tokenid = 0; for(int i = 0; i < sorted_size; i++) { int tile_off = i % MPerBlock; - if(tile_off < token_per_tile) + if(tile_off < token_per_tile && tokenid < tokens * topk) { sorted_token_ids.mData[i] = (tokenid % tokens) | ((tokenid / tokens) << 24); tokenid++; @@ -377,7 +371,7 @@ int main(int argc, char* argv[]) input[k * 2 + 1] = (i4x2 >> 0) & 0xf; } - // permute 01234567->20643175 + // permute 01234567->04152637 { int hi = input[2]; int lo = input[0]; @@ -461,9 +455,10 @@ int main(int argc, char* argv[]) float gb_per_sec = num_btype / 1.E6 / ave_time; std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec - << " GB/s" << device_op.GetTypeString() << std::endl; + << " GB/s, " << device_op.GetTypeString() << std::endl; } + bool pass = true; if(do_verification) { invoker.Run(argument, StreamConfig{nullptr, false, 0, 0, 1}); @@ -516,11 +511,11 @@ int main(int argc, char* argv[]) } e_device_buf.FromDevice(e_t_n_device_result.mData.data()); - return ck::utils::check_err( - e_t_n_device_result, e_t_n_host_result, "Error: Incorrect results!", 1e-3, 5e-2) - ? 0 - : 1; + pass &= ck::utils::check_err( + e_t_n_device_result, e_t_n_host_result, "Error: Incorrect results!", 1e-3, 5e-2) + ? 0 + : 1; } - return 0; + return pass; } 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 528503a2c4..3ddfae503f 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 @@ -63,11 +63,7 @@ struct MulABScaleExpertWeight { (void)d0; -#if CK_USE_PK4_LAYOUT_SHUFFLE - e = ck::type_convert(c * d1 * d2 * 16); -#else e = ck::type_convert(c * d1 * d2); -#endif } // for reference cpu template <> @@ -75,11 +71,7 @@ struct MulABScaleExpertWeight float& e, const float& c, const float& d0, const float& d1, const float& d2) const { // for reference cpu -#if CK_USE_PK4_LAYOUT_SHUFFLE - e = ck::type_convert(c * d0 * d1 * d2 * 16); -#else e = ck::type_convert(c * d0 * d1 * d2); -#endif } }; @@ -122,15 +114,17 @@ using AElementOp = PassThrough; using BElementOp = PassThrough; using CDEElementOp = MulABScaleExpertWeight; -static constexpr auto GemmSpec = ck::tensor_operation::device::GemmSpecialization::Default; -static constexpr ck::index_t MPerBlock = 128; +static constexpr auto GemmSpec = ck::tensor_operation::device::GemmSpecialization::Default; + +#if 0 +static constexpr ck::index_t MPerBlock = 16; static constexpr ck::index_t BLOCKSIZE = 256; -static constexpr ck::index_t MXDLPerWave = 4; -static constexpr ck::index_t NXDLPerWave = 1; -static constexpr ck::index_t NPerBlock = 128; -static constexpr ck::index_t MNPerXDL = 32; +static constexpr ck::index_t MXDLPerWave = 1; +static constexpr ck::index_t NXDLPerWave = 2; +static constexpr ck::index_t NPerBlock = 256; +static constexpr ck::index_t MNPerXDL = 16; static constexpr ck::index_t KPerBlock = 128 / sizeof(A0DataType); -static constexpr ck::index_t CShuffleNLane = 32; +static constexpr ck::index_t CShuffleNLane = 8; static constexpr ck::index_t CShuffleMLane = BLOCKSIZE / CShuffleNLane; static constexpr ck::index_t AK1 = 16 / sizeof(A0DataType); static constexpr ck::index_t BK1 = 32 / sizeof(B0DataType); @@ -148,9 +142,25 @@ using DeviceOpInstance = ck::tensor_operation::device::Devic MXDLPerWave, NXDLPerWave, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, AK1, AK1, 0, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, BK1, BK1, 0, - 1, 1, S<1, CShuffleMLane, 1, CShuffleNLane>, S, + 1, 2, S<1, CShuffleMLane, 1, CShuffleNLane>, S, ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v1, false, false, A0DataType>; // clang-format on +#else +static constexpr ck::index_t MPerBlock = 128; +using DeviceOpInstance = ck::tensor_operation::device::DeviceMoeGemm + // clang-format off + < Row, Col, DsLayout, ELayout, A0DataType, B0DataType, DsDataType, EDataType, AccDataType, CShuffleDataType, + AElementOp, BElementOp, CDEElementOp, GemmSpec, + 256, MPerBlock, 128, 128, + 16, 32, + 32, 32, + 4, 1, + S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, + S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 32, 32, 0, + 1, 1, S<1, 32, 1, 8>, S<2, 1, 1, 1>, + ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v1, false, false, A0DataType>; +// clang-format on +#endif int main(int argc, char* argv[]) { @@ -166,8 +176,8 @@ int main(int argc, char* argv[]) ck::index_t N = 4096; ck::index_t K = 14336; ck::index_t experts = 8; - ck::index_t sorted_tile_num = 19; ck::index_t valid_tile_num = 16; + ck::index_t sorted_tile_num = valid_tile_num + 3; ck::index_t sorted_size = sorted_tile_num * MPerBlock; ck::index_t valid_size = valid_tile_num * MPerBlock; ck::index_t tokens = 512; @@ -229,7 +239,7 @@ int main(int argc, char* argv[]) for(int i = 0; i < sorted_size; i++) { int tile_off = i % MPerBlock; - if(tile_off < token_per_tile) + if(tile_off < token_per_tile && tokenid < tokens * topk) { sorted_token_ids.mData[i] = (tokenid % tokens) | ((tokenid / tokens) << 24); tokenid++; @@ -429,9 +439,10 @@ int main(int argc, char* argv[]) float gb_per_sec = num_btype / 1.E6 / ave_time; std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec - << " GB/s" << device_op.GetTypeString() << std::endl; + << " GB/s, " << device_op.GetTypeString() << std::endl; } + bool pass = true; if(do_verification) { // gemm2 use atomic, so need to reinit outputs @@ -479,11 +490,11 @@ int main(int argc, char* argv[]) e_device_buf.FromDevice(e_t_n_device_result.mData.data()); - return ck::utils::check_err( - e_t_n_device_result, e_t_n_host_result, "Error: Incorrect results!", 1e-3, 5e-2) - ? 0 - : 1; + pass &= ck::utils::check_err( + e_t_n_device_result, e_t_n_host_result, "Error: Incorrect results!", 1e-3, 5e-2) + ? 0 + : 1; } - return 0; + return pass; } diff --git a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_b_preshuffle_dequant_v1.hpp b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_b_preshuffle_dequant_v1.hpp index 7a565fbaa7..c5307c4d7a 100644 --- a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_b_preshuffle_dequant_v1.hpp +++ b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_b_preshuffle_dequant_v1.hpp @@ -141,6 +141,7 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_bdequant_v1{}([&](auto i) { ignore = i; - __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA + if constexpr(MPerBlock >= 128 && NPerBlock >= 128) + { + __builtin_amdgcn_sched_group_barrier(0x008, 2 * mfma_interleave, 0); + } + else + { + __builtin_amdgcn_sched_group_barrier(0x008, mfma_interleave, 0); + } __builtin_amdgcn_sched_group_barrier(0x020, 1, 0); // VMEM read }); @@ -203,11 +211,12 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_bdequant_v1{}([&](auto i) { - ignore = i; - __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA - __builtin_amdgcn_sched_group_barrier(0x100, 2, 0); // DS read - }); + static_for<0, MPerXDL == 32 ? num_ds_read_inst_a / 2 : num_ds_read_inst_a, 1>{}( + [&](auto i) { + ignore = i; + __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA + __builtin_amdgcn_sched_group_barrier(0x100, MPerXDL == 32 ? 2 : 1, 0); // DS read + }); } template {}); static constexpr auto xdlops_gemm = XdlopsGemm{}; 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 672998d811..38456af555 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 @@ -79,23 +79,60 @@ __device__ inline half4_t i4_to_half4_scale(int q, const ck::half2_t& scale) return res.template AsType()[Number<0>{}]; } -__device__ inline f8x4_t i4_to_f8x4(int q) +__device__ inline uint32_t i4_to_f8x4(int q) { - const int LO = 0x000f000f; - const int HI = 0x00f000f0; + // register values [0, 1, 2, 3] + static constexpr uint32_t reg0 = 0x4C484000; + // register values [4, 5, 6, 7] + static constexpr uint32_t reg1 = 0x56545250; + // register values [-4, -3, -2, -1] + static constexpr uint32_t reg2 = 0xD2D4D6D8; + // register values [-8, -7, -6, -5] + static constexpr uint32_t reg3 = 0xC0C8CCD0; - int lo = amd_assembly_and_b32(q, LO); - int hi = amd_assembly_and_b32(q, HI); + uint32_t tmp_pos, tmp_neg, tmp_res, final_sel; - float f32_0 = amd_assemble_cvt_f32_i4(lo); - float f32_1 = amd_assemble_cvt_f32_i4(lo >> 16); - float f32_2 = amd_assemble_cvt_f32_i4(hi); - float f32_3 = amd_assemble_cvt_f32_i4(hi >> 16); + uint32_t dict_sel = q & 0x07070707; + uint32_t sign = q >> 1; + + #if 0 + asm volatile("v_and_or_b32 %0, %1, %2, %3" + : "=v"(final_sel) + : "v"(sign), "v"(0x04040404), "v"(0x03020100)); +#else + final_sel = (sign & 0x04040404) | 0x03020100; +#endif - return amd_assembly_cvt_f8_to_f32(f32_0, f32_1, f32_2, f32_3); + // vector_type res; + + tmp_pos = __builtin_amdgcn_perm(reg1, reg0, dict_sel); + tmp_neg = __builtin_amdgcn_perm(reg3, reg2, 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>{}]; + return tmp_res; } -__device__ inline f8x8_t i4_to_fp8x8(int q) { return amd_assembly_i4_to_fp8x8(q); } +__device__ inline f8x8_t i4_to_fp8x8(int q) +{ +#if 1 + vector_type result; + + uint32_t res_lo = i4_to_f8x4(bit_cast(q)); + uint32_t res_hi = i4_to_f8x4(bit_cast(q) >> 4); + + result.template AsType()(Number<0>{}) = + bit_cast(__builtin_amdgcn_perm(res_hi, res_lo, 0x06040200)); + result.template AsType()(Number<1>{}) = + bit_cast(__builtin_amdgcn_perm(res_hi, res_lo, 0x07050301)); + + return result.template AsType()[Number<0>{}]; +#else + return amd_assembly_i4_to_fp8x8(q); +#endif +} __device__ inline bhalf4_t i4_to_bhalf4(int q) { @@ -163,8 +200,7 @@ struct PassThroughPack8 __host__ __device__ constexpr void operator()(ck::f8x8_t& y, const ck::pk_i4x4_t& x) const { #if CK_USE_PK4_LAYOUT_SHUFFLE - y = i4_to_fp8x8(bit_cast(x)); - + y= i4_to_fp8x8(bit_cast(x)); #else // Added pk_i4_t to f8x2_fnuz_t conversion vector_type dst; diff --git a/include/ck/utility/amd_inline_asm.hpp b/include/ck/utility/amd_inline_asm.hpp index 0ed60df2c3..a775243041 100644 --- a/include/ck/utility/amd_inline_asm.hpp +++ b/include/ck/utility/amd_inline_asm.hpp @@ -62,7 +62,7 @@ inline __device__ f8x8_t amd_assembly_i4_to_fp8x8(int a) uint32_t fp8x4_0; uint32_t fp8x4_1; float tmp_0, tmp_1, tmp_2; - + asm volatile("v_cvt_off_f32_i4 %[v_tmp_0], %[v_src]\n" "v_cvt_off_f32_i4 %[v_tmp_1], %[v_src], src0_sel:BYTE_2\n" "v_cvt_pk_fp8_f32 %[v_dst_0], %[v_tmp_0], %[v_tmp_1]\n" @@ -83,8 +83,11 @@ inline __device__ f8x8_t amd_assembly_i4_to_fp8x8(int a) [v_dst_1] "+v"(fp8x4_1), [v_src] "+v"(i4x8) :); - - return bit_cast(((static_cast(fp8x4_1) << 32) | fp8x4_0)); + + vector_type out; + out.template AsType()(Number<0>{}) = bit_cast(fp8x4_0); + out.template AsType()(Number<1>{}) = bit_cast(fp8x4_1); + return out.template AsType()[Number<0>{}]; } // c0 += inner_product(a, b0) diff --git a/library/include/ck/library/reference_tensor_operation/cpu/reference_moe_gemm.hpp b/library/include/ck/library/reference_tensor_operation/cpu/reference_moe_gemm.hpp index af735925ed..bde19379bd 100644 --- a/library/include/ck/library/reference_tensor_operation/cpu/reference_moe_gemm.hpp +++ b/library/include/ck/library/reference_tensor_operation/cpu/reference_moe_gemm.hpp @@ -94,7 +94,7 @@ struct ReferenceMoeGemm : public device::BaseOperator else i4 = (i4x2 >> 4) & 0xf; #if CK_USE_PK4_LAYOUT_SHUFFLE - v_a = i4_to_f32_gfx9(i4); + v_a = i4_to_f32_gfx9(i4) * 16; #else v_a = i4 - 8; #endif @@ -113,7 +113,7 @@ struct ReferenceMoeGemm : public device::BaseOperator else i4 = (i4x2 >> 4) & 0xf; #if CK_USE_PK4_LAYOUT_SHUFFLE - v_b = i4_to_f32_gfx9(i4); + v_b = i4_to_f32_gfx9(i4) * 16; #else v_b = i4 - 8; #endif diff --git a/library/include/ck/library/reference_tensor_operation/cpu/reference_moe_gemm2.hpp b/library/include/ck/library/reference_tensor_operation/cpu/reference_moe_gemm2.hpp index 1e8a086bc4..3b928010bb 100644 --- a/library/include/ck/library/reference_tensor_operation/cpu/reference_moe_gemm2.hpp +++ b/library/include/ck/library/reference_tensor_operation/cpu/reference_moe_gemm2.hpp @@ -109,7 +109,7 @@ struct ReferenceMoeGemm2 : public device::BaseOperator else i4 = (i4x2 >> 4) & 0xf; #if CK_USE_PK4_LAYOUT_SHUFFLE - v_a = i4_to_f32_gfx9(i4); + v_a = i4_to_f32_gfx9(i4) * 16; #else v_a = i4 - 8; #endif @@ -127,7 +127,7 @@ struct ReferenceMoeGemm2 : public device::BaseOperator else i4 = (i4x2 >> 4) & 0xf; #if CK_USE_PK4_LAYOUT_SHUFFLE - v_b = i4_to_f32_gfx9(i4); + v_b = i4_to_f32_gfx9(i4) * 16; #else v_b = i4 - 8; #endif