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

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),