From 93aed9744c43662d6c0777903b8a2211087dd195 Mon Sep 17 00:00:00 2001 From: solin Date: Tue, 25 Nov 2025 16:04:47 +0000 Subject: [PATCH] draft --- .../18_flatmm/mxgemm_moe/mxfp4_moe_flatmm.hpp | 78 ----- .../flatmm/kernel/mx_moe_flatmm_kernel.hpp | 283 +++++++++--------- 2 files changed, 149 insertions(+), 212 deletions(-) delete mode 100644 example/ck_tile/18_flatmm/mxgemm_moe/mxfp4_moe_flatmm.hpp diff --git a/example/ck_tile/18_flatmm/mxgemm_moe/mxfp4_moe_flatmm.hpp b/example/ck_tile/18_flatmm/mxgemm_moe/mxfp4_moe_flatmm.hpp deleted file mode 100644 index 4151b53b4e..0000000000 --- a/example/ck_tile/18_flatmm/mxgemm_moe/mxfp4_moe_flatmm.hpp +++ /dev/null @@ -1,78 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. - -#pragma once - -#include -#include - -#include "ck_tile/core.hpp" -#include "ck_tile/host/kernel_launch.hpp" -#include "ck_tile/ops/moe_flatmm.hpp" - -// GEMM config with 16x16 warp tile for FP4×FP4 MoE -struct MXfp4_MOE_FlatmmConfig16 -{ - static constexpr ck_tile::index_t M_Tile = 64; // MOE 用更小的 M_Tile - static constexpr ck_tile::index_t N_Tile = 256; - static constexpr ck_tile::index_t K_Tile = 256; - - static constexpr ck_tile::index_t M_Warp = 1; - static constexpr ck_tile::index_t N_Warp = 4; - static constexpr ck_tile::index_t K_Warp = 1; - - static constexpr ck_tile::index_t M_Warp_Tile = 16; - static constexpr ck_tile::index_t N_Warp_Tile = 16; - static constexpr ck_tile::index_t K_Warp_Tile = 128; // FP4×FP4 使用更大的 K_Warp_Tile - - static constexpr bool kPadM = false; - static constexpr bool kPadN = false; - static constexpr bool kPadK = false; - - static constexpr bool TransposeC = false; - static constexpr bool UseStructuredSparsity = false; - - static constexpr int kBlockPerCu = 1; - static constexpr int TileParitionerGroupNum = 8; - static constexpr int TileParitionerM01 = 4; - static constexpr auto Scheduler = ck_tile::GemmPipelineScheduler::Default; - static constexpr ck_tile::index_t NumWaveGroups = 1; - static constexpr bool DoubleSmemBuffer = false; - - static constexpr int N_Repeat = N_Tile / N_Warp_Tile / N_Warp; - static constexpr bool TiledMMAPermuteN = false; - - using ComputeDataType = ck_tile::fp16_t; - static constexpr int VectorSizeC = 16; -}; - -auto create_args(int argc, char* argv[]) -{ - ck_tile::ArgParser arg_parser; - arg_parser.insert("num_experts", "8", "Num of experts - 8 by default") - .insert("num_tokens", "256", "M dimensions - 256 by default.") - .insert("topk", "2", "Top K - 2 by default.") - .insert("n", "2048", "N dimensions - 2048 by default.") - .insert("k", "1024", "K dimensions - 1024 by default.") - .insert("stride_a", "", "Tensor A strides - it is empty by default.") - .insert("stride_b", "", "Tensor B strides - it is empty by default.") - .insert("stride_c", "", "Tensor C strides - it is empty by default.") - .insert("a_layout", "R", "A tensor data layout - Row by default.") - .insert("b_layout", "C", "B tensor data layout - Col by default.") - .insert("c_layout", "R", "C tensor data layout - Row by default.") - .insert("gemm_kind", - "gemm1_gate_up", - "Gemm kind in FFN network [gemm1_gate_up | gemm2] - " - "gemm1_gate_up by default.") - .insert("v", "1", "0. No validation, 1. Validation on CPU.") - .insert("warmup", "5", "number of iterations before benchmark the kernel") - .insert("mx_prec", - "fp4xfp4", - "MX precision (fp4xfp4 for both A and B)") - .insert("init", "0", "0:random, 1:constant(1)") - .insert("warp_tile", "0", "0: 16x16") - .insert("repeat", "20", "number of iterations to benchmark the kernel."); - - bool result = arg_parser.parse(argc, argv); - return std::make_tuple(result, arg_parser); -} \ No newline at end of file diff --git a/include/ck_tile/ops/flatmm/kernel/mx_moe_flatmm_kernel.hpp b/include/ck_tile/ops/flatmm/kernel/mx_moe_flatmm_kernel.hpp index b43574dd1e..b5f29834b3 100644 --- a/include/ck_tile/ops/flatmm/kernel/mx_moe_flatmm_kernel.hpp +++ b/include/ck_tile/ops/flatmm/kernel/mx_moe_flatmm_kernel.hpp @@ -22,17 +22,16 @@ template struct MXMoeFlatmmKernel { - using TilePartitioner = remove_cvref_t; - using FlatmmPipeline = remove_cvref_t; - using BlockGemmShape = - remove_cvref_t; - using EpiloguePipeline = remove_cvref_t; - using ALayout = remove_cvref_t; - using BLayout = remove_cvref_t; - using ELayout = remove_cvref_t; - using DsLayout = remove_cvref_t; - using DsDataType = remove_cvref_t; - static constexpr index_t kBlockSize = FlatmmPipeline::BlockSize; + using TilePartitioner = remove_cvref_t; + using FlatmmPipeline = remove_cvref_t; + using BlockGemmShape = remove_cvref_t; + using EpiloguePipeline = remove_cvref_t; + using ALayout = remove_cvref_t; + using BLayout = remove_cvref_t; + using ELayout = remove_cvref_t; + using DsLayout = remove_cvref_t; + using DsDataType = remove_cvref_t; + static constexpr index_t kBlockSize = FlatmmPipeline::BlockSize; static constexpr bool UsePersistentKernel = FlatmmPipeline::UsePersistentKernel; using ADataType = remove_cvref_t; @@ -127,28 +126,27 @@ struct MXMoeFlatmmKernel CK_TILE_HOST static constexpr auto MakeKernelArgs(const MoeFlatmmHostArgs& hostArgs) { - return MXMoeFlatmmKernelArgs{ - hostArgs.p_sorted_token_ids, - hostArgs.p_sorted_expert_ids, - hostArgs.p_max_token_id, - hostArgs.p_sorted_expert_weights, - hostArgs.a_ptr, - hostArgs.b_ptr, - hostArgs.e_ptr, - hostArgs.NumTokens, - hostArgs.TopK, - hostArgs.M, - hostArgs.N, - hostArgs.K, - hostArgs.stride_A, - hostArgs.stride_B, - hostArgs.stride_C, - hostArgs.k_batch, - hostArgs.n_padded_zeros, - hostArgs.k_padded_zeros, - hostArgs.scale_m, - hostArgs.scale_n, - hostArgs.exp_bias}; + return MXMoeFlatmmKernelArgs{hostArgs.p_sorted_token_ids, + hostArgs.p_sorted_expert_ids, + hostArgs.p_max_token_id, + hostArgs.p_sorted_expert_weights, + hostArgs.a_ptr, + hostArgs.b_ptr, + hostArgs.e_ptr, + hostArgs.NumTokens, + hostArgs.TopK, + hostArgs.M, + hostArgs.N, + hostArgs.K, + hostArgs.stride_A, + hostArgs.stride_B, + hostArgs.stride_C, + hostArgs.k_batch, + hostArgs.n_padded_zeros, + hostArgs.k_padded_zeros, + hostArgs.scale_m, + hostArgs.scale_n, + hostArgs.exp_bias}; } [[nodiscard]] CK_TILE_HOST static const std::string GetName() @@ -180,8 +178,7 @@ struct MXMoeFlatmmKernel e = hipOccupancyMaxActiveBlocksPerMultiprocessor( &maxActiveBlocksPerCU, - reinterpret_cast( - kentry<1, MXMoeFlatmmKernel, MXMoeFlatmmKernelArgs>), + reinterpret_cast(kentry<1, MXMoeFlatmmKernel, MXMoeFlatmmKernelArgs>), block_size, dync_smem_size); @@ -201,7 +198,7 @@ struct MXMoeFlatmmKernel { return max(FlatmmPipeline::GetSmemSize(), EpiloguePipeline::GetSmemSize()); } - + CK_TILE_HOST_DEVICE static constexpr index_t GetSmemPongSize() { return FlatmmPipeline::GetSmemSize(); @@ -227,72 +224,7 @@ struct MXMoeFlatmmKernel index_t splitted_k; }; - // template - // CK_TILE_HOST static bool IsSupportedArgument(const KernelArgs& kargs) - // { - // if constexpr(EpiloguePipeline::GetVectorSizeC() % 2 != 0 && - // is_any_of::value) - // { - // return false; - // } - - // if constexpr(UsePersistentKernel) - // { - // if(kargs.k_batch != 1) - // { - // return false; - // } - // } - - // if constexpr(std::is_same_v) - // { - // if(kargs.stride_A < kargs.K || kargs.K % FlatmmPipeline::GetVectorSizeA() != 0) - // { - // return false; - // } - // } - // else - // { - // if(kargs.stride_A < kargs.M || kargs.M % FlatmmPipeline::GetVectorSizeA() != 0) - // { - // return false; - // } - // } - - // if constexpr(std::is_same_v) - // { - // if(kargs.stride_B < kargs.N) - // { - // return false; - // } - // } - // else - // { - // if(kargs.stride_B < kargs.K) - // { - // return false; - // } - // } - - // bool DTensorIsValid = true; - - // if constexpr(std::is_same_v) - // { - // if(kargs.stride_C < kargs.N) - // { - // return false; - // } - // } - // else - // { - // if(kargs.stride_C < kargs.M) - // { - // return false; - // } - // } - // return DTensorIsValid; - // } - template + template CK_TILE_HOST static bool IsSupportedArgument(const KernelArgs& kargs) { if constexpr(EpiloguePipeline::GetVectorSizeC() % 2 != 0 && @@ -487,7 +419,7 @@ struct MXMoeFlatmmKernel const auto& b_flat_tensor_view = [&]() { return make_naive_tensor_view( - b_flat_ptr + expert_id * kFlatN * kFlatK, + b_flat_ptr, make_tuple(kFlatN - kargs.n_padded_zeros / NPerXdl, kFlatK), make_tuple(kFlatK, 1), number{}, @@ -557,8 +489,11 @@ struct MXMoeFlatmmKernel scale_b_desc); }(); - return make_tuple( - a_tensor_view, b_flat_tensor_view, c_tensor_view, scale_a_tensor_view, scale_b_tensor_view); + return make_tuple(a_tensor_view, + b_flat_tensor_view, + c_tensor_view, + scale_a_tensor_view, + scale_b_tensor_view); } template @@ -586,17 +521,17 @@ struct MXMoeFlatmmKernel const auto& c_tensor_view = views.at(I2); if constexpr(std::is_same_v) { - return pad_tensor_view(c_tensor_view, - make_tuple(number{}, - number{}), - sequence{}); + return pad_tensor_view( + c_tensor_view, + make_tuple(number{}, number{}), + sequence{}); } else { - return pad_tensor_view(c_tensor_view, - make_tuple(number{}, - number{}), - sequence{}); + return pad_tensor_view( + c_tensor_view, + make_tuple(number{}, number{}), + sequence{}); } }(); @@ -670,39 +605,119 @@ struct MXMoeFlatmmKernel template CK_TILE_DEVICE void operator()(MXMoeFlatmmKernelArgs kargs) const { - auto tilePartitioner = TilePartitioner{kargs.M, kargs.N}; - const auto [iM, iN] = tilePartitioner.GetOutputTileIndex(blockIdx.x); - const index_t coord_m = __builtin_amdgcn_readfirstlane(iM * TilePartitioner::MPerBlock); - const index_t coord_n = __builtin_amdgcn_readfirstlane(iN * TilePartitioner::NPerBlock); + auto tilePartitioner = TilePartitioner{kargs.M, kargs.N}; + const auto [iM, iN] = tilePartitioner.GetOutputTileIndex(blockIdx.x); + const index_t coord_m = __builtin_amdgcn_readfirstlane(iM * TilePartitioner::MPerBlock); + const index_t coord_n = __builtin_amdgcn_readfirstlane(iN * TilePartitioner::NPerBlock); - this->operator()(kargs, coord_m, coord_n); + this->operator()(kargs, iM, iN); } template - CK_TILE_DEVICE void operator()(MXMoeFlatmmKernelArgs kargs, index_t coord_m, index_t coord_n) const + CK_TILE_DEVICE void operator()(MXMoeFlatmmKernelArgs kargs, index_t iM, index_t iN) const { - // Similar structure to MoeFlatmmKernel::operator() but with MX pipeline - const SplitKBatchOffset splitk_batch_offset(kargs); - const ADataType* a_ptr = static_cast(kargs.a_ptr) + - splitk_batch_offset.a_k_split_offset / APackedSize; - const BDataType* b_flat_ptr = static_cast(kargs.b_ptr) + - splitk_batch_offset.b_k_split_offset / BPackedSize; - EDataType* e_ptr = static_cast(kargs.e_ptr); + const index_t coord_m = __builtin_amdgcn_readfirstlane(iM * TilePartitioner::MPerBlock); + const index_t coord_n = __builtin_amdgcn_readfirstlane(iN * TilePartitioner::NPerBlock); + const index_t max_token_id = kargs.p_max_token_id[0]; + // Early exit if beyond valid range - CHECK THIS FIRST before any array access! + if(coord_m >= max_token_id) + return; + + // Allocate LDS __shared__ char smem_ptr_ping[GetSmemPingSize()]; __shared__ char smem_ptr_pong[GetSmemPongSize()]; + // MOE routing: Get expert ID for this tile (safe now after boundary check) + const index_t expert_id = kargs.p_sorted_expert_ids[iM]; + + // Setup A tensor gather offsets using sorted token IDs + constexpr auto a_dram_dist = FlatmmPipeline::GetADramTileDistribution(); + const auto a_coord = a_dram_dist.calculate_index(); // 2d thread offset, [i_row, i_col] + + constexpr ck_tile::index_t DramMRepeat = + decltype(a_dram_dist)::DstrEncode::hs_lengthss_[number<0>{}][number<0>{}]; + statically_indexed_array a_offsets; + + constexpr index_t token_id_offset = 24; + constexpr index_t token_id_mask = (1 << token_id_offset) - 1; + + auto row_to_token_idx = [&](auto row_idx) { + const index_t fused_token = + kargs.p_sorted_token_ids[row_idx]; // topk-idx[31:24] + token_idx[23:0] + index_t gather_token_id = fused_token & token_id_mask; + if constexpr(!IsInputGemm) + { + gather_token_id = gather_token_id * kargs.TopK + (fused_token >> token_id_offset); + } + return gather_token_id; + }; + + // Calculate gather offsets for each row in the tile + static_for<0, DramMRepeat, 1>{}([&](auto m0) { + const auto row_idx = + coord_m + m0 * (TilePartitioner::MPerBlock / DramMRepeat) + a_coord[I0]; + index_t gather_token_id = row_to_token_idx(row_idx); + a_offsets[m0] = std::is_same_v + ? gather_token_id * kargs.stride_A + : gather_token_id; + }); + + // Prepare pointers with split-K offset and expert routing + const SplitKBatchOffset splitk_batch_offset(kargs); + const long_index_t expert_stride = + __builtin_amdgcn_readfirstlane(long_index_t(kargs.N) * kargs.K); + + const ADataType* a_ptr = + static_cast(kargs.a_ptr) + splitk_batch_offset.a_k_split_offset; + const BDataType* b_flat_ptr = + static_cast(kargs.b_ptr) + + (splitk_batch_offset.b_k_split_offset + expert_stride * expert_id) / BPackedSize; + EDataType* e_ptr = static_cast(kargs.e_ptr); + + // Create MX tensor views with expert routing + const AccDataType* exp_weight_ptr = + static_cast(kargs.p_sorted_expert_weights); + const auto& gemm_tensor_views = MakeGemmTensorViews( + a_ptr, b_flat_ptr, e_ptr, exp_weight_ptr, expert_id, kargs, splitk_batch_offset); + + // Create padded views + const auto& gemm_pad_views = MakeGemmPadViews(gemm_tensor_views); + + // Create tile windows + auto gemm_tile_windows = MakeGemmTileWindows(gemm_pad_views, coord_m, coord_n); + + // Extract windows for GEMM + const auto& a_block_window = gemm_tile_windows.at(I0); + const auto& b_flat_block_window = gemm_tile_windows.at(I1); + auto& c_block_window = gemm_tile_windows.at(I2); + const auto& scale_a_block_window = gemm_tile_windows.at(I3); + const auto& scale_b_block_window = gemm_tile_windows.at(I4); + + // Calculate number of loops const index_t num_loop = TilePartitioner::GetLoopNum(splitk_batch_offset.splitted_k); - // MOE routing metadata - const auto* sorted_token_ids = kargs.p_sorted_token_ids; - const auto* sorted_expert_ids = kargs.p_sorted_expert_ids; - const auto* max_token_id = kargs.p_max_token_id; - const auto* sorted_exp_weights = static_cast(kargs.p_sorted_expert_weights); + // Create scatter-gather tile for A tensor (MOE token routing) + auto a_gather_block_tile = + ck_tile::make_tile_scatter_gather(a_block_window.get_bottom_tensor_view(), + a_block_window.get_window_lengths(), + a_block_window.get_window_origin(), + a_dram_dist, + a_offsets); - // Full MOE routing and GEMM logic would go here - // Following the pattern from moe_flatmm_kernel.hpp but using MX tensor views - // This is a placeholder for the complete implementation + // Execute GEMM with MX scales via Pipeline + const auto& c_block_tile = FlatmmPipeline{}(a_gather_block_tile, + b_flat_block_window, + scale_a_block_window, + scale_b_block_window, + num_loop, + smem_ptr_ping, + smem_ptr_pong); + + // Write output using epilogue + // For MX MOE, we pass empty ds (no bias), and the epilogue handles the shuffle + constexpr auto empty_ds_dram_windows = make_tuple(); + EpiloguePipeline{}(c_block_window, c_block_tile, empty_ds_dram_windows, smem_ptr_ping); } };