feat: implement sm90 megamoe phase4 l1 epilogue

2026-06-18 01:09:45 +08:00
parent 2bb1756787
commit f3553f976c
6 changed files with 330 additions and 10 deletions
--- a/megamoe_dev_test_scripts/phase1/interface_smoke.py
+++ b/megamoe_dev_test_scripts/phase1/interface_smoke.py
@@ -83,7 +83,7 @@ def main() -> None:
    assert buffer.l1_acts_sf.shape[1] == args.hidden // 128
    assert buffer.l1_acts_sf.dtype == torch.float32
    assert buffer.l2_acts.shape[1] == args.intermediate_hidden
-    assert buffer.l2_acts_sf.shape[1] == args.intermediate_hidden // 128
+    assert buffer.l2_acts_sf.shape[1] == args.intermediate_hidden // 64
    assert buffer.l2_acts_sf.dtype == torch.float32

    num_tokens = args.num_tokens