From f2376cdc4da852ec48ab5bcb747fdfa65a4dfd4a Mon Sep 17 00:00:00 2001 From: siyu Date: Sat, 10 Oct 2026 23:48:06 +0800 Subject: [PATCH] [Fast Start] Preload the FlashInfer autotune table on the weight cache daemon (#60085) Signed-off-by: liusy58 Signed-off-by: Isotr0py Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: Isotr0py --- docs/features/preload.md | 30 +++++++- .../model_loader/test_weight_cache.py | 73 ++++++++++++++++++- .../test_flashinfer_autotune_warmup.py | 31 ++++++++ vllm/entrypoints/cli/preload.py | 9 ++- .../model_loader/weight_cache/daemon.py | 39 ++++++++++ .../model_loader/weight_cache/utils.py | 35 ++++++++- .../warmup/flashinfer_autotune_cache.py | 16 +++- vllm/model_executor/warmup/kernel_warmup.py | 10 ++- vllm/v1/worker/gpu/model_runner.py | 12 ++- vllm/v1/worker/gpu_model_runner.py | 15 +++- 10 files changed, 253 insertions(+), 17 deletions(-) diff --git a/docs/features/preload.md b/docs/features/preload.md index 74839bf35b..6e8fad9ad0 100644 --- a/docs/features/preload.md +++ b/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 diff --git a/tests/model_executor/model_loader/test_weight_cache.py b/tests/model_executor/model_loader/test_weight_cache.py index 7960a23791..9628cb47e0 100644 --- a/tests/model_executor/model_loader/test_weight_cache.py +++ b/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 + # " (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 diff --git a/tests/model_executor/test_flashinfer_autotune_warmup.py b/tests/model_executor/test_flashinfer_autotune_warmup.py index f7ea6847f2..59046a5f3e 100644 --- a/tests/model_executor/test_flashinfer_autotune_warmup.py +++ b/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() + ) diff --git a/vllm/entrypoints/cli/preload.py b/vllm/entrypoints/cli/preload.py index 9eab7421f1..90e746fdb6 100644 --- a/vllm/entrypoints/cli/preload.py +++ b/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 " diff --git a/vllm/model_executor/model_loader/weight_cache/daemon.py b/vllm/model_executor/model_loader/weight_cache/daemon.py index 2c6944f1eb..0fc63da786 100644 --- a/vllm/model_executor/model_loader/weight_cache/daemon.py +++ b/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)) ) diff --git a/vllm/model_executor/model_loader/weight_cache/utils.py b/vllm/model_executor/model_loader/weight_cache/utils.py index 1abbf572f8..bc5de3cc97 100644 --- a/vllm/model_executor/model_loader/weight_cache/utils.py +++ b/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) diff --git a/vllm/model_executor/warmup/flashinfer_autotune_cache.py b/vllm/model_executor/warmup/flashinfer_autotune_cache.py index afc3ef3ac0..ca2c4d1421 100644 --- a/vllm/model_executor/warmup/flashinfer_autotune_cache.py +++ b/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() diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index e5751730e8..685b33a565 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/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: diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 674cf23527..e04a381b94 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/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, ) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 22ad662b90..90bea22c9d 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/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: