ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-07 15:47:05 +00:00
Code Issues Packages Projects Releases Wiki Activity

draft of mxfp4 moe

Browse Source
This commit is contained in:
solin
2025-11-25 08:43:24 +00:00
parent e95337c58c
commit 24dfa0dd48
7 changed files with 1643 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
example/ck_tile/18_flatmm/CMakeLists.txt
View File

@@ -13,6 +13,7 @@ if(has_supported_gpu)
add_executable(tile_example_mixed_prec_flatmm EXCLUDE_FROM_ALL mixed_prec/mixed_prec_flatmm.cpp)
add_executable(tile_example_moe_flatmm EXCLUDE_FROM_ALL moe_flatmm.cpp)
add_executable(tile_example_a16w4_moe_flatmm EXCLUDE_FROM_ALL mixed_prec/a16w4_moe_flatmm.cpp)
add_executable(tile_example_a4w4_moe_flatmm EXCLUDE_FROM_ALL mxgemm_moe/mx_moe_flatmm.cpp)
add_executable(tile_example_grouped_flatmm EXCLUDE_FROM_ALL grouped_flatmm.cpp)
include(mxgemm/mx_flatmm_instance.cmake)
@@ -25,6 +26,7 @@ if(has_supported_gpu)
# ... because they are auto-generated
set(EXAMPLE_FLATMM_COMPILE_OPTIONS -Wno-undefined-func-template)
set(EXAMPLE_MOE_FLATMM_COMPILE_OPTIONS)
list(APPEND EXAMPLE_FLATMM_COMPILE_OPTIONS -Wno-error -v --save-temps)
if(CK_USE_OCP_FP8)
list(APPEND EXAMPLE_FLATMM_COMPILE_OPTIONS -DCK_TILE_USE_OCP_FP8)
@@ -35,6 +37,7 @@ if(has_supported_gpu)
target_compile_options(tile_example_moe_flatmm PRIVATE ${EXAMPLE_FLATMM_COMPILE_OPTIONS})
target_compile_options(tile_example_a16w4_moe_flatmm PRIVATE ${EXAMPLE_FLATMM_COMPILE_OPTIONS})
target_compile_options(tile_example_grouped_flatmm PRIVATE ${EXAMPLE_FLATMM_COMPILE_OPTIONS})
target_compile_options(tile_example_a4w4_moe_flatmm PRIVATE ${EXAMPLE_FLATMM_COMPILE_OPTIONS})
target_compile_options(tile_example_mx_flatmm PRIVATE ${EXAMPLE_FLATMM_COMPILE_OPTIONS}) # TODO: 950 only
endif()

421
example/ck_tile/18_flatmm/mxgemm_moe/mx_moe_flatmm.cpp Normal file
View File

