chenchenghao/DeepGEMM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: implement sm90 megamoe phase4 l1 epilogue

Browse Source
This commit is contained in:
Xinyi Liu
2026-06-18 01:09:45 +08:00
parent 2bb1756787
commit f3553f976c
6 changed files with 330 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
csrc/apis/sm90_mega.hpp
View File

@@ -16,7 +16,7 @@ using SM90MegaMoEBufferViews = std::tuple<
torch::Tensor, torch::Tensor, torch::Tensor,
torch::Tensor, torch::Tensor,
torch::Tensor, torch::Tensor, torch::Tensor,
torch::Tensor>;
torch::Tensor, torch::Tensor>;
static int get_token_alignment_for_sm90_mega_moe() {
return layout::kLCMCandidateBlockM;
@@ -38,7 +38,7 @@ get_symm_buffer_size_for_sm90_mega_moe(
const auto fp8_intermediate_token_layout = layout::Data(intermediate_hidden);
const auto bf16_token_layout = layout::Data(hidden * 2);
const auto fp8_sf_layout = layout::Data(hidden / 128 * static_cast<int>(sizeof(float)), false);
const auto fp8_intermediate_sf_layout = layout::Data(intermediate_hidden / 128 * static_cast<int>(sizeof(float)), false);
const auto fp8_intermediate_sf_layout = layout::Data(intermediate_hidden / 64 * static_cast<int>(sizeof(float)), false);
const auto input_topk_idx_layout = layout::Data(num_topk * sizeof(int64_t), false);
const auto input_topk_weights_layout = layout::Data(num_topk * sizeof(float), false);
const auto l1_topk_weights_layout = layout::Data(sizeof(float), false);
@@ -141,7 +141,7 @@ get_symm_buffer_size_for_sm90_mega_moe(
torch::TensorOptions().dtype(torch::kFloat8_e4m3fn).device(buffer.device()));
auto l2_acts_sf = torch::from_blob(
math::advance_ptr(buffer.data_ptr(), reinterpret_cast<int64_t>(l2_sf_buffer.base)),
{num_max_padded_sf_pool_tokens, intermediate_hidden / 128},
{num_max_padded_sf_pool_tokens, intermediate_hidden / 64},
{1, num_max_padded_sf_pool_tokens},
torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device()));
auto expert_recv_count_sum = torch::from_blob(
@@ -152,6 +152,10 @@ get_symm_buffer_size_for_sm90_mega_moe(
runtime_workspace.get_l1_arrival_count_ptr(),
{static_cast<int>(runtime_workspace.num_max_pool_blocks)},
torch::TensorOptions().dtype(torch::kInt).device(buffer.device()));
auto l2_arrival_mask = torch::from_blob(
reinterpret_cast<int64_t*>(runtime_workspace.get_l2_arrival_mask_ptr()),
{static_cast<int>(runtime_workspace.num_max_pool_blocks)},
torch::TensorOptions().dtype(torch::kInt64).device(buffer.device()));
auto token_src_metadata = torch::from_blob(
reinterpret_cast<int32_t*>(runtime_workspace.get_token_src_metadata_ptr()),
{num_max_pool_tokens, 3},
@@ -163,7 +167,7 @@ get_symm_buffer_size_for_sm90_mega_moe(
return std::make_tuple(x, x_sf, topk_idx, topk_weights,
l1_acts, l1_acts_sf, l1_topk_weights,
l2_acts, l2_acts_sf,
expert_recv_count_sum, l1_arrival_count, token_src_metadata,
expert_recv_count_sum, l1_arrival_count, l2_arrival_mask, token_src_metadata,
l1_accum_debug);
};
return {reinterpret_cast<int64_t>(combine_token_buffer.get_end_ptr()), slice_input_buffers};
@@ -240,8 +244,9 @@ static void fp8_mega_moe(
const auto [x, x_sf, topk_idx, topk_weights,
l1_acts, l1_acts_sf, l1_topk_weights,
l2_acts, l2_acts_sf,
expert_recv_count_sum, l1_arrival_count, token_src_metadata,
expert_recv_count_sum, l1_arrival_count, l2_arrival_mask, token_src_metadata,
l1_accum_debug] = slice(sym_buffer);
(void)l2_arrival_mask;
// Dispatch into SM90 path
DG_HOST_ASSERT(arch_major == 9);

3
csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp
View File

@@ -135,6 +135,7 @@ static void sm90_fp8_mega_moe(
// Make tensormap
constexpr int kGranK = 128;
constexpr int kL2ActsGranK = 64;
const auto tensor_map_l1_acts = make_tma_2d_desc(l1_acts,
hidden, config.num_max_pool_tokens,
config.block_k, config.block_m,
@@ -162,7 +163,7 @@ static void sm90_fp8_mega_moe(
128);
const auto tensor_map_l2_acts_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l2_acts_sf,
config.num_padded_sf_pool_tokens, intermediate_hidden,
config.block_m, kGranK,
config.block_m, kL2ActsGranK,
1, 0);
const auto tensor_map_l2_weights = make_tma_2d_desc(l2_weights,
intermediate_hidden, num_experts_per_rank * hidden,