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

feat: implement sm90 megamoe phase3 l1 wgmma

Browse Source
This commit is contained in:
Xinyi Liu
2026-06-18 00:46:41 +08:00
parent ed6067c7a5
commit 842124b1c5
5 changed files with 479 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
csrc/apis/sm90_mega.hpp
View File

@@ -15,7 +15,8 @@ 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, torch::Tensor, torch::Tensor,
torch::Tensor>;
static int get_token_alignment_for_sm90_mega_moe() {
return layout::kLCMCandidateBlockM;
@@ -85,10 +86,16 @@ get_symm_buffer_size_for_sm90_mega_moe(
fp8_intermediate_sf_layout, 1, num_max_padded_sf_pool_tokens,
l2_token_buffer.get_end_ptr());
// Phase 3 verification output: one scaled FP32 L1 accumulator tile.
const auto l1_accum_debug_layout = layout::Data(128 * static_cast<int>(sizeof(float)), false);
const auto l1_accum_debug_buffer = layout::Buffer(
l1_accum_debug_layout, 1, 128,
l2_sf_buffer.get_end_ptr());
// Combine input buffer
const auto combine_token_buffer = layout::Buffer(
bf16_token_layout, num_topk, num_max_tokens_per_rank,
l2_sf_buffer.get_end_ptr());
l1_accum_debug_buffer.get_end_ptr());
// Check SF buffer requirements
DG_HOST_ASSERT(hidden % 128 == 0 and intermediate_hidden % 128 == 0);
@@ -149,10 +156,15 @@ get_symm_buffer_size_for_sm90_mega_moe(
reinterpret_cast<int32_t*>(runtime_workspace.get_token_src_metadata_ptr()),
{num_max_pool_tokens, 3},
torch::TensorOptions().dtype(torch::kInt).device(buffer.device()));
auto l1_accum_debug = torch::from_blob(
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()));
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, token_src_metadata,
l1_accum_debug);
};
return {reinterpret_cast<int64_t>(combine_token_buffer.get_end_ptr()), slice_input_buffers};
}
@@ -228,7 +240,8 @@ 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] = slice(sym_buffer);
expert_recv_count_sum, l1_arrival_count, token_src_metadata,
l1_accum_debug] = slice(sym_buffer);
// Dispatch into SM90 path
DG_HOST_ASSERT(arch_major == 9);
@@ -237,6 +250,7 @@ static void fp8_mega_moe(
l2_acts, l2_acts_sf,
l1_weights, l2_weights,
l1_weights_sf, l2_weights_sf,
l1_accum_debug,
cumulative_local_expert_recv_stats,
sym_buffer_ptrs,
rank_idx, num_max_tokens_per_rank,

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

@@ -29,6 +29,7 @@ public:
// Runtime arguments
void* y;
void* l1_accum_debug;
int* cumulative_local_expert_recv_stats;
int num_tokens;
layout::SymBuffer<> sym_buffer_ptrs;
@@ -101,7 +102,8 @@ static void __instantiate_kernel() {{
args.tensor_map_l2_acts,
args.tensor_map_l2_acts_sf,
args.tensor_map_l2_weights,
args.l2_weights_sf
args.l2_weights_sf,
args.l1_accum_debug
));
}
};
@@ -112,6 +114,7 @@ static void sm90_fp8_mega_moe(
const torch::Tensor& l2_acts, const torch::Tensor& l2_acts_sf,
const torch::Tensor& l1_weights, const torch::Tensor& l2_weights,
const torch::Tensor& l1_weights_sf, const torch::Tensor& l2_weights_sf,
const torch::Tensor& l1_accum_debug,
const std::optional<torch::Tensor> cumulative_local_expert_recv_stats,
const std::vector<int64_t>& sym_buffer_ptrs,
const int& rank_idx, const int& num_max_tokens_per_rank,
@@ -184,6 +187,7 @@ static void sm90_fp8_mega_moe(
.fast_math = fast_math,
.config = config,
.y = y.data_ptr(),
.l1_accum_debug = l1_accum_debug.data_ptr(),
.cumulative_local_expert_recv_stats = cumulative_local_expert_recv_stats_ptr,
.num_tokens = num_tokens,
.sym_buffer_ptrs = layout::SymBuffer<>(sym_buffer_ptrs, rank_idx),