From 45a0463f1f3209e17990be02e38929a4b8a07a7e Mon Sep 17 00:00:00 2001 From: root Date: Wed, 26 Mar 2025 08:14:11 +0000 Subject: [PATCH] moe gemm draft v0.1 --- example/ck_tile/XX_moe_gemm/moe_gemm.hpp | 4 +- .../ck_tile/XX_moe_gemm/moe_gemm1_xdl_fp8.cpp | 2 +- .../XX_moe_gemm/run_moe_gemm_example.inc | 10 +- .../reference_fused_single_moe_gemm.hpp | 76 +++++---- .../ops/epilogue/cshuffle_epilogue.hpp | 14 +- .../ops/gemm/kernel/moe_gemm_kernel.hpp | 153 +++++++++++++++++- 6 files changed, 215 insertions(+), 44 deletions(-) diff --git a/example/ck_tile/XX_moe_gemm/moe_gemm.hpp b/example/ck_tile/XX_moe_gemm/moe_gemm.hpp index e12ae0b253..e7c647cb5c 100644 --- a/example/ck_tile/XX_moe_gemm/moe_gemm.hpp +++ b/example/ck_tile/XX_moe_gemm/moe_gemm.hpp @@ -34,9 +34,9 @@ using moe_gemm_kargs = ck_tile::MoeGemmHostArgs; auto create_args(int argc, char* argv[]) { ck_tile::ArgParser arg_parser; - arg_parser.insert("experts", "1", "Num of experts - 8 by default") + arg_parser.insert("experts", "8", "Num of experts - 8 by default") .insert("NumTokens", "128", "M dimensions - 128 by default.") - .insert("TopK", "1", "Top K - 2 by default.") + .insert("TopK", "3", "Top K - 2 by default.") .insert("N", "4096", "N dimensions - 4096 by default.") .insert("K", "4096", "K dimensions - 4096 by default.") .insert("stride_A", "", "Tensor A strides - it is empty by default.") diff --git a/example/ck_tile/XX_moe_gemm/moe_gemm1_xdl_fp8.cpp b/example/ck_tile/XX_moe_gemm/moe_gemm1_xdl_fp8.cpp index 2c63a6b35a..df000d633b 100644 --- a/example/ck_tile/XX_moe_gemm/moe_gemm1_xdl_fp8.cpp +++ b/example/ck_tile/XX_moe_gemm/moe_gemm1_xdl_fp8.cpp @@ -28,7 +28,7 @@ struct MoeGemmKernelParam static const int kBlockPerCu = 1; static const ck_tile::index_t M_Tile = 128; static const ck_tile::index_t N_Tile = 128; - static const ck_tile::index_t K_Tile = 16; // need to ensure the M_per_thread = 1 + static const ck_tile::index_t K_Tile = 32; // need to ensure the M_per_thread = 1 static const ck_tile::index_t M_Warp = 2; static const ck_tile::index_t N_Warp = 2; diff --git a/example/ck_tile/XX_moe_gemm/run_moe_gemm_example.inc b/example/ck_tile/XX_moe_gemm/run_moe_gemm_example.inc index dc55396161..db049c9cc8 100644 --- a/example/ck_tile/XX_moe_gemm/run_moe_gemm_example.inc +++ b/example/ck_tile/XX_moe_gemm/run_moe_gemm_example.inc @@ -94,7 +94,7 @@ int run_moe_gemm_example_with_layouts(int argc, const ck_tile::index_t num_tokens = arg_parser.get_int("NumTokens"); const ck_tile::index_t topk = arg_parser.get_int("TopK"); const ck_tile::index_t repeat = arg_parser.get_int("repeat"); - // const ck_tile::index_t experts = arg_parser.get_int("experts"); + const ck_tile::index_t experts = arg_parser.get_int("experts"); // TODO: replace the magic declaration const ck_tile::index_t MPerBlock = 128; @@ -116,7 +116,9 @@ int run_moe_gemm_example_with_layouts(int argc, // TODO: add the experts' weights in b auto b_k_n_tensor = ck_tile::HostTensor( - ck_tile::host_tensor_descriptor(K, N, stride_B, is_row_major(b_layout))); + is_row_major(b_layout) + ? ck_tile::host_tensor_descriptor(experts * K, N, stride_B, is_row_major(b_layout)) + : ck_tile::host_tensor_descriptor(K, experts * N, stride_B, is_row_major(b_layout))); auto c_m_n_tensor = ck_tile::HostTensor( ck_tile::host_tensor_descriptor(M, N, stride_C, is_row_major(CLayout{}))); @@ -159,8 +161,8 @@ int run_moe_gemm_example_with_layouts(int argc, std::unique_ptr max_token_id_dev = std::make_unique( sizeof(ck_tile::index_t) * max_token_id.get_element_space_size_in_bytes()); - max_token_id.mData = {valid_tile_num * MPerBlock, 0, 1, 2, 3, 4, 5, 6, 7, 8}; - int eids[] = {0, 1, 2, 3, 4, 5, 6, 7, 3, 3, 3}; // {2, 1, 1, 2, 2, 2, 1, 2} + max_token_id.mData = {valid_tile_num * MPerBlock, 0, 1, 2, 3, 4, 6, 7, 8, 8}; + int eids[] = {0, 1, 2, 3, 4, 4, 5, 6, 3, 3, 3, 3}; // {2, 1, 1, 2, 2, 2, 1, 2} for(int i = 0; i < sorted_tile_num; i++) { expert_ids.mData[i] = eids[i]; diff --git a/include/ck_tile/host/reference/reference_fused_single_moe_gemm.hpp b/include/ck_tile/host/reference/reference_fused_single_moe_gemm.hpp index 8a227264ad..59ff88c4f2 100644 --- a/include/ck_tile/host/reference/reference_fused_single_moe_gemm.hpp +++ b/include/ck_tile/host/reference/reference_fused_single_moe_gemm.hpp @@ -77,7 +77,8 @@ template + typename LayoutC, + bool IsInputGemm = true> __global__ void naive_gemm_kernel(const ck_tile::index_t* p_sorted_token_ids_, const ck_tile::index_t* p_sorted_expert_ids_, const ck_tile::index_t* p_max_token_id_, @@ -100,15 +101,26 @@ __global__ void naive_gemm_kernel(const ck_tile::index_t* p_sorted_token_ids_, // assert(p_sorted_expert_ids_ != nullptr); // assert(TopK == 1); // assert(Num_tokens == 128); - if(Num_tokens == 128 && TopK == 1 && p_sorted_expert_ids_ != nullptr) {} + // if(Num_tokens == 128 && TopK == 1 && p_sorted_expert_ids_ != nullptr) {} // index_t max_tokens = p_max_token_id_[0]; - index_t token_id = 0; - // index_t expert_id = 0; + index_t gather_token_id = 0; + index_t scatter_token_id = 0; + index_t expert_id = 0; if(row < p_max_token_id_[0]) { - token_id = p_sorted_token_ids_[row] & 0xffffff; + expert_id = p_sorted_expert_ids_[row / 128]; + gather_token_id = p_sorted_token_ids_[row] & 0xffffff; + scatter_token_id = p_sorted_token_ids_[row] & 0xffffff; + if(!IsInputGemm) + { + gather_token_id = gather_token_id * TopK + (p_sorted_token_ids_[row] >> 24); + } + else + { + scatter_token_id = scatter_token_id * TopK + (p_sorted_token_ids_[row] >> 24); + } } else { @@ -124,13 +136,14 @@ __global__ void naive_gemm_kernel(const ck_tile::index_t* p_sorted_token_ids_, constexpr index_t packed_size_b = ck_tile::numeric_traits::PackedSize; // Adjust indexing based on matrix layout int a_index = (std::is_same_v) - ? token_id * strideA + k - : k * strideA + token_id; + ? gather_token_id * strideA + k + : k * strideA + gather_token_id; // TODO: add experts weights dispatch - int b_index = (std::is_same_v) - ? col * strideB + k - : k * strideB + col; + int b_index = + expert_id * N * K + ((std::is_same_v) + ? col * strideB + k + : k * strideB + col); AccDataType v_a; AccDataType v_b; @@ -162,8 +175,8 @@ __global__ void naive_gemm_kernel(const ck_tile::index_t* p_sorted_token_ids_, } int c_index = (std::is_same_v) - ? token_id * strideC + col - : col * strideC + token_id; + ? scatter_token_id * strideC + col + : col * strideC + scatter_token_id; C[c_index] = ck_tile::type_convert(acc); } } @@ -174,7 +187,8 @@ template + typename LayoutC, + bool IsInputGemm = true> void reference_moe_gemm_gpu(const index_t* p_sorted_token_ids_, const index_t* p_sorted_expert_ids_, const index_t* p_max_token_id_, @@ -194,21 +208,27 @@ void reference_moe_gemm_gpu(const index_t* p_sorted_token_ids_, int numThreadsPerBlock = 256; // Common choice for threads per block int numBlocks = (totalElements + numThreadsPerBlock - 1) / numThreadsPerBlock; - naive_gemm_kernel - <<>>(p_sorted_token_ids_, - p_sorted_expert_ids_, - p_max_token_id_, - a_ptr, - b_ptr, - c_ptr, - Num_tokens, - TopK, - M, - N, - K, - stride_a, - stride_b, - stride_c); + naive_gemm_kernel<<>>(p_sorted_token_ids_, + p_sorted_expert_ids_, + p_max_token_id_, + a_ptr, + b_ptr, + c_ptr, + Num_tokens, + TopK, + M, + N, + K, + stride_a, + stride_b, + stride_c); return; } diff --git a/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp b/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp index 392872ada0..6ed0de57ab 100644 --- a/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp +++ b/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp @@ -121,6 +121,7 @@ struct CShuffleEpilogue template CK_TILE_DEVICE auto operator()(ODramWindow& out_dram_window, const OAccTile& o_acc_tile, @@ -177,10 +178,17 @@ struct CShuffleEpilogue statically_indexed_array offsets; static_for<0, 2 /*CMrepeats*/, 1>{}([&](auto m0) { - auto token_id = token_pos + m0 + c_coord[0] + mIter * kMPerXdl * kMWave; - auto fused_token = p_sorted_tokens_id[token_id]; + auto token_id = token_pos + m0 + c_coord[0] + mIter * kMPerXdl * kMWave; + auto fused_token = p_sorted_tokens_id[token_id]; + index_t token_offset = fused_token & 0xffffff; - offsets[m0] = token_offset * 4096; // Problem::kN_; + + if constexpr(IsInputGemm) + { + token_offset = token_offset * 3 /*TopK*/ + (fused_token >> 24); + } + + offsets[m0] = token_offset * 4096; // Problem::kN_; }); // printf("c_coord[number<0>{}]: %d \n", coord[number<0>{}]); // printf("mIter: %d", mIter+0); diff --git a/include/ck_tile/ops/gemm/kernel/moe_gemm_kernel.hpp b/include/ck_tile/ops/gemm/kernel/moe_gemm_kernel.hpp index 9b107659a3..bff913cfb8 100644 --- a/include/ck_tile/ops/gemm/kernel/moe_gemm_kernel.hpp +++ b/include/ck_tile/ops/gemm/kernel/moe_gemm_kernel.hpp @@ -48,7 +48,10 @@ struct MoeGemmHostArgs : public ck_tile::GemmHostArgs static constexpr index_t KBatch = 1; }; -template +template struct MoeGemmKernel : public GemmKernel { using TilePartitioner = remove_cvref_t; @@ -58,6 +61,8 @@ struct MoeGemmKernel : public GemmKernel; using CLayout = remove_cvref_t; + static constexpr bool IsInputGemm = IsInputGemm_; + using ADataType = remove_cvref_t; using BDataType = remove_cvref_t; using CDataType = remove_cvref_t; @@ -66,8 +71,14 @@ struct MoeGemmKernel : public GemmKernel; using GemmKernelArgs = typename Base::GemmKernelArgs; + using SplitKBatchOffset = typename Base::SplitKBatchOffset; + static constexpr index_t KernelBlockSize = GemmPipeline::BlockSize; + static constexpr auto I0 = number<0>(); + static constexpr auto I1 = number<1>(); + static constexpr auto I2 = number<2>(); + struct MoeGemmKernelArgs : public GemmKernelArgs { const ck_tile::index_t* p_sorted_token_ids; @@ -158,6 +169,127 @@ struct MoeGemmKernel : public GemmKernel + CK_TILE_DEVICE static auto MakeGemmTensorViews(const ADataType* a_ptr, + const BDataType* b_ptr, + CDataType* c_ptr, + const MoeGemmKernelArgs& kargs, + const SplitKBatchOffset& splitk_batch_offset) + { + static_assert(!TilePartitioner::BlockGemmShape::PermuteA, "Not implemented!"); + const auto& a_tensor_view = [&]() { + if constexpr(std::is_same_v) + { + return make_naive_tensor_view( + a_ptr, + make_tuple(IsInputGemm ? kargs.NumTokens : kargs.NumTokens * kargs.TopK, + splitk_batch_offset.splitted_k), + make_tuple(kargs.stride_A, 1), + number{}, + number<1>{}); + } + else + { + return make_naive_tensor_view( + a_ptr, + make_tuple(splitk_batch_offset.splitted_k, + IsInputGemm ? kargs.NumTokens : kargs.NumTokens * kargs.TopK), + make_tuple(kargs.stride_A, 1), + number{}, + number<1>{}); + } + }(); + + const auto& b_tensor_view = [&]() { + if constexpr(std::is_same_v) + { + if constexpr(TilePartitioner::BlockGemmShape::PermuteB) + { + constexpr index_t K1 = GemmPipeline::GetSmemPackB(); + const index_t K0 = splitk_batch_offset.splitted_k / K1; + constexpr index_t VectorSizeB = std::min(K1, GemmPipeline::GetVectorSizeB()); + const auto b_k0_n_k1_desc = + make_naive_tensor_descriptor(make_tuple(K0, kargs.N, K1), + make_tuple(kargs.N * K1, K1, I1), + number{}, + number<1>{}); + const auto b_n_k_desc = transform_tensor_descriptor( + b_k0_n_k1_desc, + make_tuple(make_merge_transform(make_tuple(K0, K1)), + make_pass_through_transform(kargs.N)), + make_tuple(sequence<0, 2>{}, sequence<1>{}), + make_tuple(sequence<0>{}, sequence<1>{})); + return make_tensor_view(b_ptr, b_n_k_desc); + } + else + { + return make_naive_tensor_view( + b_ptr, + make_tuple(splitk_batch_offset.splitted_k, kargs.N), + make_tuple(kargs.stride_B, 1), + number{}, + number<1>{}); + } + } + else + { + if constexpr(TilePartitioner::BlockGemmShape::PermuteB) + { + constexpr index_t K1 = GemmPipeline::GetSmemPackB(); + const index_t K0 = splitk_batch_offset.splitted_k / K1; + constexpr index_t VectorSizeB = std::min(K1, GemmPipeline::GetVectorSizeB()); + const auto b_k0_n_k1_desc = + make_naive_tensor_descriptor(make_tuple(K0, kargs.N, K1), + make_tuple(kargs.N * K1, K1, I1), + number{}, + number<1>{}); + const auto b_n_k_desc = transform_tensor_descriptor( + b_k0_n_k1_desc, + make_tuple(make_merge_transform(make_tuple(K0, K1)), + make_pass_through_transform(kargs.N)), + make_tuple(sequence<0, 2>{}, sequence<1>{}), + make_tuple(sequence<1>{}, sequence<0>{})); + return make_tensor_view(b_ptr, b_n_k_desc); + } + else + { + return make_naive_tensor_view( + b_ptr, + make_tuple(kargs.N, splitk_batch_offset.splitted_k), + make_tuple(kargs.stride_B, 1), + number{}, + number<1>{}); + } + } + }(); + + // TODO: enable vector write for C in ColMajor + const auto& c_tensor_view = [&]() { + if constexpr(std::is_same_v) + { + return make_naive_tensor_view( + c_ptr, + make_tuple(IsInputGemm ? kargs.NumTokens * kargs.TopK : kargs.NumTokens, + kargs.N), + make_tuple(kargs.stride_C, 1), + number{}, + number<1>{}); + } + else + { + return make_naive_tensor_view( + c_ptr, + make_tuple(IsInputGemm ? kargs.NumTokens * kargs.TopK : kargs.NumToken, + kargs.N), + make_tuple(1, kargs.stride_C), + number<1>{}, + number<1>{}); + } + }(); + + return make_tuple(a_tensor_view, b_tensor_view, c_tensor_view); + } + template CK_TILE_DEVICE static auto GetATransformGemmView(const AView& view, const index_t token_id) { @@ -269,6 +401,7 @@ struct MoeGemmKernel : public GemmKernel CK_TILE_DEVICE void operator()(const MoeGemmKernelArgs gemm_desc) const { // TODO: implement C scatter store accordring to expert_id @@ -288,6 +421,10 @@ struct MoeGemmKernel : public GemmKernel{}]: %d \n",a_coord[number<0>{}]); // TODO: token_id should include topk offset depends on ffn1 or ffn2 const index_t token_id = fused_token & 0xffffff; - // const index_t expert_stride = __builtin_amdgcn_readfirstlane(problem.N * problem.K); + if constexpr(!IsInputGemm) + { + token_id = token_id * gemm_desc.TopK + (fused_token >> 24); + } + + const index_t expert_stride = __builtin_amdgcn_readfirstlane(gemm_desc.N * gemm_desc.K); const typename Base::SplitKBatchOffset splitk_batch_offset(gemm_desc); // options const ADataType* a_ptr = static_cast(gemm_desc.a_ptr) + splitk_batch_offset.a_k_split_offset; - const BDataType* b_ptr = - static_cast(gemm_desc.b_ptr) + splitk_batch_offset.b_k_split_offset; + const BDataType* b_ptr = static_cast(gemm_desc.b_ptr) + + splitk_batch_offset.b_k_split_offset + expert_stride * expert_id; CDataType* c_ptr = static_cast(gemm_desc.c_ptr); const auto& gemm_tensor_views_tuple = - Base::MakeGemmTensorViews(a_ptr, b_ptr, c_ptr, gemm_desc, splitk_batch_offset); + MakeGemmTensorViews(a_ptr, b_ptr, c_ptr, gemm_desc, splitk_batch_offset); const auto& gemm_pad_views = Base::MakeGemmPadViews(gemm_tensor_views_tuple); const auto& transformed_views = TransformGemmPadViews(gemm_pad_views, token_id); auto gemm_tile_windows = MakeGemmTileWindows(