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, 09:15 UTC
[UX] Add a persistent cache for FlashInfer autotuning (#42537)
Signed-off-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>
Mohammad Miadh Angkad committed
May 19, 2026 at 11:25 UTC
da03e549b34685c4e63a091e973d907aee48a68c
4 files changed
+103
-16
docs/usage/security.md
+1
@@ -325,6 +325,7 @@ Most cache paths default to subdirectories under a single root. Changing `VLLM_C
325
| --- | --- | --- |
326
| `VLLM_CACHE_ROOT` | `~/.cache/vllm` | Base cache directory. Respects `XDG_CACHE_HOME` if set. All paths below inherit from this unless explicitly overridden. |
327
| *(torch.compile)* | `$VLLM_CACHE_ROOT/torch_compile_cache/` | Compilation cache for AOT-compiled models, Inductor graphs, and Triton kernels. Controlled by `VLLM_DISABLE_COMPILE_CACHE` (set to `1` to disable). |
328
+| `VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR` | `$VLLM_CACHE_ROOT/flashinfer_autotune_cache/<flashinfer-version>/<arch>/<cache-hash>/` | FlashInfer autotune config cache. |
329
| `VLLM_ASSETS_CACHE` | `$VLLM_CACHE_ROOT/assets/` | Downloaded assets (e.g., tokenizer files). |
330
| `VLLM_XLA_CACHE_PATH` | `$VLLM_CACHE_ROOT/xla_cache/` | XLA/TPU compilation cache. |
331
| `VLLM_MEDIA_CACHE` | *(disabled)* | Optional cache for downloaded media (images, video, audio). Not enabled unless explicitly set. |
tests/model_executor/test_flashinfer_autotune_cache.py
new
+60
@@ -0,0 +1,60 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import sys
5
+from hashlib import sha256
6
+from pathlib import Path
7
+from types import SimpleNamespace
8
+
9
+from vllm.model_executor.warmup import kernel_warmup
10
+
11
+
12
+def test_resolve_flashinfer_autotune_file_default_layout(
13
+ monkeypatch, tmp_path: Path
14
+) -> None:
15
+ fake_jit = SimpleNamespace(
16
+ env=SimpleNamespace(
17
+ FLASHINFER_WORKSPACE_DIR=Path("/flashinfer-cache/0.6.11.post2/103a")
18
+ )
19
+ )
20
+ fake_flashinfer = SimpleNamespace(jit=fake_jit)
21
+ monkeypatch.setitem(sys.modules, "flashinfer", fake_flashinfer)
22
+ monkeypatch.setitem(sys.modules, "flashinfer.jit", fake_jit)
23
+ monkeypatch.setattr(
24
+ kernel_warmup, "aot_compile_hash_factors", lambda _: ["env-hash", "config-hash"]
25
+ )
26
+ monkeypatch.setattr(kernel_warmup.envs, "VLLM_CACHE_ROOT", str(tmp_path))
27
+ monkeypatch.setattr(kernel_warmup.envs, "VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR", None)
28
+
29
+ runner = SimpleNamespace(vllm_config=SimpleNamespace())
30
+ cache_hash = sha256(str(["env-hash", "config-hash"]).encode()).hexdigest()
31
+
32
+ path = kernel_warmup._resolve_flashinfer_autotune_file(runner)
33
+
34
+ assert path == (
35
+ tmp_path
36
+ / "flashinfer_autotune_cache"
37
+ / "0.6.11.post2"
38
+ / "103a"
39
+ / cache_hash
40
+ / "autotune_configs.json"
41
+ )
42
+ assert path.parent.is_dir()
43
+
44
+
45
+def test_resolve_flashinfer_autotune_file_uses_override_dir(
46
+ monkeypatch, tmp_path: Path
47
+) -> None:
48
+ monkeypatch.setattr(
49
+ kernel_warmup.envs, "VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR", str(tmp_path)
50
+ )
51
+ monkeypatch.setattr(
52
+ kernel_warmup, "aot_compile_hash_factors", lambda _: ["env-hash", "config-hash"]
53
+ )
54
+
55
+ runner = SimpleNamespace(vllm_config=SimpleNamespace())
56
+ cache_hash = sha256(str(["env-hash", "config-hash"]).encode()).hexdigest()
57
+
58
+ path = kernel_warmup._resolve_flashinfer_autotune_file(runner)
59
+
60
+ assert path == tmp_path / cache_hash / "autotune_configs.json"
vllm/envs.py
+6
@@ -182,6 +182,7 @@ if TYPE_CHECKING:
182
VLLM_FLASHINFER_MOE_BACKEND: Literal["throughput", "latency", "masked_gemm"] = (
183
"latency"
184
)
185
+ VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR: str | None = None
186
VLLM_FLASHINFER_ALLREDUCE_BACKEND: Literal["auto", "trtllm", "mnnvl"] = "auto"
187
VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE: int = 394 * 1024 * 1024
188
VLLM_XGRAMMAR_CACHE_MB: int = 0
@@ -1419,6 +1420,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
1420
"latency",
1421
["throughput", "latency", "masked_gemm"],
1422
),
1423
+ # Override the directory for the FlashInfer autotune config cache.
1424
+ "VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR": lambda: os.getenv(
1425
+ "VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR", None
1426
+ ),
1427
# Flashinfer fused allreduce backend.
1428
"VLLM_FLASHINFER_ALLREDUCE_BACKEND": env_with_choices(
1429
"VLLM_FLASHINFER_ALLREDUCE_BACKEND",
@@ -1933,6 +1938,7 @@ def compile_factors() -> dict[str, object]:
1938
"VLLM_LOG_STATS_INTERVAL",
1939
"VLLM_DEBUG_LOG_API_SERVER_RESPONSE",
1940
"VLLM_TUNED_CONFIG_FOLDER",
1941
+ "VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR",
1942
"VLLM_ENGINE_ITERATION_TIMEOUT_S",
1943
"VLLM_HTTP_TIMEOUT_KEEP_ALIVE",
1944
"VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS",
vllm/model_executor/warmup/kernel_warmup.py
+36
-16
@@ -6,11 +6,14 @@ This is useful specifically for JIT'ed kernels as we don't want JIT'ing to
6
happen during model execution.
7
"""
8
9
+import hashlib
10
+from pathlib import Path
11
from typing import TYPE_CHECKING
12
13
import torch
14
15
import vllm.envs as envs
16
+from vllm.compilation.caching import aot_compile_hash_factors
17
from vllm.logger import init_logger
18
from vllm.model_executor.warmup.deep_gemm_warmup import deep_gemm_warmup
19
from vllm.platforms import current_platform
@@ -24,6 +27,31 @@ if TYPE_CHECKING:
27
logger = init_logger(__name__)
28
29
30
+def _flashinfer_autotune_cache_hash(runner: "GPUModelRunner") -> str:
31
+ factors = aot_compile_hash_factors(runner.vllm_config)
32
+ return hashlib.sha256(str(factors).encode()).hexdigest()
33
+
34
+
35
+def _resolve_flashinfer_autotune_file(runner: "GPUModelRunner") -> Path:
36
+ override_dir = envs.VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR
37
+ if override_dir:
38
+ root = Path(override_dir).expanduser()
39
+ else:
40
+ from flashinfer.jit import env as flashinfer_jit_env
41
+
42
+ flashinfer_workspace = flashinfer_jit_env.FLASHINFER_WORKSPACE_DIR
43
+ root = (
44
+ Path(envs.VLLM_CACHE_ROOT)
45
+ / "flashinfer_autotune_cache"
46
+ / flashinfer_workspace.parent.name
47
+ / flashinfer_workspace.name
48
+ )
49
+
50
+ output_dir = root / _flashinfer_autotune_cache_hash(runner)
51
+ output_dir.mkdir(parents=True, exist_ok=True)
52
+ return output_dir / "autotune_configs.json"
53
+
54
+
55
def kernel_warmup(worker: "Worker"):
56
# Deep GEMM warmup
57
do_deep_gemm_warmup = (
@@ -91,17 +119,15 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:
119
Tuning is performed only on rank 0. The resulting cache is broadcast
120
to every rank so all ranks dispatch the same kernel tactic.
121
"""
94
- import os
95
- import tempfile
96
-
122
import vllm.utils.flashinfer as fi_utils
123
from vllm.distributed.parallel_state import get_world_group
124
125
world = get_world_group()
126
is_leader = world.rank_in_group == 0
127
103
- cache_dir = tempfile.mkdtemp(prefix="vllm_flashinfer_autotune_")
104
- cache_path = os.path.join(cache_dir, "autotune_cache.json")
128
+ cache_path = _resolve_flashinfer_autotune_file(runner)
129
+ if is_leader:
130
+ logger.info("Using FlashInfer autotune cache file: %s", cache_path)
131
132
# We skip EPLB here since we don't want to record dummy metrics.
133
# When autotuning with number of tokens m, flashinfer will autotune
@@ -115,7 +141,7 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:
141
142
with torch.inference_mode():
143
if is_leader:
118
- with fi_utils.autotune(tune_mode=True, cache=cache_path):
144
+ with fi_utils.autotune(tune_mode=True, cache=str(cache_path)):
145
runner._dummy_run(**dummy_run_kwargs)
146
else:
147
runner._dummy_run(**dummy_run_kwargs)
@@ -123,7 +149,7 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:
149
# Broadcast autotune cache from rank 0 to all other ranks so every
150
# rank loads the same set of chosen tactics.
151
tune_results: bytes | None = None
126
- if is_leader and os.path.exists(cache_path):
152
+ if is_leader and cache_path.exists():
153
with open(cache_path, "rb") as f:
154
tune_results = f.read()
155
@@ -135,21 +161,15 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:
161
"Falling back to default tactics."
162
)
163
else:
138
- if not is_leader:
164
+ if not is_leader and world.local_rank == 0:
165
with open(cache_path, "wb") as f:
166
f.write(tune_results)
167
+ world.barrier()
168
from flashinfer.autotuner import AutoTuner
169
143
- AutoTuner.get().load_configs(cache_path)
170
+ AutoTuner.get().load_configs(str(cache_path))
171
logger.info(
172
"FlashInfer autotune cache loaded on rank %d from %s.",
173
world.rank_in_group,
174
cache_path,
175
)
149
-
150
- try:
151
- if os.path.exists(cache_path):
152
- os.unlink(cache_path)
153
- os.rmdir(cache_dir)
154
- except OSError:
155
- pass