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

[Quantization] Select linear backends per quantization (#51204)

Browse Source 
Signed-off-by: Netanel Haber <58652339+netanel-haber@users.noreply.github.com>
This commit is contained in:
Netanel Haber
2026-09-14 16:24:18 -04:00
committed by GitHub
parent 995e8581f4
commit d4ee7fe7a9
6 changed files with 154 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/features/quantization/README.md
+17
View File
@@ -23,6 +23,23 @@ The following are the supported quantization formats for vLLM:
- [TorchAO](torchao.md)
- [FP8 ViT Encoder Attention](fp8_vit_attn.md)
## Selecting Linear Backends per Quantization
`--linear-backend` selects one backend for all quantized linear layers. For
mixed-precision models that use more than one linear quantization scheme, use
`linear_backend_per_quant` to override the backend for individual schemes:
```bash
vllm serve <model> \
--linear-backend cutlass \
--kernel-config '{"linear_backend_per_quant":{"nvfp4_w4a16":"humming"}}'
```
Here, NVFP4 W4A16 linear layers use Humming, while all other quantized linear
layers use CUTLASS. Per-quantization overrides take precedence over
`--linear-backend`; schemes without an override continue to use the global
setting, including automatic selection when it is `auto`.
## Supported Hardware
The table below shows the compatibility of various quantization implementations with different hardware platforms in vLLM:
tests/kernels/quantization/test_scaled_mm_kernel_selection.py
+70
View File
@@ -12,20 +12,90 @@ from unittest.mock import patch
import pytest
import torch
from vllm.config import KernelConfig, VllmConfig, set_current_vllm_config
from vllm.model_executor.kernels.linear import (
AiterInt8ScaledMMLinearKernel,
CPUInt8ScaledMMLinearKernel,
HummingFP8ScaledMMLinearKernel,
Int8ScaledMMLinearKernel,
Int8ScaledMMLinearLayerConfig,
ScaledMMLinearKernel,
_get_linear_backend,
_resolve_backend_kernels,
init_fp8_linear_kernel,
init_int8_linear_kernel,
register_linear_kernel,
)
from vllm.model_executor.layers.quantization.utils.quant_utils import (
kFp8StaticTensorSym,
)
from vllm.platforms import PlatformEnum
pytestmark = pytest.mark.cpu_test
def test_auto_linear_backend_is_unchanged_without_overrides():
kernels = [object]
quantizations = (
"fp8_block_w8a8",
"fp8_w8a8",
"int8_w8a8",
"mixed_precision",
"mxfp8",
"mxfp4",
"mxfp6",
"w8a16_fp8",
"nvfp4_w4a4",
"nvfp4_w4a16",
)
config = VllmConfig(kernel_config=KernelConfig(linear_backend="auto"))
with set_current_vllm_config(config):
for quantization in quantizations:
assert _get_linear_backend(quantization=quantization) == "auto"
assert (
_resolve_backend_kernels(kernels, "test", quantization=quantization)
is kernels
)
def test_linear_backend_override_is_quantization_specific():
config = VllmConfig(
kernel_config=KernelConfig(
linear_backend="cutlass",
linear_backend_per_quant={"nvfp4_w4a16": "humming"},
)
)
with set_current_vllm_config(config):
assert _get_linear_backend(quantization="nvfp4_w4a16") == "humming"
assert _get_linear_backend(quantization="nvfp4_w4a4") == "cutlass"
assert _get_linear_backend(quantization="fp8_w8a8") == "cutlass"
@patch.object(HummingFP8ScaledMMLinearKernel, "is_supported", return_value=(True, None))
@patch("vllm.model_executor.kernels.linear.current_platform")
def test_fp8_linear_backend_override(platform_mock, _):
platform_mock._enum = PlatformEnum.CUDA
config = VllmConfig(
kernel_config=KernelConfig(
linear_backend="cutlass",
linear_backend_per_quant={"fp8_w8a8": "humming"},
)
)
with set_current_vllm_config(config):
kernel = init_fp8_linear_kernel(
activation_quant_key=kFp8StaticTensorSym,
weight_quant_key=kFp8StaticTensorSym,
input_dtype=torch.bfloat16,
out_dtype=torch.bfloat16,
weight_shape=(128, 128),
)
assert type(kernel) is HummingFP8ScaledMMLinearKernel
def test_is_supported_is_abstract():
"""Test that is_supported() is properly defined as abstract."""
assert issubclass(ScaledMMLinearKernel, ABC)
tests/model_executor/kernels/test_b12x_linear.py
+6 -2
View File
@@ -115,7 +115,9 @@ def test_b12x_backend_registration_priority_and_selection(
assert names.index(before) < names.index(kernel_cls.__name__) < names.index(after)
monkeypatch.setattr(linear_mod.current_platform, "_enum", PlatformEnum.CUDA)
monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "b12x")
monkeypatch.setattr(
linear_mod, "_get_linear_backend", lambda *, quantization: "b12x"
)
monkeypatch.setattr(
kernel_cls,
"is_supported",
@@ -784,7 +786,9 @@ def test_b12x_backend_preserves_w4a16_fallback(monkeypatch) -> None:
import vllm.model_executor.kernels.linear as linear_mod
monkeypatch.setattr(linear_mod.current_platform, "_enum", PlatformEnum.CUDA)
monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "b12x")
monkeypatch.setattr(
linear_mod, "_get_linear_backend", lambda *, quantization: "b12x"
)
monkeypatch.setitem(
linear_mod._POSSIBLE_NVFP4_KERNELS,
PlatformEnum.CUDA,
vllm/config/kernel.py
+9
View File
@@ -326,6 +326,13 @@ class KernelConfig:
- "xpu_woq": Use XPU kernels for weight-only quantization (e.g. W8A16)
"""
linear_backend_per_quant: dict[str, LinearBackend] | None = Field(
default=None, min_length=1
)
"""Backend overrides keyed by linear quantization scheme. Overrides take
precedence over ``linear_backend``; for example,
``{"nvfp4_w4a16": "humming"}``."""
@field_validator("moe_backend", mode="before")
@classmethod
def _normalize_moe_backend(cls, value: Any) -> Any:
@@ -359,6 +366,8 @@ class KernelConfig:
"enable_flashinfer_autotune",
"ir_op_priority", # handled separately below
}
if self.linear_backend_per_quant is None:
ignored_factors.add("linear_backend_per_quant")
factors = get_hash_factors(self, ignored_factors)
factors["ir_op_priority"] = self.ir_op_priority.compute_hash()
return hash_factors(factors)
vllm/model_executor/kernels/linear/__init__.py
+51 -14
View File
@@ -235,14 +235,19 @@ from vllm.platforms import PlatformEnum, current_platform
logger = init_logger(__name__)
def _get_linear_backend() -> str:
def _get_linear_backend(*, quantization: str) -> str:
"""Get the linear_backend setting from the current vllm config."""
from vllm.config import get_current_vllm_config_or_none
config = get_current_vllm_config_or_none()
if config is not None:
return config.kernel_config.linear_backend
return "auto"
if (config := get_current_vllm_config_or_none()) is None:
return "auto"
overrides = config.kernel_config.linear_backend_per_quant
if overrides is not None and (override := overrides.get(quantization)):
logger.info_once(
"Applied linear backend override for %r: %r", quantization, override
)
return override
return config.kernel_config.linear_backend
# Kernel classes covered by each --linear-backend value.
@@ -370,6 +375,8 @@ def _filter_kernels_by_backend(
def _resolve_backend_kernels(
kernels: list[type],
layer_desc: str,
*,
quantization: str,
) -> list[type]:
"""Apply --linear-backend filtering to one layer type's kernel list.
@@ -380,7 +387,7 @@ def _resolve_backend_kernels(
layer types (e.g. NVFP4 MoE projections next to FP8 attention
projections).
"""
linear_backend = _get_linear_backend()
linear_backend = _get_linear_backend(quantization=quantization)
if linear_backend == "auto":
return kernels
filtered = _filter_kernels_by_backend(linear_backend, kernels)
@@ -625,6 +632,8 @@ def choose_scaled_mm_linear_kernel(
possible_kernels: dict[PlatformEnum, list[type[_KernelT]]],
compute_capability: int | None = None,
force_kernel: type[_KernelT] | None = None,
*,
quantization: str,
) -> type[_KernelT]:
"""
Choose a _KernelT that can implement the given config for the
@@ -642,6 +651,7 @@ def choose_scaled_mm_linear_kernel(
force_kernel (Optional[type[_KernelT]]): An Optional forced kernel to override
the possible_kernels if it can be implemented. If None, it will only try the
possible kernels.
quantization: Quantization scheme used to select a backend override.
Raises:
ValueError: If no kernel can implement the given config.
@@ -668,7 +678,11 @@ def choose_scaled_mm_linear_kernel(
platform_kernels = possible_kernels.get(current_platform._enum, [])
# Apply --linear-backend filtering when set.
platform_kernels = _resolve_backend_kernels(platform_kernels, "scaled-mm")
platform_kernels = _resolve_backend_kernels(
platform_kernels,
"scaled-mm",
quantization=quantization,
)
for kernel in platform_kernels:
is_supported_and_can_implement, failure_reason = (
@@ -705,6 +719,7 @@ def init_fp8_linear_kernel(
kernel_type = choose_scaled_mm_linear_kernel(
config=scaled_mm_linear_kernel_config,
possible_kernels=_POSSIBLE_FP8_BLOCK_KERNELS, # type: ignore[misc]
quantization="fp8_block_w8a8",
force_kernel=force_kernel,
)
if module_name:
@@ -736,6 +751,7 @@ def init_fp8_linear_kernel(
kernel_type = choose_scaled_mm_linear_kernel(
config=scaled_mm_linear_kernel_config,
possible_kernels=_POSSIBLE_FP8_KERNELS, # type: ignore[arg-type]
quantization="fp8_w8a8",
force_kernel=force_kernel,
)
if module_name:
@@ -772,6 +788,7 @@ def init_int8_linear_kernel(
kernel_type = choose_scaled_mm_linear_kernel(
config,
_POSSIBLE_INT8_KERNELS,
quantization="int8_w8a8",
)
logger.info_once(
@@ -824,7 +841,11 @@ def choose_mp_linear_kernel(
platform_kernels = _POSSIBLE_KERNELS.get(current_platform._enum, [])
# Apply --linear-backend filtering when set.
platform_kernels = _resolve_backend_kernels(platform_kernels, "mixed-precision")
platform_kernels = _resolve_backend_kernels(
platform_kernels,
"mixed-precision",
quantization="mixed_precision",
)
failure_reasons = []
for kernel in platform_kernels:
@@ -875,7 +896,11 @@ def init_mxfp8_linear_kernel(*, bmm_batch_size: int | None = None) -> Mxfp8Linea
possible = list(_POSSIBLE_MXFP8_KERNELS.get(platform, []))
# Apply --linear-backend filtering when set.
possible = _resolve_backend_kernels(possible, "MXFP8")
possible = _resolve_backend_kernels(
possible,
"MXFP8",
quantization="mxfp8",
)
failure_reasons = []
for kernel_cls in possible:
@@ -917,7 +942,11 @@ def init_mxfp4_linear_kernel(
possible = list(_POSSIBLE_MXFP4_KERNELS.get(platform, []))
# Apply --linear-backend filtering when set.
possible = _resolve_backend_kernels(possible, "MXFP4")
possible = _resolve_backend_kernels(
possible,
"MXFP4",
quantization="mxfp4",
)
failure_reasons = []
for kernel_cls in possible:
@@ -957,7 +986,7 @@ def init_mxfp6_linear_kernel(
activation_quant_key=activation_quant_key,
)
linear_backend = _get_linear_backend()
linear_backend = _get_linear_backend(quantization="mxfp6")
platform = current_platform._enum
possible = list(_POSSIBLE_MXFP6_KERNELS.get(platform, []))
@@ -1017,7 +1046,10 @@ def init_wfp8_a16_linear_kernel(
)
kernel_type = choose_scaled_mm_linear_kernel(
config, _POSSIBLE_WFP8A16_KERNELS, force_kernel=force_kernel
config,
_POSSIBLE_WFP8A16_KERNELS,
quantization="w8a16_fp8",
force_kernel=force_kernel,
)
if module_name:
@@ -1048,7 +1080,8 @@ def init_nvfp4_linear_kernel(use_a16: bool = False) -> NvFp4LinearKernel:
# batch-invariant CUTLASS implementation when available, otherwise fall
# back to emulation. It overrides --linear-backend.
force_kernel: type[NvFp4LinearKernel] | None = None
linear_backend = _get_linear_backend()
quantization = "nvfp4_w4a16" if use_a16 else "nvfp4_w4a4"
linear_backend = _get_linear_backend(quantization=quantization)
if envs.VLLM_BATCH_INVARIANT:
bi_supported, reason = CutlassNvFp4LinearKernel.is_supported()
if bi_supported:
@@ -1110,7 +1143,11 @@ def init_nvfp4_linear_kernel(use_a16: bool = False) -> NvFp4LinearKernel:
possible = [kernel for kernel in possible if kernel in a16_kernels]
# Apply --linear-backend filtering when set.
possible = _resolve_backend_kernels(possible, "NVFP4")
possible = _resolve_backend_kernels(
possible,
"NVFP4",
quantization=quantization,
)
failure_reasons = []
for kernel_cls in possible:
vllm/model_executor/kernels/linear/mxfp4/aiter.py
+1 -1
View File
@@ -158,7 +158,7 @@ class AiterMxfp4LinearKernel(MxFp4LinearKernel):
from vllm._aiter_ops import is_aiter_found_and_supported
from vllm.model_executor.kernels.linear import _get_linear_backend
linear_backend = _get_linear_backend()
linear_backend = _get_linear_backend(quantization="mxfp4")
if (
current_platform.is_rocm()