biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: moe_pipeline.py gate/up split — L1 output is 2*intermediate, not intermediate

Browse Source
This commit is contained in:
biondizzle
2026-05-16 19:28:15 +00:00
parent da31ce7e1a
commit 4c06c51ec3
1 changed files with 5 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
cutedsl/moe_pipeline.py
View File

@@ -199,10 +199,11 @@ def run_nvfp4_moe(
# ════════════════════════════════════════════════════════════════
# SiLU(gate) * up (BF16 — nonlinear requires BF16)
# ════════════════════════════════════════════════════════════════
intermediate = l1_out.shape[1]
gate = l1_out[:, :intermediate]
up = l1_out[:, intermediate:]
activated = torch.nn.functional.silu(gate) * up # (num_slots, half) BF16
# L1 output is (tokens, 2*intermediate) — gate and up fused
intermediate_size = l1_out.shape[1] // 2
gate = l1_out[:, :intermediate_size]
up = l1_out[:, intermediate_size:]
activated = torch.nn.functional.silu(gate) * up # (num_slots, intermediate) BF16
# ════════════════════════════════════════════════════════════════
# L2: down projection (NVFP4 × NVFP4 → BF16)