feat: implement sm90 megamoe phase5 l2 scatter

2026-06-18 15:17:20 +08:00
parent fc8218750c
commit 9bd0519605
5 changed files with 379 additions and 72 deletions
--- a/megamoe_dev_test_scripts/phase1/interface_smoke.py
+++ b/megamoe_dev_test_scripts/phase1/interface_smoke.py
@@ -85,6 +85,8 @@ def main() -> None:
    assert buffer.l2_acts.shape[1] == args.intermediate_hidden
    assert buffer.l2_acts_sf.shape[1] == args.intermediate_hidden // 64
    assert buffer.l2_acts_sf.dtype == torch.float32
+    assert buffer.combine_acts.shape == (args.num_topk, buffer.num_max_tokens_per_rank, args.hidden)
+    assert buffer.combine_acts.dtype == torch.bfloat16

    num_tokens = args.num_tokens
    buffer.x[:num_tokens].copy_(torch.randn((num_tokens, args.hidden), device='cuda').to(torch.float8_e4m3fn))