@@ -0,0 +1,421 @@
// SPDX-License-Identifier: MIT
// Copyright (c) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
#include <hip/hip_runtime.h>
#include <cstring>
#include <iostream>
#include <ostream>
#include <string>
#include <tuple>
#include <memory>
#include "mx_moe_flatmm.hpp"
#include "ck_tile/core.hpp"
#include "ck_tile/ops/epilogue.hpp"
#include "ck_tile/ops/gemm.hpp"
#include "ck_tile/ops/flatmm.hpp"
#include "ck_tile/ops/moe_flatmm.hpp"
#include "ck_tile/host.hpp"
#include "ck_tile/host/reference/reference_moe_gemm.hpp"
#include "ck_tile/ops/flatmm/kernel/mx_moe_flatmm_kernel.hpp"
template <typename Layout>
static constexpr inline auto is_row_major(Layout layout_)
{
return ck_tile::bool_constant<std::is_same_v<ck_tile::remove_cvref_t<decltype(layout_)>,
ck_tile::tensor_layout::gemm::RowMajor>>{};
}
// ==================== Kernel Implementation ====================
template <typename FlatmmConfig,
typename ADataType,
typename BDataType,
typename DsDatatype,
typename AccDataType,
typename CDataType,
typename ALayout,
typename BLayout,
typename DsLayout,
typename ELayout,
ck_tile::MoeFlatmmKind moe_kind = ck_tile::MoeFlatmmKind::kFFN_gemm1_gate_only,
typename CDEElementWise = ck_tile::element_wise::PassThrough,
typename MoeFlatmmHostArgs>
float mx_moe_flatmm(const MoeFlatmmHostArgs& args, const ck_tile::stream_config& s)
{
using CodegenFlatmmShape = ck_tile::TileGemmShape<
ck_tile::sequence<FlatmmConfig::M_Tile, FlatmmConfig::N_Tile, FlatmmConfig::K_Tile>,
ck_tile::sequence<FlatmmConfig::M_Warp, FlatmmConfig::N_Warp, FlatmmConfig::K_Warp>,
ck_tile::sequence<FlatmmConfig::M_Warp_Tile,
FlatmmConfig::N_Warp_Tile,
FlatmmConfig::K_Warp_Tile>>;
using TilePartitioner =
ck_tile::GemmSpatiallyLocalTilePartitioner<CodegenFlatmmShape,
FlatmmConfig::TileParitionerGroupNum,
FlatmmConfig::TileParitionerM01>;
using CodegenGemmTraits = ck_tile::TileGemmUniversalTraits<FlatmmConfig::kPadM,
FlatmmConfig::kPadN,
FlatmmConfig::kPadK,
FlatmmConfig::DoubleSmemBuffer,
ALayout,
BLayout,
ELayout,
FlatmmConfig::TransposeC,
FlatmmConfig::UseStructuredSparsity,
false, // UsePersistentKernel_
FlatmmConfig::NumWaveGroups,
true>; // Preshuffle_
// ⭐ FP4×FP4 always uses MX pipeline
constexpr bool MXFP4_Pipeline = true;
static_assert(std::is_same_v<ADataType, ck_tile::pk_fp4_t> &&
std::is_same_v<BDataType, ck_tile::pk_fp4_t>,
"mx_moe_flatmm requires FP4×FP4");
if constexpr(moe_kind == ck_tile::MoeFlatmmKind::kFFN_gemm1_gate_up)
{
static_assert(
FlatmmConfig::N_Tile % (FlatmmConfig::N_Warp * FlatmmConfig::N_Warp_Tile * 2) == 0,
"requires NRepeat is multiple of 2 for FFN_gemm1_gate_up");
}
using ComputeDataType = ck_tile::pk_fp4_t; // ⭐ FP4→FP16 dequantize
using GemmPipelineProblem = ck_tile::GemmPipelineProblem<ComputeDataType,
ComputeDataType,
AccDataType,
CodegenFlatmmShape,
ck_tile::TileGemmTraits<FlatmmConfig::kPadM,
FlatmmConfig::kPadN,
FlatmmConfig::kPadK,
ALayout,
BLayout,
ELayout,
FlatmmConfig::NumWaveGroups>>;
using BaseGemmPipeline = ck_tile::BaseFlatmmPipelineAGmemBGmemCRegV1<GemmPipelineProblem>;
const ck_tile::index_t k_grain = args.k_batch * FlatmmConfig::K_Tile;
const ck_tile::index_t K_split = (args.K + k_grain - 1) / k_grain * FlatmmConfig::K_Tile;
const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split);
const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop);
const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop);
float ave_time{0};
const auto Run = [&](const auto has_hot_loop_,
const auto tail_number_,
const auto memory_operation_) {
constexpr bool has_hot_loop_v = has_hot_loop_.value;
constexpr auto tail_number_v = tail_number_.value;
constexpr auto scheduler = FlatmmConfig::Scheduler;
constexpr auto memory_operation = memory_operation_.value;
// ⭐ 使用 MXF4 Pipeline (FP4×FP4)
using CodegenPipelineProblem =
ck_tile::MXFlatmmPipelineProblem<ADataType,
BDataType,
AccDataType,
CodegenFlatmmShape,
CodegenGemmTraits,
scheduler,
has_hot_loop_v,
tail_number_v>;
constexpr int BlockedXDLN_PerWarp = 2;
using GemmEpilogue = ck_tile::CShuffleEpilogue<
ck_tile::CShuffleEpilogueProblem<ADataType,
BDataType,
DsDatatype,
AccDataType,
CDataType,
DsLayout,
ELayout,
CDEElementWise,
TilePartitioner::MPerBlock,
TilePartitioner::NPerBlock,
FlatmmConfig::M_Warp,
FlatmmConfig::N_Warp,
FlatmmConfig::M_Warp_Tile,
FlatmmConfig::N_Warp_Tile,
FlatmmConfig::K_Warp_Tile,
CodegenPipelineProblem::TransposeC,
memory_operation,
FlatmmConfig::NumWaveGroups,
false,
1,
FlatmmConfig::TiledMMAPermuteN,
BlockedXDLN_PerWarp>>;
// ⭐ 使用 MXF4MoeFlatmmPipeline (FP4×FP4 专用)
using CodegenFlatmmPipeline =
ck_tile::MXF4FlatmmPipelineAGmemBGmemCRegV1<CodegenPipelineProblem>;
// ⭐ FP4×FP4 使用 MoeSilu (不是 Swiglu,因为没有 bias)
using FusedAct = ck_tile::moe::MoeSilu;
using Kernel = ck_tile::MXMoeFlatmmKernel<TilePartitioner,
CodegenFlatmmPipeline,
GemmEpilogue,
moe_kind,
FusedAct>;
auto kargs = Kernel::MakeKernelArgs(args);
const dim3 grids = Kernel::GridSize(kargs);
constexpr dim3 blocks = Kernel::BlockSize();
if(!Kernel::IsSupportedArgument(kargs))
{
throw std::runtime_error("Wrong! Arguments not supported! Skipping gemm!\n");
}
if(s.log_level_ > 0)
{
std::cout << "Launching kernel with args:" << CodegenFlatmmShape::GetName() << "\n"
<< "Shape: " << CodegenFlatmmShape::GetName() << "\n"
<< "problem: " << CodegenPipelineProblem::GetName() << "\n"
<< "pipeline: " << CodegenFlatmmPipeline::GetName() << "\n"
<< "grid: {" << grids.x << ", " << grids.y << ", " << grids.z << "}"
<< ", blocks: {" << blocks.x << ", " << blocks.y << ", " << blocks.z << "}"
<< std::endl;
}
ave_time = ck_tile::launch_kernel(
s,
ck_tile::make_kernel<FlatmmConfig::kBlockPerCu>(Kernel{}, grids, blocks, 0, kargs));
return ave_time;
};
const auto RunSplitk = [&](const auto has_hot_loop_, const auto tail_number_) {
if(args.k_batch == 1)
{
Run(has_hot_loop_,
tail_number_,
ck_tile::integral_constant<ck_tile::memory_operation_enum,
ck_tile::memory_operation_enum::set>{});
}
else
{
Run(has_hot_loop_,
tail_number_,
ck_tile::integral_constant<ck_tile::memory_operation_enum,
ck_tile::memory_operation_enum::atomic_add>{});
}
};
BaseGemmPipeline::TailHandler(RunSplitk, has_hot_loop, tail_num);
return ave_time;
}
// ==================== Weight Shuffle ====================
template <class FlatmmConfig, ck_tile::MoeFlatmmKind moe_kind, class IterSrc, class IterDst>
void shuffle_mx_moe_weight(const IterSrc src, IterDst dst, int experts_cnt, int N, int K)
{
int KPack = 16;
int NLane = FlatmmConfig::N_Warp_Tile;
int KLane = 64 / NLane;
int K_pk = K / 2; // FP4 packed
int K0 = K_pk / (KLane * KPack);
int tempk;
if constexpr(moe_kind == ck_tile::MoeFlatmmKind::kFFN_gemm1_gate_up)
{
int up_stride = N / 2 / NLane;
for(long eid = 0; eid < experts_cnt; ++eid)
{
for(int n = 0; n < N; ++n)
{
for(int k = 0; k < K_pk; ++k)
{
int n0 = n / NLane;
int n1 = n % NLane;
int n0_interleave = n >= N / 2 ? (n0 - up_stride) * 2 + 1 : n0 * 2;
int k0 = k / (KLane * KPack);
tempk = k % (KLane * KPack);
int k1 = tempk / KPack;
int k2 = tempk % KPack;
long outputIndex = eid * N * K_pk + n0_interleave * KPack * NLane * KLane * K0 +
k0 * KPack * NLane * KLane + k1 * KPack * NLane +
n1 * KPack + k2;
dst[outputIndex] = src[eid * N * K_pk + n * K_pk + k];
}
}
}
}
else
{
for(long eid = 0; eid < experts_cnt; ++eid)
{
for(int n = 0; n < N; ++n)
{
for(int k = 0; k < K_pk; ++k)
{
int n0 = n / NLane;
int n1 = n % NLane;
int k0 = k / (KLane * KPack);
tempk = k % (KLane * KPack);
int k1 = tempk / KPack;
int k2 = tempk % KPack;
long outputIndex = eid * N * K_pk + n0 * KPack * NLane * KLane * K0 +
k0 * KPack * NLane * KLane + k1 * KPack * NLane +
n1 * KPack + k2;
dst[outputIndex] = src[eid * N * K_pk + n * K_pk + k];
}
}
}
}
}
// ==================== Scale Shuffle ====================
template <typename FlatmmConfig, ck_tile::MoeFlatmmKind moe_kind, typename T>
auto shuffle_mx_moe_scale(const ck_tile::HostTensor<T>& scale, int experts_cnt)
{
assert(scale.get_lengths().size() == 2);
int n_ = scale.get_lengths()[1];
int k_ = scale.get_lengths()[0];
int k_per_expert = k_ / experts_cnt;
constexpr int K_Pack = 2;
constexpr int N_Pack = 2;
constexpr int GranularityK = 32;
constexpr int K_Lane = 64 / FlatmmConfig::N_Warp_Tile;
static_assert(FlatmmConfig::N_Warp_Tile == 16, "only support XDL_N == 16");
static_assert(FlatmmConfig::N_Repeat % N_Pack == 0);
static_assert(FlatmmConfig::K_Tile % (K_Pack * K_Lane * GranularityK) == 0);
if constexpr(moe_kind == ck_tile::MoeFlatmmKind::kFFN_gemm1_gate_up)
{
ck_tile::HostTensor<T> shfl_scale({
experts_cnt,
k_per_expert / K_Pack / K_Lane,
K_Pack,
K_Lane,
N_Pack,
n_ / FlatmmConfig::N_Warp_Tile / N_Pack,
FlatmmConfig::N_Warp_Tile,
});
std::copy(scale.begin(), scale.end(), shfl_scale.begin());
return ck_tile::reference_permute(shfl_scale, {0, 5, 1, 3, 6, 2, 4});
}
else
{
ck_tile::HostTensor<T> shfl_scale({
experts_cnt,
k_per_expert / K_Pack / K_Lane,
K_Pack,
K_Lane,
n_ / FlatmmConfig::N_Warp_Tile / N_Pack,
N_Pack,
FlatmmConfig::N_Warp_Tile,
});
std::copy(scale.begin(), scale.end(), shfl_scale.begin());
return ck_tile::reference_permute(shfl_scale, {0, 4, 1, 3, 6, 2, 5});
}
}
// ==================== Include Implementation ====================
#include "run_mx_moe_flatmm.inc"
// ==================== Wrapper Function ====================
template <typename FlatmmConfig>
int run_mx_moe_flatmm_example(int argc, char* argv[])
{
auto [result, arg_parser] = create_args(argc, argv);
if(!result)
{
return -1;
}
const std::string a_layout = arg_parser.get_str("a_layout");
const std::string b_layout = arg_parser.get_str("b_layout");
const std::string mx_prec = arg_parser.get_str("mx_prec");
using Row = ck_tile::tensor_layout::gemm::RowMajor;
using Col = ck_tile::tensor_layout::gemm::ColumnMajor;
if(a_layout == "R" && b_layout == "C")
{
const std::string gemm_kind = arg_parser.get_str("gemm_kind");
if(gemm_kind == "gemm1_gate_up")
{
if(mx_prec == "fp4xfp4")
{
return run_mx_moe_flatmm_with_layouts<
ck_tile::pk_fp4_t,
ck_tile::pk_fp4_t,
ck_tile::fp16_t,
FlatmmConfig,
ck_tile::MoeFlatmmKind::kFFN_gemm1_gate_up>(argc, argv, Row{}, Col{}, Row{});
}
else
{
throw std::runtime_error("Only support fp4xfp4 for gemm1_gate_up!");
}
}
else if(gemm_kind == "gemm2")
{
if(mx_prec == "fp4xfp4")
{
return run_mx_moe_flatmm_with_layouts<ck_tile::pk_fp4_t,
ck_tile::pk_fp4_t,
ck_tile::fp16_t,
FlatmmConfig,
ck_tile::MoeFlatmmKind::kFFN_gemm2>(
argc, argv, Row{}, Col{}, Row{});
}
else
{
throw std::runtime_error("Only support fp4xfp4 for gemm2!");
}
}
else
{
throw std::runtime_error("Unrecognized gemm_kind parameter, only accept value "
"[gemm1_gate_up | gemm2]");
}
}
else
{
throw std::runtime_error("Unsupported data layout configuration for A,B and C tensors!");
}
return -1;
}
// ==================== Main Entry ====================
int main(int argc, char* argv[])
{
auto [result, arg_parser] = create_args(argc, argv);
if(!result)
return EXIT_FAILURE;
try
{
int warp_tile = arg_parser.get_int("warp_tile");
if(warp_tile == 0)
{
return !run_mx_moe_flatmm_example<MXfp4_MOE_FlatmmConfig16>(argc, argv);
}
else
{
throw std::runtime_error("Only warp_tile=0 (16x16) is supported now!");
}
}
catch(const std::runtime_error& e)
{
std::cerr << "Runtime error: " << e.what() << '\n';
return EXIT_FAILURE;
}
}

