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

feat: implement sm90 megamoe phase2 dispatch-only

Browse Source
This commit is contained in:
Xinyi Liu
2026-06-18 00:00:36 +08:00
parent 74eb59dfaa
commit 540e5aeadc
4 changed files with 622 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
deep_gemm/mega/__init__.py
+18 -8
View File
@@ -60,14 +60,24 @@ class SymmBuffer:
torch.cuda.synchronize()
# Create input buffer views
(self.x, self.x_sf,
self.topk_idx, self.topk_weights,
self.l1_acts, self.l1_acts_sf,
self.l2_acts, self.l2_acts_sf) = slice_input_buffers(self.buffer)
self.l1_topk_weights = None
self.expert_recv_count_sum = None
self.l1_arrival_count = None
self.token_src_metadata = None
buffer_views = slice_input_buffers(self.buffer)
if _is_sm90():
(self.x, self.x_sf,
self.topk_idx, self.topk_weights,
self.l1_acts, self.l1_acts_sf, self.l1_topk_weights,
self.l2_acts, self.l2_acts_sf,
self.expert_recv_count_sum,
self.l1_arrival_count,
self.token_src_metadata) = buffer_views
else:
(self.x, self.x_sf,
self.topk_idx, self.topk_weights,
self.l1_acts, self.l1_acts_sf,
self.l2_acts, self.l2_acts_sf) = buffer_views
self.l1_topk_weights = None
self.expert_recv_count_sum = None
self.l1_arrival_count = None
self.token_src_metadata = None
def destroy(self):
self.handle = None