ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-16 08:44:55 +00:00
Code Issues Packages Projects Releases Wiki Activity

moe gemm draft v0.1

Browse Source
This commit is contained in:
root
2025-03-26 08:14:11 +00:00
parent 938f4234f6
commit 45a0463f1f
6 changed files with 215 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
example/ck_tile/XX_moe_gemm/moe_gemm.hpp
View File

@@ -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.")

2
example/ck_tile/XX_moe_gemm/moe_gemm1_xdl_fp8.cpp
View File

@@ -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;

10
example/ck_tile/XX_moe_gemm/run_moe_gemm_example.inc
View File

@@ -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<BDataType>(
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<CDataType>(
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<ck_tile::DeviceMem> max_token_id_dev = std::make_unique<ck_tile::DeviceMem>(
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];

76
include/ck_tile/host/reference/reference_fused_single_moe_gemm.hpp
View File

@@ -77,7 +77,8 @@ template <typename ADataType,
typename CDataType,
typename LayoutA,
typename LayoutB,
typename LayoutC>
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<BDataType>::PackedSize;
// Adjust indexing based on matrix layout
int a_index = (std::is_same_v<LayoutA, tensor_layout::gemm::RowMajor>)
? 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<LayoutB, tensor_layout::gemm::ColumnMajor>)
? col * strideB + k
: k * strideB + col;
int b_index =
expert_id * N * K + ((std::is_same_v<LayoutB, tensor_layout::gemm::ColumnMajor>)
? 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<LayoutC, tensor_layout::gemm::RowMajor>)
? 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<CDataType>(acc);
}
}
@@ -174,7 +187,8 @@ template <typename ADataType,
typename CDataType,
typename LayoutA,
typename LayoutB,
typename LayoutC>
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<ADataType, BDataType, AccDataType, CDataType, LayoutA, LayoutB, LayoutC>
<<<numBlocks, numThreadsPerBlock>>>(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<ADataType,
BDataType,
AccDataType,
CDataType,
LayoutA,
LayoutB,
LayoutC,
IsInputGemm><<<numBlocks, numThreadsPerBlock>>>(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;
}

14
include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp
View File

@@ -121,6 +121,7 @@ struct CShuffleEpilogue
template <typename ODramWindow,
typename OAccTile,
bool IsInputGemm = true,
memory_operation_enum out_memory_data_op = memory_operation_enum::set>
CK_TILE_DEVICE auto operator()(ODramWindow& out_dram_window,
const OAccTile& o_acc_tile,
@@ -177,10 +178,17 @@ struct CShuffleEpilogue
statically_indexed_array<index_t, 2> 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);

153
include/ck_tile/ops/gemm/kernel/moe_gemm_kernel.hpp
View File

@@ -48,7 +48,10 @@ struct MoeGemmHostArgs : public ck_tile::GemmHostArgs
static constexpr index_t KBatch = 1;
};
template <typename TilePartitioner_, typename GemmPipeline_, typename EpiloguePipeline_>
template <typename TilePartitioner_,
typename GemmPipeline_,
typename EpiloguePipeline_,
bool IsInputGemm_ = true>
struct MoeGemmKernel : public GemmKernel<TilePartitioner_, GemmPipeline_, EpiloguePipeline_>
{
using TilePartitioner = remove_cvref_t<TilePartitioner_>;
@@ -58,6 +61,8 @@ struct MoeGemmKernel : public GemmKernel<TilePartitioner_, GemmPipeline_, Epilog
using BLayout = remove_cvref_t<typename GemmPipeline::BLayout>;
using CLayout = remove_cvref_t<typename GemmPipeline::CLayout>;
static constexpr bool IsInputGemm = IsInputGemm_;
using ADataType = remove_cvref_t<typename GemmPipeline::ADataType>;
using BDataType = remove_cvref_t<typename GemmPipeline::BDataType>;
using CDataType = remove_cvref_t<typename EpiloguePipeline::ODataType>;
@@ -66,8 +71,14 @@ struct MoeGemmKernel : public GemmKernel<TilePartitioner_, GemmPipeline_, Epilog
using Base = GemmKernel<TilePartitioner_, GemmPipeline_, EpiloguePipeline_>;
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<TilePartitioner_, GemmPipeline_, Epilog
return max(GemmPipeline::GetSmemSize(), EpiloguePipeline::GetSmemSize());
}
template <memory_operation_enum DstInMemOp = memory_operation_enum::set>
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<ALayout, tensor_layout::gemm::RowMajor>)
{
return make_naive_tensor_view<address_space_enum::global>(
a_ptr,
make_tuple(IsInputGemm ? kargs.NumTokens : kargs.NumTokens * kargs.TopK,
splitk_batch_offset.splitted_k),
make_tuple(kargs.stride_A, 1),
number<GemmPipeline::GetVectorSizeA()>{},
number<1>{});
}
else
{
return make_naive_tensor_view<address_space_enum::global>(
a_ptr,
make_tuple(splitk_batch_offset.splitted_k,
IsInputGemm ? kargs.NumTokens : kargs.NumTokens * kargs.TopK),
make_tuple(kargs.stride_A, 1),
number<GemmPipeline::GetVectorSizeA()>{},
number<1>{});
}
}();
const auto& b_tensor_view = [&]() {
if constexpr(std::is_same_v<BLayout, tensor_layout::gemm::RowMajor>)
{
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<VectorSizeB>{},
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<address_space_enum::global>(b_ptr, b_n_k_desc);
}
else
{
return make_naive_tensor_view<address_space_enum::global>(
b_ptr,
make_tuple(splitk_batch_offset.splitted_k, kargs.N),
make_tuple(kargs.stride_B, 1),
number<GemmPipeline::GetVectorSizeB()>{},
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<VectorSizeB>{},
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<address_space_enum::global>(b_ptr, b_n_k_desc);
}
else
{
return make_naive_tensor_view<address_space_enum::global>(
b_ptr,
make_tuple(kargs.N, splitk_batch_offset.splitted_k),
make_tuple(kargs.stride_B, 1),
number<GemmPipeline::GetVectorSizeB()>{},
number<1>{});
}
}
}();
// TODO: enable vector write for C in ColMajor
const auto& c_tensor_view = [&]() {
if constexpr(std::is_same_v<CLayout, tensor_layout::gemm::RowMajor>)
{
return make_naive_tensor_view<address_space_enum::global, DstInMemOp>(
c_ptr,
make_tuple(IsInputGemm ? kargs.NumTokens * kargs.TopK : kargs.NumTokens,
kargs.N),
make_tuple(kargs.stride_C, 1),
number<EpiloguePipeline::GetVectorSizeC()>{},
number<1>{});
}
else
{
return make_naive_tensor_view<address_space_enum::global, DstInMemOp>(
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 <typename AView>
CK_TILE_DEVICE static auto GetATransformGemmView(const AView& view, const index_t token_id)
{
@@ -269,6 +401,7 @@ struct MoeGemmKernel : public GemmKernel<TilePartitioner_, GemmPipeline_, Epilog
return make_tuple(a_block_window, b_block_window, c_block_window);
}
template <bool IsInputGemm = true>
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<TilePartitioner_, GemmPipeline_, Epilog
const index_t prefix_blk_m = gemm_desc.p_max_token_id[1 + expert_id];
const index_t blk_cnt_of_eid = gemm_desc.p_max_token_id[2 + expert_id];
// printf("expert_blk_id: %d, expert_id: %d \n",expert_blk_id, expert_id);
// expert_id = expert_blk_id;
const index_t block_start = prefix_blk_m * NBlocks;
const index_t ecnt = blk_cnt_of_eid - prefix_blk_m;
@@ -312,23 +449,27 @@ struct MoeGemmKernel : public GemmKernel<TilePartitioner_, GemmPipeline_, Epilog
// });
const index_t fused_token = gemm_desc.p_sorted_token_ids[sorted_token_id];
// printf("a_coord[number<0>{}]: %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<const ADataType*>(gemm_desc.a_ptr) + splitk_batch_offset.a_k_split_offset;
const BDataType* b_ptr =
static_cast<const BDataType*>(gemm_desc.b_ptr) + splitk_batch_offset.b_k_split_offset;
const BDataType* b_ptr = static_cast<const BDataType*>(gemm_desc.b_ptr) +
splitk_batch_offset.b_k_split_offset + expert_stride * expert_id;
CDataType* c_ptr = static_cast<CDataType*>(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(