78
example/ck_tile/18_flatmm/mxgemm_moe/mx_moe_flatmm.hpp Normal file
View File

@@ -0,0 +1,78 @@
// SPDX-License-Identifier: MIT
// Copyright (c) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
#pragma once
#include <string>
#include <tuple>
#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 = 128; // MOE 用更小的 M_Tile
static constexpr ck_tile::index_t N_Tile = 512;
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("experts", "8", "Num of experts - 8 by default")
.insert("NumTokens", "128", "M dimensions - 128 by default.")
.insert("TopK", "3", "Top K - 2 by default.")
.insert("N", "4096", "N dimensions - 2048 by default.")
.insert("K", "4096", "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("validate", "1", "0. No validation, 1. Validation on CPU.")
.insert("warmup", "50", "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);
}

78
example/ck_tile/18_flatmm/mxgemm_moe/mxfp4_moe_flatmm.hpp Normal file
View File

@@ -0,0 +1,78 @@
// SPDX-License-Identifier: MIT
// Copyright (c) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
#pragma once
#include <string>
#include <tuple>
#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);
}

353
example/ck_tile/18_flatmm/mxgemm_moe/run_mx_moe_flatmm.inc Normal file
View File

@@ -0,0 +1,353 @@
// SPDX-License-Identifier: MIT
// Copyright (c) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
template <typename FlatmmConfig,
typename ADataType,
typename BDataType,
typename DsDatatype,
typename AccDataType,
typename CDataType,
typename ALayout,
typename BLayout,
typename DsLayout,
typename ELayout,
ck_tile::MoeFlatmmKind kind,
typename CDEElementWise = ck_tile::element_wise::PassThrough,
typename MoeHostArgs>
float invoke_mx_moe_flatmm(int n_warmup, int n_repeat, const MoeHostArgs& args)
{
float ave_time = mx_moe_flatmm<FlatmmConfig,
ADataType,
BDataType,
DsDatatype,
AccDataType,
CDataType,
ALayout,
BLayout,
DsLayout,
ELayout,
kind,
CDEElementWise>(
args, ck_tile::stream_config{nullptr, true, 1, n_warmup, n_repeat, true, true, 50});
std::string op_name{"MX MoE Gemm"};
constexpr int PackedSize = ck_tile::numeric_traits<BDataType>::PackedSize;
std::size_t flop = std::size_t(2) * args.M * args.N * args.K;
std::size_t num_byte = sizeof(ADataType) * args.M * args.K / PackedSize +
sizeof(BDataType) * args.N * args.K / PackedSize +
sizeof(CDataType) * args.M * args.N;
float tflops = static_cast<float>(flop) / 1.E9 / ave_time;
float gb_per_sec = num_byte / 1.E6 / ave_time;
std::cout << "Perf: " << std::setw(10) << ave_time << " ms, " << tflops << " TFlops, "
<< gb_per_sec << " GB/s, " << op_name << std::endl;
return ave_time;
}
template <typename PrecActType,
typename PrecWeightType,
typename CDataType,
typename FlatmmConfig,
ck_tile::MoeFlatmmKind kind,
typename ALayout,
typename BLayout,
typename CLayout>
int run_mx_moe_flatmm_with_layouts(int argc,
char* argv[],
const ALayout a_layout = ALayout{},
const BLayout b_layout = BLayout{},
[[maybe_unused]] const CLayout c_layout = CLayout{})
{
auto [result, arg_parser] = create_args(argc, argv);
if(!result)
return -1;
using ADataType = PrecActType;
using BDataType = PrecWeightType;
using AccDataType = float;
using ScaleType = ck_tile::e8m0_t;
constexpr int ScaleGranularityM = 1;
constexpr int ScaleGranularityN = 1;
constexpr int ScaleGranularityK = 32;
// MOE parameters
const ck_tile::index_t num_tokens = arg_parser.get_int("NumTokens");
const ck_tile::index_t experts = arg_parser.get_int("experts");
const ck_tile::index_t topk = arg_parser.get_int("TopK");
const ck_tile::index_t N = arg_parser.get_int("N");
const ck_tile::index_t K = arg_parser.get_int("K");
ck_tile::index_t stride_A = arg_parser.get_int("stride_A");
ck_tile::index_t stride_B = arg_parser.get_int("stride_B");
ck_tile::index_t stride_C = arg_parser.get_int("stride_C");
ck_tile::index_t init_method = arg_parser.get_int("init");
const ck_tile::index_t warmup = arg_parser.get_int("warmup");
const ck_tile::index_t repeat = arg_parser.get_int("repeat");
const ck_tile::index_t MPerBlock = FlatmmConfig::M_Tile;
ck_tile::index_t sorted_tile_num = (num_tokens + MPerBlock - 1) / MPerBlock * MPerBlock * topk;
ck_tile::index_t valid_tile_num = sorted_tile_num;
ck_tile::index_t sorted_size = sorted_tile_num * MPerBlock;
const ck_tile::index_t M = sorted_tile_num * MPerBlock;
const ck_tile::index_t outputN = kind == ck_tile::MoeFlatmmKind::kFFN_gemm1_gate_up ? N / 2 : N;
static_assert(std::is_same_v<BLayout, ck_tile::tensor_layout::gemm::ColumnMajor>);
constexpr bool IsInputGemm = kind != ck_tile::MoeFlatmmKind::kFFN_gemm2;
stride_A = ck_tile::get_default_stride(
IsInputGemm ? num_tokens : num_tokens * topk, K, stride_A, is_row_major(a_layout));
stride_B = ck_tile::get_default_stride(K, N, stride_B, is_row_major(b_layout));
stride_C = ck_tile::get_default_stride(
IsInputGemm ? num_tokens * topk : num_tokens, outputN, stride_C, is_row_major(CLayout{}));
auto a_m_k_tensor = ck_tile::HostTensor<ADataType>(ck_tile::host_tensor_descriptor(
IsInputGemm ? num_tokens : num_tokens * topk, K, stride_A, is_row_major(a_layout)));
auto b_k_n_tensor = ck_tile::HostTensor<BDataType>(
is_row_major(b_layout)
? ck_tile::host_tensor_descriptor(experts * N, K, 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<CDataType>(ck_tile::host_tensor_descriptor(
IsInputGemm ? num_tokens * topk : num_tokens, outputN, stride_C, is_row_major(CLayout{})));
ck_tile::HostTensor<ScaleType> scale_a(ck_tile::HostTensorDescriptor(
{(IsInputGemm ? num_tokens : M) / ScaleGranularityM, K / ScaleGranularityK},
{K / ScaleGranularityK, 1}));
ck_tile::HostTensor<ScaleType> scale_b(ck_tile::HostTensorDescriptor(
{K * experts / ScaleGranularityK, N / ScaleGranularityN}, {N / ScaleGranularityN, 1}));
if(init_method == 0)
{
ck_tile::FillUniformDistribution<ADataType>{0.0f, 1.0f}(a_m_k_tensor);
ck_tile::FillUniformDistribution<BDataType>{-.5f, .5f}(b_k_n_tensor);
ck_tile::FillUniformDistribution<ScaleType>{0.f, 1.f}(scale_a);
ck_tile::FillUniformDistribution<ScaleType>{0.f, 1.f}(scale_b);
}
else
{
ck_tile::FillUniformDistribution<ADataType>{1.0f, 1.0f}(a_m_k_tensor);
ck_tile::FillUniformDistribution<BDataType>{1.0f, 1.0f}(b_k_n_tensor);
ck_tile::FillUniformDistribution<ScaleType>{1.0f, 1.0f}(scale_a);
ck_tile::FillUniformDistribution<ScaleType>{1.0f, 1.0f}(scale_b);
}
ck_tile::HostTensor<BDataType> b_shuffle_host(
ck_tile::host_tensor_descriptor(K, experts * N, stride_B, is_row_major(b_layout)));
shuffle_mx_moe_weight<FlatmmConfig, kind>(
b_k_n_tensor.begin(), b_shuffle_host.begin(), experts, N, K);
ck_tile::HostTensor<ScaleType> scale_a_shuffle =
shuffle_mx_moe_scale<FlatmmConfig, kind>(scale_a, 1);
ck_tile::HostTensor<ScaleType> scale_b_shuffle =
shuffle_mx_moe_scale<FlatmmConfig, kind>(scale_b, experts);
ck_tile::DeviceMem scale_a_shuffle_dev_buf(scale_a_shuffle.get_element_space_size_in_bytes());
ck_tile::DeviceMem scale_b_shuffle_dev_buf(scale_b_shuffle.get_element_space_size_in_bytes());
std::cout << "mx_moe_flatmm:" << "\n num_experts: " << experts << "\n num_tokens: " << num_tokens
<< "\n topk: " << topk << "\n sorted_tile_num: " << sorted_tile_num
<< "\n problem_n: " << N << "\n problem_k: " << K
<< "\n a_m_k: " << a_m_k_tensor.mDesc << "\n b_k_n: " << b_k_n_tensor.mDesc
<< "\n b_shuffle: " << b_shuffle_host.mDesc << "\n c_m_n: " << c_m_n_tensor.mDesc
<< std::endl;
ck_tile::HostTensor<ck_tile::index_t> expert_ids(
ck_tile::HostTensorDescriptor({sorted_tile_num}, {1}));
ck_tile::HostTensor<ck_tile::index_t> sorted_token_ids(
ck_tile::HostTensorDescriptor({sorted_size}, {1}));
ck_tile::HostTensor<AccDataType> expert_weight(
ck_tile::HostTensorDescriptor({sorted_size}, {1}));
ck_tile::HostTensor<ck_tile::index_t> max_token_id(
ck_tile::HostTensorDescriptor({1 + sorted_tile_num}));
if(init_method == 0)
{
ck_tile::FillUniformDistribution<AccDataType>{0.0f, 1.0f}(expert_weight);
}
else
{
ck_tile::FillUniformDistribution<AccDataType>{1.0f, 1.0f}(expert_weight);
}
max_token_id.mData = {valid_tile_num * MPerBlock, 0, 1, 2, 3, 4, 6, 7, 8, 8};
for(int i = 0; i < sorted_tile_num; i++)
{
expert_ids.mData[i] = i / ((valid_tile_num + experts - 1) / experts);
}
int token_per_tile = (num_tokens * topk + valid_tile_num - 1) / valid_tile_num;
int tokenid = 0;
for(int i = 0; i < sorted_tile_num * MPerBlock; i++)
{
int tile_off = i % MPerBlock;
if(tile_off < token_per_tile && tokenid < num_tokens * topk)
{
sorted_token_ids.mData[i] = (tokenid % num_tokens) | ((tokenid / num_tokens) << 24);
tokenid++;
}
else
{
sorted_token_ids.mData[i] = num_tokens;
}
}
ck_tile::DeviceMem a_m_k_dev_buf{a_m_k_tensor.get_element_space_size_in_bytes()};
ck_tile::DeviceMem b_origin_dev_buf{b_k_n_tensor.get_element_space_size_in_bytes()};
ck_tile::DeviceMem b_shuffle_dev_buf{b_shuffle_host.get_element_space_size_in_bytes()};
ck_tile::DeviceMem c_m_n_dev_buf{c_m_n_tensor.get_element_space_size_in_bytes()};
a_m_k_dev_buf.ToDevice(a_m_k_tensor.data());
b_origin_dev_buf.ToDevice(b_k_n_tensor.data());
b_shuffle_dev_buf.ToDevice(b_shuffle_host.data());
c_m_n_dev_buf.SetZero();
c_m_n_tensor.SetZero();
ck_tile::DeviceMem sorted_token_ids_dev{sorted_token_ids.get_element_space_size_in_bytes()};
ck_tile::DeviceMem expert_ids_dev{expert_ids.get_element_space_size_in_bytes()};
ck_tile::DeviceMem max_token_id_dev{max_token_id.get_element_space_size_in_bytes()};
ck_tile::DeviceMem expert_weight_dev{expert_weight.get_element_space_size_in_bytes()};
sorted_token_ids_dev.ToDevice(sorted_token_ids.data());
expert_ids_dev.ToDevice(expert_ids.data());
max_token_id_dev.ToDevice(max_token_id.data());
expert_weight_dev.ToDevice(expert_weight.data());
scale_a_shuffle_dev_buf.ToDevice(scale_a_shuffle.data());
scale_b_shuffle_dev_buf.ToDevice(scale_b_shuffle.data());
const ck_tile::index_t* p_sorted_token_ids_dev =
static_cast<ck_tile::index_t*>(sorted_token_ids_dev.GetDeviceBuffer());
const ck_tile::index_t* p_expert_ids_dev =
static_cast<ck_tile::index_t*>(expert_ids_dev.GetDeviceBuffer());
const ck_tile::index_t* p_max_token_id_dev =
static_cast<ck_tile::index_t*>(max_token_id_dev.GetDeviceBuffer());
const AccDataType* p_sorted_expert_weight_dev =
static_cast<AccDataType*>(expert_weight_dev.GetDeviceBuffer());
auto scale_a_shuffle_dev_ptr =
ck_tile::FlatmmScalePointer<ScaleGranularityM, ScaleGranularityK>{
static_cast<float*>(scale_a_shuffle_dev_buf.GetDeviceBuffer()),
(IsInputGemm ? num_tokens : M) / ScaleGranularityM};
auto scale_b_shuffle_dev_ptr =
ck_tile::FlatmmScalePointer<ScaleGranularityN, ScaleGranularityK>{
static_cast<float*>(scale_b_shuffle_dev_buf.GetDeviceBuffer()), N / ScaleGranularityN};
using MoeFlatmmArgs = ck_tile::MoeFlatmmHostArgs<
ck_tile::FlatmmScalePointer<ScaleGranularityM, ScaleGranularityK>,
ck_tile::FlatmmScalePointer<ScaleGranularityN, ScaleGranularityK>,
ck_tile::FlatmmScalePointer<-1>>;
MoeFlatmmArgs gemm_desc{p_sorted_token_ids_dev,
p_sorted_expert_weight_dev,
p_expert_ids_dev,
p_max_token_id_dev,
a_m_k_dev_buf.GetDeviceBuffer(),
b_shuffle_dev_buf.GetDeviceBuffer(),
c_m_n_dev_buf.GetDeviceBuffer(),
num_tokens,
experts,
topk,
1, // k_batch
M,
N,
K,
stride_A,
stride_B,
stride_C,
scale_a_shuffle_dev_ptr,
scale_b_shuffle_dev_ptr};
invoke_mx_moe_flatmm<FlatmmConfig,
ADataType,
BDataType,
ck_tile::tuple<>,
AccDataType,
CDataType,
ALayout,
BLayout,
ck_tile::tuple<>,
CLayout,
kind>(warmup, repeat, gemm_desc);
c_m_n_dev_buf.FromDevice(c_m_n_tensor.data());
bool pass{true};
if(arg_parser.get_int("validate"))
{
ck_tile::HostTensor<CDataType> c_m_n_host_ref(
ck_tile::host_tensor_descriptor(IsInputGemm ? num_tokens * topk : num_tokens,
outputN,
stride_C,
is_row_major(CLayout{})));
c_m_n_host_ref.SetZero();
// Convert scale_a from e8m0 to float
ck_tile::HostTensor<AccDataType> scale_a_float(ck_tile::HostTensorDescriptor(
{(IsInputGemm ? num_tokens : M) / ScaleGranularityM, K / ScaleGranularityK},
{K / ScaleGranularityK, 1}));
std::copy(scale_a.begin(), scale_a.end(), scale_a_float.begin());
ck_tile::DeviceMem scale_a_float_dev_buf(scale_a_float.get_element_space_size_in_bytes());
scale_a_float_dev_buf.ToDevice(scale_a_float.data());
// Convert scale_b from e8m0 to float
ck_tile::HostTensor<AccDataType> scale_b_float(ck_tile::HostTensorDescriptor(
{K * experts / ScaleGranularityK, N / ScaleGranularityN}, {N / ScaleGranularityN, 1}));
std::copy(scale_b.begin(), scale_b.end(), scale_b_float.begin());
ck_tile::DeviceMem scale_b_float_dev_buf(scale_b_float.get_element_space_size_in_bytes());
scale_b_float_dev_buf.ToDevice(scale_b_float.data());
std::unique_ptr<ck_tile::DeviceMem> c_m_n_ref_buf =
std::make_unique<ck_tile::DeviceMem>(c_m_n_tensor.get_element_space_size_in_bytes());
c_m_n_ref_buf->SetZero();
ck_tile::reference_moe_gemm_gpu<ADataType,
BDataType,
AccDataType,
CDataType,
ALayout,
BLayout,
CLayout,
static_cast<int>(kind),
ck_tile::moe::MoeSilu>(
p_sorted_token_ids_dev,
p_expert_ids_dev,
p_max_token_id_dev,
static_cast<const ADataType*>(a_m_k_dev_buf.GetDeviceBuffer()),
static_cast<const BDataType*>(b_origin_dev_buf.GetDeviceBuffer()),
static_cast<CDataType*>(c_m_n_ref_buf->GetDeviceBuffer()),
p_sorted_expert_weight_dev,
num_tokens,
MPerBlock,
topk,
M,
N,
K,
stride_A,
stride_B,
stride_C,
(IsInputGemm ? num_tokens : M),
1,
ScaleGranularityK,
static_cast<float*>(scale_a_float_dev_buf.GetDeviceBuffer()),
static_cast<float*>(scale_b_float_dev_buf.GetDeviceBuffer()),
nullptr);
c_m_n_ref_buf->FromDevice(c_m_n_host_ref.data());
const float rtol = std::is_same_v<ADataType, ck_tile::pk_fp4_t> && IsInputGemm ? 1e-2 : 1e-2;
const float atol = std::is_same_v<ADataType, ck_tile::pk_fp4_t> && IsInputGemm ? 1e-2 : 1e-2;
pass = ck_tile::check_err(
c_m_n_tensor, c_m_n_host_ref, "Error: Incorrect results!", rtol, atol);
std::cout << "Relative error threshold: " << rtol << " Absolute error threshold: " << atol
<< std::endl;
std::cout << "The CPU verification result is:" << (pass ? "correct" : "fail") << std::endl;
}
return pass;
}

2
include/ck_tile/host/reference/reference_moe_gemm.hpp
View File

@@ -144,7 +144,7 @@ __global__ void moe_gemm_kernel(const ck_tile::index_t* p_sorted_token_ids_,
}
else if constexpr(std::is_same_v<ADataType, pk_fp4_t>)
{
const fp32x2_t fp32_val = pk_fp4_to_fp32x2(A[a_index / packed_size_a]);
const fp32x2_t fp32_val = pk_fp4_to_fp32x2(A[a_index / packed_size_a], 1.0f);
if(k % 2 == 1)
v_a = fp32_val.hi;
else

709
include/ck_tile/ops/flatmm/kernel/mx_moe_flatmm_kernel.hpp Normal file
View File

@@ -0,0 +1,709 @@
// SPDX-License-Identifier: MIT
// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
#pragma once
#include "ck_tile/core/numeric/math.hpp"
#include "ck_tile/core/utility/literals.hpp"
#include "ck_tile/ops/elementwise/unary_element_wise_operation.hpp"
#include "ck_tile/ops/flatmm/kernel/flatmm_kernel.hpp"
#include "ck_tile/ops/flatmm/kernel/moe_flatmm_kernel.hpp"
#include "ck_tile/ops/gemm/kernel/gemm_tile_partitioner.hpp"
#include "ck_tile/host.hpp"
namespace ck_tile {
// MX MOE FlatMM Kernel - combines MX (FP4xFP4) with MOE routing
// Based on MXFlatmmKernel structure with MOE extensions from MoeFlatmmKernel
template <typename TilePartitioner_,
typename MXFlatmmPipeline_,
typename EpiloguePipeline_,
MoeFlatmmKind kind,
typename FusedActivation = moe::MoeSilu>
struct MXMoeFlatmmKernel
{
using TilePartitioner = remove_cvref_t<TilePartitioner_>;
using FlatmmPipeline = remove_cvref_t<MXFlatmmPipeline_>;
using BlockGemmShape =
remove_cvref_t<typename MXFlatmmPipeline_::BlockGemmShape>;
using EpiloguePipeline = remove_cvref_t<EpiloguePipeline_>;
using ALayout = remove_cvref_t<typename FlatmmPipeline::ALayout>;
using BLayout = remove_cvref_t<typename FlatmmPipeline::BLayout>;
using ELayout = remove_cvref_t<typename FlatmmPipeline::CLayout>;
using DsLayout = remove_cvref_t<typename EpiloguePipeline::DsLayout>;
using DsDataType = remove_cvref_t<typename EpiloguePipeline::DsDataType>;
static constexpr index_t kBlockSize = FlatmmPipeline::BlockSize;
static constexpr bool UsePersistentKernel = FlatmmPipeline::UsePersistentKernel;
using ADataType = remove_cvref_t<typename FlatmmPipeline::ADataType>;
using BDataType = remove_cvref_t<typename FlatmmPipeline::BDataType>;
using EDataType = remove_cvref_t<typename EpiloguePipeline::ODataType>;
using AccDataType = float;
using ActivationOp = FusedActivation;
// MX-specific packing parameters (from MXFlatmmKernel)
static constexpr int MThreadPerXdl = BlockGemmShape::WarpTile::at(number<0>{});
static constexpr int NThreadPerXdl = BlockGemmShape::WarpTile::at(number<1>{});
static constexpr int KThreadPerXdl = 64 / MThreadPerXdl;
static constexpr int APackedSize = numeric_traits<ADataType>::PackedSize;
static constexpr int BPackedSize = numeric_traits<BDataType>::PackedSize;
static constexpr int MXdlPack = FlatmmPipeline::MXdlPack;
static constexpr int NXdlPack = FlatmmPipeline::NXdlPack;
static constexpr int KXdlPack = FlatmmPipeline::KXdlPack;
static constexpr index_t NumDTensor = DsDataType::size();
static constexpr auto I0 = number<0>();
static constexpr auto I1 = number<1>();
static constexpr auto I2 = number<2>();
static constexpr auto I3 = number<3>();
static constexpr auto I4 = number<4>();
static constexpr auto I5 = number<5>();
static_assert(DsLayout::size() == DsDataType::size(),
"The size of DsLayout and DsDataType should be the same");
// MOE-specific parameters (from MoeFlatmmKernel)
static constexpr bool IsInputGemm = kind != MoeFlatmmKind::kFFN_gemm2;
static constexpr bool IsGateUp = kind == MoeFlatmmKind::kFFN_gemm1_gate_up;
static constexpr index_t kMPerBlock = EpiloguePipeline::kMPerBlock;
static constexpr index_t kNPerBlock = EpiloguePipeline::kNPerBlock;
static constexpr index_t MWave = EpiloguePipeline::MWave;
static constexpr index_t NWave = EpiloguePipeline::NWave;
static constexpr index_t MPerXdl = EpiloguePipeline::MPerXdl;
static constexpr index_t NPerXdl = EpiloguePipeline::NPerXdl;
static constexpr index_t KPerXdl = EpiloguePipeline::KPerXdl;
static constexpr index_t isCTransposed = EpiloguePipeline::isCTransposed;
static constexpr index_t kMPerIteration = MPerXdl * MWave;
static constexpr index_t kNPerIteration = NPerXdl * NWave;
static constexpr index_t kNRepeat = kNPerBlock / kNPerIteration;
static constexpr int OutputNPerBlock =
IsGateUp ? TilePartitioner::NPerBlock / 2 : TilePartitioner::NPerBlock;
// MX always uses FP4 for both A and B
static constexpr bool MXFP4_Pipeline = true;
static constexpr int MXFP4N_Pack = 2;
static constexpr int MXFP4K_Pack = 2;
static constexpr int N_Pack = MXFP4N_Pack;
static constexpr int K_Pack = MXFP4K_Pack;
// Kernel arguments structure
template <class ScaleM = FlatmmScalePointer<-1>,
class ScaleN = FlatmmScalePointer<-1>,
class ExpertBias = FlatmmScalePointer<-1>>
struct MXMoeFlatmmKernelArgs
{
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;
const void* p_sorted_expert_weights;
const void* a_ptr;
const void* b_ptr;
void* e_ptr;
ck_tile::index_t NumTokens;
ck_tile::index_t TopK;
ck_tile::index_t M;
ck_tile::index_t N;
ck_tile::index_t K;
ck_tile::index_t stride_A;
ck_tile::index_t stride_B;
ck_tile::index_t stride_C;
ck_tile::index_t k_batch;
ck_tile::index_t n_padded_zeros;
ck_tile::index_t k_padded_zeros;
ScaleM scale_m;
ScaleN scale_n;
ExpertBias exp_bias;
};
template <class ScaleM = FlatmmScalePointer<-1>,
class ScaleN = FlatmmScalePointer<-1>,
class ExpertBias = FlatmmScalePointer<-1>>
CK_TILE_HOST static constexpr auto
MakeKernelArgs(const MoeFlatmmHostArgs<ScaleM, ScaleN, ExpertBias>& hostArgs)
{
return MXMoeFlatmmKernelArgs<ScaleM, ScaleN, ExpertBias>{
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()
{
return concat(
'_', "mx_moe_flatmm", gemm_prec_str<ADataType, BDataType>, FlatmmPipeline::GetName());
}
static constexpr auto BlockSize() -> dim3 { return dim3(kBlockSize); }
static constexpr auto GridSize(index_t M, index_t N, index_t KBatch)
{
return dim3(TilePartitioner::GridSize(M, N), 1, KBatch);
}
template <class MXMoeFlatmmKernelArgs>
static constexpr auto GridSize(const MXMoeFlatmmKernelArgs& kargs)
{
if constexpr(UsePersistentKernel)
{
hipDeviceProp_t prop;
int deviceId = 0;
constexpr int block_size = MXMoeFlatmmKernel::BlockSize().x;
int dync_smem_size = 0;
int maxActiveBlocksPerCU = 0;
[[maybe_unused]] auto e = hipGetDeviceProperties(&prop, deviceId);
e = hipOccupancyMaxActiveBlocksPerMultiprocessor(
&maxActiveBlocksPerCU,
reinterpret_cast<void*>(
kentry<1, MXMoeFlatmmKernel, MXMoeFlatmmKernelArgs>),
block_size,
dync_smem_size);
const int persistent_block_size = prop.multiProcessorCount * maxActiveBlocksPerCU;
const int total_work_tile_cnt = TilePartitioner::GridSize(kargs.M, kargs.N);
assert(kargs.k_batch == 1);
return dim3(min(persistent_block_size, total_work_tile_cnt), 1, kargs.k_batch);
}
else
{
return dim3(TilePartitioner::GridSize(kargs.M, kargs.N), 1, kargs.k_batch);
}
}
CK_TILE_HOST_DEVICE static constexpr index_t GetSmemPingSize()
{
return max(FlatmmPipeline::GetSmemSize(), EpiloguePipeline::GetSmemSize());
}
CK_TILE_HOST_DEVICE static constexpr index_t GetSmemPongSize()
{
return FlatmmPipeline::GetSmemSize();
}
struct SplitKBatchOffset
{
template <class KernelArgs>
__device__ SplitKBatchOffset(const KernelArgs& kargs, const std::size_t k_id = blockIdx.z)
{
const auto k_total = kargs.K;
const auto k_split = integer_divide_ceil(k_total, kargs.k_batch);
const std::size_t splitted_k_start = k_id * k_split;
splitted_k = min(k_split, static_cast<index_t>(k_total - splitted_k_start));
a_k_split_offset = splitted_k_start;
b_k_split_offset = splitted_k_start;
}
index_t a_k_split_offset;
index_t b_k_split_offset;
index_t splitted_k;
};
// template <typename KernelArgs>
// CK_TILE_HOST static bool IsSupportedArgument(const KernelArgs& kargs)
// {
// if constexpr(EpiloguePipeline::GetVectorSizeC() % 2 != 0 &&
// is_any_of<EDataType, fp16_t, bf16_t>::value)
// {
// return false;
// }
// if constexpr(UsePersistentKernel)
// {
// if(kargs.k_batch != 1)
// {
// return false;
// }
// }
// if constexpr(std::is_same_v<ALayout, tensor_layout::gemm::RowMajor>)
// {
// 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<BLayout, tensor_layout::gemm::RowMajor>)
// {
// 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<ELayout, tensor_layout::gemm::RowMajor>)
// {
// if(kargs.stride_C < kargs.N)
// {
// return false;
// }
// }
// else
// {
// if(kargs.stride_C < kargs.M)
// {
// return false;
// }
// }
// return DTensorIsValid;
// }
template <typename KernelArgs>
CK_TILE_HOST static bool IsSupportedArgument(const KernelArgs& kargs)
{
if constexpr(EpiloguePipeline::GetVectorSizeC() % 2 != 0 &&
is_any_of<EDataType, fp16_t, bf16_t>::value)
{
if(kargs.k_batch != 1)
{
std::cerr << "Conditions not met for Kbatch >1 !" << std::endl;
return false;
}
}
if constexpr(UsePersistentKernel)
{
if(kargs.k_batch != 1)
{
std::cerr << "Persistent mode doesn't support Kbatch >1 !" << std::endl;
return false;
}
}
if constexpr(std::is_same_v<ALayout, tensor_layout::gemm::RowMajor>)
{
if(kargs.K % TilePartitioner::KPerBlock != 0 && FlatmmPipeline::kPadK == false)
{
std::cerr << "Can't support K that is not a multiple of KPerBlock"
" without padding!"
<< std::endl;
return false;
}
if(kargs.K % FlatmmPipeline::GetVectorSizeA() != 0)
{
std::cerr << "K is not a multiple of vector load size for A tensor!" << std::endl;
return false;
}
}
else
{
if(kargs.M % TilePartitioner::MPerBlock != 0 && FlatmmPipeline::kPadM == false)
{
std::cerr << "Can't support M that is not a multiple of MPerBlock"
" without padding!"
<< std::endl;
return false;
}
if(kargs.M % FlatmmPipeline::GetVectorSizeA() != 0)
{
std::cerr << "M is not a multiple of vector load size for A tensor!" << std::endl;
return false;
}
}
if constexpr(std::is_same_v<BLayout, tensor_layout::gemm::RowMajor>)
{
// if(kargs.N % TilePartitioner::NPerBlock != 0 && FlatmmPipeline::kPadN == false)
// {
// std::cerr << "Can't support N that is not a multiple of NPerBlock"
// " without padding!"
// << std::endl;
// return false;
// }
if(kargs.N % FlatmmPipeline::GetVectorSizeB() != 0)
{
std::cerr << "N is not a multiple of vector load size for B tensor!" << std::endl;
return false;
}
}
else
{
if(kargs.K % TilePartitioner::KPerBlock != 0 && FlatmmPipeline::kPadK == false)
{
std::cerr << "Can't support K that is not a multiple of KPerBlock"
" without padding!"
<< std::endl;
return false;
}
if(kargs.K % FlatmmPipeline::GetVectorSizeB() != 0)
{
std::cerr << "K is not a multiple of vector load size for B tensor!" << std::endl;
return false;
}
}
bool DTensorIsValid = {true};
static_for<0, NumDTensor, 1>{}([&](auto index) {
using DiLayout = remove_cvref_t<std::tuple_element_t<index.value, DsLayout>>;
if(std::is_same_v<DiLayout, ELayout> == false)
{
DTensorIsValid = false;
}
if constexpr(std::is_same_v<DiLayout, tensor_layout::gemm::RowMajor>)
{
if(kargs.N % TilePartitioner::NPerBlock != 0 && FlatmmPipeline::kPadN == false)
{
CK_TILE_ERROR("Can't support N for tensor D that is not a multiple of "
"NPerBlock without padding!");
DTensorIsValid = false;
}
if(kargs.N % EpiloguePipeline::GetVectorSizeD(index) != 0)
{
CK_TILE_ERROR("N is not a multiple of vector load size for D tensor!");
DTensorIsValid = false;
}
}
else
{
if(kargs.M % TilePartitioner::MPerBlock != 0 && FlatmmPipeline::kPadM == false)
{
CK_TILE_ERROR("Can't support M for tensor D that is not a multiple of "
"MPerBlock without padding!");
DTensorIsValid = false;
}
if(kargs.M % EpiloguePipeline::GetVectorSizeD(index) != 0)
{
CK_TILE_ERROR("M is not a multiple of vector load size for D tensor!");
DTensorIsValid = false;
}
}
});
if constexpr(std::is_same_v<ELayout, tensor_layout::gemm::RowMajor>)
{
if(kargs.stride_C % TilePartitioner::NPerBlock != 0 && FlatmmPipeline::kPadN == false)
{
std::cerr << "Can't support N that is not a multiple of NPerBlock"
" without padding!"
<< std::endl;
return false;
}
if(kargs.N % EpiloguePipeline::GetVectorSizeC() != 0)
{
std::cerr << "N is not a multiple of vector load size for C tensor!" << std::endl;
return false;
}
}
else
{
if(kargs.M % TilePartitioner::MPerBlock != 0 && FlatmmPipeline::kPadM == false)
{
std::cerr << "Can't support M that is not a multiple of MPerBlock"
" without padding!"
<< std::endl;
return false;
}
if(kargs.M % EpiloguePipeline::GetVectorSizeC() != 0)
{
std::cerr << "M is not a multiple of vector load size for C tensor!" << std::endl;
return false;
}
}
return DTensorIsValid;
}
// Tensor view creation with MOE expert routing
template <memory_operation_enum DstInMemOp = IsInputGemm ? memory_operation_enum::set
: memory_operation_enum::atomic_add,
typename KernelArgs>
CK_TILE_DEVICE static auto
MakeGemmTensorViews(const ADataType* a_ptr,
const BDataType* b_flat_ptr,
EDataType* e_ptr,
[[maybe_unused]] const AccDataType* exp_weight_ptr,
const int expert_id,
const KernelArgs& kargs,
const SplitKBatchOffset& splitk_batch_offset)
{
// A tensor view (token activations)
const auto& a_tensor_view = [&]() {
if constexpr(std::is_same_v<ALayout, tensor_layout::gemm::RowMajor>)
{
return make_naive_tensor_view<address_space_enum::global>(
a_ptr,
make_tuple(kargs.M, splitk_batch_offset.splitted_k),
make_tuple(kargs.stride_A, 1),
number<FlatmmPipeline::GetVectorSizeA()>{},
number<1>{});
}
else
{
return make_naive_tensor_view<address_space_enum::global>(
a_ptr,
make_tuple(splitk_batch_offset.splitted_k, kargs.M),
make_tuple(kargs.stride_A, 1),
number<FlatmmPipeline::GetVectorSizeA()>{},
number<1>{});
}
}();
// B tensor view (expert weights) - per-expert offset
index_t kFlatK = kargs.K * BlockGemmShape::WarpTile::at(I1);
index_t kFlatN = kargs.N * kargs.K / kFlatK;
const auto& b_flat_tensor_view = [&]() {
return make_naive_tensor_view<address_space_enum::global>(
b_flat_ptr + expert_id * kFlatN * kFlatK,
make_tuple(kFlatN - kargs.n_padded_zeros / NPerXdl, kFlatK),
make_tuple(kFlatK, 1),
number<FlatmmPipeline::GetVectorSizeB()>{},
number<1>{});
}();
// C tensor view (output)
const auto& c_tensor_view = [&]() {
if constexpr(std::is_same_v<ELayout, tensor_layout::gemm::RowMajor>)
{
return make_naive_tensor_view<address_space_enum::global, DstInMemOp>(
e_ptr,
make_tuple(kargs.M, kargs.N),
make_tuple(kargs.stride_C, 1),
number<EpiloguePipeline::GetVectorSizeC()>{},
number<1>{});
}
else
{
return make_naive_tensor_view<address_space_enum::global, DstInMemOp>(
e_ptr,
make_tuple(kargs.N, kargs.M),
make_tuple(kargs.stride_C, 1),
number<1>{},
number<1>{});
}
}();
// MX scale tensors (from MXFlatmmKernel)
auto scale_a = kargs.scale_m;
auto scale_b = kargs.scale_n;
static constexpr int BlockScaleSize = 32;
const auto&& scale_packs_m = integer_divide_ceil(kargs.M, (MXdlPack * MThreadPerXdl));
const auto&& scale_packs_n = integer_divide_ceil(kargs.N, (NXdlPack * NThreadPerXdl));
const auto&& scale_packs_k = kargs.K / BlockScaleSize / (KXdlPack * KThreadPerXdl);
// A scale tensor view
const auto& scale_a_tensor_view = [&]() {
const auto scale_a_naive_desc = make_naive_tensor_descriptor_packed(
make_tuple(scale_packs_m, scale_packs_k, KThreadPerXdl, MThreadPerXdl));
const auto scale_a_desc = transform_tensor_descriptor(
scale_a_naive_desc,
make_tuple(make_merge_transform(make_tuple(scale_packs_m, MThreadPerXdl)),
make_merge_transform(make_tuple(scale_packs_k, KThreadPerXdl))),
make_tuple(sequence<0, 3>{}, sequence<1, 2>{}),
make_tuple(sequence<0>{}, sequence<1>{}));
return make_tensor_view<address_space_enum::global>(
reinterpret_cast<const int32_t*>(scale_a.ptr), scale_a_desc);
}();
// B scale tensor view - per-expert offset
const auto& scale_b_tensor_view = [&]() {
const auto scale_b_navie_desc = make_naive_tensor_descriptor_packed(
make_tuple(scale_packs_n, scale_packs_k, KThreadPerXdl, NThreadPerXdl));
const auto scale_b_desc = transform_tensor_descriptor(
scale_b_navie_desc,
make_tuple(make_merge_transform(make_tuple(scale_packs_n, NThreadPerXdl)),
make_merge_transform(make_tuple(scale_packs_k, KThreadPerXdl))),
make_tuple(sequence<0, 3>{}, sequence<1, 2>{}),
make_tuple(sequence<0>{}, sequence<1>{}));
return make_tensor_view<address_space_enum::global>(
reinterpret_cast<const int32_t*>(scale_b.ptr) +
expert_id * scale_packs_n * scale_packs_k,
scale_b_desc);
}();
return make_tuple(
a_tensor_view, b_flat_tensor_view, c_tensor_view, scale_a_tensor_view, scale_b_tensor_view);
}
template <typename TensorView>
CK_TILE_DEVICE static auto MakeGemmPadViews(const TensorView& views)
{
const auto& a_pad_view = [&]() {
const auto& a_tensor_view = views.at(I0);
if constexpr(std::is_same_v<ALayout, tensor_layout::gemm::RowMajor>)
{
return pad_tensor_view(a_tensor_view,
make_tuple(number<TilePartitioner::MPerBlock>{},
number<TilePartitioner::KPerBlock>{}),
sequence<false, FlatmmPipeline::kPadK>{});
}
else
{
return pad_tensor_view(a_tensor_view,
make_tuple(number<TilePartitioner::KPerBlock>{},
number<TilePartitioner::MPerBlock>{}),
sequence<false, FlatmmPipeline::kPadM>{});
}
}();
const auto& c_pad_view = [&]() {
const auto& c_tensor_view = views.at(I2);
if constexpr(std::is_same_v<ELayout, tensor_layout::gemm::RowMajor>)
{
return pad_tensor_view(c_tensor_view,
make_tuple(number<TilePartitioner::MPerBlock>{},
number<OutputNPerBlock>{}),
sequence<false, FlatmmPipeline::kPadN>{});
}
else
{
return pad_tensor_view(c_tensor_view,
make_tuple(number<OutputNPerBlock>{},
number<TilePartitioner::MPerBlock>{}),
sequence<FlatmmPipeline::kPadN, false>{});
}
}();
return make_tuple(a_pad_view, views.at(I1), c_pad_view, views.at(I3), views.at(I4));
}
template <typename PadView>
CK_TILE_DEVICE static auto MakeGemmTileWindows(const PadView& views,
[[maybe_unused]] const index_t coord_m,
const index_t coord_n)
{
const auto& a_pad_view = views.at(I0);
const auto& b_flat_pad_view = views.at(I1);
const auto& c_pad_view = views.at(I2);
const auto& a_block_window = [&]() {
if constexpr(std::is_same_v<ALayout, tensor_layout::gemm::RowMajor>)
{
return make_tile_window(a_pad_view,
make_tuple(number<TilePartitioner::MPerBlock>{},
number<TilePartitioner::KPerBlock>{}),
{0, 0});
}
else
{
return make_tile_window(a_pad_view,
make_tuple(number<TilePartitioner::KPerBlock>{},
number<TilePartitioner::MPerBlock>{}),
{0, 0});
}
}();
constexpr bool isNonInterleaveGateUp = !IsGateUp || MXFP4_Pipeline;
const auto& b_flat_block_window =
make_tile_window(b_flat_pad_view,
make_tuple(number<FlatmmPipeline::flatNPerWarp>{},
number<FlatmmPipeline::flatKPerWarp>{}),
{static_cast<int>(coord_n / BlockGemmShape::WarpTile::at(I1) /
(isNonInterleaveGateUp ? 1 : 2)),
0});
const int output_N_offset = IsGateUp ? coord_n / 2 : coord_n;
auto c_block_window = make_tile_window(
c_pad_view,
make_tuple(number<TilePartitioner::MPerBlock>{}, number<OutputNPerBlock>{}),
{0, output_N_offset});
static constexpr int BlockScaleSize = 32;
auto scale_a_block_window = make_tile_window(
views.at(I3),
make_tuple(number<TilePartitioner::MPerBlock / MXdlPack>{},
number<TilePartitioner::KPerBlock / (BlockScaleSize * KXdlPack)>{}),
{0, 0});
auto scale_b_block_window = make_tile_window(
views.at(I4),
make_tuple(number<TilePartitioner::NPerBlock / NXdlPack>{},
number<TilePartitioner::KPerBlock / (BlockScaleSize * KXdlPack)>{}),
{coord_n / NXdlPack, 0});
return make_tuple(a_block_window,
b_flat_block_window,
c_block_window,
scale_a_block_window,
scale_b_block_window);
}
template <class MXMoeFlatmmKernelArgs>
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);
this->operator()(kargs, coord_m, coord_n);
}
template <class MXMoeFlatmmKernelArgs>
CK_TILE_DEVICE void operator()(MXMoeFlatmmKernelArgs kargs, index_t coord_m, index_t coord_n) const
{
// Similar structure to MoeFlatmmKernel::operator() but with MX pipeline
const SplitKBatchOffset splitk_batch_offset(kargs);
const ADataType* a_ptr = static_cast<const ADataType*>(kargs.a_ptr) +
splitk_batch_offset.a_k_split_offset / APackedSize;
const BDataType* b_flat_ptr = static_cast<const BDataType*>(kargs.b_ptr) +
splitk_batch_offset.b_k_split_offset / BPackedSize;
EDataType* e_ptr = static_cast<EDataType*>(kargs.e_ptr);
__shared__ char smem_ptr_ping[GetSmemPingSize()];
__shared__ char smem_ptr_pong[GetSmemPongSize()];
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<const AccDataType*>(kargs.p_sorted_expert_weights);
// 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
}
};
} // namespace ck_tile