feat: implement sm90 megamoe phase5 l2 scatter

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
+4 -1
View File
@@ -70,7 +70,8 @@ class SymmBuffer:
self.l1_arrival_count,
self.l2_arrival_mask,
self.token_src_metadata,
self.l1_accum_debug) = buffer_views
self.l1_accum_debug,
self.combine_acts) = buffer_views
else:
(self.x, self.x_sf,
self.topk_idx, self.topk_weights,
@@ -82,6 +83,7 @@ class SymmBuffer:
self.l2_arrival_mask = None
self.token_src_metadata = None
self.l1_accum_debug = None
self.combine_acts = None
def destroy(self):
self.handle = None
@@ -101,6 +103,7 @@ class SymmBuffer:
self.l2_arrival_mask = None
self.token_src_metadata = None
self.l1_accum_debug = None
self.combine_acts = None
def get_symm_buffer_for_mega_moe(group: dist.ProcessGroup,