From 172448514ca2e04ed8974d69869738d521e5545b Mon Sep 17 00:00:00 2001 From: biondizzle Date: Mon, 1 Jun 2026 00:10:50 +0000 Subject: [PATCH] fix: fold weight_scale_2 into global_scale_b for NVFP4 GEMM Critical bug fix: weight_scale_2 (the second-level NVFP4 scale) was being dropped entirely in the production pipeline. The dequant formula is lut[w] * weight_scale * weight_scale_2, so weight_scale_2 must be folded into the GEMM's global_scale_b parameter. Fixes in: - Nvfp4Linear: ws2 field, folded in finalize_weights() - Nvfp4MoE: l1_ws2/l2_ws2 lists, folded in _ensure_stacked() - Nvfp4SharedExpert: l1_ws2/l2_ws2 lists, folded in finalize_weights() - single_shot_inference.py: pass weight_scale_2 through all loading paths - Also fix missing o_a_prod key fallback in attention output --- dsv4/layers/linear.py | 10 ++++++++++ dsv4/layers/moe.py | 14 ++++++++++++++ dsv4/layers/shared_expert.py | 16 ++++++++++++++++ single_shot_inference.py | 32 +++++++++++++++++++++++--------- 4 files changed, 63 insertions(+), 9 deletions(-) diff --git a/dsv4/layers/linear.py b/dsv4/layers/linear.py index 9d7e9d0a..e708b598 100644 --- a/dsv4/layers/linear.py +++ b/dsv4/layers/linear.py @@ -52,6 +52,7 @@ class Nvfp4Linear: self.fp4 = None # list of 1 tensor self.sf = None # list of 1 tensor self.gs = None # list of 1 float + self.ws2 = None # list of 1 tensor — weight_scale_2 (scalar, folded into global_scale_b) # Processed weights self._mat_b = None @@ -73,10 +74,19 @@ class Nvfp4Linear: self._scale_b = assemble_scales_3d_side(self.sf) self._gsb = torch.tensor(self.gs, dtype=torch.float32, device=self.device) + # Fold weight_scale_2 into global_scale_b + # Dequant formula: w = lut[w_packed] * weight_scale * weight_scale_2 + # Production GEMM: y = (x * scale_a * gsa) @ (w * scale_b * gsb) + # So gsb = input_scale * weight_scale_2 + if self.ws2 is not None and len(self.ws2) > 0 and self.ws2[0] is not None: + ws2_val = self.ws2[0].float().item() + self._gsb = self._gsb * ws2_val + # Free raw weights self.fp4 = None self.sf = None self.gs = None + self.ws2 = None # Eagerly JIT-compile the GEMM kernel for this (K, N) shape. # Uses num_groups=1 since this is a single linear layer. diff --git a/dsv4/layers/moe.py b/dsv4/layers/moe.py index 63d3a146..038ce14e 100644 --- a/dsv4/layers/moe.py +++ b/dsv4/layers/moe.py @@ -273,8 +273,22 @@ class Nvfp4MoE: self._l1_gsb = torch.tensor(self.l1_gs, dtype=torch.float32, device=self.device) self._l2_gsb = torch.tensor(self.l2_gs, dtype=torch.float32, device=self.device) + + # Fold weight_scale_2 into global_scale_b + # gsb = input_scale * weight_scale_2 + if self.l1_ws2 is not None: + for i, ws2 in enumerate(self.l1_ws2): + if ws2 is not None: + self._l1_gsb[i] *= ws2.float().item() + if self.l2_ws2 is not None: + for i, ws2 in enumerate(self.l2_ws2): + if ws2 is not None: + self._l2_gsb[i] *= ws2.float().item() + self.l1_gs = None self.l2_gs = None + self.l1_ws2 = None + self.l2_ws2 = None # Allocate buffers and eagerly warmup JIT compilation. # cute.compile does NOT corrupt GPU memory (verified 2026-05-20). diff --git a/dsv4/layers/shared_expert.py b/dsv4/layers/shared_expert.py index 5d51cbba..aa0fb90a 100644 --- a/dsv4/layers/shared_expert.py +++ b/dsv4/layers/shared_expert.py @@ -71,6 +71,9 @@ class Nvfp4SharedExpert: self.l2_fp4 = None self.l2_sf = None self.l2_gs = None + # weight_scale_2 per layer (scalar, folded into global_scale_b in finalize_weights) + self.l1_ws2 = None + self.l2_ws2 = None # Processed weights (set by finalize_weights) self._l1_mat_b = None @@ -108,6 +111,17 @@ class Nvfp4SharedExpert: self._l1_gsb = torch.tensor(self.l1_gs, dtype=torch.float32, device=self.device) self._l2_gsb = torch.tensor(self.l2_gs, dtype=torch.float32, device=self.device) + # Fold weight_scale_2 into global_scale_b + # gsb = input_scale * weight_scale_2 + if self.l1_ws2 is not None: + for i, ws2 in enumerate(self.l1_ws2): + if ws2 is not None: + self._l1_gsb[i] *= ws2.float().item() + if self.l2_ws2 is not None: + for i, ws2 in enumerate(self.l2_ws2): + if ws2 is not None: + self._l2_gsb[i] *= ws2.float().item() + # Free raw weights self.l1_fp4 = None self.l1_sf = None @@ -115,6 +129,8 @@ class Nvfp4SharedExpert: self.l2_fp4 = None self.l2_sf = None self.l2_gs = None + self.l1_ws2 = None + self.l2_ws2 = None def _allocate_buffers(self): """Pre-allocate all buffers at max size for cudagraph compatibility.""" diff --git a/single_shot_inference.py b/single_shot_inference.py index d456bda8..308dda1b 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -133,6 +133,7 @@ def make_nvfp4_linear(in_features, out_features, device, all_w, pfx, proj_name): weight, ws, ws2, isc = get_nvfp4_weight(all_w, pfx, proj_name) assert weight is not None, f"{pfx}.{proj_name}.weight not found" lin.fp4 = [weight.to(d)]; lin.sf = [ws.to(d)] + lin.ws2 = [ws2.to(d) if ws2 is not None else None] # weight_scale_2 gs = isc.float().item() if isc is not None else 1.0 / (6.0 * 448.0) lin.gs = [gs]; lin.finalize_weights(); return lin @@ -357,7 +358,13 @@ def forward_attention(x_normed, w, li, cfg, rope_cos, rope_sin, g_flat = g_out.permute(1, 0, 2).reshape(T, o_groups * o_rank) F_attn = prod_lin['o_b'](g_flat) else: - F_attn = prod_lin['o_a'](attn_out.reshape(T, n_h * hd)) + # o_a_proj as full-rank BF16 linear (no low-rank) + oa_full = w.get(f"{pfx}.o_a_proj.weight") + if oa_full is not None: + F_attn = F.linear(attn_out.reshape(T, n_h * hd), oa_full.bfloat16().to(dev)) + else: + log.warning(f"L{li}: No o_a_proj weight, zero attention output") + F_attn = torch.zeros(T, cfg["hidden_size"], dtype=torch.bfloat16, device=dev) return F_attn, q_a # ===================================================================== @@ -394,23 +401,26 @@ def forward_layer(X_l, w, li, cfg, rope_cos, rope_sin, # ===================================================================== def _load_moe_weights_stacked(all_w, li, pfx, dev, moe, cfg): n_e = cfg["n_routed_experts"] - l1_fp4_list, l1_sf_list, l1_gs_list = [], [], [] - l2_fp4_list, l2_sf_list, l2_gs_list = [], [], [] + l1_fp4_list, l1_sf_list, l1_gs_list, l1_ws2_list = [], [], [], [] + l2_fp4_list, l2_sf_list, l2_gs_list, l2_ws2_list = [], [], [], [] for eid in range(n_e): ep = f"{pfx}.experts.{eid}" - gw, gws, _, gisc = get_nvfp4_weight(all_w, ep, 'gate_proj') - uw, uws, _, uisc = get_nvfp4_weight(all_w, ep, 'up_proj') + gw, gws, gws2, gisc = get_nvfp4_weight(all_w, ep, 'gate_proj') + uw, uws, uws2, uisc = get_nvfp4_weight(all_w, ep, 'up_proj') if gw is not None and uw is not None: l1_fp4_list.append(torch.cat([gw, uw], dim=0).to(dev)) if gws is not None and uws is not None: l1_sf_list.append(torch.cat([gws, uws], dim=0).to(dev)) gs = gisc.float().item() if gisc is not None else 1.0 / (6.0 * 448.0) l1_gs_list.append(gs) - dw, dws, _, disc = get_nvfp4_weight(all_w, ep, 'down_proj') + # weight_scale_2: scalar, folded into global_scale_b + l1_ws2_list.append(gws2.to(dev) if gws2 is not None else None) + dw, dws, dws2, disc = get_nvfp4_weight(all_w, ep, 'down_proj') if dw is not None: l2_fp4_list.append(dw.to(dev)) if dws is not None: l2_sf_list.append(dws.to(dev)) gs2 = disc.float().item() if disc is not None else 1.0 / (6.0 * 448.0) l2_gs_list.append(gs2) + l2_ws2_list.append(dws2.to(dev) if dws2 is not None else None) if not l1_fp4_list: log.warning(f"L{li}: No expert weights found"); return l1_stacked = torch.stack(l1_fp4_list).to(dev) l1_sf_stacked = torch.stack(l1_sf_list).to(dev) if l1_sf_list else None @@ -418,19 +428,23 @@ def _load_moe_weights_stacked(all_w, li, pfx, dev, moe, cfg): l2_sf_stacked = torch.stack(l2_sf_list).to(dev) if l2_sf_list else None del l1_fp4_list, l1_sf_list, l2_fp4_list, l2_sf_list moe.prepare_weights_from_stacked(l1_stacked, l1_sf_stacked, l1_gs_list, l2_stacked, l2_sf_stacked, l2_gs_list) + moe.l1_ws2 = l1_ws2_list + moe.l2_ws2 = l2_ws2_list def _load_shared_expert_weights(all_w, li, pfx, dev, se, cfg): - gw, gws, _, gisc = get_nvfp4_weight(all_w, f"{pfx}.shared_experts", 'gate_proj') - uw, uws, _, uisc = get_nvfp4_weight(all_w, f"{pfx}.shared_experts", 'up_proj') - dw, dws, _, disc = get_nvfp4_weight(all_w, f"{pfx}.shared_experts", 'down_proj') + gw, gws, gws2, gisc = get_nvfp4_weight(all_w, f"{pfx}.shared_experts", 'gate_proj') + uw, uws, uws2, uisc = get_nvfp4_weight(all_w, f"{pfx}.shared_experts", 'up_proj') + dw, dws, dws2, disc = get_nvfp4_weight(all_w, f"{pfx}.shared_experts", 'down_proj') if gw is not None and uw is not None: se.l1_fp4 = [torch.cat([gw, uw], dim=0).to(dev)] se.l1_sf = [torch.cat([gws, uws], dim=0).to(dev)] if gws is not None and uws is not None else [torch.zeros(1, device=dev, dtype=torch.float8_e4m3fn)] se.l1_gs = [gisc.float().item() if gisc is not None else 1.0 / (6.0 * 448.0)] + se.l1_ws2 = [gws2.to(dev) if gws2 is not None else None] # weight_scale_2 if dw is not None: se.l2_fp4 = [dw.to(dev)] se.l2_sf = [dws.to(dev)] if dws is not None else [torch.zeros(1, device=dev, dtype=torch.float8_e4m3fn)] se.l2_gs = [disc.float().item() if disc is not None else 1.0 / (6.0 * 448.0)] + se.l2_ws2 = [dws2.to(dev) if dws2 is not None else None] # weight_scale_2 def _cache_layer_weights_no_experts(all_w, n_layers, devices): cached = {}