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

Switch router gate from NVFP4 to BF16 (dequantize)

Browse Source 
Dequantize NVFP4 gate weight to BF16 for router. Wrong top-k expert
selection from NVFP4 quantization noise is much worse than a small GEMM
error — one wrong expert poisons the whole token.

Also: lm_head already BF16 from previous commit.
This commit is contained in:
biondizzle
2026-06-03 12:30:34 +00:00
parent cfea22cd6f
commit dc5a24687e
1 changed files with 15 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

48
single_shot_inference.py
View File

@@ -1306,50 +1306,32 @@ def main():
router.load_weights(hash_lut=all_w[f"{pfx}.gate.tid2eid"].to(dev, torch.int32))
else:
eb = all_w.get(f"{pfx}.gate.e_score_correction_bias")
# NVFP4 production GEMM for router gate
# Custom CuTeDSL fused kernel crashes MLIR optimizer,
# so we use Nvfp4Linear (proven production path).
from dsv4.layers.linear import Nvfp4Linear
# Router gate: BF16 path — dequantize NVFP4 weight to BF16.
# NVFP4 router gate can produce wrong top-k experts; BF16 preserves
# the full logit distribution for accurate expert selection.
gate_w, gate_ws, gate_ws2, gate_isc = get_nvfp4_weight(all_w, pfx, 'gate')
E = cfg["n_routed_experts"]
if gate_w is not None and gate_ws is not None:
# Checkpoint has NVFP4 gate weight (N_packed, K_packed) — correct layout
gate_lin = Nvfp4Linear(in_features=H, out_features=E, device=dev)
gate_w_view = gate_w.to(dev).view(torch.float4_e2m1fn_x2) if gate_w.dtype == torch.uint8 else gate_w.to(dev)
gate_lin.fp4 = [gate_w_view]
gate_lin.sf = [gate_ws.to(dev)]
ws2_v = gate_ws2.float().item() if gate_ws2 is not None else 1.0
isc_v = gate_isc.float().item() if gate_isc is not None else 1.0/(6.0*448.0)
gate_lin.gs = [1.0]
gate_lin.ws2 = [torch.tensor([ws2_v], device=dev, dtype=torch.float32)]
gate_lin._activation_global_scale = isc_v # placeholder — runtime gsa overrides this
gate_lin._use_runtime_gsa = True # compute gsa from actual input to avoid E4M3 overflow
gate_lin.finalize_weights()
router.load_nvfp4_gate(gate_lin)
# Dequantize NVFP4 gate weight → BF16, store as W_gate
gate_bf16 = dequant_nvfp4(gate_w.to(dev), gate_ws.to(dev),
gate_ws2.to(dev) if gate_ws2 is not None else None,
gate_isc.to(dev) if gate_isc is not None else None)
# W_gate shape: (E, H) for F.linear(x, W_gate)
router.W_gate = gate_bf16
router.gate_lin = None # force BF16 dispatch path
router.load_weights(e_bias=eb.to(dev, torch.float32))
if li < 5: print(f" L{li}: NVFP4 router gate (checkpoint)", flush=True)
if li < 5: print(f" L{li}: BF16 router gate (dequantized from NVFP4)", flush=True)
else:
# BF16 gate weight: quantize to NVFP4
# BF16 gate weight from checkpoint
gw = all_w.get(f"{pfx}.gate.weight")
if gw is not None:
g_bf16 = gw if gw.shape == (E, H) else gw.T.contiguous()
g_bf16 = g_bf16.bfloat16().to(dev)
from dsv4.ops.quantize import quantize_to_nvfp4
g_fp4, g_sf, g_gs = quantize_to_nvfp4(g_bf16)
gate_lin = Nvfp4Linear(in_features=H, out_features=E, device=dev)
gate_lin.fp4 = [g_fp4]
gate_lin.sf = [g_sf]
gate_lin.gs = [g_gs]
gate_lin.ws2 = [torch.tensor([g_gs], device=dev, dtype=torch.float32)]
gate_lin._activation_global_scale = 1.0 / (6.0 * 448.0) # placeholder — runtime gsa overrides
gate_lin._use_runtime_gsa = True # compute gsa from actual input to avoid E4M3 overflow
gate_lin.finalize_weights()
router.load_nvfp4_gate(gate_lin)
router.W_gate = g_bf16.bfloat16().to(dev)
router.gate_lin = None
router.load_weights(e_bias=eb.to(dev, torch.float32))
if li < 5: print(f" L{li}: NVFP4 router gate (quantized, gs={g_gs:.6f})", flush=True)
if li < 5: print(f" L{li}: BF16 router gate (checkpoint BF16)", flush=True)
else:
router.load_weights(e_bias=eb.to(dev, torch.float32))
router.load_weights(e_bias=eb.to(dev, torch.float32))
router.finalize_weights(); routers[li] = router
moe = Nvfp4MoE(num_experts=cfg["n_routed_experts"], hidden_size=H,