From 70bd10930bbb91af38fc243c53df4b94da80a8c9 Mon Sep 17 00:00:00 2001 From: Netanel Haber <58652339+netanel-haber@users.noreply.github.com> Date: Thu, 30 Jul 2026 23:19:56 +0300 Subject: [PATCH] [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 --- tests/quantization/test_modelopt.py | 27 +++++++++++++--- .../model_executor/kernels/linear/__init__.py | 5 +++ .../layers/quantization/modelopt.py | 31 +++++++++---------- 3 files changed, 42 insertions(+), 21 deletions(-) diff --git a/tests/quantization/test_modelopt.py b/tests/quantization/test_modelopt.py index c8f8a429e3..cac49a5a80 100644 --- a/tests/quantization/test_modelopt.py +++ b/tests/quantization/test_modelopt.py @@ -13,7 +13,12 @@ import pytest import torch from tests.quantization.utils import is_quant_method_supported +from vllm.config import VllmConfig, set_current_vllm_config from vllm.config.model import ModelConfig +from vllm.model_executor.kernels.linear import ( + HummingNvFp4LinearKernel, + MarlinNvFp4LinearKernel, +) from vllm.model_executor.layers.linear import UnquantizedLinearMethod from vllm.model_executor.layers.quantization.modelopt import ( ModelOptFp8Config, @@ -21,6 +26,7 @@ from vllm.model_executor.layers.quantization.modelopt import ( ModelOptMxFp8Config, ModelOptNvFp4Config, ModelOptNvFp4LinearMethod, + ModelOptNvFp4W4A16LinearMethod, ) from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, @@ -468,15 +474,13 @@ def test_modelopt_nvfp4_config_dispatches_w4a4_method(): def test_modelopt_nvfp4_config_dispatches_w4a16_method(): - """``quant_method="W4A16_NVFP4"`` routes to the new + """``quant_method="W4A16_NVFP4"`` routes to ``ModelOptNvFp4W4A16LinearMethod`` instead of the W4A4 sibling. Mirrors the FP8 dispatch precedent (``ModelOptFp8Config`` selects one of three FP8 LinearMethods on ``quant_method``); a regression here would mean a W4A16 NVFP4 checkpoint silently loaded under the - W4A4 method, which would try to register an ``input_scale`` runtime - parameter and (more importantly) call the cutlass W4A4 NVFP4 GEMM - instead of FP4 Marlin. + W4A4 activation-quantization path. """ from vllm.model_executor.layers.quantization.modelopt import ( ModelOptNvFp4Config, @@ -495,6 +499,21 @@ def test_modelopt_nvfp4_config_dispatches_w4a16_method(): assert config.quant_method == "W4A16_NVFP4" +@pytest.mark.parametrize( + ("linear_backend", "kernel_cls"), + [("auto", MarlinNvFp4LinearKernel), ("humming", HummingNvFp4LinearKernel)], +) +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") +def test_modelopt_w4a16_respects_linear_backend(linear_backend, kernel_cls): + vllm_config = VllmConfig() + vllm_config.kernel_config.linear_backend = linear_backend + with set_current_vllm_config(vllm_config): + method = ModelOptNvFp4W4A16LinearMethod( + ModelOptNvFp4Config(quant_method="W4A16_NVFP4") + ) + assert isinstance(method.kernel, kernel_cls) + + @pytest.mark.parametrize( "quant_method, expected_use_a16, act_key_is_none", [ diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 39ed4309ca..2892f1aa71 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -891,6 +891,7 @@ def init_nvfp4_linear_kernel(use_a16: bool = False) -> NvFp4LinearKernel: """Select and instantiate the best NVFP4 linear kernel for the current platform.""" config = NvFp4LinearLayerConfig() + a16_kernels = (MarlinNvFp4LinearKernel, HummingNvFp4LinearKernel) # VLLM_BATCH_INVARIANT forces deterministic execution. Prefer the # batch-invariant CUTLASS implementation when available, otherwise fall @@ -931,6 +932,8 @@ def init_nvfp4_linear_kernel(use_a16: bool = False) -> NvFp4LinearKernel: force_kernel = MarlinNvFp4LinearKernel if force_kernel is not None: + if use_a16 and force_kernel not in a16_kernels: + raise ValueError(f"{force_kernel.__name__} does not support W4A16") is_supported, reason = force_kernel.is_supported() if not is_supported: raise ValueError( @@ -943,6 +946,8 @@ def init_nvfp4_linear_kernel(use_a16: bool = False) -> NvFp4LinearKernel: # Auto-select from registry (or --linear-backend filtered). platform = current_platform._enum possible = list(_POSSIBLE_NVFP4_KERNELS.get(platform, [])) + if use_a16: + possible = [kernel for kernel in possible if kernel in a16_kernels] # Apply --linear-backend filtering when set. if linear_backend != "auto": diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index e11097b7de..48ebf0fef5 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -12,8 +12,6 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm.config import get_current_vllm_config from vllm.logger import init_logger from vllm.model_executor.kernels.linear import ( - MarlinNvFp4LinearKernel, - NvFp4LinearLayerConfig, init_fp8_linear_kernel, init_mxfp8_linear_kernel, init_nvfp4_linear_kernel, @@ -1247,16 +1245,16 @@ class ModelOptNvFp4W4A16LinearMethod(LinearMethodBase): """Linear method for ModelOpt NVFP4 W4A16. 4-bit NVFP4 weights, fp16/bf16 activations. Loads ModelOpt-style names - directly (no on-disk conversion) and dispatches to the FP4 Marlin GEMM: + directly (no on-disk conversion) and dispatches to a W4A16 GEMM: weight uint8 packed NVFP4 (2 nibbles/byte along input dim) weight_scale fp8-e4m3 per 16-elem group along input dim weight_scale_2 fp32 per-tensor global scale = amax / (6.0 * 448.0) - No activation quantization. Marlin expects the global scale in the same - form ModelOpt stores (amax/2688), so we rename weight_scale_2 -> - weight_global_scale **without reciprocation** -- the CT W4A16 path - reciprocates only because CT stores the inverse on disk. + No activation quantization. ModelOpt stores the global scale as + amax/2688, so we rename weight_scale_2 -> weight_global_scale without + reciprocation. The selected kernel converts it to its runtime format. + The CT W4A16 path reciprocates because CT stores the inverse on disk. We also register a placeholder input_scale parameter so that W4A4-shaped checkpoints (which contain *_proj.input_scale tensors) can be loaded @@ -1267,17 +1265,12 @@ class ModelOptNvFp4W4A16LinearMethod(LinearMethodBase): def __init__(self, quant_config: ModelOptNvFp4Config) -> None: self.quant_config = quant_config - # Vestigial slot mirrored from ModelOptNvFp4LinearMethod: the parent - # config's get_quant_method only fills marlin_input_dtype when - # backend == "marlin"; we don't set that since we pin the kernel - # below, but we keep the attribute for shape parity. self.marlin_input_dtype = None - # Direct-instantiate the Marlin NVFP4 adapter rather than going through - # init_nvfp4_linear_kernel(): the latter's priority list returns a - # cutlass W4A4 kernel as first-pick on this hardware, which would - # silently try to quantize activations (we have no input_scale). For - # W4A16 there is exactly one valid kernel, so we pin it. - self.kernel = MarlinNvFp4LinearKernel(NvFp4LinearLayerConfig()) + # `init_nvfp4_linear_kernel(use_a16=True)` is best of both worlds: + # 1. `use_a16=True` forces `Marlin`: https://github.com/vllm-project/vllm/commit/e68988a#diff-7135ab92aa94dfacb1ad3c77fc13f9c4ffe0b977f8eac5d86c2afe243e5f92a6R842-R889 + # for `--linear-backend=auto`, avoiding a W4A4 kernel that requires input_scale. + # 2. Specifying e.g. `--linear-backend=humming` will override. + self.kernel = init_nvfp4_linear_kernel(use_a16=True) def create_weights( self, @@ -1300,6 +1293,7 @@ class ModelOptNvFp4W4A16LinearMethod(LinearMethodBase): layer.logical_widths = output_partition_sizes layer.input_size_per_partition = input_size_per_partition layer.output_size_per_partition = output_size_per_partition + layer.output_partition_sizes = output_partition_sizes if input_size_per_partition % 16 != 0: raise ValueError( @@ -1355,6 +1349,9 @@ class ModelOptNvFp4W4A16LinearMethod(LinearMethodBase): layer.register_parameter("input_scale", input_scale) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + if not hasattr(layer, "has_bias"): + layer.has_bias = getattr(layer, "bias", None) is not None + # Discard the input_scale placeholder. Whether it carries values # (W4A4 ckpt loaded as W4A16) or is uninitialized (native W4A16 # ckpt), W4A16 mode does not quantize activations, so this is unused.