From b146191da27513a52d0e8eeed8c267a28ba266eb Mon Sep 17 00:00:00 2001 From: OscarXu Date: Tue, 20 May 2025 16:34:07 +0800 Subject: [PATCH] v2 tok_weight WIP --- .../moe_gemm1_xdl_fp8_blockscale.cpp | 6 +-- .../moe_gemm2_xdl_fp8_blockscale.cpp | 46 +++++++++++++------ .../cpu/reference_moe_gemm2_blockscale.hpp | 10 +++- 3 files changed, 44 insertions(+), 18 deletions(-) diff --git a/example/65_gemm_multiply_multiply/moe_gemm1_xdl_fp8_blockscale.cpp b/example/65_gemm_multiply_multiply/moe_gemm1_xdl_fp8_blockscale.cpp index fee16b5472..dec048199b 100644 --- a/example/65_gemm_multiply_multiply/moe_gemm1_xdl_fp8_blockscale.cpp +++ b/example/65_gemm_multiply_multiply/moe_gemm1_xdl_fp8_blockscale.cpp @@ -124,10 +124,10 @@ static constexpr ck::index_t Scale_Block_N = 128; static constexpr ck::index_t Scale_Block_K = 128; static constexpr ck::index_t Nswizzle = false; -static constexpr ck::index_t ActOP = 1; // 0: gelu_and_mul, 1: silu_and_mul +static constexpr ck::index_t ActOP = 0; // 0: gelu_and_mul, 1: silu_and_mul static constexpr bool MulRoutedWeight = false; -#if 1 +#if 0 static constexpr ck::index_t MPerBlock = 32; static constexpr ck::index_t NPerBlock = 128; static constexpr ck::index_t MNPerXDL = 16; @@ -179,7 +179,7 @@ static constexpr ck::index_t MPerBlock = 64; using DeviceOpInstance = ck::tensor S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, 4, 2, S<1, 32, 1, 8>, S<2, 1, 1, 1>, - ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v1, ActOP, Nswizzle, true, MulRoutedWeight, int32_t, A0DataType>; + ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v3, ActOP, Nswizzle, true, MulRoutedWeight, int32_t, A0DataType>; #endif // clang-format on diff --git a/example/65_gemm_multiply_multiply/moe_gemm2_xdl_fp8_blockscale.cpp b/example/65_gemm_multiply_multiply/moe_gemm2_xdl_fp8_blockscale.cpp index dd6f58f678..b466153706 100644 --- a/example/65_gemm_multiply_multiply/moe_gemm2_xdl_fp8_blockscale.cpp +++ b/example/65_gemm_multiply_multiply/moe_gemm2_xdl_fp8_blockscale.cpp @@ -123,6 +123,7 @@ static constexpr auto GemmSpec = ck::tensor_operation::device::GemmSpecializatio static constexpr ck::index_t Scale_Block_M = 1; static constexpr ck::index_t Scale_Block_N = 128; static constexpr ck::index_t Scale_Block_K = 128; +static constexpr bool MulRoutedWeight = true; #if 0 static constexpr ck::index_t MPerBlock = 32; @@ -156,22 +157,22 @@ using DeviceOpInstance = ck::tensor_operation::device::DeviceMoeGemmBlockScale< S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, AK1, AK1, 0, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, AK1, AK1, 0, 2, 2, S<1, CShuffleMLane, 1, CShuffleNLane>, S, - ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v1, 0, false, false, true, int32_t, A0DataType>; + ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v1, 0, false, false, MulRoutedWeight, int32_t, A0DataType>; #else -static constexpr ck::index_t MPerBlock = 128; using DeviceOpInstance = ck::tensor_operation::device::DeviceMoeGemmBlockScale< +static constexpr ck::index_t MPerBlock = 64; using DeviceOpInstance = ck::tensor_operation::device::DeviceMoeGemmBlockScale< Row, Col, DsLayout, ELayout, A0DataType, A1DataType, B0DataType, B1DataType, DsDataType, EDataType, AccDataType, CShuffleDataType, AElementOp, BElementOp, CDEElementOp, GemmSpec, - 256, Scale_Block_M, Scale_Block_N, Scale_Block_K, - MPerBlock, 128, 128, + 64, Scale_Block_M, Scale_Block_N, Scale_Block_K, + MPerBlock, 16, 128, 16, 16, 16, 16, - 4, 4, - S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, - S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, - 2, 2, S<1, 32, 1, 8>, S<2, 1, 1, 1>, - ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v3, 0, false, false, true, int32_t, A0DataType>; + 4, 1, + S<8, 8, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, + S<8, 8, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, + 2, 1, S<1, 8, 1, 8>, S<2, 1, 1, 1>, + ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v3, 0, false, false, MulRoutedWeight, int32_t, A0DataType>; #endif // clang-format on @@ -187,7 +188,7 @@ int main(int argc, char* argv[]) // per expert: constexpr ck::index_t valid_tile_num = - 52; // 13 for 128; 52 for 32; 4096 for ds // > token * topk / MPerBlock + 2; // 13 for 128; 52 for 32; 4096 for ds // > token * topk / MPerBlock constexpr 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; @@ -195,7 +196,7 @@ int main(int argc, char* argv[]) // GEMM shape ck::index_t N = 6144; ck::index_t K = 4096; - ck::index_t experts = 8; + ck::index_t experts = 2; ck::index_t tokens = 832; ck::index_t topk = 2; #else @@ -349,6 +350,22 @@ int main(int argc, char* argv[]) b1_e_n_k.GenerateTensorValue(GeneratorTensor_3{0, 1.0}); d2_e_n.GenerateTensorValue(GeneratorTensor_3{0, 1.0}); break; + case 6: + a0_t_k_k.GenerateTensorValue(GeneratorTensor_3{0.5, 0.5}); + a1_t_k_k.GenerateTensorValue(GeneratorTensor_3{1.0, 1.0}); + b0_e_n_k.GenerateTensorValue(GeneratorTensor_3{0.5, 0.5}); + b1_e_n_k.GenerateTensorValue(GeneratorTensor_3{1.0, 1.0}); + d2_e_n.GenerateTensorValue(GeneratorTensor_3{0.0, 1.0}); + for (auto i = 0; i < 128; i++){ + a0_t_k_k.mData[i] = ck::type_convert(0.01 * i); + a0_t_k_k.mData[i + 128] = ck::type_convert(0.01 * i + 1.28); + for (auto j = 0; j < 32; j++){ + b0_e_n_k.mData[i * j] = ck::type_convert(0.01 * i * j); + } + d2_e_n.mData[i] = ck::type_convert(0.01 * i); + d2_e_n.mData[i*2] = ck::type_convert(0.01 * i + 1.28); + } + break; default: a0_t_k_k.GenerateTensorValue(GeneratorTensor_3{0.0, 1.0}); a1_t_k_k.GenerateTensorValue(GeneratorTensor_3{0, 1.0}); @@ -487,7 +504,8 @@ int main(int argc, char* argv[]) AccDataType, PassThrough, PassThrough, - CDEElementOp>; + CDEElementOp, + MulRoutedWeight>; auto ref_moe_gemm = ReferenceGemmInstance{}; auto ref_invoker = ref_moe_gemm.MakeInvoker(); auto ref_argument = ref_moe_gemm.MakeArgument(sorted_token_ids, @@ -518,7 +536,7 @@ int main(int argc, char* argv[]) printf("e_t_n_device_result: \n"); for(int t = 0; t < tokens; ++t) { - for(int n = 0; n < N; ++n) + for(int n = 0; n < 5; ++n) { printf("%.2f ", ck::type_convert(e_t_n_device_result(t, n))); } @@ -528,7 +546,7 @@ int main(int argc, char* argv[]) printf("e_t_n_host_result: \n"); for(int t = 0; t < tokens; ++t) { - for(int n = 0; n < N; ++n) + for(int n = 0; n < 5; ++n) { printf("%.2f ", ck::type_convert(e_t_n_host_result(t, n))); } diff --git a/library/include/ck/library/reference_tensor_operation/cpu/reference_moe_gemm2_blockscale.hpp b/library/include/ck/library/reference_tensor_operation/cpu/reference_moe_gemm2_blockscale.hpp index 3972403d8e..3899dffc5c 100644 --- a/library/include/ck/library/reference_tensor_operation/cpu/reference_moe_gemm2_blockscale.hpp +++ b/library/include/ck/library/reference_tensor_operation/cpu/reference_moe_gemm2_blockscale.hpp @@ -23,6 +23,7 @@ template struct ReferenceMoeGemm2BlockScale : public device::BaseOperator @@ -133,7 +134,14 @@ struct ReferenceMoeGemm2BlockScale : public device::BaseOperator ck::type_convert(v_a) * ck::type_convert(v_b); } CDataType v_c{0}; - arg.c_element_op_(v_c, v_acc, v_topk_w); + if constexpr(MulRoutedWeight) + { + arg.c_element_op_(v_c, v_acc, v_topk_w); + } + else + { + arg.c_element_op_(v_c, v_acc, 1.f); + } arg.c_t_n_(t, n) += v_c; } };