@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 70bd10930b

[Quantization] Honor `--linear-backend` for ModelOpt W4A16 (#50273)

Signed-off-by: Netanel Haber <58652339+netanel-haber@users.noreply.github.com> Co-authored-by: OpenAI Codex <noreply@openai.com>

Netanel Haber committed Jul 30, 2026 at 23:19 UTC 70bd10930bbb91af38fc243c53df4b94da80a8c9

3 files changed +42 -21

tests/quantization/test_modelopt.py

+23 -4

@@ -13,7 +13,12 @@ import pytest

13 import torch

15 from tests.quantization.utils import is_quant_method_supported

16 +from vllm.config import VllmConfig, set_current_vllm_config

17 from vllm.config.model import ModelConfig

18 +from vllm.model_executor.kernels.linear import (

19 + HummingNvFp4LinearKernel,

20 + MarlinNvFp4LinearKernel,

21 +)

22 from vllm.model_executor.layers.linear import UnquantizedLinearMethod

23 from vllm.model_executor.layers.quantization.modelopt import (

24 ModelOptFp8Config,

@@ -21,6 +26,7 @@ from vllm.model_executor.layers.quantization.modelopt import (

26 ModelOptMxFp8Config,

27 ModelOptNvFp4Config,

28 ModelOptNvFp4LinearMethod,

29 + ModelOptNvFp4W4A16LinearMethod,

30 )

31 from vllm.model_executor.layers.vocab_parallel_embedding import (

32 ParallelLMHead,

@@ -468,15 +474,13 @@ def test_modelopt_nvfp4_config_dispatches_w4a4_method():

474

475

476 def test_modelopt_nvfp4_config_dispatches_w4a16_method():

471 - """``quant_method="W4A16_NVFP4"`` routes to the new

477 + """``quant_method="W4A16_NVFP4"`` routes to

478 ``ModelOptNvFp4W4A16LinearMethod`` instead of the W4A4 sibling.

479

480 Mirrors the FP8 dispatch precedent (``ModelOptFp8Config`` selects

481 one of three FP8 LinearMethods on ``quant_method``); a regression

482 here would mean a W4A16 NVFP4 checkpoint silently loaded under the

477 - W4A4 method, which would try to register an ``input_scale`` runtime

478 - parameter and (more importantly) call the cutlass W4A4 NVFP4 GEMM

479 - instead of FP4 Marlin.

483 + W4A4 activation-quantization path.

484 """

485 from vllm.model_executor.layers.quantization.modelopt import (

486 ModelOptNvFp4Config,

@@ -495,6 +499,21 @@ def test_modelopt_nvfp4_config_dispatches_w4a16_method():

499 assert config.quant_method == "W4A16_NVFP4"

500

501

502 +@pytest.mark.parametrize(

503 + ("linear_backend", "kernel_cls"),

504 + [("auto", MarlinNvFp4LinearKernel), ("humming", HummingNvFp4LinearKernel)],

505 +)

506 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only")

507 +def test_modelopt_w4a16_respects_linear_backend(linear_backend, kernel_cls):

508 + vllm_config = VllmConfig()

509 + vllm_config.kernel_config.linear_backend = linear_backend

510 + with set_current_vllm_config(vllm_config):

511 + method = ModelOptNvFp4W4A16LinearMethod(

512 + ModelOptNvFp4Config(quant_method="W4A16_NVFP4")

513 + )

514 + assert isinstance(method.kernel, kernel_cls)

515 +

516 +

517 @pytest.mark.parametrize(

518 "quant_method, expected_use_a16, act_key_is_none",

519 [

vllm/model_executor/kernels/linear/__init__.py

@@ -891,6 +891,7 @@ def init_nvfp4_linear_kernel(use_a16: bool = False) -> NvFp4LinearKernel:

891 """Select and instantiate the best NVFP4 linear kernel for the

892 current platform."""

893 config = NvFp4LinearLayerConfig()

894 + a16_kernels = (MarlinNvFp4LinearKernel, HummingNvFp4LinearKernel)

895

896 # VLLM_BATCH_INVARIANT forces deterministic execution. Prefer the

897 # batch-invariant CUTLASS implementation when available, otherwise fall

@@ -931,6 +932,8 @@ def init_nvfp4_linear_kernel(use_a16: bool = False) -> NvFp4LinearKernel:

932 force_kernel = MarlinNvFp4LinearKernel

933

934 if force_kernel is not None:

935 + if use_a16 and force_kernel not in a16_kernels:

936 + raise ValueError(f"{force_kernel.__name__} does not support W4A16")

937 is_supported, reason = force_kernel.is_supported()

938 if not is_supported:

939 raise ValueError(

@@ -943,6 +946,8 @@ def init_nvfp4_linear_kernel(use_a16: bool = False) -> NvFp4LinearKernel:

946 # Auto-select from registry (or --linear-backend filtered).

947 platform = current_platform._enum

948 possible = list(_POSSIBLE_NVFP4_KERNELS.get(platform, []))

949 + if use_a16:

950 + possible = [kernel for kernel in possible if kernel in a16_kernels]

951

952 # Apply --linear-backend filtering when set.

953 if linear_backend != "auto":

vllm/model_executor/layers/quantization/modelopt.py

+14 -17

@@ -12,8 +12,6 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk

12 from vllm.config import get_current_vllm_config

13 from vllm.logger import init_logger

14 from vllm.model_executor.kernels.linear import (

15 - MarlinNvFp4LinearKernel,

16 - NvFp4LinearLayerConfig,

15 init_fp8_linear_kernel,

16 init_mxfp8_linear_kernel,

17 init_nvfp4_linear_kernel,

@@ -1247,16 +1245,16 @@ class ModelOptNvFp4W4A16LinearMethod(LinearMethodBase):

1245 """Linear method for ModelOpt NVFP4 W4A16.

1246

1247 4-bit NVFP4 weights, fp16/bf16 activations. Loads ModelOpt-style names

1250 - directly (no on-disk conversion) and dispatches to the FP4 Marlin GEMM:

1248 + directly (no on-disk conversion) and dispatches to a W4A16 GEMM:

1249

1250 weight uint8 packed NVFP4 (2 nibbles/byte along input dim)

1251 weight_scale fp8-e4m3 per 16-elem group along input dim

1252 weight_scale_2 fp32 per-tensor global scale = amax / (6.0 * 448.0)

1253

1256 - No activation quantization. Marlin expects the global scale in the same

1257 - form ModelOpt stores (amax/2688), so we rename weight_scale_2 ->

1258 - weight_global_scale **without reciprocation** -- the CT W4A16 path

1259 - reciprocates only because CT stores the inverse on disk.

1254 + No activation quantization. ModelOpt stores the global scale as

1255 + amax/2688, so we rename weight_scale_2 -> weight_global_scale without

1256 + reciprocation. The selected kernel converts it to its runtime format.

1257 + The CT W4A16 path reciprocates because CT stores the inverse on disk.

1258

1259 We also register a placeholder input_scale parameter so that W4A4-shaped

1260 checkpoints (which contain *_proj.input_scale tensors) can be loaded

@@ -1267,17 +1265,12 @@ class ModelOptNvFp4W4A16LinearMethod(LinearMethodBase):

1265

1266 def __init__(self, quant_config: ModelOptNvFp4Config) -> None:

1267 self.quant_config = quant_config

1270 - # Vestigial slot mirrored from ModelOptNvFp4LinearMethod: the parent

1271 - # config's get_quant_method only fills marlin_input_dtype when

1272 - # backend == "marlin"; we don't set that since we pin the kernel

1273 - # below, but we keep the attribute for shape parity.

1268 self.marlin_input_dtype = None

1275 - # Direct-instantiate the Marlin NVFP4 adapter rather than going through

1276 - # init_nvfp4_linear_kernel(): the latter's priority list returns a

1277 - # cutlass W4A4 kernel as first-pick on this hardware, which would

1278 - # silently try to quantize activations (we have no input_scale). For

1279 - # W4A16 there is exactly one valid kernel, so we pin it.

1280 - self.kernel = MarlinNvFp4LinearKernel(NvFp4LinearLayerConfig())

1269 + # `init_nvfp4_linear_kernel(use_a16=True)` is best of both worlds:

1270 + # 1. `use_a16=True` forces `Marlin`: https://github.com/vllm-project/vllm/commit/e68988a#diff-7135ab92aa94dfacb1ad3c77fc13f9c4ffe0b977f8eac5d86c2afe243e5f92a6R842-R889

1271 + # for `--linear-backend=auto`, avoiding a W4A4 kernel that requires input_scale.

1272 + # 2. Specifying e.g. `--linear-backend=humming` will override.

1273 + self.kernel = init_nvfp4_linear_kernel(use_a16=True)

1274

1275 def create_weights(

1276 self,

@@ -1300,6 +1293,7 @@ class ModelOptNvFp4W4A16LinearMethod(LinearMethodBase):

1293 layer.logical_widths = output_partition_sizes

1294 layer.input_size_per_partition = input_size_per_partition

1295 layer.output_size_per_partition = output_size_per_partition

1296 + layer.output_partition_sizes = output_partition_sizes

1297

1298 if input_size_per_partition % 16 != 0:

1299 raise ValueError(

@@ -1355,6 +1349,9 @@ class ModelOptNvFp4W4A16LinearMethod(LinearMethodBase):

1349 layer.register_parameter("input_scale", input_scale)

1350

1351 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:

1352 + if not hasattr(layer, "has_bias"):

1353 + layer.has_bias = getattr(layer, "bias", None) is not None

1354 +

1355 # Discard the input_scale placeholder. Whether it carries values

1356 # (W4A4 ckpt loaded as W4A16) or is uninitialized (native W4A16

1357 # ckpt), W4A16 mode does not quantize activations, so this is unused.