[Fast Start] Preload the FlashInfer autotune table on the weight cache daemon (#60085)

Signed-off-by: liusy58 <mg21330037@smail.nju.edu.cn>
Signed-off-by: Isotr0py <Isotr0py@outlook.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: Isotr0py <Isotr0py@outlook.com>

docs/features/preload.md

@@ -50,6 +50,26 @@ vllm serve meta-llama/Llama-3.1-8B-Instruct \
 The daemon itself must load from disk; passing `--load-format ipc_cache` to
 `vllm preload` is an error.
 
+## Preloading the FlashInfer autotune cache
+
+FlashInfer has several implementations of each operation and chooses between
+them by benchmarking. vLLM runs that pass during kernel warmup and persists the
+result per rank in the on-disk FlashInfer autotune cache (under
+`VLLM_CACHE_ROOT`, or `VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR`), so a later engine
+with the same configuration loads the tuned configs instead of profiling again.
+Normally the first serving engine pays for the pass, and on a large MoE model it
+dominates what is left of the startup time once the weights come from the
+daemon.
+
+Unless FlashInfer autotune is disabled (`--no-enable-flashinfer-autotune`),
+`vllm preload` pays it at preload time: each daemon runs the autotune pass in
+place — profiling dummy runs against the shard it already holds — after
+loading and before binding its socket, then writes the cache:
+
+```bash
+vllm preload --model /path/to/model --tensor-parallel-size 4
+```
+
 ## How it works
 
 1. `vllm preload` spawns one daemon process per GPU. Each daemon loads its
@@ -164,8 +184,14 @@ Socket paths are derived from the GPU UUID, so they are stable regardless of
 - **Platform**: CUDA and ROCm only; other platforms raise
   `UnsupportedPlatformForIPCError` even when `fallback` is enabled, since it
   is a permanent misconfiguration rather than a transient daemon outage.
-- **Parallelism**: tensor, expert and data parallelism are supported;
-  launching the daemon with pipeline parallelism is rejected.
+- **Parallelism**: tensor, expert, pipeline and data parallelism are
+  supported, including across nodes.
+- **Autotune preloading**: tuning runs across the daemon's world group,
+  including across nodes and DP ranks. With model-based speculative decoding
+  both daemon groups skip it (the target cannot run the drafter's dummy
+  passes, and a table tuned on the draft model would key on the draft's
+  config, which no engine looks up), so those deployments' first engine
+  tunes and fills the cache instead.
 - **Quantization**: every quantization method in the model must declare
   support for pre-processed weights (the daemon transfers weights *after*
   quantization post-processing). Unsupported methods raise

tests/model_executor/model_loader/test_weight_cache.py

@@ -12,12 +12,14 @@ import subprocess
 import sys
 import tempfile
 import threading
+import time
 import urllib.error
 import urllib.request
 from dataclasses import dataclass, field
 from typing import Any
 
 import pytest
+import regex as re
 
 from vllm import SamplingParams
 from vllm.assets.image import ImageAsset
@@ -81,6 +83,21 @@ class WeightCacheDaemon:
                 f"Weight cache daemon is not ready: HTTP {error.code}"
             ) from error
 
+    def wait_until_ready(self, timeout_s: float = 900) -> None:
+        """Poll until READY (which covers the autotune tuners) or the preload
+        process dies."""
+        deadline = time.monotonic() + timeout_s
+        while True:
+            try:
+                self.check_health()
+                return
+            except (AssertionError, urllib.error.URLError):
+                if time.monotonic() > deadline:
+                    raise
+                if self._proc is not None and self._proc.poll() is not None:
+                    raise AssertionError("vllm preload exited before ready") from None
+                time.sleep(2)
+
     def _stop(self) -> None:
         assert self._proc is not None
         self._proc.terminate()
@@ -99,6 +116,10 @@ class ModelCase:
     images: list | None = None
     llm_kwargs: dict[str, Any] = field(default_factory=dict)
     daemon_args: list[str] = field(default_factory=list)
+    # Additionally assert the daemons tune in place before serving and write
+    # the on-disk FlashInfer autotune cache, and the restart engine reuses it
+    # (needs SM90+/FlashInfer).
+    check_flashinfer_cache: bool = False
 
 
 def generate(
@@ -157,6 +178,9 @@ K3_CASE = ModelCase(
         max_model_len=4096,
     ),
     daemon_args=["--trust-remote-code"],
+    # Kimi K3 has tunable FlashInfer ops (a dense model like Qwen2.5 tunes
+    # nothing and upstream never writes the cache file for it).
+    check_flashinfer_cache=True,
 )
 
 # Qwen3.5-0.8B ships one MTP layer in the target checkpoint, so method="mtp"
