@tduncan / Nvidia-Optimization-Stack-VLLM / commits / d4ee7fe7a9

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

Signed-off-by: Netanel Haber <58652339+netanel-haber@users.noreply.github.com>

Netanel Haber committed Sep 14, 2026 at 23:24 UTC d4ee7fe7a98b3012f64309e18dd425e565069a01

6 files changed +154 -17

docs/features/quantization/README.md

+17

@@ -23,6 +23,23 @@ The following are the supported quantization formats for vLLM:

23 - [TorchAO](torchao.md)

24 - [FP8 ViT Encoder Attention](fp8_vit_attn.md)

26 +## Selecting Linear Backends per Quantization

27 +

28 +`--linear-backend` selects one backend for all quantized linear layers. For

29 +mixed-precision models that use more than one linear quantization scheme, use

30 +`linear_backend_per_quant` to override the backend for individual schemes:

31 +

32 +```bash

33 +vllm serve <model> \

34 + --linear-backend cutlass \

35 + --kernel-config '{"linear_backend_per_quant":{"nvfp4_w4a16":"humming"}}'

36 +```

37 +

38 +Here, NVFP4 W4A16 linear layers use Humming, while all other quantized linear

39 +layers use CUTLASS. Per-quantization overrides take precedence over

40 +`--linear-backend`; schemes without an override continue to use the global

41 +setting, including automatic selection when it is `auto`.

42 +

43 ## Supported Hardware

45 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

@@ -12,20 +12,90 @@ from unittest.mock import patch

12 import pytest

13 import torch

15 +from vllm.config import KernelConfig, VllmConfig, set_current_vllm_config

16 from vllm.model_executor.kernels.linear import (

17 AiterInt8ScaledMMLinearKernel,

18 CPUInt8ScaledMMLinearKernel,

19 + HummingFP8ScaledMMLinearKernel,

20 Int8ScaledMMLinearKernel,

21 Int8ScaledMMLinearLayerConfig,

22 ScaledMMLinearKernel,

23 + _get_linear_backend,

24 + _resolve_backend_kernels,

25 + init_fp8_linear_kernel,

26 init_int8_linear_kernel,

27 register_linear_kernel,

28 )

29 +from vllm.model_executor.layers.quantization.utils.quant_utils import (

30 + kFp8StaticTensorSym,

31 +)

32 from vllm.platforms import PlatformEnum

34 pytestmark = pytest.mark.cpu_test

37 +def test_auto_linear_backend_is_unchanged_without_overrides():

38 + kernels = [object]

39 + quantizations = (

40 + "fp8_block_w8a8",

41 + "fp8_w8a8",

42 + "int8_w8a8",

43 + "mixed_precision",

44 + "mxfp8",

45 + "mxfp4",

46 + "mxfp6",

47 + "w8a16_fp8",

48 + "nvfp4_w4a4",

49 + "nvfp4_w4a16",

50 + )

51 + config = VllmConfig(kernel_config=KernelConfig(linear_backend="auto"))

52 +

53 + with set_current_vllm_config(config):

54 + for quantization in quantizations:

55 + assert _get_linear_backend(quantization=quantization) == "auto"

56 + assert (

57 + _resolve_backend_kernels(kernels, "test", quantization=quantization)

58 + is kernels

59 + )

60 +

61 +

62 +def test_linear_backend_override_is_quantization_specific():

63 + config = VllmConfig(

64 + kernel_config=KernelConfig(

65 + linear_backend="cutlass",

66 + linear_backend_per_quant={"nvfp4_w4a16": "humming"},

67 + )

68 + )

69 +

70 + with set_current_vllm_config(config):

71 + assert _get_linear_backend(quantization="nvfp4_w4a16") == "humming"

72 + assert _get_linear_backend(quantization="nvfp4_w4a4") == "cutlass"

73 + assert _get_linear_backend(quantization="fp8_w8a8") == "cutlass"

74 +

75 +

