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

feat: implement sm90 megamoe phase5 l2 scatter

Browse Source
This commit is contained in:
Xinyi Liu
2026-06-18 15:17:20 +08:00
parent fc8218750c
commit 9bd0519605
5 changed files with 379 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
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, torch::Tensor, torch::Tensor>;
static int get_token_alignment_for_sm90_mega_moe() {
return layout::kLCMCandidateBlockM;
@@ -164,11 +164,15 @@ get_symm_buffer_size_for_sm90_mega_moe(
math::advance_ptr(buffer.data_ptr(), reinterpret_cast<int64_t>(l1_accum_debug_buffer.base)),
{128, 128},
torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device()));
auto combine_acts = torch::from_blob(
math::advance_ptr(buffer.data_ptr(), reinterpret_cast<int64_t>(combine_token_buffer.base)),
{num_topk, num_max_tokens_per_rank, hidden},
torch::TensorOptions().dtype(torch::kBFloat16).device(buffer.device()));
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, l2_arrival_mask, token_src_metadata,
l1_accum_debug);
l1_accum_debug, combine_acts);
};
return {reinterpret_cast<int64_t>(combine_token_buffer.get_end_ptr()), slice_input_buffers};
}
@@ -245,8 +249,9 @@ static void fp8_mega_moe(
l1_acts, l1_acts_sf, l1_topk_weights,
l2_acts, l2_acts_sf,
expert_recv_count_sum, l1_arrival_count, l2_arrival_mask, token_src_metadata,
l1_accum_debug] = slice(sym_buffer);
l1_accum_debug, combine_acts] = slice(sym_buffer);
(void)l2_arrival_mask;
(void)combine_acts;
// Dispatch into SM90 path
DG_HOST_ASSERT(arch_major == 9);