From a35bc01f279e320a76dc6d8b0a994006931a3557 Mon Sep 17 00:00:00 2001 From: felix Date: Sun, 26 Oct 2025 07:59:49 +0000 Subject: [PATCH] add instance for tokens<=8 --- .../moe_gemm1_xdl_mx_fp4_bns.cpp | 43 +++++++++++-------- .../moe_gemm1_xdl_mx_fp4_bpreshuffle.cpp | 9 ++-- 2 files changed, 28 insertions(+), 24 deletions(-) diff --git a/example/67_gemm_microscaling/moe_gemm1_xdl_mx_fp4_bns.cpp b/example/67_gemm_microscaling/moe_gemm1_xdl_mx_fp4_bns.cpp index b6d5d8f211..4e53be5f33 100644 --- a/example/67_gemm_microscaling/moe_gemm1_xdl_mx_fp4_bns.cpp +++ b/example/67_gemm_microscaling/moe_gemm1_xdl_mx_fp4_bns.cpp @@ -148,11 +148,13 @@ constexpr ck::index_t ScaleBlockSize = 32; // scaling block constexpr ck::index_t KPerBlock = 256 / DataPackedSize; // 256 f4 = 128 fp4x2 static constexpr ck::index_t Nswizzle = false; static constexpr ck::index_t ActOP = 0; // 0: gelu_and_mul, 1: silu_and_mul -static constexpr ck::index_t MPerBlock = 128; +static constexpr ck::index_t MPerBlock = 32; static constexpr ck::index_t NPerBlock = 64; -static constexpr ck::index_t BlockSize = 256; +static constexpr ck::index_t BlockSize = 128; static constexpr bool MulRoutedWeight = true; +// return ck_moe_stage1_gemm<{A0DataType}, {B0DataType}, {AccDataType}, {EDataType}, {CDEElementOp}, V3, 256, +// 64, 128, 128/sizeof({A0DataType}), 1, 4, {Nswizzle}, {Quant} == static_cast(QuantType::per_Tensor), {MulRoutedWeight}, {ActOP}>; // clang-format off using DeviceOpInstance = ck::tensor_operation::device::DeviceMoeGemmMXBNS< A0Layout, B0Layout, DsLayout, ELayout, @@ -162,10 +164,10 @@ using DeviceOpInstance = ck::tensor_operation::device::Devic MPerBlock, NPerBlock, KPerBlock, 16, 16, 16, 16, - 4, 2, - 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<8, 1, 1, 1>, + 2, 2, + S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, + S<8, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, + 2, 2, S<1, 32, 1, 4>, S<8, 1, 1, 1>, ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v3, ActOP, Nswizzle, true, MulRoutedWeight, ck::index_t, A0DataType>; // clang-format on @@ -178,16 +180,12 @@ int main(int argc, char* argv[]) // per expert: // GEMM shape - constexpr ck::index_t sorted_tile_num = 13; - constexpr ck::index_t valid_tile_num = sorted_tile_num; - ck::index_t sorted_size = sorted_tile_num * MPerBlock; - ck::index_t valid_size = valid_tile_num * MPerBlock; - ck::index_t N = 4096; - ck::index_t K = 6144; - ck::index_t experts = 8; - ck::index_t tokens = 832; - ck::index_t topk = 2; + ck::index_t N = 256; + ck::index_t K = 7168; + ck::index_t experts = 256; + ck::index_t tokens = 64; + ck::index_t topk = 8; if(argc == 1) { @@ -223,6 +221,10 @@ int main(int argc, char* argv[]) throw std::runtime_error("wrong! K must be multiple of ScaleBlockSize."); }; + ck::index_t sorted_tile_num = experts > tokens * topk ? experts : tokens * topk; + ck::index_t valid_tile_num = sorted_tile_num; + ck::index_t sorted_size = sorted_tile_num * MPerBlock; + ck::index_t valid_size = valid_tile_num * MPerBlock; ck::index_t StrideA = K; ck::index_t StrideB = K; ck::index_t StrideE = N; @@ -246,14 +248,16 @@ int main(int argc, char* argv[]) for(int i = 0; i < sorted_tile_num; i++) { - expert_ids.mData[i] = i / ck::math::integer_divide_ceil(valid_tile_num, experts); + expert_ids.mData[i] = i / (valid_tile_num / experts); } + 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++; @@ -461,10 +465,11 @@ int main(int argc, char* argv[]) std::size_t(2) * tokens * N * 2 * topk * K + std::size_t(2) * tokens * N * 2 * topk * K / ScaleBlockSize; + int valid_expert = tokens * topk < experts? tokens * topk : experts; std::size_t num_btype = sizeof(A0DataType) / 2 * tokens * topk * K + - sizeof(B0DataType) / 2 * K * N * 2 * experts + + sizeof(B0DataType) * K * N * 2 / 2 * valid_expert + sizeof(XDataType) * tokens * topk * K / ScaleBlockSize + - sizeof(XDataType) * K / ScaleBlockSize * N * 2 * experts + + sizeof(XDataType) * K / ScaleBlockSize * N * 2 * valid_expert + sizeof(EDataType) * tokens * topk * N; float tflops = static_cast(flop) / 1.E9 / ave_time; diff --git a/example/67_gemm_microscaling/moe_gemm1_xdl_mx_fp4_bpreshuffle.cpp b/example/67_gemm_microscaling/moe_gemm1_xdl_mx_fp4_bpreshuffle.cpp index 10b66795b3..f6e4dee394 100644 --- a/example/67_gemm_microscaling/moe_gemm1_xdl_mx_fp4_bpreshuffle.cpp +++ b/example/67_gemm_microscaling/moe_gemm1_xdl_mx_fp4_bpreshuffle.cpp @@ -208,11 +208,6 @@ int main(int argc, char* argv[]) // per expert: // GEMM shape - constexpr ck::index_t sorted_tile_num = 256; - constexpr ck::index_t valid_tile_num = sorted_tile_num; - ck::index_t sorted_size = sorted_tile_num * MPerBlock; - ck::index_t valid_size = valid_tile_num * MPerBlock; - ck::index_t N = 256; ck::index_t K = 7168; ck::index_t experts = 256; @@ -253,6 +248,10 @@ int main(int argc, char* argv[]) throw std::runtime_error("wrong! K must be multiple of ScaleBlockSize."); }; + ck::index_t sorted_tile_num = experts > tokens * topk ? experts : tokens * topk; + ck::index_t valid_tile_num = sorted_tile_num; + ck::index_t sorted_size = sorted_tile_num * MPerBlock; + ck::index_t valid_size = valid_tile_num * MPerBlock; ck::index_t StrideA = K; ck::index_t StrideB = K; ck::index_t StrideE = N;