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

fix: fold weight_scale_2 into global_scale_b for NVFP4 GEMM

Browse Source 
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
This commit is contained in:
biondizzle
2026-06-01 00:10:50 +00:00
parent 563df02aef
commit 172448514c
4 changed files with 63 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
dsv4/layers/linear.py
View File

@@ -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.

14
dsv4/layers/moe.py
View File

@@ -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).

16
dsv4/layers/shared_expert.py
View File

@@ -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."""

32
single_shot_inference.py
View File

@@ -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 = {}