76 +@patch.object(HummingFP8ScaledMMLinearKernel, "is_supported", return_value=(True, None))

77 +@patch("vllm.model_executor.kernels.linear.current_platform")

78 +def test_fp8_linear_backend_override(platform_mock, _):

79 + platform_mock._enum = PlatformEnum.CUDA

80 + config = VllmConfig(

81 + kernel_config=KernelConfig(

82 + linear_backend="cutlass",

83 + linear_backend_per_quant={"fp8_w8a8": "humming"},

84 + )

85 + )

86 +

87 + with set_current_vllm_config(config):

88 + kernel = init_fp8_linear_kernel(

89 + activation_quant_key=kFp8StaticTensorSym,

90 + weight_quant_key=kFp8StaticTensorSym,

91 + input_dtype=torch.bfloat16,

92 + out_dtype=torch.bfloat16,

93 + weight_shape=(128, 128),

94 + )

95 +

96 + assert type(kernel) is HummingFP8ScaledMMLinearKernel

97 +

98 +

99 def test_is_supported_is_abstract():

100 """Test that is_supported() is properly defined as abstract."""

101 assert issubclass(ScaledMMLinearKernel, ABC)

tests/model_executor/kernels/test_b12x_linear.py

+6 -2

@@ -115,7 +115,9 @@ def test_b12x_backend_registration_priority_and_selection(

115 assert names.index(before) < names.index(kernel_cls.__name__) < names.index(after)

116

117 monkeypatch.setattr(linear_mod.current_platform, "_enum", PlatformEnum.CUDA)

118 - monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "b12x")

118 + monkeypatch.setattr(

119 + linear_mod, "_get_linear_backend", lambda *, quantization: "b12x"

120 + )

121 monkeypatch.setattr(

122 kernel_cls,

123 "is_supported",

@@ -784,7 +786,9 @@ def test_b12x_backend_preserves_w4a16_fallback(monkeypatch) -> None:

786 import vllm.model_executor.kernels.linear as linear_mod

787

788 monkeypatch.setattr(linear_mod.current_platform, "_enum", PlatformEnum.CUDA)

787 - monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "b12x")

789 + monkeypatch.setattr(

790 + linear_mod, "_get_linear_backend", lambda *, quantization: "b12x"

791 + )

