From f4b444b45616c5c1f2568a047332fa22bb1492e1 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Mon, 1 Jun 2026 02:19:35 +0000 Subject: [PATCH] =?UTF-8?q?fix:=20NVFP4=20global=20scale=20bug=20=E2=80=94?= =?UTF-8?q?=20gsb=3Dweight=5Fscale=5F2=20(not=20input=5Fscale*ws2),=20gsa?= =?UTF-8?q?=3Dinput=5Fscale?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- single_shot_inference.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/single_shot_inference.py b/single_shot_inference.py index aa3e49ee..2b4cd9be 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -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]