From d4ee7fe7a98b3012f64309e18dd425e565069a01 Mon Sep 17 00:00:00 2001 From: Netanel Haber <58652339+netanel-haber@users.noreply.github.com> Date: Mon, 14 Sep 2026 23:24:18 +0300 Subject: [PATCH] [Quantization] Select linear backends per quantization (#51204) Signed-off-by: Netanel Haber <58652339+netanel-haber@users.noreply.github.com> --- docs/features/quantization/README.md | 17 +++++ .../test_scaled_mm_kernel_selection.py | 70 +++++++++++++++++++ .../kernels/test_b12x_linear.py | 8 ++- vllm/config/kernel.py | 9 +++ .../model_executor/kernels/linear/__init__.py | 65 +++++++++++++---- .../kernels/linear/mxfp4/aiter.py | 2 +- 6 files changed, 154 insertions(+), 17 deletions(-) diff --git a/docs/features/quantization/README.md b/docs/features/quantization/README.md index 69ece36076..ba1a5b86fe 100644 --- a/docs/features/quantization/README.md +++ b/docs/features/quantization/README.md @@ -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 \ + --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: diff --git a/tests/kernels/quantization/test_scaled_mm_kernel_selection.py b/tests/kernels/quantization/test_scaled_mm_kernel_selection.py index bedebdb59b..8275714744 100644 --- a/tests/kernels/quantization/test_scaled_mm_kernel_selection.py +++ b/tests/kernels/quantization/test_scaled_mm_kernel_selection.py @@ -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) diff --git a/tests/model_executor/kernels/test_b12x_linear.py b/tests/model_executor/kernels/test_b12x_linear.py index 27368ead27..9d4f9beeb9 100644 --- a/tests/model_executor/kernels/test_b12x_linear.py +++ b/tests/model_executor/kernels/test_b12x_linear.py @@ -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, diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index 4505e0cfd9..0ac5bf4391 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -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) diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 1df73273ca..d3c41f97d8 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -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: diff --git a/vllm/model_executor/kernels/linear/mxfp4/aiter.py b/vllm/model_executor/kernels/linear/mxfp4/aiter.py index 13fd5db275..420e7f9493 100644 --- a/vllm/model_executor/kernels/linear/mxfp4/aiter.py +++ b/vllm/model_executor/kernels/linear/mxfp4/aiter.py @@ -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()