final scatter

2026-05-15 08:57:43 +00:00
parent f2cacfc2f2
commit fae418c3a3
1 changed files with 5 additions and 1 deletions
--- a/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py
+++ b/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py
@@ -360,4 +360,8 @@ def nvfp4_mega_moe_full(

    # Step 6: Final scatter — routing weights applied ONCE
    y.zero_()
-    y.index_add_(0, slot_token, slot_weight[:, None] * l2_slots)
+    y.index_add_(
+        0,
+        slot_token,
+        l2_slots * slot_weight.to(l2_slots.dtype).unsqueeze(1),
+    )