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

5
deep_gemm/mega/__init__.py
View File

@@ -68,7 +68,8 @@ class SymmBuffer:
self.l2_acts, self.l2_acts_sf,
self.expert_recv_count_sum,
self.l1_arrival_count,
self.token_src_metadata) = buffer_views
self.token_src_metadata,
self.l1_accum_debug) = buffer_views
else:
(self.x, self.x_sf,
self.topk_idx, self.topk_weights,
@@ -78,6 +79,7 @@ class SymmBuffer:
self.expert_recv_count_sum = None
self.l1_arrival_count = None
self.token_src_metadata = None
self.l1_accum_debug = None
def destroy(self):
self.handle = None
@@ -95,6 +97,7 @@ class SymmBuffer:
self.expert_recv_count_sum = None
self.l1_arrival_count = None
self.token_src_metadata = None
self.l1_accum_debug = None
def get_symm_buffer_for_mega_moe(group: dist.ProcessGroup,