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

NVFP4 router gate: use Nvfp4Linear for both checkpoint and quantized paths

Browse Source 
- Checkpoint path: load NVFP4 gate weight directly into Nvfp4Linear
- BF16 path: quantize and load into Nvfp4Linear
- Both paths use proven production GEMM (no custom kernel)
- load_nvfp4_fused_gate now creates Nvfp4Linear from BF16 weight
This commit is contained in:
biondizzle
2026-06-01 11:25:50 +00:00
parent e5dbe1ed22
commit 7a05d3d3af
2 changed files with 53 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
dsv4/layers/router.py
View File

@@ -163,17 +163,29 @@ class Router:
self.gate_lin = gate_lin
def load_nvfp4_fused_gate(self, gate_weight, gate_weight_scale,
gate_ws2, gate_input_scale) -> None:
"""Set raw NVFP4 gate tensors for the fused single-kernel path.
Preferred over load_nvfp4_gate (2-kernel) when available.
The fused kernel handles activation quantization + GEMM +
router epilogue in a single kernel launch.
"""
gate_ws2, gate_input_scale,
gate_weight_bf16=None) -> None:
"""Set raw NVFP4 gate tensors and create Nvfp4Linear for production GEMM."""
self.gate_weight = gate_weight.to(device=self.device)
self.gate_weight_scale = gate_weight_scale.to(device=self.device)
self.gate_ws2 = gate_ws2.to(device=self.device) if gate_ws2 is not None else None
self.gate_input_scale = gate_input_scale.to(device=self.device)
self.gate_input_scale = gate_input_scale.to(self.device)
# Create Nvfp4Linear from BF16 weight (handles layout correctly)
if gate_weight_bf16 is not None:
from dsv4.layers.linear import Nvfp4Linear
from dsv4.ops.quantize import quantize_to_nvfp4
E = gate_weight_bf16.shape[0]
gate_lin = Nvfp4Linear(in_features=self.hidden_size, out_features=E, device=self.device)
g_fp4, g_sf, g_gs = quantize_to_nvfp4(gate_weight_bf16.bfloat16().to(self.device))
gate_lin.fp4 = [g_fp4]
gate_lin.sf = [g_sf]
gate_lin.gs = [g_gs]
ws2_val = gate_ws2.float().item() if gate_ws2.numel() == 1 else gate_ws2.float().mean().item()
gate_lin.ws2 = [torch.tensor([ws2_val], device=self.device, dtype=torch.float32)]
gate_lin._activation_global_scale = gate_input_scale.float().item() if gate_input_scale.numel() == 1 else gate_input_scale.float().mean().item()
gate_lin.finalize_weights()
self.gate_lin = gate_lin
def finalize_weights(self) -> None:
"""Allocate output buffers and JIT-compile the routing kernel.

73
single_shot_inference.py
View File

@@ -694,55 +694,48 @@ 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")
# Try NVFP4 gate weights first (production path)
# 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
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:
# NVFP4 gate: load raw tensors for fused single-kernel path
# 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
gate_lin.finalize_weights()
router.load_nvfp4_gate(gate_lin)
router.load_weights(e_bias=eb.to(dev, torch.float32))
router.load_nvfp4_fused_gate(
gate_weight=gate_w.to(dev),
gate_weight_scale=gate_ws.to(dev),
gate_ws2=gate_ws2.to(dev) if gate_ws2 is not None else torch.tensor(1.0, device=dev),
gate_input_scale=gate_isc.to(dev) if gate_isc is not None else torch.tensor(1.0 / (6.0 * 448.0), device=dev),
)
if li < 5: print(f" L{li}: NVFP4 router gate (checkpoint)", flush=True)
else:
# BF16 gate weight: quantize to NVFP4 for fused kernel
# BF16 gate weight: quantize to NVFP4
gw = all_w.get(f"{pfx}.gate.weight")
if gw is not None:
if gw.shape == (cfg["n_routed_experts"], H): gw = gw.T.contiguous()
gw = gw.bfloat16().to(dev)
# Quantize BF16 → NVFP4 for fused router kernel
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
try:
gw_fp4, gw_sf, gw_gs = quantize_to_nvfp4(gw)
router.load_weights(e_bias=eb.to(dev, torch.float32))
router.load_nvfp4_fused_gate(
gate_weight=gw_fp4,
gate_weight_scale=gw_sf,
gate_ws2=torch.tensor([gw_gs], device=dev, dtype=torch.float32),
gate_input_scale=torch.tensor([1.0 / (6.0 * 448.0)], device=dev, dtype=torch.float32),
)
if li < 5: print(f" L{li}: Fused NVFP4 gate OK (gs={gw_gs:.6f})", flush=True)
except Exception as e:
print(f" L{li}: Fused NVFP4 gate FAILED: {e}", flush=True)
import traceback; traceback.print_exc()
# Fallback: create Nvfp4Linear from BF16 weight
from dsv4.layers.linear import Nvfp4Linear
gate_lin = Nvfp4Linear(in_features=H, out_features=cfg["n_routed_experts"], device=dev)
gate_lin.fp4 = None; gate_lin.sf = None # will quantize from BF16
from dsv4.ops.quantize import quantize_to_nvfp4
g_bf16 = gw if gw.shape == (cfg["n_routed_experts"], H) else gw.T.contiguous()
g_fp4, g_sf, g_gs = quantize_to_nvfp4(g_bf16.bfloat16().to(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)]
gate_lin._activation_global_scale = 1.0 / (6.0 * 448.0)
gate_lin.finalize_weights()
router.load_nvfp4_gate(gate_lin)
router.load_weights(e_bias=eb.to(dev, torch.float32))
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)
gate_lin.finalize_weights()
router.load_nvfp4_gate(gate_lin)
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)
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,