[Quantization] Preserve precision in online NVFP4 expert packing (#50029)
Signed-off-by: S1ro1 <matej.sirovatka@gmail.com>
Matej Sirovatka committed
Aug 6, 2026 at 02:23 UTC
9c22668436a4d94aab87ea74a220e060415cf1d8
2 files changed
+58
-10
tests/quantization/test_online.py
+37
@@ -9,6 +9,7 @@ from tests.quantization.utils import (
9
_test_online_quant_peak_mem_impl,
10
is_quant_method_supported,
11
)
12
+from vllm._custom_ops import scaled_fp4_quant
13
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
14
from vllm.model_executor.layers.quantization.online.fp8 import (
15
Fp8PerBlockOnlineLinearMethod,
@@ -18,6 +19,7 @@ from vllm.model_executor.layers.quantization.online.fp8 import (
19
)
20
from vllm.model_executor.layers.quantization.online.nvfp4 import (
21
Nvfp4OnlineMoEMethod,
22
+ _quantize_moe_weight_to_nvfp4,
23
)
24
from vllm.platforms import current_platform
25
from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
@@ -184,6 +186,41 @@ def test_online_nvfp4_per_token_moe(vllm_runner, monkeypatch) -> None:
186
print(outputs[0][1])
187
188
189
+@pytest.mark.skipif(
190
+ not (
191
+ current_platform.is_cuda() and current_platform.is_device_capability_family(100)
192
+ ),
193
+ reason="NVFP4 weight quantization needs a Blackwell (SM100) GPU.",
194
+)
195
+def test_online_nvfp4_quantizes_original_expert_weights() -> None:
196
+ torch.manual_seed(0)
197
+ weight = torch.randn(2, 32, 32, device="cuda", dtype=torch.bfloat16)
198
+
199
+ quantized, block_scale, global_decode_scale = _quantize_moe_weight_to_nvfp4(weight)
200
+ global_encode_scale = 1.0 / global_decode_scale
201
+ expected = [
202
+ scaled_fp4_quant(
203
+ expert_weight,
204
+ expert_scale,
205
+ is_sf_swizzled_layout=False,
206
+ )
207
+ for expert_weight, expert_scale in zip(
208
+ weight,
209
+ global_encode_scale,
210
+ strict=True,
211
+ )
212
+ ]
213
+
214
+ assert torch.equal(
215
+ quantized,
216
+ torch.stack([expert_weight for expert_weight, _ in expected]),
217
+ )
218
+ assert torch.equal(
219
+ block_scale,
220
+ torch.stack([expert_scale for _, expert_scale in expected]),
221
+ )
222
+
223
+
224
@pytest.mark.skipif(
225
not is_quant_method_supported("fp8"),
226
reason="FP8 is not supported on this GPU type.",
vllm/model_executor/layers/quantization/online/nvfp4.py
+21
-10
@@ -40,23 +40,34 @@ def _quantize_moe_weight_to_nvfp4(
40
global scale ``(E,)`` stored as ``amax / (fp4_max * fp8_max)``.
41
"""
42
assert weight.dim() == 3, f"expected 3D expert weights, got {weight.shape}"
43
- num_experts, n, k = weight.shape
43
+ k = weight.shape[-1]
44
assert k % 16 == 0, f"last dim must be a multiple of 16, got {k}"
45
46
amax = weight.abs().amax(dim=(1, 2)).to(torch.float32).clamp_min(1e-8)
47
global_scale = (FLOAT4_E2M1_MAX * FLOAT8_E4M3_MAX) / amax
48
weight_scale_2 = (1.0 / global_scale).to(torch.float32)
49
50
- # scaled_fp4_quant(w, g) == scaled_fp4_quant(w * g, 1), so fold each
51
- # expert's scale in and quantize all experts in one call (fp32 to keep the
52
- # large scale precise), rather than looping per expert.
53
- scaled = (weight.float() * global_scale[:, None, None]).to(weight.dtype)
54
- scaled = scaled.reshape(-1, k)
55
- one = torch.ones((), device=weight.device, dtype=torch.float32)
56
- qweight, block_scale = scaled_fp4_quant(scaled, one, is_sf_swizzled_layout=False)
50
+ # Keep the original BF16/FP16 values as the quantizer input. Folding each
51
+ # expert's FP32 global scale into the weight would add a BF16/FP16 rounding
52
+ # before the group-16 scale and E2M1 values are selected.
53
+ weight = weight.contiguous()
54
+ quantized_experts = [
55
+ scaled_fp4_quant(
56
+ expert_weight,
57
+ expert_scale,
58
+ is_sf_swizzled_layout=False,
59
+ )
60
+ for expert_weight, expert_scale in zip(
61
+ weight,
62
+ global_scale,
63
+ strict=True,
64
+ )
65
+ ]
66
+ qweight = torch.stack([quantized for quantized, _ in quantized_experts])
67
+ block_scale = torch.stack([block_scale for _, block_scale in quantized_experts])
68
return (
58
- qweight.reshape(num_experts, n, k // 2),
59
- block_scale.reshape(num_experts, n, k // 16),
69
+ qweight,
70
+ block_scale,
71
weight_scale_2,
72
)
73