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

fix: NVFP4 global scale bug — gsb=weight_scale_2 (not input_scale*ws2), gsa=input_scale

Browse Source
This commit is contained in:
biondizzle
2026-06-01 02:19:35 +00:00
parent 1eed28dd09
commit f4b444b456
1 changed files with 13 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
single_shot_inference.py
View File

@@ -141,9 +141,19 @@ def make_nvfp4_linear(in_features, out_features, device, all_w, pfx, proj_name):
actual_in = weight.shape[1] * 2 # K_packed * 2 = BF16 input dim (for buffer allocation)
lin = Nvfp4Linear(actual_in, actual_out, max_num_tokens=8192, device=d)
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
# Global scales for NVFP4 GEMM:
# gsb (weight global scale) = weight_scale_2 (NOT input_scale * weight_scale_2)
# gsa (activation global scale) = input_scale from checkpoint
# Dequant: w = lut[w_packed] * weight_scale * weight_scale_2
# GEMM: y = (x * scale_a * gsa) @ (w * scale_b * gsb)
# Nvfp4Linear.finalize_weights does: gsb = gs * ws2_val
# So to get gsb = ws2_val, set gs = 1.0 and let ws2 do its job
lin.gs = [1.0] # base gs — finalize_weights will multiply by ws2
lin.ws2 = [ws2.to(d) if ws2 is not None else None]
# Set activation global scale from checkpoint input_scale
isc_val = isc.float().item() if isc is not None else 1.0 / (6.0 * 448.0)
lin._activation_global_scale = isc_val # gsa = input_scale
lin.finalize_weights(); return lin
# =====================================================================
# Compressor — CSA (ratio=4) and HCA (ratio=128) [PyTorch ref]