792 monkeypatch.setitem(

793 linear_mod._POSSIBLE_NVFP4_KERNELS,

794 PlatformEnum.CUDA,

vllm/config/kernel.py

@@ -326,6 +326,13 @@ class KernelConfig:

326 - "xpu_woq": Use XPU kernels for weight-only quantization (e.g. W8A16)

327 """

328

329 + linear_backend_per_quant: dict[str, LinearBackend] | None = Field(

330 + default=None, min_length=1

331 + )

332 + """Backend overrides keyed by linear quantization scheme. Overrides take

333 + precedence over ``linear_backend``; for example,

334 + ``{"nvfp4_w4a16": "humming"}``."""

335 +

336 @field_validator("moe_backend", mode="before")

337 @classmethod

338 def _normalize_moe_backend(cls, value: Any) -> Any:

@@ -359,6 +366,8 @@ class KernelConfig:

366 "enable_flashinfer_autotune",

367 "ir_op_priority", # handled separately below

368 }

369 + if self.linear_backend_per_quant is None:

370 + ignored_factors.add("linear_backend_per_quant")

371 factors = get_hash_factors(self, ignored_factors)

372 factors["ir_op_priority"] = self.ir_op_priority.compute_hash()

373 return hash_factors(factors)

vllm/model_executor/kernels/linear/__init__.py

+51 -14

@@ -235,14 +235,19 @@ from vllm.platforms import PlatformEnum, current_platform

235 logger = init_logger(__name__)

236

237

238 -def _get_linear_backend() -> str:

238 +def _get_linear_backend(*, quantization: str) -> str:

239 """Get the linear_backend setting from the current vllm config."""

240 from vllm.config import get_current_vllm_config_or_none

241

242 - config = get_current_vllm_config_or_none()

243 - if config is not None:

244 - return config.kernel_config.linear_backend

245 - return "auto"

242 + if (config := get_current_vllm_config_or_none()) is None:

243 + return "auto"

244 + overrides = config.kernel_config.linear_backend_per_quant

245 + if overrides is not None and (override := overrides.get(quantization)):

246 + logger.info_once(

247 + "Applied linear backend override for %r: %r", quantization, override

248 + )

249 + return override

250 + return config.kernel_config.linear_backend

251

252

253 # Kernel classes covered by each --linear-backend value.

@@ -370,6 +375,8 @@ def _filter_kernels_by_backend(

375 def _resolve_backend_kernels(

376 kernels: list[type],

377 layer_desc: str,

378 + *,

379 + quantization: str,

380 ) -> list[type]:

381 """Apply --linear-backend filtering to one layer type's kernel list.

382

@@ -380,7 +387,7 @@ def _resolve_backend_kernels(

387 layer types (e.g. NVFP4 MoE projections next to FP8 attention

388 projections).

389 """

383 - linear_backend = _get_linear_backend()

390 + linear_backend = _get_linear_backend(quantization=quantization)

391 if linear_backend == "auto":

392 return kernels

393 filtered = _filter_kernels_by_backend(linear_backend, kernels)

@@ -625,6 +632,8 @@ def choose_scaled_mm_linear_kernel(

632 possible_kernels: dict[PlatformEnum, list[type[_KernelT]]],

633 compute_capability: int | None = None,

634 force_kernel: type[_KernelT] | None = None,

635 + *,

636 + quantization: str,

637 ) -> type[_KernelT]:

638 """

639 Choose a _KernelT that can implement the given config for the

@@ -642,6 +651,7 @@ def choose_scaled_mm_linear_kernel(

651 force_kernel (Optional[type[_KernelT]]): An Optional forced kernel to override

652 the possible_kernels if it can be implemented. If None, it will only try the

653 possible kernels.

654 + quantization: Quantization scheme used to select a backend override.

655

656 Raises:

657 ValueError: If no kernel can implement the given config.

@@ -668,7 +678,11 @@ def choose_scaled_mm_linear_kernel(

678 platform_kernels = possible_kernels.get(current_platform._enum, [])

679

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

671 - platform_kernels = _resolve_backend_kernels(platform_kernels, "scaled-mm")

681 + platform_kernels = _resolve_backend_kernels(

682 + platform_kernels,

683 + "scaled-mm",

684 + quantization=quantization,

685 + )

686

687 for kernel in platform_kernels:

688 is_supported_and_can_implement, failure_reason = (

@@ -705,6 +719,7 @@ def init_fp8_linear_kernel(

719 kernel_type = choose_scaled_mm_linear_kernel(

720 config=scaled_mm_linear_kernel_config,

721 possible_kernels=_POSSIBLE_FP8_BLOCK_KERNELS, # type: ignore[misc]

722 + quantization="fp8_block_w8a8",

723 force_kernel=force_kernel,

724 )

725 if module_name:

@@ -736,6 +751,7 @@ def init_fp8_linear_kernel(

751 kernel_type = choose_scaled_mm_linear_kernel(

752 config=scaled_mm_linear_kernel_config,

753 possible_kernels=_POSSIBLE_FP8_KERNELS, # type: ignore[arg-type]

754 + quantization="fp8_w8a8",

755 force_kernel=force_kernel,

756 )

757 if module_name:

@@ -772,6 +788,7 @@ def init_int8_linear_kernel(

788 kernel_type = choose_scaled_mm_linear_kernel(

789 config,

790 _POSSIBLE_INT8_KERNELS,

791 + quantization="int8_w8a8",

792 )

793

794 logger.info_once(

@@ -824,7 +841,11 @@ def choose_mp_linear_kernel(

841 platform_kernels = _POSSIBLE_KERNELS.get(current_platform._enum, [])

842

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

827 - platform_kernels = _resolve_backend_kernels(platform_kernels, "mixed-precision")

844 + platform_kernels = _resolve_backend_kernels(

845 + platform_kernels,

846 + "mixed-precision",

847 + quantization="mixed_precision",

848 + )

849

850 failure_reasons = []

851 for kernel in platform_kernels:

@@ -875,7 +896,11 @@ def init_mxfp8_linear_kernel(*, bmm_batch_size: int | None = None) -> Mxfp8Linea

896 possible = list(_POSSIBLE_MXFP8_KERNELS.get(platform, []))

897

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

878 - possible = _resolve_backend_kernels(possible, "MXFP8")

899 + possible = _resolve_backend_kernels(

900 + possible,

901 + "MXFP8",

902 + quantization="mxfp8",

903 + )

904

905 failure_reasons = []

906 for kernel_cls in possible:

@@ -917,7 +942,11 @@ def init_mxfp4_linear_kernel(

942 possible = list(_POSSIBLE_MXFP4_KERNELS.get(platform, []))

943

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

920 - possible = _resolve_backend_kernels(possible, "MXFP4")

945 + possible = _resolve_backend_kernels(

946 + possible,

947 + "MXFP4",

948 + quantization="mxfp4",

949 + )

950

951 failure_reasons = []

952 for kernel_cls in possible:

@@ -957,7 +986,7 @@ def init_mxfp6_linear_kernel(

986 activation_quant_key=activation_quant_key,

987 )

988

960 - linear_backend = _get_linear_backend()

989 + linear_backend = _get_linear_backend(quantization="mxfp6")

990

991 platform = current_platform._enum

992 possible = list(_POSSIBLE_MXFP6_KERNELS.get(platform, []))

@@ -1017,7 +1046,10 @@ def init_wfp8_a16_linear_kernel(

1046 )

1047

1048 kernel_type = choose_scaled_mm_linear_kernel(

1020 - config, _POSSIBLE_WFP8A16_KERNELS, force_kernel=force_kernel

1049 + config,

1050 + _POSSIBLE_WFP8A16_KERNELS,

1051 + quantization="w8a16_fp8",

1052 + force_kernel=force_kernel,

1053 )

1054

1055 if module_name:

@@ -1048,7 +1080,8 @@ def init_nvfp4_linear_kernel(use_a16: bool = False) -> NvFp4LinearKernel:

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

1081 # back to emulation. It overrides --linear-backend.

1082 force_kernel: type[NvFp4LinearKernel] | None = None

1051 - linear_backend = _get_linear_backend()

1083 + quantization = "nvfp4_w4a16" if use_a16 else "nvfp4_w4a4"

1084 + linear_backend = _get_linear_backend(quantization=quantization)

1085 if envs.VLLM_BATCH_INVARIANT:

1086 bi_supported, reason = CutlassNvFp4LinearKernel.is_supported()

1087 if bi_supported:

@@ -1110,7 +1143,11 @@ def init_nvfp4_linear_kernel(use_a16: bool = False) -> NvFp4LinearKernel:

1143 possible = [kernel for kernel in possible if kernel in a16_kernels]

1144

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

1113 - possible = _resolve_backend_kernels(possible, "NVFP4")

1146 + possible = _resolve_backend_kernels(

1147 + possible,

1148 + "NVFP4",

1149 + quantization=quantization,

1150 + )

1151

1152 failure_reasons = []

1153 for kernel_cls in possible:

vllm/model_executor/kernels/linear/mxfp4/aiter.py

+1 -1

@@ -158,7 +158,7 @@ class AiterMxfp4LinearKernel(MxFp4LinearKernel):

158 from vllm._aiter_ops import is_aiter_found_and_supported

159 from vllm.model_executor.kernels.linear import _get_linear_backend

160

161 - linear_backend = _get_linear_backend()

161 + linear_backend = _get_linear_backend(quantization="mxfp4")

162

163 if (

164 current_platform.is_rocm()