[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>

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()
+    )