A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 6, 13:15 UTC
[Bugfix] Rescale NVFP4 weight scales to fix BF16 dequant underflow (#34577)
Signed-off-by: ricky-chaoju <ricky.chen@infinirc.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
Chao-Ju Chen committed
Mar 18, 2026 at 04:48 UTC
245758992ed74fbaaffcdb4e607ad817627455fc
1 file changed
+71
-6
vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py
+71
-6
@@ -27,7 +27,44 @@ def is_fp4_marlin_supported():
27
return current_platform.has_device_capability(75)
28
29
30
-def nvfp4_marlin_process_scales(marlin_scales):
30
+def _nvfp4_compute_scale_factor(marlin_scales: torch.Tensor) -> float:
31
+ """Compute the power-of-2 scale_factor needed so that all non-zero
32
+ values in marlin_scales * 2^7 are >= 2 after rescaling.
33
+ Returns a Python float (power of 2, >= 1.0)."""
34
+ ws_float = marlin_scales.float() * (2**7)
35
+ nonzero_mask = ws_float > 0
36
+ if nonzero_mask.any():
37
+ min_val = ws_float[nonzero_mask].min()
38
+ if min_val < 2:
39
+ sf = (2 / min_val).log2().ceil().exp2()
40
+ assert (ws_float[nonzero_mask] * sf <= 448 * (2**7)).all(), (
41
+ "NVFP4 scale dynamic range too large for rescaling"
42
+ )
43
+ return sf.item()
44
+ return 1.0
45
+
46
+
47
+def nvfp4_marlin_process_scales(
48
+ marlin_scales: torch.Tensor,
49
+ scale_factor: float | None = None,
50
+) -> tuple[torch.Tensor, float]:
51
+ """Process NVFP4 weight scales into the special S0E5M3 format for Marlin.
52
+
53
+ Args:
54
+ marlin_scales: Weight scales tensor in half precision, already
55
+ permuted for the Marlin kernel layout.
56
+ scale_factor: Optional power-of-2 rescaling factor. If None, the
57
+ factor is computed automatically so that every non-zero scale
58
+ satisfies ``scale * 2^7 >= 2`` (i.e., the MSB of the S0E5M3
59
+ representation is always 1). When provided (e.g., for MoE
60
+ layers where all experts must share the same factor), the
61
+ given value is used directly. The caller is responsible for
62
+ dividing ``global_scale`` by the returned ``scale_factor`` to
63
+ preserve numerical correctness.
64
+
65
+ Returns:
66
+ A tuple of (processed_scales, scale_factor).
67
+ """
68
if not (marlin_scales >= 0).all():
69
logger.warning_once(
70
"NVFP4 Marlin assumes the scales to be >=0, but has encountered "
@@ -51,11 +88,21 @@ def nvfp4_marlin_process_scales(marlin_scales):
88
# when weight_scale > 0. This allows us to have an exponent bias
89
# closer to zero after dequantization.
90
91
+ # Rescale weight_scale so that all non-zero values have MSB=1
92
+ # after multiplying by 2^7 (i.e., weight_scale * 2^7 >= 2).
93
+ # This is needed for models whose E4M3 scales were not normalized
94
+ # to fully utilize the E4M3 dynamic range (e.g., global_scale=1).
95
+ # The caller must compensate by dividing global_scale by scale_factor.
96
+ if scale_factor is None:
97
+ scale_factor = _nvfp4_compute_scale_factor(marlin_scales)
98
+ if scale_factor > 1.0:
99
+ marlin_scales = (marlin_scales.float() * scale_factor).to(torch.half)
100
+
101
marlin_scales = (marlin_scales * (2**7)).view(torch.int16) << 1
102
marlin_scales = marlin_scales.view(torch.float8_e4m3fn)
103
marlin_scales = marlin_scales[:, 1::2].contiguous()
104
58
- return marlin_scales
105
+ return marlin_scales, scale_factor
106
107
108
def mxfp4_marlin_process_scales(marlin_scales, input_dtype=None):
@@ -200,11 +247,12 @@ def prepare_fp4_layer_for_marlin(
247
)
248
249
if is_nvfp4:
203
- weight_scale = nvfp4_marlin_process_scales(weight_scale)
250
+ weight_scale, scale_factor = nvfp4_marlin_process_scales(weight_scale)
251
layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False)
252
253
weight_global_scale = layer.weight_global_scale.to(param_dtype)
254
weight_global_scale = nvfp4_marlin_process_global_scale(weight_global_scale)
255
+ weight_global_scale = weight_global_scale / scale_factor
256
layer.weight_global_scale = torch.nn.Parameter(
257
weight_global_scale, requires_grad=False
258
)
@@ -303,6 +351,10 @@ def prepare_nvfp4_moe_layer_for_marlin(
351
else:
352
size_n, size_k = K, N
353
354
+ # All experts share one global_scale, so compute the max
355
+ # scale_factor across all experts first, then apply uniformly.
356
+ combined_scale_factor = _nvfp4_compute_scale_factor(scales)
357
+
358
for i in range(E):
359
scale = scales[i].T
360
marlin_scales = marlin_permute_scales(
@@ -312,11 +364,14 @@ def prepare_nvfp4_moe_layer_for_marlin(
364
group_size=GROUP_SIZE,
365
is_a_8bit=is_a_8bit,
366
)
315
- marlin_scales = nvfp4_marlin_process_scales(marlin_scales)
367
+ marlin_scales, _ = nvfp4_marlin_process_scales(
368
+ marlin_scales, scale_factor=combined_scale_factor
369
+ )
370
tensor_list.append(marlin_scales)
371
372
scales = torch.cat([x.unsqueeze(0) for x in tensor_list], 0)
373
g_scales = nvfp4_marlin_process_global_scale(g_scales)
374
+ g_scales = g_scales / combined_scale_factor
375
return scales, g_scales
376
377
w13_scale, w13_scale_2 = premute_scales(w13_scale, w13_scale_2, "w13")
@@ -394,6 +449,11 @@ def prepare_moe_fp4_layer_for_marlin(
449
else:
450
size_n, size_k = k, n
451
452
+ # For NVFP4: compute unified scale_factor across all experts
453
+ combined_scale_factor = None
454
+ if is_nvfp4:
455
+ combined_scale_factor = _nvfp4_compute_scale_factor(scales)
456
+
457
for i in range(e):
458
scale = scales[i].T
459
@@ -405,7 +465,9 @@ def prepare_moe_fp4_layer_for_marlin(
465
is_a_8bit=is_a_8bit,
466
)
467
if is_nvfp4:
408
- marlin_scales = nvfp4_marlin_process_scales(marlin_scales)
468
+ marlin_scales, _ = nvfp4_marlin_process_scales(
469
+ marlin_scales, scale_factor=combined_scale_factor
470
+ )
471
else:
472
marlin_scales = mxfp4_marlin_process_scales(
473
marlin_scales, input_dtype=input_dtype
@@ -417,7 +479,9 @@ def prepare_moe_fp4_layer_for_marlin(
479
setattr(layer, name + "_weight_scale", scales)
480
481
if is_nvfp4:
482
+ assert combined_scale_factor is not None
483
global_scale = nvfp4_marlin_process_global_scale(global_scale)
484
+ global_scale = global_scale / combined_scale_factor
485
global_scale = torch.nn.Parameter(global_scale, requires_grad=False)
486
setattr(layer, name + "_weight_scale_2", global_scale)
487
@@ -488,9 +552,10 @@ def rand_marlin_weight_nvfp4_like(weight, group_size, input_dtype=None):
552
group_size=group_size,
553
is_a_8bit=is_a_8bit,
554
)
491
- marlin_scales = nvfp4_marlin_process_scales(marlin_scales)
555
+ marlin_scales, scale_factor = nvfp4_marlin_process_scales(marlin_scales)
556
557
global_scale = nvfp4_marlin_process_global_scale(global_scale)
558
+ global_scale = global_scale / scale_factor
559
560
return weight_ref.T, marlin_qweight, marlin_scales, global_scale
561