From ec8d00d58db371331b63a3c0d4b9bec1fa6307af Mon Sep 17 00:00:00 2001 From: OscarXu Date: Fri, 16 May 2025 04:09:26 -0500 Subject: [PATCH] mx_moe_fp4 ready for aiter with clang-format. --- .../moe_gemm1_xdl_mx_fp4.cpp | 115 +++++++------- ...xdlops_b_preshuffle_mx_moe_gufusion_v1.hpp | 142 ++++++++--------- ...xdlops_b_preshuffle_mx_moe_gufusion_v3.hpp | 146 +++++++++--------- ...ne_xdlops_b_preshuffle_mx_moe_selector.hpp | 3 +- .../gpu/grid/gridwise_moe_mx_gemm.hpp | 125 ++++++++------- .../cpu/reference_moe_mx_gemm1.hpp | 27 ++-- 6 files changed, 283 insertions(+), 275 deletions(-) diff --git a/example/67_gemm_microscaling/moe_gemm1_xdl_mx_fp4.cpp b/example/67_gemm_microscaling/moe_gemm1_xdl_mx_fp4.cpp index c6633033ee..f683c5075f 100644 --- a/example/67_gemm_microscaling/moe_gemm1_xdl_mx_fp4.cpp +++ b/example/67_gemm_microscaling/moe_gemm1_xdl_mx_fp4.cpp @@ -126,10 +126,10 @@ using CDEElementOp = MulABScaleExpertWeight; static constexpr auto GemmSpec = ck::tensor_operation::device::GemmSpecialization::Default; -constexpr ck::index_t ScaleBlockSize = 32; // scaling block size -static constexpr ck::index_t Nswizzle = false; -static constexpr ck::index_t ActOP = 0; // 0: gelu_and_mul, 1: silu_and_mul -static constexpr bool MulRoutedWeight = false; +constexpr ck::index_t ScaleBlockSize = 32; // scaling block size +static constexpr ck::index_t Nswizzle = false; +static constexpr ck::index_t ActOP = 0; // 0: gelu_and_mul, 1: silu_and_mul +static constexpr bool MulRoutedWeight = false; #if 0 static constexpr ck::index_t MPerBlock = 128; @@ -269,12 +269,12 @@ int main(int argc, char* argv[]) } Tensor a0_t_k(HostTensorDescriptor({tokens, K}, {K, 1})); - Tensor a1_t_k(HostTensorDescriptor({tokens, (K + ScaleBlockSize - 1) / ScaleBlockSize}, - {Scale_Stride_AM, 1})); + Tensor a1_t_k(HostTensorDescriptor( + {tokens, (K + ScaleBlockSize - 1) / ScaleBlockSize}, {Scale_Stride_AM, 1})); Tensor b0_e_n_k(HostTensorDescriptor({experts, K, N * 2}, {N * 2 * K, 1, K})); Tensor b1_e_n_k( HostTensorDescriptor({experts, (K + ScaleBlockSize - 1) / ScaleBlockSize, N * 2}, - {(N * 2 *Scale_Stride_BN), 1, Scale_Stride_BN})); + {(N * 2 * Scale_Stride_BN), 1, Scale_Stride_BN})); Tensor b0_preshuffled(HostTensorDescriptor({experts, K, N * 2}, {N * 2 * K, 1, K})); Tensor d0_t_n(HostTensorDescriptor({tokens, N}, {StrideDs[0], 0})); Tensor d1_e_n( @@ -375,7 +375,7 @@ int main(int argc, char* argv[]) int n1 = n % NLane; int k0 = k / (KLane * KPack); - tempk = k % (KLane * KPack); + tempk = k % (KLane * KPack); int k1 = tempk / KPack; int k2 = tempk % KPack; @@ -390,32 +390,32 @@ int main(int argc, char* argv[]) b0_device_buf.ToDevice(b0_preshuffled.mData.data()); - auto invoker = device_op.MakeInvoker(); - auto argument = - device_op.MakeArgument(sorted_token_ids_dev.GetDeviceBuffer(), - expert_ids_dev.GetDeviceBuffer(), - max_token_id_dev.GetDeviceBuffer(), - a0_device_buf.GetDeviceBuffer(), - a1_device_buf.GetDeviceBuffer(), - b0_device_buf.GetDeviceBuffer(), - b1_device_buf.GetDeviceBuffer(), - std::array{nullptr, nullptr, d2_device_buf.GetDeviceBuffer()}, - e_device_buf.GetDeviceBuffer(), - tokens, - topk, - sorted_size, - N, - K, - StrideA, - Scale_Stride_AM, - StrideB, - Scale_Stride_BN, - StrideDs, - StrideE, - KBatch, - a_element_op, - b_element_op, - cde_element_op); + auto invoker = device_op.MakeInvoker(); + auto argument = device_op.MakeArgument( + sorted_token_ids_dev.GetDeviceBuffer(), + expert_ids_dev.GetDeviceBuffer(), + max_token_id_dev.GetDeviceBuffer(), + a0_device_buf.GetDeviceBuffer(), + a1_device_buf.GetDeviceBuffer(), + b0_device_buf.GetDeviceBuffer(), + b1_device_buf.GetDeviceBuffer(), + std::array{nullptr, nullptr, d2_device_buf.GetDeviceBuffer()}, + e_device_buf.GetDeviceBuffer(), + tokens, + topk, + sorted_size, + N, + K, + StrideA, + Scale_Stride_AM, + StrideB, + Scale_Stride_BN, + StrideDs, + StrideE, + KBatch, + a_element_op, + b_element_op, + cde_element_op); if(!device_op.IsSupportedArgument(argument)) { @@ -443,7 +443,8 @@ 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\n" << device_op.GetTypeString() << std::endl; + << " GB/s\n" + << device_op.GetTypeString() << std::endl; } if(do_verification) @@ -454,20 +455,21 @@ int main(int argc, char* argv[]) Tensor c_t_k_n({tokens, topk, N}, {topk * N, N, 1}); - using ReferenceGemmInstance = ck::tensor_operation::host::ReferenceMoeMXGemm1; - auto ref_moe_gemm = ReferenceGemmInstance{}; - auto ref_invoker = ref_moe_gemm.MakeInvoker(); + using ReferenceGemmInstance = + ck::tensor_operation::host::ReferenceMoeMXGemm1; + auto ref_moe_gemm = ReferenceGemmInstance{}; + auto ref_invoker = ref_moe_gemm.MakeInvoker(); auto ref_argument = ref_moe_gemm.MakeArgument(sorted_token_ids, expert_ids, @@ -497,16 +499,19 @@ int main(int argc, char* argv[]) } for(int n = 0; n < N; ++n) { - e_t_n_host_result(t, topk_id, n) = ck::type_convert(c_t_k_n(t, topk_id, n)); + e_t_n_host_result(t, topk_id, n) = + ck::type_convert(c_t_k_n(t, topk_id, n)); } } e_device_buf.FromDevice(e_t_n_device_result.mData.data()); - auto status = ck::utils::check_err( - e_t_n_device_result, e_t_n_host_result, "Error: Incorrect results!", 1e-3, 5e-1) - ? 0 - : 1; - if (status == 0){ + auto status = + ck::utils::check_err( + e_t_n_device_result, e_t_n_host_result, "Error: Incorrect results!", 1e-3, 5e-1) + ? 0 + : 1; + if(status == 0) + { printf("Validation Pass.\n"); } return status; diff --git a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_b_preshuffle_mx_moe_gufusion_v1.hpp b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_b_preshuffle_mx_moe_gufusion_v1.hpp index 3cb0f7b791..11bf96cb53 100644 --- a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_b_preshuffle_mx_moe_gufusion_v1.hpp +++ b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_b_preshuffle_mx_moe_gufusion_v1.hpp @@ -58,44 +58,44 @@ template -struct BlockwiseGemmXdlops_pipeline_bpreshuffle_mx_moe_gufusion_v1 - : BlockwiseGemmXdlops_mx_pipeline_base +struct BlockwiseGemmXdlops_pipeline_bpreshuffle_mx_moe_gufusion_v1< + BlockGemmPipelineScheduler::Intrawave, + ThreadBlockSize, + ScaleBlockSize, + ADataType, + AScaleDataType, + BDataType, + BScaleDataType, + ATileDesc, + BTileDesc, + AMmaTileDesc, + BMmaTileDesc, + ABlockTransferSrcScalarPerVector, + BBlockTransferSrcScalarPerVector, + MPerBlock, + NPerBlock, + KPerBlock, + MPerXDL, + NPerXDL, + MRepeat, + NRepeat, + KPack> : BlockwiseGemmXdlops_mx_pipeline_base { @@ -283,10 +283,10 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_mx_moe_gufusion_v1( b_scale_thread_desc_copy.GetElementSpaceSize()); b_scale_thread_copy_up.Run(b_scale_grid_desc, - b_scale_grid_buf_up, - b_scale_thread_desc_copy, - make_tuple(I0, I0), - b_scale_thread_buf_copy_up); + b_scale_grid_buf_up, + b_scale_thread_desc_copy, + make_tuple(I0, I0), + b_scale_thread_buf_copy_up); b_scale_thread_bufs_up(I0)(Number{}) = b_scale_thread_buf_copy_up[Number<0>{}]; @@ -350,8 +350,8 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_mx_moe_gufusion_v1( b_scale_thread_desc_copy.GetElementSpaceSize()); b_scale_thread_copy_up.Run(b_scale_grid_desc, - b_scale_grid_buf_up, - b_scale_thread_desc_copy, - make_tuple(I0, I0), - b_scale_thread_buf_copy_up); + b_scale_grid_buf_up, + b_scale_thread_desc_copy, + make_tuple(I0, I0), + b_scale_thread_buf_copy_up); b_scale_thread_bufs_up(I1)(Number{}) = b_scale_thread_buf_copy_up[Number<0>{}]; @@ -418,8 +418,8 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_mx_moe_gufusion_v1{}]; b_thread_vec_up.template AsType()(ik) = b_thread_bufs_up[mfma_reg_buf] - [Number{}]; + [Number{}]; }); constexpr index_t a_scale_offset = @@ -518,7 +518,7 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_mx_moe_gufusion_v1{}]; b_scale_thread_vec_up.template AsType()(s) = b_scale_thread_bufs_up[mfma_reg_buf] - [Number{}]; + [Number{}]; }); using mfma_input_type_a = @@ -610,10 +610,10 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_mx_moe_gufusion_v1( b_scale_thread_desc_copy.GetElementSpaceSize()); b_scale_thread_copy_up.Run(b_scale_grid_desc, - b_scale_grid_buf_up, - b_scale_thread_desc_copy, - make_tuple(I0, I0), - b_scale_thread_buf_copy_up); + b_scale_grid_buf_up, + b_scale_thread_desc_copy, + make_tuple(I0, I0), + b_scale_thread_buf_copy_up); b_scale_thread_bufs_up(mfma_reg_buf)(Number{}) = b_scale_thread_buf_copy_up[Number<0>{}]; @@ -687,7 +687,8 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_mx_moe_gufusion_v1 a_scale_thread_vec; vector_type b_scale_thread_vec; - vector_type b_scale_thread_vec_up; + vector_type + b_scale_thread_vec_up; // Pack b_scale_thread_buf into b_scale_thread_vec static_for<0, ScalesPerXdlopsRunPerThread, 1>{}([&](auto s) { @@ -761,7 +762,7 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_mx_moe_gufusion_v1()(ik) = b_thread_bufs[I1][Number{}]; - b_thread_vec_up.template AsType()(ik) = + b_thread_vec_up.template AsType()(ik) = b_thread_bufs_up[I1][Number{}]; }); @@ -774,7 +775,8 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_mx_moe_gufusion_v1 a_scale_thread_vec; vector_type b_scale_thread_vec; - vector_type b_scale_thread_vec_up; + vector_type + b_scale_thread_vec_up; // Pack b_scale_thread_buf into b_scale_thread_vec static_for<0, ScalesPerXdlopsRunPerThread, 1>{}([&](auto s) { @@ -821,7 +823,6 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_mx_moe_gufusion_v1 a_thread_vec; vector_type b_thread_vec; vector_type b_thread_vec_up; - static_for<0, KPack / ComputePackedSize, 1>{}([&](auto ik) { a_thread_vec.template AsType()(ik) = @@ -843,7 +844,8 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_mx_moe_gufusion_v1 a_scale_thread_vec; vector_type b_scale_thread_vec; - vector_type b_scale_thread_vec_up; + vector_type + b_scale_thread_vec_up; // Pack b_scale_thread_buf into b_scale_thread_vec static_for<0, ScalesPerXdlopsRunPerThread, 1>{}([&](auto s) { diff --git a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_b_preshuffle_mx_moe_gufusion_v3.hpp b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_b_preshuffle_mx_moe_gufusion_v3.hpp index b1d1e1692d..94e034654d 100644 --- a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_b_preshuffle_mx_moe_gufusion_v3.hpp +++ b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_b_preshuffle_mx_moe_gufusion_v3.hpp @@ -58,44 +58,44 @@ template -struct BlockwiseGemmXdlops_pipeline_bpreshuffle_mx_moe_gufusion_v3 - : BlockwiseGemmXdlops_mx_pipeline_base +struct BlockwiseGemmXdlops_pipeline_bpreshuffle_mx_moe_gufusion_v3< + BlockGemmPipelineScheduler::Intrawave, + ThreadBlockSize, + ScaleBlockSize, + ADataType, + AScaleDataType, + BDataType, + BScaleDataType, + ATileDesc, + BTileDesc, + AMmaTileDesc, + BMmaTileDesc, + ABlockTransferSrcScalarPerVector, + BBlockTransferSrcScalarPerVector, + MPerBlock, + NPerBlock, + KPerBlock, + MPerXDL, + NPerXDL, + MRepeat, + NRepeat, + KPack> : BlockwiseGemmXdlops_mx_pipeline_base { @@ -389,10 +389,10 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_mx_moe_gufusion_v3( b_scale_thread_desc_copy.GetElementSpaceSize()); b_scale_thread_copy_up.Run(b_scale_grid_desc, - b_scale_grid_buf_up, - b_scale_thread_desc_copy, - make_tuple(I0, I0), - b_scale_thread_buf_copy_up); + b_scale_grid_buf_up, + b_scale_thread_desc_copy, + make_tuple(I0, I0), + b_scale_thread_buf_copy_up); b_scale_thread_bufs_up(I0)(Number{}) = b_scale_thread_buf_copy_up[Number<0>{}]; @@ -455,8 +455,8 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_mx_moe_gufusion_v3vgpr-> lds0 @@ -532,10 +532,10 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_mx_moe_gufusion_v3( b_scale_thread_desc_copy.GetElementSpaceSize()); b_scale_thread_copy_up.Run(b_scale_grid_desc, - b_scale_grid_buf_up, - b_scale_thread_desc_copy, - make_tuple(I0, I0), - b_scale_thread_buf_copy_up); + b_scale_grid_buf_up, + b_scale_thread_desc_copy, + make_tuple(I0, I0), + b_scale_thread_buf_copy_up); b_scale_thread_bufs_up(local_read_buf)(Number{}) = b_scale_thread_buf_copy_up[Number<0>{}]; @@ -574,10 +574,10 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_mx_moe_gufusion_v3{}]; b_thread_vec_up.template AsType()(ik) = b_thread_bufs_up[mfma_reg_buf] - [Number{}]; + [Number{}]; }); constexpr index_t a_scale_offset = @@ -624,7 +624,7 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_mx_moe_gufusion_v3{}]; b_scale_thread_vec_up.template AsType()(s) = b_scale_thread_bufs_up[mfma_reg_buf] - [Number{}]; + [Number{}]; }); using mfma_input_type_a = @@ -725,10 +725,10 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_mx_moe_gufusion_v3( b_scale_thread_desc_copy.GetElementSpaceSize()); b_scale_thread_copy_up.Run(b_scale_grid_desc, - b_scale_grid_buf_up, - b_scale_thread_desc_copy, - make_tuple(I0, I0), - b_scale_thread_buf_copy_up); + b_scale_grid_buf_up, + b_scale_thread_desc_copy, + make_tuple(I0, I0), + b_scale_thread_buf_copy_up); b_scale_thread_bufs_up(I1)(Number{}) = b_scale_thread_buf_copy_up[Number<0>{}]; @@ -755,10 +755,10 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_mx_moe_gufusion_v3{}([&](auto m0) { @@ -788,7 +788,8 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_mx_moe_gufusion_v3 a_scale_thread_vec; vector_type b_scale_thread_vec; - vector_type b_scale_thread_vec_up; + vector_type + b_scale_thread_vec_up; // Pack b_scale_thread_buf into b_scale_thread_vec static_for<0, ScalesPerXdlopsRunPerThread, 1>{}([&](auto s) { @@ -862,7 +863,7 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_mx_moe_gufusion_v3()(ik) = b_thread_bufs[I1][Number{}]; - b_thread_vec_up.template AsType()(ik) = + b_thread_vec_up.template AsType()(ik) = b_thread_bufs_up[I1][Number{}]; }); @@ -875,7 +876,8 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_mx_moe_gufusion_v3 a_scale_thread_vec; vector_type b_scale_thread_vec; - vector_type b_scale_thread_vec_up; + vector_type + b_scale_thread_vec_up; // Pack b_scale_thread_buf into b_scale_thread_vec static_for<0, ScalesPerXdlopsRunPerThread, 1>{}([&](auto s) { @@ -922,7 +924,6 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_mx_moe_gufusion_v3 a_thread_vec; vector_type b_thread_vec; vector_type b_thread_vec_up; - static_for<0, KPack / ComputePackedSize, 1>{}([&](auto ik) { a_thread_vec.template AsType()(ik) = @@ -944,7 +945,8 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_mx_moe_gufusion_v3 a_scale_thread_vec; vector_type b_scale_thread_vec; - vector_type b_scale_thread_vec_up; + vector_type + b_scale_thread_vec_up; // Pack b_scale_thread_buf into b_scale_thread_vec static_for<0, ScalesPerXdlopsRunPerThread, 1>{}([&](auto s) { diff --git a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_b_preshuffle_mx_moe_selector.hpp b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_b_preshuffle_mx_moe_selector.hpp index 2adf0915f2..e4b7009c8f 100644 --- a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_b_preshuffle_mx_moe_selector.hpp +++ b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_b_preshuffle_mx_moe_selector.hpp @@ -95,7 +95,8 @@ constexpr auto BlockGemmMXBPreshufflePipeline_Selector() NPerXDL, MRepeat, NRepeat, - KPack>{};; + KPack>{}; + ; } else { diff --git a/include/ck/tensor_operation/gpu/grid/gridwise_moe_mx_gemm.hpp b/include/ck/tensor_operation/gpu/grid/gridwise_moe_mx_gemm.hpp index 76fdcb3f22..8f3748d15f 100644 --- a/include/ck/tensor_operation/gpu/grid/gridwise_moe_mx_gemm.hpp +++ b/include/ck/tensor_operation/gpu/grid/gridwise_moe_mx_gemm.hpp @@ -1317,8 +1317,9 @@ struct GridwiseMoeGemmMX }); const index_t expert_stride = __builtin_amdgcn_readfirstlane(problem.N * problem.K * (IsInputGemm ? 2 : 1)); - const index_t expert_scale_stride = __builtin_amdgcn_readfirstlane( - problem.N * (IsInputGemm ? 2 : 1) * math::integer_divide_ceil(problem.K, ScaleBlockSize)); + const index_t expert_scale_stride = + __builtin_amdgcn_readfirstlane(problem.N * (IsInputGemm ? 2 : 1) * + math::integer_divide_ceil(problem.K, ScaleBlockSize)); // N0, K0, Blocksize*KPack const index_t n_block_data_idx_on_grid = @@ -1506,50 +1507,49 @@ struct GridwiseMoeGemmMX 0, KPack / KGroup * (get_thread_local_1d_id() % warpSize))); const BScaleDataType* p_b_scale_grid_up = p_b_scale_grid + expert_scale_stride / 2; - const auto b_scale_grid_buf_up = make_dynamic_buffer( + const auto b_scale_grid_buf_up = make_dynamic_buffer( p_b_scale_grid_up + expert_id * expert_scale_stride, b_scale_grid_desc_bn_ak.GetElementSpaceSize()); - auto b_scale_thread_copy_up = - ThreadwiseTensorSliceTransfer_v2, // SliceLengths - Sequence<0, 1>, // DimAccessOrder - 1, // SrcVectorDim - 1, // SrcScalarPerVector - 1, - true>( - b_scale_grid_desc_bn_ak, + auto b_scale_thread_copy_up = ThreadwiseTensorSliceTransfer_v2< + BScaleDataType, + BScaleDataType, + decltype(b_scale_grid_desc_bn_ak), + decltype(BlockwiseGemmPipe::b_scale_thread_desc_copy), + Sequence<1, 1>, // SliceLengths + Sequence<0, 1>, // DimAccessOrder + 1, // SrcVectorDim + 1, // SrcScalarPerVector + 1, + true>( + b_scale_grid_desc_bn_ak, make_multi_index(block_n_id * NPerBlock + b_thread_offset_n, thread_offset_k)); - - blockwise_gemm_pipeline.template Run( - a_grid_desc_ak0_m_ak1, - a_block_desc_ak0_m_ak1, - a_blockwise_copy, - a_grid_buf, - a_block_buf, - a_block_slice_copy_step, - b_grid_desc_bpreshuffled, - b_block_desc_bk0_n_bk1, - b_blockwise_copy, - b_blockwise_copy_up, - b_grid_buf, - b_grid_buf_up, - b_block_buf, - b_block_slice_copy_step, - c_thread_buf, - c_thread_buf_up, - a_scale_grid_desc_am_ak, - a_scale_thread_copy, - a_scale_grid_buf, - b_scale_grid_desc_bn_ak, - b_scale_thread_copy, - b_scale_thread_copy_up, - b_scale_grid_buf, - b_scale_grid_buf_up, - num_k_block_main_loop); + blockwise_gemm_pipeline.template Run( + a_grid_desc_ak0_m_ak1, + a_block_desc_ak0_m_ak1, + a_blockwise_copy, + a_grid_buf, + a_block_buf, + a_block_slice_copy_step, + b_grid_desc_bpreshuffled, + b_block_desc_bk0_n_bk1, + b_blockwise_copy, + b_blockwise_copy_up, + b_grid_buf, + b_grid_buf_up, + b_block_buf, + b_block_slice_copy_step, + c_thread_buf, + c_thread_buf_up, + a_scale_grid_desc_am_ak, + a_scale_thread_copy, + a_scale_grid_buf, + b_scale_grid_desc_bn_ak, + b_scale_thread_copy, + b_scale_thread_copy_up, + b_scale_grid_buf, + b_scale_grid_buf_up, + num_k_block_main_loop); } else { @@ -1611,7 +1611,7 @@ struct GridwiseMoeGemmMX static_for<0, MXdlPerWave, 1>{}([&](auto m0) { // MXDLPerWave static_for<0, M2, 1>{}([&](auto m2) { // m_inst_num_groups_per_blk const index_t m_pos = block_m_id * MPerBlock + m0 * M1 * M2 * M3 * M4 + - m1 * M2 * M3 * M4 + m2 * M3 * M4 + m3 * M4; + m1 * M2 * M3 * M4 + m2 * M3 * M4 + m3 * M4; if constexpr(MulRoutedWeight) { topk_weights = *c_style_pointer_cast*>( @@ -1655,8 +1655,8 @@ struct GridwiseMoeGemmMX c_thread_buf_fp32(cidx) = c_thread_buf[cidx]; if constexpr(MulRoutedWeight) { - c_thread_buf_fp32(cidx) = topk_weights.AsType()[m4] * - c_thread_buf_fp32[cidx]; + c_thread_buf_fp32(cidx) = + topk_weights.AsType()[m4] * c_thread_buf_fp32[cidx]; } } }); @@ -2227,21 +2227,21 @@ struct GridwiseMoeGemmMX 0, KPack / KGroup * (get_thread_local_1d_id() % warpSize))); const BScaleDataType* p_b_scale_grid_up = p_b_scale_grid + expert_scale_stride / 2; - const auto b_scale_grid_buf_up = make_dynamic_buffer( + const auto b_scale_grid_buf_up = make_dynamic_buffer( p_b_scale_grid_up + expert_id * expert_scale_stride, b_scale_grid_desc_bn_ak.GetElementSpaceSize()); - auto b_scale_thread_copy_up = - ThreadwiseTensorSliceTransfer_v2, // SliceLengths - Sequence<0, 1>, // DimAccessOrder - 1, // SrcVectorDim - 1, // SrcScalarPerVector - 1, - true>( - b_scale_grid_desc_bn_ak, + auto b_scale_thread_copy_up = ThreadwiseTensorSliceTransfer_v2< + BScaleDataType, + BScaleDataType, + decltype(b_scale_grid_desc_bn_ak), + decltype(BlockwiseGemmPipe::b_scale_thread_desc_copy), + Sequence<1, 1>, // SliceLengths + Sequence<0, 1>, // DimAccessOrder + 1, // SrcVectorDim + 1, // SrcScalarPerVector + 1, + true>( + b_scale_grid_desc_bn_ak, make_multi_index(block_n_id * NPerBlock + b_thread_offset_n, thread_offset_k)); blockwise_gemm_pipeline.template Run( a_grid_desc_ak0_m_ak1, @@ -2320,19 +2320,18 @@ struct GridwiseMoeGemmMX constexpr auto N2 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I7); // mul scales - + static_assert(M0 * M1 * M2 * M3 * M4 == MPerBlock); static_assert(M4 == 4); const index_t m1 = get_warp_local_1d_id() / NWave; const index_t m3 = threadIdx.x % get_warp_size() / MPerXdl; - vector_type topk_weights; // for gemm2 only static_for<0, NXdlPerWave, 1>{}([&](auto n0) { static_for<0, MXdlPerWave, 1>{}([&](auto m0) { // MXDLPerWave static_for<0, M2, 1>{}([&](auto m2) { // m_inst_num_groups_per_blk const index_t m_pos = block_m_id * MPerBlock + m0 * M1 * M2 * M3 * M4 + - m1 * M2 * M3 * M4 + m2 * M3 * M4 + m3 * M4; + m1 * M2 * M3 * M4 + m2 * M3 * M4 + m3 * M4; if constexpr(MulRoutedWeight) { topk_weights = *c_style_pointer_cast*>( @@ -2376,8 +2375,8 @@ struct GridwiseMoeGemmMX c_thread_buf_fp32(cidx) = c_thread_buf[cidx]; if constexpr(MulRoutedWeight) { - c_thread_buf_fp32(cidx) = topk_weights.AsType()[m4] * - c_thread_buf_fp32[cidx]; + c_thread_buf_fp32(cidx) = + topk_weights.AsType()[m4] * c_thread_buf_fp32[cidx]; } } }); diff --git a/library/include/ck/library/reference_tensor_operation/cpu/reference_moe_mx_gemm1.hpp b/library/include/ck/library/reference_tensor_operation/cpu/reference_moe_mx_gemm1.hpp index 4a40f4c66c..b4673a447e 100644 --- a/library/include/ck/library/reference_tensor_operation/cpu/reference_moe_mx_gemm1.hpp +++ b/library/include/ck/library/reference_tensor_operation/cpu/reference_moe_mx_gemm1.hpp @@ -26,9 +26,9 @@ template + bool MulRoutedWeight = true, + typename ComputeTypeA = CDataType, + typename ComputeTypeB = ComputeTypeA> struct ReferenceMoeMXGemm1 : public device::BaseOperator { // Argument @@ -91,7 +91,7 @@ struct ReferenceMoeMXGemm1 : public device::BaseOperator const int full_n = arg.c_t_k_n_.mDesc.GetLengths()[2]; arg.c_t_k_n_.SetZero(); auto f_mk_kn_mn = [&](auto m, auto n) { - const int K = arg.a_t_k_.mDesc.GetLengths()[1]; + const int K = arg.a_t_k_.mDesc.GetLengths()[1]; const ck::index_t SCALE_BLOCK = K / arg.b_e_n_k_scale_.mDesc.GetLengths()[1]; AccDataType v_acc{0}; AccDataType v_acc_up{0}; @@ -108,7 +108,7 @@ struct ReferenceMoeMXGemm1 : public device::BaseOperator { for(int k = 0; k < K; ++k) { - auto a_f4x2 = arg.a_t_k_(t, k).data; + auto a_f4x2 = arg.a_t_k_(t, k).data; auto a_scale = arg.a_t_k_scale_(t, k / SCALE_BLOCK); if constexpr(is_same_v) { @@ -127,16 +127,15 @@ struct ReferenceMoeMXGemm1 : public device::BaseOperator type_convert(a_scale); arg.a_element_op_(v_a, v_a); } - auto b_f4x2 = arg.b_e_n_k_(e, k, n).data; - auto b_f4x2_up = arg.b_e_n_k_(e, k, n + full_n).data; - auto b_scale = arg.b_e_n_k_scale_(e, k / SCALE_BLOCK, n); - auto b_scale_up = arg.b_e_n_k_scale_(e, k / SCALE_BLOCK, n + full_n); + auto b_f4x2 = arg.b_e_n_k_(e, k, n).data; + auto b_f4x2_up = arg.b_e_n_k_(e, k, n + full_n).data; + auto b_scale = arg.b_e_n_k_scale_(e, k / SCALE_BLOCK, n); + auto b_scale_up = arg.b_e_n_k_scale_(e, k / SCALE_BLOCK, n + full_n); if constexpr(is_same_v) { - - f4_t f4 = 0; - f4_t f4_up = 0; + f4_t f4 = 0; + f4_t f4_up = 0; if(k % 2 == 1) { f4 = (b_f4x2 >> 0) & 0xf; @@ -150,14 +149,14 @@ struct ReferenceMoeMXGemm1 : public device::BaseOperator v_b = type_convert(f4) * type_convert(b_scale); v_b_up = type_convert(f4_up) * - type_convert(b_scale_up); + type_convert(b_scale_up); } else { v_b = type_convert(b_f4x2) * type_convert(b_scale); v_b_up = type_convert(b_f4x2_up) * - type_convert(b_scale_up); + type_convert(b_scale_up); arg.b_element_op_(v_b, v_b); arg.b_element_op_(v_b_up, v_b_up); }