A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 13:00 UTC
[Kernel][UX] Add `--linear-backend` arg for linear kernel selection (#39538)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
May 15, 2026 at 20:07 UTC
b2a27b82d970efa0203c06be6dc0d94526edaab0
4 files changed
+265
-33
tests/models/quantization/test_nvfp4.py
+8
-9
@@ -90,9 +90,9 @@ def test_models(example_prompts, model_name) -> None:
90
EAGER = [True, False]
91
92
SM_100_NVFP4_BACKENDS = [
93
- "flashinfer-cudnn",
94
- "flashinfer-trtllm",
95
- "flashinfer-cutlass",
93
+ "flashinfer_cudnn",
94
+ "flashinfer_trtllm",
95
+ "flashinfer_cutlass",
96
]
97
98
@@ -102,12 +102,12 @@ SM_100_NVFP4_BACKENDS = [
102
"backend",
103
[
104
"emulation",
105
- "flashinfer-cudnn",
106
- "flashinfer-trtllm", # the small seq_len ensures trtllm_8x4_layout backend is used
107
- "flashinfer-cutlass",
105
+ "flashinfer_cudnn",
106
+ "flashinfer_trtllm", # the small seq_len ensures trtllm_8x4_layout backend is used
107
+ "flashinfer_cutlass",
108
],
109
)
110
-def test_nvfp4(vllm_runner, model, eager, backend, monkeypatch):
110
+def test_nvfp4(vllm_runner, model, eager, backend):
111
if (
112
not current_platform.has_device_capability(100)
113
and backend in SM_100_NVFP4_BACKENDS
@@ -116,8 +116,7 @@ def test_nvfp4(vllm_runner, model, eager, backend, monkeypatch):
116
f"The backend {backend} is not supported with current_platform.has_device_capability(100) == False"
117
)
118
119
- monkeypatch.setenv("VLLM_NVFP4_GEMM_BACKEND", backend)
120
- with vllm_runner(model, enforce_eager=eager) as llm:
119
+ with vllm_runner(model, enforce_eager=eager, linear_backend=backend) as llm:
120
output = llm.generate_greedy(["1 2 3 4 5"], max_tokens=2)
121
assert output[0][1] == "1 2 3 4 5 6"
122
vllm/config/kernel.py
+44
@@ -135,6 +135,24 @@ MoEBackend = Literal[
135
"emulation",
136
]
137
138
+LinearBackend = Literal[
139
+ "auto",
140
+ "cutlass",
141
+ "flashinfer_cutlass",
142
+ "flashinfer_trtllm",
143
+ "flashinfer_cudnn",
144
+ "marlin",
145
+ "triton",
146
+ "deep_gemm",
147
+ "torch",
148
+ "aiter",
149
+ "machete",
150
+ "fbgemm",
151
+ "conch",
152
+ "exllama",
153
+ "emulation",
154
+]
155
+
156
157
@config
158
class KernelConfig:
@@ -168,6 +186,25 @@ class KernelConfig:
186
running QDQ on activations.
187
"""
188
189
+ linear_backend: LinearBackend = "auto"
190
+ """Backend for quantized linear layer GEMM kernels. Available options:
191
+
192
+ - "auto": Automatically select the best backend based on model and hardware
193
+ - "cutlass": Use CUTLASS-based kernels
194
+ - "flashinfer_cutlass": Use FlashInfer with CUTLASS kernels
195
+ - "flashinfer_trtllm": Use FlashInfer with TensorRT-LLM kernels
196
+ - "flashinfer_cudnn": Use FlashInfer with cuDNN kernels
197
+ - "marlin": Use Marlin kernels
198
+ - "triton": Use Triton-based kernels
199
+ - "deep_gemm": Use DeepGEMM kernels
200
+ - "torch": Use PyTorch native scaled_mm kernels
201
+ - "aiter": Use AMD AITer kernels (ROCm only)
202
+ - "machete": Use Machete kernels (mixed-precision)
203
+ - "fbgemm": Use FBGEMM kernels
204
+ - "conch": Use Conch mixed-precision kernels
205
+ - "exllama": Use Exllama mixed-precision kernels
206
+ - "emulation": Use slow dequant-to-BF16 emulation (for testing only)"""
207
+
208
@field_validator("moe_backend", mode="before")
209
@classmethod
210
def _normalize_moe_backend(cls, value: Any) -> Any:
@@ -175,6 +212,13 @@ class KernelConfig:
212
return value.lower().replace("-", "_")
213
return value
214
215
+ @field_validator("linear_backend", mode="before")
216
+ @classmethod
217
+ def _normalize_linear_backend(cls, value: Any) -> Any:
218
+ if isinstance(value, str):
219
+ return value.lower().replace("-", "_")
220
+ return value
221
+
222
def compute_hash(self) -> str:
223
"""
224
Produces a hash unique to the pass configuration.
vllm/engine/arg_utils.py
+7
-1
@@ -71,7 +71,7 @@ from vllm.config.cache import (
71
PrefixCachingHashAlgo,
72
)
73
from vllm.config.device import Device
74
-from vllm.config.kernel import IrOpPriorityConfig, MoEBackend
74
+from vllm.config.kernel import IrOpPriorityConfig, LinearBackend, MoEBackend
75
from vllm.config.lora import MaxLoRARanks
76
from vllm.config.mamba import MambaBackendEnum
77
from vllm.config.model import (
@@ -477,6 +477,7 @@ class EngineArgs:
477
enable_expert_parallel: bool = ParallelConfig.enable_expert_parallel
478
enable_ep_weight_filter: bool = ParallelConfig.enable_ep_weight_filter
479
moe_backend: MoEBackend = KernelConfig.moe_backend
480
+ linear_backend: LinearBackend = KernelConfig.linear_backend
481
all2all_backend: All2AllBackend = ParallelConfig.all2all_backend
482
enable_elastic_ep: bool = ParallelConfig.enable_elastic_ep
483
enable_dbo: bool = ParallelConfig.enable_dbo
@@ -1412,6 +1413,9 @@ class EngineArgs:
1413
moe_backend_kwargs = kernel_kwargs["moe_backend"]
1414
moe_backend_kwargs["type"] = lambda s: s.lower().replace("-", "_")
1415
kernel_group.add_argument("--moe-backend", **moe_backend_kwargs)
1416
+ linear_backend_kwargs = kernel_kwargs["linear_backend"]
1417
+ linear_backend_kwargs["type"] = lambda s: s.lower().replace("-", "_")
1418
+ kernel_group.add_argument("--linear-backend", **linear_backend_kwargs)
1419
1420
# vLLM arguments
1421
vllm_kwargs = get_kwargs(VllmConfig)
@@ -2085,6 +2089,8 @@ class EngineArgs:
2089
kernel_config.enable_flashinfer_autotune = self.enable_flashinfer_autotune
2090
if self.moe_backend != "auto":
2091
kernel_config.moe_backend = self.moe_backend
2092
+ if self.linear_backend != "auto":
2093
+ kernel_config.linear_backend = self.linear_backend
2094
2095
# Transfer top-level ir_op_priority into KernelConfig.ir_op_priority
2096
for op_name, op_priority in asdict(self.ir_op_priority).items():
vllm/model_executor/kernels/linear/__init__.py
+206
-23
@@ -13,6 +13,7 @@ or kernel implementation, add it to this __init__.py to maintain
13
import stability.
14
"""
15
16
+import warnings
17
from typing import TypeVar
18
19
import torch
@@ -161,6 +162,96 @@ from vllm.platforms import PlatformEnum, current_platform
162
163
logger = init_logger(__name__)
164
165
+
166
+def _get_linear_backend() -> str:
167
+ """Get the linear_backend setting from the current vllm config."""
168
+ from vllm.config import get_current_vllm_config_or_none
169
+
170
+ config = get_current_vllm_config_or_none()
171
+ if config is not None:
172
+ return config.kernel_config.linear_backend
173
+ return "auto"
174
+
175
+
176
+# Mapping from linear_backend name to the set of kernel classes it covers.
177
+# When a user sets --linear-backend <name>, only kernels in the corresponding
178
+# set are considered candidates. If none can implement the layer config,
179
+# an error is raised to respect the user's explicit intent.
180
+_LINEAR_BACKEND_KERNEL_MAP: dict[str, set[type]] = {
181
+ "cutlass": {
182
+ CutlassInt8ScaledMMLinearKernel,
183
+ CutlassFP8ScaledMMLinearKernel,
184
+ CutlassFp8BlockScaledMMKernel,
185
+ CutlassW4A8LinearKernel,
186
+ CutlassNvFp4LinearKernel,
187
+ },
188
+ "flashinfer_cutlass": {
189
+ FlashInferFP8ScaledMMLinearKernel,
190
+ FlashInferFp8DeepGEMMDynamicBlockScaledKernel,
191
+ FlashInferCutlassMxfp8LinearKernel,
192
+ FlashInferCutlassNvFp4LinearKernel,
193
+ FlashInferMxFp4LinearKernel,
194
+ },
195
+ "flashinfer_trtllm": {
196
+ FlashInferTrtllmNvFp4LinearKernel,
197
+ },
198
+ "flashinfer_cudnn": {
199
+ FlashInferCudnnNvFp4LinearKernel,
200
+ },
201
+ "marlin": {
202
+ MarlinFP8ScaledMMLinearKernel,
203
+ MarlinLinearKernel,
204
+ MarlinMxfp8LinearKernel,
205
+ MarlinNvFp4LinearKernel,
206
+ MarlinMxFp4LinearKernel,
207
+ },
208
+ "triton": {
209
+ TritonInt8ScaledMMLinearKernel,
210
+ TritonFp8BlockScaledMMKernel,
211
+ TritonW4A16LinearKernel,
212
+ },
213
+ "deep_gemm": {
214
+ DeepGemmFp8BlockScaledMMKernel,
215
+ },
216
+ "torch": {
217
+ PerTensorTorchFP8ScaledMMLinearKernel,
218
+ ChannelWiseTorchFP8ScaledMMLinearKernel,
219
+ RowWiseTorchFP8ScaledMMLinearKernel,
220
+ },
221
+ "aiter": {
222
+ AiterInt8ScaledMMLinearKernel,
223
+ AiterFp8BlockScaledMMKernel,
224
+ AiterPerTokenFp8ScaledMMLinearKernel,
225
+ AiterPreshuffledPerTokenFp8ScaledMMLinearKernel,
226
+ },
227
+ "machete": {
228
+ MacheteLinearKernel,
229
+ },
230
+ "fbgemm": {
231
+ FbgemmNvFp4LinearKernel,
232
+ },
233
+ "conch": {
234
+ ConchLinearKernel,
235
+ },
236
+ "exllama": {
237
+ ExllamaLinearKernel,
238
+ },
239
+ "emulation": {
240
+ EmulationMxfp8LinearKernel,
241
+ EmulationNvFp4LinearKernel,
242
+ },
243
+}
244
+
245
+
246
+def _filter_kernels_by_backend(
247
+ backend: str,
248
+ kernels: list[type],
249
+) -> list[type]:
250
+ """Filter a kernel priority list to only those matching the backend."""
251
+ backend_kernels = _LINEAR_BACKEND_KERNEL_MAP.get(backend, set())
252
+ return [k for k in kernels if k in backend_kernels]
253
+
254
+
255
# in priority/performance order (when available)
256
_POSSIBLE_INT8_KERNELS: dict[PlatformEnum, list[type[Int8ScaledMMLinearKernel]]] = {
257
PlatformEnum.CPU: [CPUInt8ScaledMMLinearKernel],
@@ -375,7 +466,20 @@ def choose_scaled_mm_linear_kernel(
466
scope="global",
467
)
468
378
- for kernel in possible_kernels[current_platform._enum]:
469
+ platform_kernels = possible_kernels[current_platform._enum]
470
+
471
+ # Apply --linear-backend filtering when set.
472
+ linear_backend = _get_linear_backend()
473
+ if linear_backend != "auto":
474
+ filtered = _filter_kernels_by_backend(linear_backend, platform_kernels)
475
+ if not filtered:
476
+ raise ValueError(
477
+ f"--linear-backend={linear_backend} was requested but no "
478
+ f"'{linear_backend}' kernel exists for this layer type."
479
+ )
480
+ platform_kernels = filtered
481
+
482
+ for kernel in platform_kernels:
483
is_supported_and_can_implement, failure_reason = (
484
is_supported_and_can_implement_kernel(kernel, config, compute_capability)
485
)
@@ -526,8 +630,21 @@ def choose_mp_linear_kernel(
630
if _cc is not None:
631
compute_capability = _cc[0] * 10 + _cc[1]
632
633
+ platform_kernels = _POSSIBLE_KERNELS[current_platform._enum]
634
+
635
+ # Apply --linear-backend filtering when set.
636
+ linear_backend = _get_linear_backend()
637
+ if linear_backend != "auto":
638
+ filtered = _filter_kernels_by_backend(linear_backend, platform_kernels)
639
+ if not filtered:
640
+ raise ValueError(
641
+ f"--linear-backend={linear_backend} was requested but no "
642
+ f"'{linear_backend}' kernel exists for mixed-precision layers."
643
+ )
644
+ platform_kernels = filtered
645
+
646
failure_reasons = []
530
- for kernel in _POSSIBLE_KERNELS[current_platform._enum]:
647
+ for kernel in platform_kernels:
648
if kernel.__name__ in envs.VLLM_DISABLED_KERNELS:
649
failure_reasons.append(
650
f" {kernel.__name__} disabled by environment variable"
@@ -564,7 +681,18 @@ def init_mxfp8_linear_kernel() -> Mxfp8LinearKernel:
681
config = Mxfp8LinearLayerConfig()
682
683
platform = current_platform._enum
567
- possible = _POSSIBLE_MXFP8_KERNELS.get(platform, [])
684
+ possible = list(_POSSIBLE_MXFP8_KERNELS.get(platform, []))
685
+
686
+ # Apply --linear-backend filtering when set.
687
+ linear_backend = _get_linear_backend()
688
+ if linear_backend != "auto":
689
+ filtered = _filter_kernels_by_backend(linear_backend, possible)
690
+ if not filtered:
691
+ raise ValueError(
692
+ f"--linear-backend={linear_backend} was requested but no "
693
+ f"'{linear_backend}' kernel exists for MXFP8 layers."
694
+ )
695
+ possible = filtered
696
697
failure_reasons = []
698
for kernel_cls in possible:
@@ -596,8 +724,10 @@ def init_mxfp8_linear_kernel() -> Mxfp8LinearKernel:
724
def init_mxfp4_linear_kernel() -> MxFp4LinearKernel:
725
"""Select and instantiate the best MXFP4 linear kernel for the
726
current platform."""
727
+ linear_backend = _get_linear_backend()
728
+
729
force_kernel: type[MxFp4LinearKernel] | None = None
600
- if envs.VLLM_MXFP4_USE_MARLIN:
730
+ if linear_backend == "auto" and envs.VLLM_MXFP4_USE_MARLIN:
731
force_kernel = MarlinMxFp4LinearKernel
732
733
if force_kernel is not None:
@@ -611,7 +741,17 @@ def init_mxfp4_linear_kernel() -> MxFp4LinearKernel:
741
return force_kernel(MxFp4LinearLayerConfig())
742
743
platform = current_platform._enum
614
- possible = _POSSIBLE_MXFP4_KERNELS.get(platform, [])
744
+ possible = list(_POSSIBLE_MXFP4_KERNELS.get(platform, []))
745
+
746
+ # Apply --linear-backend filtering when set.
747
+ if linear_backend != "auto":
748
+ filtered = _filter_kernels_by_backend(linear_backend, possible)
749
+ if not filtered:
750
+ raise ValueError(
751
+ f"--linear-backend={linear_backend} was requested but no "
752
+ f"'{linear_backend}' kernel exists for MXFP4 layers."
753
+ )
754
+ possible = filtered
755
756
failure_reasons = []
757
for kernel_cls in possible:
@@ -686,26 +826,59 @@ def init_nvfp4_linear_kernel() -> NvFp4LinearKernel:
826
current platform."""
827
config = NvFp4LinearLayerConfig()
828
689
- # Env-var overrides.
829
+ # VLLM_BATCH_INVARIANT unconditionally forces emulation for deterministic
830
+ # execution. It overrides both --linear-backend and the deprecated env
831
+ # vars below.
832
force_kernel: type[NvFp4LinearKernel] | None = None
833
+ linear_backend = _get_linear_backend()
834
if envs.VLLM_BATCH_INVARIANT:
692
- logger.info_once(
693
- "VLLM_BATCH_INVARIANT forces NVFP4 linear to use the "
694
- "emulation backend for deterministic execution."
695
- )
696
- force_kernel = EmulationNvFp4LinearKernel
697
- elif envs.VLLM_USE_FBGEMM:
698
- force_kernel = FbgemmNvFp4LinearKernel
699
- elif envs.VLLM_USE_NVFP4_CT_EMULATIONS:
835
+ if linear_backend not in ("auto", "emulation"):
836
+ logger.warning_once(
837
+ "VLLM_BATCH_INVARIANT overrides --linear-backend=%s; using "
838
+ "the emulation backend for deterministic execution.",
839
+ linear_backend,
840
+ )
841
+ else:
842
+ logger.info_once(
843
+ "VLLM_BATCH_INVARIANT forces NVFP4 linear to use the "
844
+ "emulation backend for deterministic execution."
845
+ )
846
force_kernel = EmulationNvFp4LinearKernel
701
- elif envs.VLLM_NVFP4_GEMM_BACKEND is not None:
702
- backend_name = envs.VLLM_NVFP4_GEMM_BACKEND
703
- force_kernel = _NVFP4_BACKEND_TO_KERNEL.get(backend_name)
704
- if force_kernel is None:
705
- raise ValueError(
706
- f"Unknown VLLM_NVFP4_GEMM_BACKEND={backend_name!r}. "
707
- f"Valid choices: {list(_NVFP4_BACKEND_TO_KERNEL.keys())}"
847
+ elif linear_backend == "auto":
848
+ # Deprecated env-var overrides — only honoured when --linear-backend
849
+ # is "auto". Will be removed in v0.21; users should migrate to
850
+ # --linear-backend.
851
+ if envs.VLLM_USE_FBGEMM:
852
+ warnings.warn(
853
+ "VLLM_USE_FBGEMM is deprecated and will be removed in "
854
+ "v0.21. Use --linear-backend fbgemm instead.",
855
+ DeprecationWarning,
856
+ stacklevel=2,
857
+ )
858
+ force_kernel = FbgemmNvFp4LinearKernel
859
+ elif envs.VLLM_USE_NVFP4_CT_EMULATIONS:
860
+ warnings.warn(
861
+ "VLLM_USE_NVFP4_CT_EMULATIONS is deprecated and will be "
862
+ "removed in v0.21. Use --linear-backend emulation instead.",
863
+ DeprecationWarning,
864
+ stacklevel=2,
865
)
866
+ force_kernel = EmulationNvFp4LinearKernel
867
+ elif envs.VLLM_NVFP4_GEMM_BACKEND is not None:
868
+ warnings.warn(
869
+ "VLLM_NVFP4_GEMM_BACKEND is deprecated and will be "
870
+ "removed in v0.21. Use --linear-backend instead.",
871
+ DeprecationWarning,
872
+ stacklevel=2,
873
+ )
874
+ backend_name = envs.VLLM_NVFP4_GEMM_BACKEND
875
+ force_kernel = _NVFP4_BACKEND_TO_KERNEL.get(backend_name)
876
+ if force_kernel is None:
877
+ raise ValueError(
878
+ f"Unknown VLLM_NVFP4_GEMM_BACKEND={backend_name!r}. "
879
+ f"Valid choices: "
880
+ f"{list(_NVFP4_BACKEND_TO_KERNEL.keys())}"
881
+ )
882
883
if force_kernel is not None:
884
is_supported, reason = force_kernel.is_supported()
@@ -717,9 +890,19 @@ def init_nvfp4_linear_kernel() -> NvFp4LinearKernel:
890
logger.info_once("Using %s for NVFP4 GEMM", force_kernel.__name__)
891
return force_kernel(config)
892
720
- # Auto-select from registry.
893
+ # Auto-select from registry (or --linear-backend filtered).
894
platform = current_platform._enum
722
- possible = _POSSIBLE_NVFP4_KERNELS.get(platform, [])
895
+ possible = list(_POSSIBLE_NVFP4_KERNELS.get(platform, []))
896
+
897
+ # Apply --linear-backend filtering when set.
898
+ if linear_backend != "auto":
899
+ filtered = _filter_kernels_by_backend(linear_backend, possible)
900
+ if not filtered:
901
+ raise ValueError(
902
+ f"--linear-backend={linear_backend} was requested but no "
903
+ f"'{linear_backend}' kernel exists for NVFP4 layers."
904
+ )
905
+ possible = filtered
906
907
failure_reasons = []
908
for kernel_cls in possible: