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 3, 00:30 UTC
[Quantization][ModelOpt] W4A16 NVFP4 fused MoE + mixed-precision dispatch (#42566)
Signed-off-by: Juhi Mittal <juhim@nvidia.com>
Juhi Mittal committed
May 22, 2026 at 13:51 UTC
e203006a8b05b6380f7deb261c28f384711674c1
2 files changed
+177
-5
tests/quantization/test_modelopt.py
+140
-1
@@ -6,7 +6,8 @@ Run `pytest tests/quantization/test_modelopt.py`.
6
"""
7
8
import os
9
-from typing import NoReturn
9
+from typing import Any, NoReturn
10
+from unittest.mock import MagicMock, patch
11
12
import pytest
13
import torch
@@ -285,3 +286,141 @@ def test_modelopt_nvfp4_config_dispatches_w4a16_method():
286
assert config.LinearMethodCls is ModelOptNvFp4W4A16LinearMethod
287
assert config.LinearMethodCls is not ModelOptNvFp4LinearMethod
288
assert config.quant_method == "W4A16_NVFP4"
289
+
290
+
291
+@pytest.mark.parametrize(
292
+ "quant_method, expected_use_a16, act_key_is_none",
293
+ [
294
+ ("NVFP4", False, False), # W4A4 default
295
+ ("W4A16_NVFP4", True, True), # native W4A16 ckpt
296
+ ],
297
+)
298
+def test_modelopt_nvfp4_moe_dispatches_to_marlin_when_w4a16(
299
+ quant_method, expected_use_a16, act_key_is_none
300
+):
301
+ """``ModelOptNvFp4FusedMoE``: when the ckpt's ``quant_method`` is
302
+ ``W4A16_NVFP4``, the MoE class must pass ``activation_key=None`` to
303
+ ``select_nvfp4_moe_backend``. That filters out every W4A4 backend
304
+ (their ``_supports_quant_scheme`` requires
305
+ ``(kNvfp4Static, kNvfp4Dynamic)`` exactly); Marlin survives because
306
+ it only checks ``weight_key``. A regression here would mean a W4A16
307
+ ckpt silently went to the cutlass W4A4 path.
308
+ """
309
+ from vllm.model_executor.layers.quantization.modelopt import (
310
+ ModelOptNvFp4Config,
311
+ ModelOptNvFp4FusedMoE,
312
+ )
313
+ from vllm.model_executor.layers.quantization.utils.quant_utils import (
314
+ kNvfp4Dynamic,
315
+ kNvfp4Static,
316
+ )
317
+
318
+ config = ModelOptNvFp4Config(
319
+ quant_method=quant_method,
320
+ is_checkpoint_nvfp4_serialized=True,
321
+ kv_cache_quant_algo=None,
322
+ exclude_modules=[],
323
+ group_size=16,
324
+ )
325
+
326
+ mock_select = MagicMock(return_value=(MagicMock(), MagicMock()))
327
+ with (
328
+ patch(
329
+ "vllm.model_executor.layers.quantization.modelopt.select_nvfp4_moe_backend",
330
+ mock_select,
331
+ ),
332
+ patch(
333
+ "vllm.model_executor.layers.quantization.modelopt."
334
+ "is_global_sf_supported_for_nvfp4_backend",
335
+ return_value=False,
336
+ ),
337
+ ):
338
+ moe = ModelOptNvFp4FusedMoE(config, MagicMock())
339
+
340
+ assert moe.use_a16 is expected_use_a16
341
+ _, kwargs = mock_select.call_args
342
+ assert kwargs["weight_key"] is kNvfp4Static
343
+ if act_key_is_none:
344
+ assert kwargs["activation_key"] is None
345
+ else:
346
+ assert kwargs["activation_key"] is kNvfp4Dynamic
347
+
348
+
349
+@pytest.mark.parametrize(
350
+ "per_layer_algo, expected_linear_cls_name",
351
+ [
352
+ ("NVFP4", "ModelOptNvFp4LinearMethod"),
353
+ ("W4A16_NVFP4", "ModelOptNvFp4W4A16LinearMethod"),
354
+ ],
355
+)
356
+def test_modelopt_mixed_precision_dispatches_w4a16_layer(
357
+ per_layer_algo, expected_linear_cls_name
358
+):
359
+ """``ModelOptMixedPrecisionConfig.get_quant_method`` must route a Linear
360
+ layer to the right LinearMethod based on its per-layer ``quant_algo``
361
+ entry in ``quantized_layers``. Verifies the new ``W4A16_NVFP4`` branch
362
+ coexists with the existing ``NVFP4`` branch without regression. A
363
+ regression here would mean a W4A16 layer in a mixed-precision ckpt
364
+ silently fell through to ``UnquantizedLinearMethod``.
365
+
366
+ NOTE: FP8 dispatch (the third branch of get_quant_method) is not
367
+ covered here because ``ModelOptFp8LinearMethod.__init__`` reads
368
+ ``get_current_vllm_config().model_config.dtype``, which requires a
369
+ fully constructed ``ModelConfig`` (real model path). FP8 routing in
370
+ mixed-precision is exercised by the existing integration tests
371
+ above that use the ``vllm_runner`` fixture (e.g.
372
+ ``test_modelopt_fp8_checkpoint_setup``). Our PR doesn't change the
373
+ FP8 branch, so this isn't a coverage gap.
374
+ """
375
+ from vllm.model_executor.layers.linear import LinearBase
376
+ from vllm.model_executor.layers.quantization import modelopt as m
377
+
378
+ hf_quant_config: dict[str, Any] = {
379
+ "quantization": {
380
+ "quant_algo": "MIXED_PRECISION",
381
+ "kv_cache_quant_algo": None,
382
+ "exclude_modules": [],
383
+ "group_size": 16,
384
+ "quantized_layers": {
385
+ "model.layers.0.fake_proj": {"quant_algo": per_layer_algo},
386
+ },
387
+ }
388
+ }
389
+ config = m.ModelOptMixedPrecisionConfig.from_config(hf_quant_config)
390
+
391
+ fake_layer = MagicMock(spec=LinearBase)
392
+ method = config.get_quant_method(fake_layer, "model.layers.0.fake_proj")
393
+
394
+ expected_cls = getattr(m, expected_linear_cls_name)
395
+ assert isinstance(method, expected_cls), (
396
+ f"Expected {expected_linear_cls_name}, got {type(method).__name__}"
397
+ )
398
+
399
+
400
+def test_modelopt_mixed_precision_builds_w4a16_sibling_config():
401
+ """Sanity: ``ModelOptMixedPrecisionConfig._from_config`` builds **two**
402
+ NVFP4 sub-configs — one for W4A4 (default) and one tagged
403
+ ``quant_method='W4A16_NVFP4'`` — so per-layer dispatch can hand
404
+ Marlin-bound layers the right config without re-instantiating it on
405
+ every call.
406
+ """
407
+ from vllm.model_executor.layers.quantization import modelopt as m
408
+
409
+ hf_quant_config: dict[str, Any] = {
410
+ "quantization": {
411
+ "quant_algo": "MIXED_PRECISION",
412
+ "kv_cache_quant_algo": None,
413
+ "exclude_modules": [],
414
+ "group_size": 16,
415
+ "quantized_layers": {
416
+ "model.layers.0.a": {"quant_algo": "NVFP4"},
417
+ "model.layers.0.b": {"quant_algo": "W4A16_NVFP4"},
418
+ },
419
+ }
420
+ }
421
+ config = m.ModelOptMixedPrecisionConfig.from_config(hf_quant_config)
422
+
423
+ assert config.nvfp4_config.quant_method == "NVFP4"
424
+ assert config.nvfp4_config.LinearMethodCls is m.ModelOptNvFp4LinearMethod
425
+ assert config.w4a16_nvfp4_config.quant_method == "W4A16_NVFP4"
426
+ assert config.w4a16_nvfp4_config.LinearMethodCls is m.ModelOptNvFp4W4A16LinearMethod
vllm/model_executor/layers/quantization/modelopt.py
+37
-4
@@ -1392,11 +1392,17 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1392
) -> None:
1393
super().__init__(moe_config)
1394
self.quant_config = quant_config
1395
- # Select experts implementation.
1395
+ # W4A16 mode fires for W4A16_NVFP4 on-disk checkpoints. With
1396
+ # activation_key=None every W4A4 backend's _supports_quant_scheme
1397
+ # rejects itself (they all require (kNvfp4Static, kNvfp4Dynamic)
1398
+ # exactly); only Marlin survives. Marlin's MoE path drops
1399
+ # activation scales in convert_to_nvfp4_moe_kernel_format, so no
1400
+ # other change is needed.
1401
+ self.use_a16 = quant_config.quant_method == "W4A16_NVFP4"
1402
self.nvfp4_backend, self.experts_cls = select_nvfp4_moe_backend(
1403
config=self.moe,
1404
weight_key=kNvfp4Static,
1399
- activation_key=kNvfp4Dynamic,
1405
+ activation_key=None if self.use_a16 else kNvfp4Dynamic,
1406
)
1407
1408
self.use_global_sf = is_global_sf_supported_for_nvfp4_backend(
@@ -2211,12 +2217,14 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
2217
quantized_layers: dict[str, dict[str, Any]],
2218
fp8_config: ModelOptFp8Config,
2219
nvfp4_config: ModelOptNvFp4Config,
2220
+ w4a16_nvfp4_config: ModelOptNvFp4Config,
2221
) -> None:
2222
super().__init__(exclude_modules)
2223
self.kv_cache_quant_method = kv_cache_quant_method
2224
self.quantized_layers = quantized_layers
2225
self.fp8_config = fp8_config
2226
self.nvfp4_config = nvfp4_config
2227
+ self.w4a16_nvfp4_config = w4a16_nvfp4_config
2228
2229
def get_name(self) -> QuantizationMethods:
2230
return "modelopt_mixed"
@@ -2261,10 +2269,15 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
2269
"'quantized_layers' mapping in the quantization config."
2270
)
2271
2264
- # Determine group_size from the first NVFP4 entry if not provided.
2272
+ # Determine group_size from the first NVFP4-family entry if not
2273
+ # provided. Both NVFP4 (W4A4) and W4A16_NVFP4 share the same packing
2274
+ # + group-size convention; either entry resolves the value.
2275
if group_size is None:
2276
for layer_info in quantized_layers.values():
2267
- if layer_info.get("quant_algo", "").upper() == "NVFP4":
2277
+ if layer_info.get("quant_algo", "").upper() in (
2278
+ "NVFP4",
2279
+ "W4A16_NVFP4",
2280
+ ):
2281
group_size = layer_info.get("group_size", 16)
2282
break
2283
if group_size is None:
@@ -2282,6 +2295,18 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
2295
exclude_modules=[],
2296
group_size=group_size,
2297
)
2298
+ # Sibling config for layers that declare quant_algo: "W4A16_NVFP4".
2299
+ # ModelOptNvFp4Config.__init__ keys LinearMethodCls off quant_method,
2300
+ # so this instance auto-selects ModelOptNvFp4W4A16LinearMethod. The
2301
+ # MoE side reads quant_config.quant_method == "W4A16_NVFP4" to set
2302
+ # use_a16 → Marlin backend in ModelOptNvFp4FusedMoE.__init__.
2303
+ w4a16_nvfp4_config = ModelOptNvFp4Config(
2304
+ quant_method="W4A16_NVFP4",
2305
+ is_checkpoint_nvfp4_serialized=True,
2306
+ kv_cache_quant_algo=kv_cache_quant_method,
2307
+ exclude_modules=[],
2308
+ group_size=group_size,
2309
+ )
2310
2311
return cls(
2312
kv_cache_quant_method=kv_cache_quant_method,
@@ -2289,6 +2314,7 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
2314
quantized_layers=quantized_layers,
2315
fp8_config=fp8_config,
2316
nvfp4_config=nvfp4_config,
2317
+ w4a16_nvfp4_config=w4a16_nvfp4_config,
2318
)
2319
2320
def _resolve_quant_algo(self, prefix: str) -> str | None:
@@ -2355,6 +2381,8 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
2381
return ModelOptFp8LinearMethod(self.fp8_config)
2382
if quant_algo == "NVFP4":
2383
return ModelOptNvFp4LinearMethod(self.nvfp4_config)
2384
+ if quant_algo == "W4A16_NVFP4":
2385
+ return ModelOptNvFp4W4A16LinearMethod(self.w4a16_nvfp4_config)
2386
# Layer not in quantized_layers — leave unquantized
2387
return UnquantizedLinearMethod()
2388
@@ -2369,6 +2397,11 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
2397
quant_config=self.nvfp4_config,
2398
moe_config=layer.moe_config,
2399
)
2400
+ if quant_algo == "W4A16_NVFP4":
2401
+ return ModelOptNvFp4FusedMoE(
2402
+ quant_config=self.w4a16_nvfp4_config,
2403
+ moe_config=layer.moe_config,
2404
+ )
2405
return None
2406
2407
return None