From fae418c3a30219ae12b6fa4924c0f0a61bcee20d Mon Sep 17 00:00:00 2001 From: biondizzle Date: Fri, 15 May 2026 08:57:43 +0000 Subject: [PATCH] final scatter --- src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py b/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py index a6c8b9a8..95ff2bda 100644 --- 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), + )