From 7a05d3d3afacf42a223c461233257553a2c3f584 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Mon, 1 Jun 2026 11:25:50 +0000 Subject: [PATCH] NVFP4 router gate: use Nvfp4Linear for both checkpoint and quantized paths - 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 --- dsv4/layers/router.py | 28 ++++++++++----- single_shot_inference.py | 73 ++++++++++++++++++---------------------- 2 files changed, 53 insertions(+), 48 deletions(-) diff --git a/dsv4/layers/router.py b/dsv4/layers/router.py index 6fa8c7d2..7f4cc879 100644 --- a/dsv4/layers/router.py +++ b/dsv4/layers/router.py @@ -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. diff --git a/single_shot_inference.py b/single_shot_inference.py index 2f6ecee3..8bab5879 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -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,