Add VLLM_FLASHINFER_AUTOTUNE_SKIP_OPS and skip CuTeDSL fp4_gemm autotuning by default (#48268)
Signed-off-by: mgoin <mgoin64@gmail.com> Signed-off-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Michael Goin committed
Jul 10, 2026 at 21:13 UTC
29fd6888922139e33980b81ec0c50477233eea00
2 files changed
+46
-2
vllm/envs.py
+14
@@ -197,6 +197,7 @@ if TYPE_CHECKING:
197
VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER: bool = True
198
VLLM_USE_FLASHINFER_MOE_INT4: bool = False
199
VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR: str | None = None
200
+ VLLM_FLASHINFER_AUTOTUNE_SKIP_OPS: list[str] | None = None
201
VLLM_FLASHINFER_ALLREDUCE_BACKEND: Literal["auto", "trtllm", "mnnvl"] = "auto"
202
VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE: int = 394 * 1024 * 1024
203
VLLM_XGRAMMAR_CACHE_MB: int = 0
@@ -1594,6 +1595,18 @@ environment_variables: dict[str, Callable[[], Any]] = {
1595
"VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR": lambda: os.getenv(
1596
"VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR", None
1597
),
1598
+ # Comma-separated FlashInfer op names to exclude from autotuning, using
1599
+ # the heuristic fallback tactic instead. Unset: skip "fp4_gemm" when the
1600
+ # CuTe-DSL NVFP4 linear kernel is selected. Empty: skip nothing.
1601
+ "VLLM_FLASHINFER_AUTOTUNE_SKIP_OPS": lambda: (
1602
+ None
1603
+ if "VLLM_FLASHINFER_AUTOTUNE_SKIP_OPS" not in os.environ
1604
+ else [
1605
+ v.strip()
1606
+ for v in os.environ["VLLM_FLASHINFER_AUTOTUNE_SKIP_OPS"].split(",")
1607
+ if v.strip()
1608
+ ]
1609
+ ),
1610
# Flashinfer fused allreduce backend.
1611
"VLLM_FLASHINFER_ALLREDUCE_BACKEND": env_with_choices(
1612
"VLLM_FLASHINFER_ALLREDUCE_BACKEND",
@@ -2110,6 +2123,7 @@ def compile_factors() -> dict[str, object]:
2123
"VLLM_DEBUG_LOG_API_SERVER_RESPONSE",
2124
"VLLM_TUNED_CONFIG_FOLDER",
2125
"VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR",
2126
+ "VLLM_FLASHINFER_AUTOTUNE_SKIP_OPS",
2127
"VLLM_ENGINE_ITERATION_TIMEOUT_S",
2128
"VLLM_HTTP_TIMEOUT_KEEP_ALIVE",
2129
"VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS",
vllm/model_executor/warmup/kernel_warmup.py
+32
-2
@@ -130,6 +130,25 @@ def kernel_warmup(worker: "Worker"):
130
cutedsl_warmup()
131
132
133
+def _flashinfer_autotune_skip_ops(runner: "GPUModelRunner") -> set[str] | None:
134
+ if envs.VLLM_FLASHINFER_AUTOTUNE_SKIP_OPS is not None:
135
+ return set(envs.VLLM_FLASHINFER_AUTOTUNE_SKIP_OPS) or None
136
+
137
+ from vllm.model_executor.kernels.linear import (
138
+ FlashInferCuteDslNvFp4LinearKernel,
139
+ )
140
+
141
+ for module in runner.get_model().modules():
142
+ for holder_name in ("quant_method", "scheme"):
143
+ kernel = getattr(getattr(module, holder_name, None), "kernel", None)
144
+ # CuTe-DSL mm_fp4 tuning JIT-compiles every tactic and its
145
+ # fallback is already the heuristic; all mm_fp4 backends share
146
+ # the "fp4_gemm" op name, so skip only when cute-dsl is selected.
147
+ if isinstance(kernel, FlashInferCuteDslNvFp4LinearKernel):
148
+ return {"fp4_gemm"}
149
+ return None
150
+
151
+
152
def flashinfer_autotune(runner: "GPUModelRunner") -> None:
153
"""
154
Autotune FlashInfer operations.
@@ -146,6 +165,15 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:
165
import vllm.utils.flashinfer as fi_utils
166
from vllm.distributed.parallel_state import get_world_group
167
168
+ autotune_kwargs: dict = {}
169
+ skip_ops = _flashinfer_autotune_skip_ops(runner)
170
+ if skip_ops:
171
+ logger.info(
172
+ "Skipping FlashInfer autotuning for ops %s",
173
+ sorted(skip_ops),
174
+ )
175
+ autotune_kwargs["skip_ops"] = skip_ops
176
+
177
use_persistent_cache = True
178
179
# When distributed, tune on every rank so the collectives stay synchronized.
@@ -153,7 +181,7 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:
181
use_persistent_cache = False
182
183
if not use_persistent_cache:
156
- with torch.inference_mode(), fi_utils.autotune():
184
+ with torch.inference_mode(), fi_utils.autotune(**autotune_kwargs):
185
runner._dummy_run(
186
num_tokens=runner.scheduler_config.max_num_batched_tokens,
187
skip_eplb=True,
@@ -181,7 +209,9 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:
209
210
with torch.inference_mode():
211
if is_leader:
184
- with fi_utils.autotune(tune_mode=True, cache=str(cache_path)):
212
+ with fi_utils.autotune(
213
+ tune_mode=True, cache=str(cache_path), **autotune_kwargs
214
+ ):
215
runner._dummy_run(**dummy_run_kwargs)
216
else:
217
runner._dummy_run(**dummy_run_kwargs)