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 8, 17:45 UTC
[Kernel] Add KernelConfig flag to enable/disable FlashInfer autotune (#34006)
Signed-off-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com> Signed-off-by: Luka Govedič <ProExpertProg@users.noreply.github.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
Mohammad Miadh Angkad committed
Feb 7, 2026 at 21:24 UTC
dd6a6e119062027fae8964335c525796f71cdc2b
5 files changed
+104
-1
vllm/config/__init__.py
+3
@@ -11,6 +11,7 @@ from vllm.config.compilation import (
11
)
12
from vllm.config.device import DeviceConfig
13
from vllm.config.ec_transfer import ECTransferConfig
14
+from vllm.config.kernel import KernelConfig
15
from vllm.config.kv_events import KVEventsConfig
16
from vllm.config.kv_transfer import KVTransferConfig
17
from vllm.config.load import LoadConfig
@@ -65,6 +66,8 @@ __all__ = [
66
"DeviceConfig",
67
# From vllm.config.ec_transfer
68
"ECTransferConfig",
69
+ # From vllm.config.kernel
70
+ "KernelConfig",
71
# From vllm.config.kv_events
72
"KVEventsConfig",
73
# From vllm.config.kv_transfer
vllm/config/kernel.py
new
+44
@@ -0,0 +1,44 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from collections.abc import Callable
5
+from typing import Any
6
+
7
+from pydantic import Field, field_validator
8
+
9
+from vllm.config.utils import config
10
+from vllm.utils.hashing import safe_hash
11
+
12
+
13
+@config
14
+class KernelConfig:
15
+ """Configuration for kernel selection and warmup behavior."""
16
+
17
+ enable_flashinfer_autotune: bool = Field(default=None)
18
+ """If True, run FlashInfer autotuning during kernel warmup."""
19
+
20
+ def compute_hash(self) -> str:
21
+ """
22
+ WARNING: Whenever a new field is added to this config,
23
+ ensure that it is included in the factors list if
24
+ it affects the computation graph.
25
+
26
+ Provide a hash that uniquely identifies all the configs
27
+ that affect the structure of the computation
28
+ graph from input ids/embeddings to the final hidden states,
29
+ excluding anything before input ids/embeddings and after
30
+ the final hidden states.
31
+ """
32
+ # no factors to consider.
33
+ # this config will not affect the computation graph.
34
+ factors: list[Any] = []
35
+ hash_str = safe_hash(str(factors).encode(), usedforsecurity=False).hexdigest()
36
+ return hash_str
37
+
38
+ @field_validator("enable_flashinfer_autotune", mode="wrap")
39
+ @classmethod
40
+ def _skip_none_validation(cls, value: Any, handler: Callable) -> Any:
41
+ """Skip validation if the value is `None` when initialization is delayed."""
42
+ if value is None:
43
+ return value
44
+ return handler(value)
vllm/config/vllm.py
+20
@@ -30,6 +30,7 @@ from .cache import CacheConfig
30
from .compilation import CompilationConfig, CompilationMode, CUDAGraphMode
31
from .device import DeviceConfig
32
from .ec_transfer import ECTransferConfig
33
+from .kernel import KernelConfig
34
from .kv_events import KVEventsConfig
35
from .kv_transfer import KVTransferConfig
36
from .load import LoadConfig
@@ -129,6 +130,9 @@ OPTIMIZATION_LEVEL_00 = {
130
"cudagraph_mode": CUDAGraphMode.NONE,
131
"use_inductor_graph_partition": False,
132
},
133
+ "kernel_config": {
134
+ "enable_flashinfer_autotune": False,
135
+ },
136
}
137
OPTIMIZATION_LEVEL_01 = {
138
"compilation_config": {
@@ -145,6 +149,9 @@ OPTIMIZATION_LEVEL_01 = {
149
"cudagraph_mode": CUDAGraphMode.PIECEWISE,
150
"use_inductor_graph_partition": False,
151
},
152
+ "kernel_config": {
153
+ "enable_flashinfer_autotune": True,
154
+ },
155
}
156
OPTIMIZATION_LEVEL_02 = {
157
"compilation_config": {
@@ -161,6 +168,9 @@ OPTIMIZATION_LEVEL_02 = {
168
"cudagraph_mode": CUDAGraphMode.FULL_AND_PIECEWISE,
169
"use_inductor_graph_partition": False,
170
},
171
+ "kernel_config": {
172
+ "enable_flashinfer_autotune": True,
173
+ },
174
}
175
OPTIMIZATION_LEVEL_03 = {
176
"compilation_config": {
@@ -177,6 +187,9 @@ OPTIMIZATION_LEVEL_03 = {
187
"cudagraph_mode": CUDAGraphMode.FULL_AND_PIECEWISE,
188
"use_inductor_graph_partition": False,
189
},
190
+ "kernel_config": {
191
+ "enable_flashinfer_autotune": True,
192
+ },
193
}
194
195
OPTIMIZATION_LEVEL_TO_CONFIG = {
@@ -211,6 +224,8 @@ class VllmConfig:
224
"""Load configuration."""
225
attention_config: AttentionConfig = Field(default_factory=AttentionConfig)
226
"""Attention configuration."""
227
+ kernel_config: KernelConfig = Field(default_factory=KernelConfig)
228
+ """Kernel configuration."""
229
lora_config: LoRAConfig | None = None
230
"""LoRA configuration."""
231
speculative_config: SpeculativeConfig | None = None
@@ -756,6 +771,11 @@ class VllmConfig:
771
772
default_config = OPTIMIZATION_LEVEL_TO_CONFIG[self.optimization_level]
773
self._apply_optimization_level_defaults(default_config)
774
+ if self.kernel_config.enable_flashinfer_autotune is None:
775
+ raise ValueError(
776
+ "KernelConfig.enable_flashinfer_autotune must be set after applying "
777
+ "optimization level defaults."
778
+ )
779
780
if (
781
self.compilation_config.cudagraph_mode.requires_piecewise_compilation()
vllm/engine/arg_utils.py
+31
@@ -40,6 +40,7 @@ from vllm.config import (
40
DeviceConfig,
41
ECTransferConfig,
42
EPLBConfig,
43
+ KernelConfig,
44
KVEventsConfig,
45
KVTransferConfig,
46
LoadConfig,
@@ -536,6 +537,10 @@ class EngineArgs:
537
pooler_config: PoolerConfig | None = ModelConfig.pooler_config
538
compilation_config: CompilationConfig = get_field(VllmConfig, "compilation_config")
539
attention_config: AttentionConfig = get_field(VllmConfig, "attention_config")
540
+ kernel_config: KernelConfig = get_field(VllmConfig, "kernel_config")
541
+ enable_flashinfer_autotune: bool = get_field(
542
+ KernelConfig, "enable_flashinfer_autotune"
543
+ )
544
worker_cls: str = ParallelConfig.worker_cls
545
worker_extension_cls: str = ParallelConfig.worker_extension_cls
546
@@ -595,6 +600,8 @@ class EngineArgs:
600
self.compilation_config = CompilationConfig(**self.compilation_config)
601
if isinstance(self.attention_config, dict):
602
self.attention_config = AttentionConfig(**self.attention_config)
603
+ if isinstance(self.kernel_config, dict):
604
+ self.kernel_config = KernelConfig(**self.kernel_config)
605
if isinstance(self.eplb_config, dict):
606
self.eplb_config = EPLBConfig(**self.eplb_config)
607
if isinstance(self.weight_transfer_config, dict):
@@ -1163,6 +1170,17 @@ class EngineArgs:
1170
**compilation_kwargs["max_cudagraph_capture_size"],
1171
)
1172
1173
+ # Kernel arguments
1174
+ kernel_kwargs = get_kwargs(KernelConfig)
1175
+ kernel_group = parser.add_argument_group(
1176
+ title="KernelConfig",
1177
+ description=KernelConfig.__doc__,
1178
+ )
1179
+ kernel_group.add_argument(
1180
+ "--enable-flashinfer-autotune",
1181
+ **kernel_kwargs["enable_flashinfer_autotune"],
1182
+ )
1183
+
1184
# vLLM arguments
1185
vllm_kwargs = get_kwargs(VllmConfig)
1186
vllm_group = parser.add_argument_group(
@@ -1189,6 +1207,7 @@ class EngineArgs:
1207
vllm_group.add_argument(
1208
"--attention-config", "-ac", **vllm_kwargs["attention_config"]
1209
)
1210
+ vllm_group.add_argument("--kernel-config", **vllm_kwargs["kernel_config"])
1211
vllm_group.add_argument(
1212
"--additional-config", **vllm_kwargs["additional_config"]
1213
)
@@ -1717,6 +1736,17 @@ class EngineArgs:
1736
else:
1737
attention_config.backend = self.attention_backend
1738
1739
+ # Kernel config overrides
1740
+ kernel_config = copy.deepcopy(self.kernel_config)
1741
+ if self.enable_flashinfer_autotune is not None:
1742
+ if kernel_config.enable_flashinfer_autotune is not None:
1743
+ raise ValueError(
1744
+ "enable_flashinfer_autotune and "
1745
+ "kernel_config.enable_flashinfer_autotune "
1746
+ "are mutually exclusive"
1747
+ )
1748
+ kernel_config.enable_flashinfer_autotune = self.enable_flashinfer_autotune
1749
+
1750
load_config = self.create_load_config()
1751
1752
# Pass reasoning_parser into StructuredOutputsConfig
@@ -1767,6 +1797,7 @@ class EngineArgs:
1797
device_config=device_config,
1798
load_config=load_config,
1799
attention_config=attention_config,
1800
+ kernel_config=kernel_config,
1801
lora_config=lora_config,
1802
speculative_config=speculative_config,
1803
structured_outputs_config=self.structured_outputs_config,
vllm/model_executor/warmup/kernel_warmup.py
+6
-1
@@ -36,8 +36,13 @@ def kernel_warmup(worker: "Worker"):
36
max_tokens = worker.scheduler_config.max_num_batched_tokens
37
deep_gemm_warmup(model, max_tokens)
38
39
+ enable_flashinfer_autotune = (
40
+ worker.vllm_config.kernel_config.enable_flashinfer_autotune
41
+ )
42
# FlashInfer autotune for Hopper (SM 9.0) and Blackwell (SM 10.0) GPUs
40
- if has_flashinfer() and current_platform.has_device_capability(90):
43
+ if enable_flashinfer_autotune is False:
44
+ logger.info("Skipping FlashInfer autotune because it is disabled.")
45
+ elif has_flashinfer() and current_platform.has_device_capability(90):
46
flashinfer_autotune(worker.model_runner)
47
48
# FlashInfer attention warmup