@@ -186,18 +210,38 @@ QWEN_MTP_CASE = ModelCase(
     [QWEN_CASE, K3_CASE, QWEN_MTP_CASE],
     ids=["qwen3.5", "kimi-k3", "qwen3.5-mtp"],
 )
-def test_ipc_cache_cold_start_and_warm_restart(vllm_runner, case: ModelCase):
+def test_ipc_cache_cold_start_and_warm_restart(
+    vllm_runner, case: ModelCase, tmp_path, monkeypatch, capfd
+):
     """Cold start falls back to disk; warm restarts load weights via CUDA IPC.
 
     All runs must produce outputs identical to a default-loader baseline. The
     warm runs disable the disk fallback, so they only pass if the weights
     really came from the daemon — for the MTP case, both the target's and the
     draft's daemon groups.
+
+    Cases with ``check_flashinfer_cache`` additionally cover in-daemon
+    autotune (default-on with `--enable-flashinfer-autotune`): each daemon
+    tunes before binding its socket and must write the on-disk FlashInfer
+    autotune cache, and the restart engine must load tactics from it instead
+    of re-profiling. The daemon's cache file is keyed with the
+    OPENAI_API_SERVER batch defaults while these engines run under LLM_CLASS,
+    so they never hit the daemon's file by design; reuse is proven against
+    the file the earlier in-process engines wrote.
     """
     if not current_platform.is_cuda_alike():
         pytest.skip("Weight cache IPC sharing requires CUDA or ROCm")
     if case is K3_CASE and not current_platform.is_device_capability_family(100):
         pytest.skip("Kimi K3 IPC weight cache requires an SM100 MXFP4 backend")
