iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: Netanel Haber <58652339+netanel-haber@users.noreply.github.com>
Co-authored-by: OpenAI Codex <noreply@openai.com>
This commit is contained in:
Netanel Haber
2026-07-30 13:19:56 -07:00
committed by GitHub
co-authored by OpenAI Codex
parent c27b080d26
commit 70bd10930b
3 changed files with 42 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/quantization/test_modelopt.py
+23 -4
View File
@@ -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",
[
vllm/model_executor/kernels/linear/__init__.py
+5
View File
@@ -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":
vllm/model_executor/layers/quantization/modelopt.py
+14 -17
View File
@@ -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.