fix: same gate/up split fix in moe_pipeline.py

2026-05-16 04:04:53 +00:00
parent 6d17988b51
commit 174ad70dca
1 changed files with 2 additions and 3 deletions
--- a/cutedsl/moe_pipeline.py
+++ b/cutedsl/moe_pipeline.py
@@ -200,9 +200,8 @@ def run_nvfp4_moe(
    # SiLU(gate) * up (BF16 — nonlinear requires BF16)
    # ════════════════════════════════════════════════════════════════
    intermediate = l1_out.shape[1]
-    half = intermediate // 2  # 3072
-    gate = l1_out[:, :half]
-    up = l1_out[:, half:]
+    gate = l1_out[:, :intermediate]
+    up = l1_out[:, intermediate:]
    activated = torch.nn.functional.silu(gate) * up  # (num_slots, half) BF16
    
    # ════════════════════════════════════════════════════════════════