+    if case.check_flashinfer_cache:
+        from vllm.utils.flashinfer import has_flashinfer
+
+        if not (has_flashinfer() and current_platform.has_device_capability(90)):
+            pytest.skip("FlashInfer autotune requires FlashInfer and SM90+")
+        # Isolate the on-disk autotune cache. The engines are in-process
+        # (envs read lazily); the daemon subprocesses spawn later and
+        # inherit it.
+        monkeypatch.setenv("VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR", str(tmp_path))
 
     # Baseline: plain disk loading with the default loader.
     baseline_outputs = generate(
@@ -212,15 +256,40 @@ def test_ipc_cache_cold_start_and_warm_restart(vllm_runner, case: ModelCase):
     with tempfile.TemporaryDirectory(prefix="vllm_ipc_empty_") as empty_socket_dir:
         cold_outputs = generate(vllm_runner, case, empty_socket_dir, fallback=True)
 
+    if case.check_flashinfer_cache:
+        # The baseline engine already wrote its own (LLM_CLASS-keyed) table.
+        engine_cache_files = set(tmp_path.rglob("autotune_configs*.json"))
+
     with WeightCacheDaemon(
         case.model,
         tp_size=1,
         extra_args=case.daemon_args,
     ) as d:
         warm_outputs = generate(vllm_runner, case, d.socket_dir, fallback=False)
-        d.check_health()
+        d.wait_until_ready()
+        if case.check_flashinfer_cache:
+            # Each daemon tunes before its ready report; exactly one new file
+            # (their OPENAI_API_SERVER-keyed table) must appear.
+            tuner_files = set(tmp_path.rglob("autotune_configs*.json"))
+            tuner_files -= engine_cache_files
+            assert len(tuner_files) == 1, (
+                f"expected the tuner's tuned table, got {tuner_files}"
+            )
         # Warm restart: a second engine lifetime against the same daemon.
+        if case.check_flashinfer_cache:
+            capfd.readouterr()  # drain: the assertion reads the increment
         restart_outputs = generate(vllm_runner, case, d.socket_dir, fallback=False)
+        if case.check_flashinfer_cache:
+            # The restart's EngineCore subprocess inherits pytest's fds; the
+            # flashinfer autotuner logs "[Autotuner]: Loaded N configs from
+            # <path>" (INFO, not once-deduplicated) on a disk cache hit.
+            out = capfd.readouterr()
+            m = re.search(
+                r"\[Autotuner\]: Loaded (\d+) configs from", out.out + out.err
+            )
+            assert m and int(m.group(1)) > 0, (
+                "restart engine did not load the on-disk autotune cache"
+            )
 
     assert cold_outputs == baseline_outputs
     assert warm_outputs == baseline_outputs

tests/model_executor/test_flashinfer_autotune_warmup.py

@@ -386,3 +386,34 @@ def test_rejected_cache_retunes_every_rank_in_its_tuning_group(autotune_run, pp,
     rerun.assert_collectives_match()
     assert set(rerun.profile_groups) == set(range(tp))
     assert {rank for rank, _, _ in rerun.saves} == set(range(tp))
+
+
+def test_autotune_cache_key_normalizes_kv_init_fields():
+    """mamba_block_size/kv_cache_layout only resolve at engine KV-init time;
+    they must not feed the autotune cache key, or nothing can compute the key
+    before startup (the daemon tunes before any KV init). block_size must
+    stay in the key: the daemon replays the engine's block-size resolution
+    before tuning, and different block sizes may select different ops."""
+    from vllm.config.cache import CacheConfig
+    from vllm.model_executor.warmup.flashinfer_autotune_cache import (
+        _normalize_cache_config_for_hash,
+    )
+
+    fresh = CacheConfig()
+    # Simulate the engine's KV-init write-backs (hybrid model values).
+    mutated = CacheConfig()
+    mutated.mamba_block_size = 512
+    mutated.kv_cache_layout = "LBHNC"
+
+    # Negative control: the write-backs are hashed factors.
+    assert mutated.compute_hash() != fresh.compute_hash()
+    normalized = _normalize_cache_config_for_hash(mutated)
+    assert normalized.compute_hash() == fresh.compute_hash()
+    # Factory state is returned untouched.
+    assert _normalize_cache_config_for_hash(fresh) is fresh
+    # block_size stays significant (the daemon keys on the resolved value).
+    resolved = CacheConfig(block_size=64)
+    assert (
+        _normalize_cache_config_for_hash(resolved).compute_hash()
+        != fresh.compute_hash()
+    )

vllm/entrypoints/cli/preload.py

@@ -39,6 +39,7 @@ from vllm.logger import init_logger
 from vllm.model_executor.model_loader.weight_cache.protocol import (
     check_ipc_platform_support,
 )
+from vllm.usage.usage_lib import UsageContext
 from vllm.utils.network_utils import get_distributed_init_method, get_open_port
 
 if typing.TYPE_CHECKING:
@@ -157,7 +158,13 @@ class PreloadSubcommand(CLISubcommand):
         )
 
         engine_args = EngineArgs.from_cli_args(args)
-        vllm_config = engine_args.create_engine_config()
+        # Warmup context for flashinfer autotune.
+        usage_context = (
+            UsageContext.OPENAI_API_SERVER
+            if engine_args.kernel_config.enable_flashinfer_autotune is not False
+            else None
+        )
+        vllm_config = engine_args.create_engine_config(usage_context=usage_context)
         if vllm_config.load_config.load_format == "ipc_cache":
             raise ValueError(
                 "The weight cache daemon itself must load from disk; use the "

vllm/model_executor/model_loader/weight_cache/daemon.py

@@ -112,11 +112,14 @@ from vllm.model_executor.model_loader.weight_cache.protocol import (
     verify_peer_is_owner,
 )
 from vllm.model_executor.model_loader.weight_cache.utils import (
+    build_warmup_runner,
     export_model_attrs,
     format_daemon_role,
     is_draft_model_cacheable,
 )
+from vllm.model_executor.warmup.kernel_warmup import flashinfer_autotune
 from vllm.platforms import current_platform
+from vllm.utils.flashinfer import has_flashinfer
 from vllm.utils.mem_utils import format_gib
 from vllm.utils.torch_utils import set_default_torch_dtype
 from vllm.v1.worker.workspace import init_workspace_manager
@@ -236,6 +239,41 @@ class WeightCacheDaemon:
             is_draft=is_draft,
         )
 
+    def warmup(self) -> None:
+        """Run the FlashInfer autotune pass against the cached model."""
+        vllm_config = self.vllm_config
+        # Mirrors the gate in the engine's kernel warmup.
+        if (
+            vllm_config.kernel_config.enable_flashinfer_autotune is False
+            or not has_flashinfer()
+            or not current_platform.has_device_capability(90)
+            or is_draft_model_cacheable(vllm_config.speculative_config)
+        ):
+            return
+        try:
+            current_platform.update_block_size_for_backend(vllm_config)
+            runner = build_warmup_runner(
+                vllm_config,
+                self.local_rank,
+                is_draft=self.is_draft,
+                model_config=self.model_config,
+            )
+            assert self.model is not None, "warmup ran before load_model"
+            runner.load_model(model=self.model)
+            # The V2 runner requires skip_attn when no KV cache exists
+            flashinfer_autotune(runner, skip_attn=vllm_config.use_v2_model_runner)
+            logger.info(
+                "Weight cache %s daemon rank %d tuned FlashInfer; the tuned "
+                "configs are in the on-disk autotune cache",
+                self.role,
+                self.global_rank,
+            )
+        finally:
+            # Free the runner and the dummy-run activations; the daemon's
+            # weights are untouched.
+            gc.collect()
+            torch.accelerator.empty_cache()
+
     def load_model(self) -> None:
         torch.accelerator.set_device_index(self.local_rank)
         # Outside the config context so the daemon keeps its own rendezvous
@@ -457,6 +495,7 @@ def _run_daemon(
         pp_rank,
     )
     daemon.load_model()
+    daemon.warmup()
     daemon.serve_forever(
         ready_callback=lambda: ready_queue.put((daemon.role, daemon.global_rank))
     )

vllm/model_executor/model_loader/weight_cache/utils.py

@@ -7,10 +7,17 @@ draft is cached, or how a daemon group is named, do not have to import the
 wire format.
 """
 
-from typing import Any
+from typing import TYPE_CHECKING, Any, cast
 
-from vllm.config import SpeculativeConfig
+import torch
+
+from vllm.config import SpeculativeConfig, VllmConfig, replace, set_current_vllm_config
 from vllm.model_executor.models.interfaces import SupportsEagleBase
+from vllm.platforms import current_platform
+
+if TYPE_CHECKING:
+    from vllm.config import ModelConfig
+    from vllm.v1.worker.gpu_model_runner import GPUModelRunner
 
 # Speculative methods whose draft model the daemon caches in its own group.
 # Other drafts keep loading from disk in the engine.
@@ -46,3 +53,27 @@ def format_daemon_role(is_draft: bool) -> str:
 def format_socket_role_suffix(is_draft: bool) -> str:
     """Socket-name suffix keeping the draft group distinct from the target."""
     return "_draft" if is_draft else ""
+
+
+def build_warmup_runner(
+    vllm_config: VllmConfig,
+    local_rank: int,
+    *,
+    is_draft: bool,
+    model_config: "ModelConfig",
+) -> "GPUModelRunner":
+    """Create a model runner for the daemon's warmup dummy runs."""
+    if is_draft:
+        vllm_config = replace(
+            vllm_config, model_config=model_config, speculative_config=None
+        )
+    device = torch.device(current_platform.device_type, local_rank)
+    with set_current_vllm_config(vllm_config):
+        if vllm_config.use_v2_model_runner:
+            from vllm.v1.worker.gpu.model_runner import GPUModelRunner as Runner
+        else:
+            from vllm.v1.worker.gpu_model_runner import GPUModelRunner as Runner
+        runner = Runner(vllm_config, device)
+    # The two runner classes share the warmup surface; the nominal type
+    # stays V1's GPUModelRunner.
+    return cast("GPUModelRunner", runner)

vllm/model_executor/warmup/flashinfer_autotune_cache.py

@@ -10,14 +10,28 @@ from pathlib import Path
 from typing import TYPE_CHECKING
 
 import vllm.envs as envs
+from vllm.config import CacheConfig, VllmConfig, replace
 
 if TYPE_CHECKING:
     from vllm.distributed.parallel_state import GroupCoordinator
     from vllm.v1.worker.gpu_model_runner import GPUModelRunner
 
 
+def _normalize_cache_config_for_hash(cache_config: CacheConfig) -> CacheConfig:
+    """Normalize cache config for hashing"""
+    if cache_config.mamba_block_size is None and cache_config.kv_cache_layout is None:
+        return cache_config
+    normalized = replace(cache_config, mamba_block_size=None)
+    assert normalized.kv_cache_layout is None
+    return normalized
+
+
 def flashinfer_autotune_cache_hash(runner: "GPUModelRunner") -> str:
-    config_hash = runner.vllm_config.compute_hash(include_version=False)
+    vllm_config: VllmConfig = runner.vllm_config
+    cache_config = _normalize_cache_config_for_hash(vllm_config.cache_config)
+    if cache_config is not vllm_config.cache_config:
+        vllm_config = replace(vllm_config, cache_config=cache_config)
+    config_hash = vllm_config.compute_hash(include_version=False)
     return hashlib.sha256(config_hash.encode()).hexdigest()
 
 

vllm/model_executor/warmup/kernel_warmup.py

@@ -406,7 +406,7 @@ def _all_ranks_have_matching_cache(path: Path, group) -> bool:
     return fingerprint is not None and all(f == fingerprint for f in gathered)
 
 
-def flashinfer_autotune(runner: "GPUModelRunner") -> None:
+def flashinfer_autotune(runner: "GPUModelRunner", *, skip_attn: bool = False) -> None:
     """Autotune FlashInfer operations.
     FlashInfer have many implementations for the same operation,
     autotuning runs benchmarks for each implementation and stores
@@ -415,6 +415,10 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:
     Without autotuning, FlashInfer will rely on heuristics, which may
     be significantly slower.
 
+    Args:
+        runner: The model runner to tune.
+        skip_attn: Additionally skip attention in the tuning dummy runs.
+
     With PP > 1, stages run different layers and may profile different ops,
     so each stage's TP group tunes separately with its own cache file;
     otherwise the world group tunes together. Per-tactic timings are
@@ -425,6 +429,7 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:
     rank. A rank with a cache hit skips the per-tactic reduce the others block
     in, so ranks keep loaded configs only if every rank in the tuning group
     has a matching file and successfully loads it.
+
     """
     from flashinfer.autotuner import AutoTuner, set_autotune_process_group
 
@@ -485,11 +490,12 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:
             hisparse_enabled = (
                 runner.vllm_config.attention_config.hisparse_config is not None
             )
+            skip_attn = skip_attn or hisparse_enabled
             if hisparse_enabled:
                 # HiSparse hot-buffer attention is bounded by decode batch
                 # size, not the prefill-sized batch used for the full model.
                 autotune_hisparse_flashinfer_attention(runner)
-            _run_flashinfer_autotune_dummy_runs(runner, skip_attn=hisparse_enabled)
+            _run_flashinfer_autotune_dummy_runs(runner, skip_attn=skip_attn)
             replayssm_autotune_warmup(runner)
             _autotune_kimi_k3_kda_qkvg(runner.get_model())
     finally:

vllm/v1/worker/gpu/model_runner.py

@@ -373,9 +373,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             tasks.extend(PoolingRunner.get_supported_tasks(self.model))
         return tuple(tasks)
 
-    def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None:
+    def load_model(
+        self,
+        load_dummy_weights: bool = False,
+        *args,
+        model: nn.Module | None = None,
+        **kwargs,
+    ) -> None:
         time_before_load = time.perf_counter()
-        if load_dummy_weights:
+        if model is None and load_dummy_weights:
             self.load_config.load_format = "dummy"
         self.eplb.prepare_load()
         eplb_models_added = False
@@ -385,7 +391,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
 
             # Capture warmup providers selected while constructing the model.
             with self.jit_warmup_registry.activate():
-                self.model = model_loader.load_model(
+                self.model = model or model_loader.load_model(
                     vllm_config=self.vllm_config,
                     model_config=self.vllm_config.model_config,
                 )

vllm/v1/worker/gpu_model_runner.py

@@ -5217,9 +5217,15 @@ class GPUModelRunner(
             setattr(self, config_name, new_config)
 
     @instrument(span_name="Loading (GPU)")
-    def load_model(self, load_dummy_weights: bool = False) -> None:
+    def load_model(
+        self,
+        load_dummy_weights: bool = False,
+        *,
+        model: nn.Module | None = None,
+    ) -> None:
         """Args:
         load_dummy_weights: load dummy weights instead of real weights.
+        model: a pre-loaded model to skip the model loader.
 
         """
         logger.info_once(
@@ -5235,13 +5241,14 @@ class GPUModelRunner(
         try:
             with DeviceMemoryProfiler() as m:
                 time_before_load = time.perf_counter()
-                if load_dummy_weights:
+                if model is None and load_dummy_weights:
                     self.load_config.load_format = "dummy"
                 model_loader = get_model_loader(self.load_config)
                 # Capture warmup providers selected while constructing the model.
                 with self.jit_warmup_registry.activate():  # type: ignore[attr-defined]
-                    self.model = model_loader.load_model(
-                        vllm_config=self.vllm_config, model_config=self.model_config
+                    self.model = model or model_loader.load_model(
+                        vllm_config=self.vllm_config,
+                        model_config=self.model_config,
                     )
                 lookback_depth = getattr(self.model, "token_lookback_depth", 0)
                 if lookback_depth > 0: