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
csrc/apis/sm90_mega.hpp
+10 -5
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);