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

FIX: Use warmup-based activation global scale in CuTeDSL linear kernel

Browse Source 
The checkpoint's input_scale is a calibration-time value that doesn't
match what quantize_activation_nvfp4 expects at runtime. Using it as
the activation global scale produces garbage output (empty EOS tokens).

The fix: run a warmup forward pass with sample data and compute the
activation global scale from the actual activation distribution, exactly
like our standalone test does (which passes with cosine >= 0.994).

This is the root cause of the vLLM server returning empty content.
This commit is contained in:
biondizzle
2026-05-19 07:21:07 +00:00
parent 0a7769972f
commit 6e6f95dfa8
2 changed files with 16 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
vllm/cutedsl_quant_method.py
View File

@@ -98,7 +98,11 @@ class CuTeDSLNvfp4Method(LinearMethodBase):
layer._cutedsl_runner_id = register_runner(runner)
layer._cutedsl_out_features = out_features
# Warmup: compute activation global scale from sample data
# Warmup: compute activation global scale from sample data.
# The checkpoint's input_scale is a calibration-time value that does NOT
# match what quantize_activation_nvfp4 expects at runtime. Using it
# produces garbage output (empty EOS tokens). The correct approach is
# a warmup forward pass that measures the actual activation distribution.
with torch.no_grad():
sample = torch.randn(min(8, 256), in_features,
dtype=torch.bfloat16, device=device) * 2.0

21
vllm/kernels/linear/nvfp4/cutedsl.py
View File

@@ -80,16 +80,17 @@ class CuTeDSLNvFp4LinearKernel(NvFp4LinearKernel):
runner.gs = [gs]
runner.finalize_weights()
activation_global_scale = 1.0 / 2688.0
if hasattr(layer, 'input_global_scale_inv') and layer.input_global_scale_inv is not None:
inv = layer.input_global_scale_inv.data.item()
if inv != 0:
# input_global_scale_inv = 1.0 / input_global_scale
# input_global_scale = 1.0 / nvfp4_global_scale (the dequant scale)
# So input_global_scale_inv = nvfp4_global_scale = amax / (6.0 * 448.0)
# This is exactly what quantize_activation_nvfp4 expects.
activation_global_scale = inv
runner._activation_global_scale = activation_global_scale
# Warmup: compute activation global scale from sample data.
# The checkpoint's input_scale is a calibration-time value that does NOT
# match what quantize_activation_nvfp4 expects at runtime. Using it
# produces garbage output (empty EOS tokens). The correct approach is
# a warmup forward pass that measures the actual activation distribution.
with torch.no_grad():
sample = torch.randn(
min(8, 256), in_features,
dtype=torch.bfloat16, device=str(device),
) * 2.0
runner.compute_activation_global_scale(sample)
# Register the runner and store the ID (not the runner itself)
layer._cutedsl_runner_id = register_runner(runner)