From 71c20f73f2f1764576fed799c152319c9fd1c820 Mon Sep 17 00:00:00 2001 From: Yang Xu <38851819+YangXu1990uiuc@users.noreply.github.com> Date: Fri, 16 Oct 2026 19:08:28 -0700 Subject: [PATCH] [FlashInfer] Migrate autotune warmup to flashinfer.autotune_v2 managed cache (#55671) Signed-off-by: Yang Xu Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Signed-off-by: YangXu Co-authored-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Co-authored-by: Kimi Code Co-authored-by: Codex Co-authored-by: Misha Goin Co-authored-by: vllm-agent --- docs/usage/security.md | 4 +- .../test_flashinfer_autotune_warmup.py | 271 +++++++++++++++++- .../warmup/flashinfer_autotune_cache.py | 50 ++++ .../warmup/flashinfer_sparse_mla_warmup.py | 62 +++- vllm/model_executor/warmup/kernel_warmup.py | 72 +++-- vllm/utils/flashinfer.py | 14 + 6 files changed, 436 insertions(+), 37 deletions(-) diff --git a/docs/usage/security.md b/docs/usage/security.md index f43ad0a6e3..4a704578fa 100644 --- a/docs/usage/security.md +++ b/docs/usage/security.md @@ -423,9 +423,9 @@ Most cache paths default to subdirectories under a single root. Changing `VLLM_C | Environment Variable | Default | Description | | --- | --- | --- | -| `VLLM_CACHE_ROOT` | `~/.cache/vllm` | Base cache directory. Respects `XDG_CACHE_HOME` if set. All paths below inherit from this unless explicitly overridden. | +| `VLLM_CACHE_ROOT` | `~/.cache/vllm` | Base directory for vLLM-owned caches. Respects `XDG_CACHE_HOME` if set. | | *(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). | -| `VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR` | `$VLLM_CACHE_ROOT/flashinfer_autotune_cache////` | FlashInfer autotune config cache. | +| `VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR` | Managed: `~/.cache/flashinfer/autotune/`; legacy: `$VLLM_CACHE_ROOT/flashinfer_autotune_cache////` | Overrides cache placement for both FlashInfer autotune cache APIs. | | `VLLM_ASSETS_CACHE` | `$VLLM_CACHE_ROOT/assets/` | Downloaded assets (e.g., tokenizer files). | | `VLLM_XLA_CACHE_PATH` | `$VLLM_CACHE_ROOT/xla_cache/` | XLA/TPU compilation cache. | | `VLLM_MEDIA_CACHE` | *(disabled)* | Optional cache for downloaded media (images, video, audio). Not enabled unless explicitly set. | diff --git a/tests/model_executor/test_flashinfer_autotune_warmup.py b/tests/model_executor/test_flashinfer_autotune_warmup.py index 59046a5f3e..f8c86cd1b5 100644 --- a/tests/model_executor/test_flashinfer_autotune_warmup.py +++ b/tests/model_executor/test_flashinfer_autotune_warmup.py @@ -4,7 +4,7 @@ import json import sys from collections import defaultdict -from contextlib import nullcontext +from contextlib import contextmanager, nullcontext from inspect import signature from pathlib import Path from types import ModuleType, SimpleNamespace @@ -12,8 +12,10 @@ from typing import Any from unittest.mock import Mock, call, patch import pytest +import torch -from vllm.model_executor.warmup import kernel_warmup as warmup +import vllm.model_executor.warmup.kernel_warmup as warmup +import vllm.utils.flashinfer as fi_utils from vllm.model_executor.warmup.kernel_warmup import ( _flashinfer_autotune_token_counts, _run_flashinfer_autotune_dummy_runs, @@ -45,7 +47,12 @@ def _make_runner(modules, *, max_tokens: int = 8192): scheduler_config=SimpleNamespace(max_num_batched_tokens=max_tokens), vllm_config=SimpleNamespace( attention_config=SimpleNamespace(hisparse_config=None), - parallel_config=SimpleNamespace(data_parallel_rank=0), + parallel_config=SimpleNamespace( + enable_elastic_ep=False, + nnodes=1, + pipeline_parallel_size=1, + data_parallel_rank=0, + ), ), get_model=Mock( return_value=SimpleNamespace(modules=Mock(return_value=modules)) @@ -133,6 +140,258 @@ def test_flashinfer_autotune_uses_token_buckets_for_each_dummy_run(skip_attn): ] +@pytest.mark.parametrize( + "has_v2,elastic_ep,node_count", + [(False, False, 1), (True, False, 1), (True, True, 1), (True, False, 2)], + ids=["v1", "v2", "elastic-v1", "multi-node-v1"], +) +@pytest.mark.parametrize( + "hisparse,world_size,rank,skip_ops,failure", + [ + (False, 1, 0, [], None), + (False, 2, 0, [], None), + (False, 2, 1, [], None), + (True, 2, 0, ["fp4_gemm", "bmm_fp8"], None), + (False, 2, 1, ["fp4_gemm"], "dummy"), + (False, 2, 1, [], "replayssm"), + (True, 2, 0, [], "kimi"), + ], +) +def test_flashinfer_autotune_lifecycle( + monkeypatch, + tmp_path, + has_v2, + elastic_ep, + node_count, + hisparse, + world_size, + rank, + skip_ops, + failure, +): + use_v2 = has_v2 and not elastic_ep and node_count == 1 and world_size == 1 + runner = _make_runner([]) + runner.vllm_config.parallel_config.enable_elastic_ep = elastic_ep + runner.vllm_config.attention_config.hisparse_config = object() if hisparse else None + model = runner.get_model() + cache_path = tmp_path / "autotune.json" + rank_cache_path = cache_path.with_name(f"autotune_dp0_rank{rank}.json") + rank_cache_path.write_text(json.dumps({"_metadata": {}, "op": 1})) + events = Mock() + events.tuner.load_configs.return_value = True + events.tuner._dirty = True + world = SimpleNamespace( + rank_in_group=rank, + world_size=world_size, + cpu_group=object(), + broadcast_object=events.broadcast, + barrier=events.barrier, + ) + + @contextmanager + def autotune_context(**kwargs): + assert torch.is_inference_mode_enabled() + events.enter() + try: + yield + finally: + events.exit() + + events.autotune.side_effect = autotune_context + if failure is not None: + getattr(events, failure).side_effect = RuntimeError("warmup failed") + monkeypatch.setitem( + sys.modules, + "flashinfer.autotuner", + SimpleNamespace( + AutoTuner=SimpleNamespace(get=lambda: events.tuner), + set_autotune_process_group=events.set_group, + ), + ) + monkeypatch.setitem( + sys.modules, + "flashinfer", + SimpleNamespace(autotune_v2=events.autotune, autotune_v2_reload=events.reload), + ) + monkeypatch.setattr( + "vllm.distributed.parallel_state.get_world_group", lambda: world + ) + monkeypatch.setattr( + "vllm.distributed.parallel_state.get_node_count", lambda: node_count + ) + monkeypatch.setattr( + "vllm.distributed.parallel_state.get_pp_group", + lambda: SimpleNamespace(world_size=1), + ) + monkeypatch.setattr(fi_utils, "has_flashinfer_autotune_v2", lambda: has_v2) + monkeypatch.setattr(fi_utils, "autotune", events.autotune) + monkeypatch.setattr(warmup.envs, "VLLM_FLASHINFER_AUTOTUNE_SKIP_OPS", skip_ops) + resolve_file = Mock(return_value=cache_path) + resolve_root = Mock(return_value=str(tmp_path)) + monkeypatch.setattr(warmup, "resolve_flashinfer_autotune_file", resolve_file) + monkeypatch.setattr(warmup, "resolve_flashinfer_autotune_v2_root", resolve_root) + + def gather(out, value, *, group): + out[:] = [value] * world_size + + monkeypatch.setattr(torch.distributed, "all_gather_object", gather) + monkeypatch.setattr(warmup, "_run_flashinfer_autotune_dummy_runs", events.dummy) + monkeypatch.setattr(warmup, "replayssm_autotune_warmup", events.replayssm) + monkeypatch.setattr(warmup, "_autotune_kimi_k3_kda_qkvg", events.kimi) + monkeypatch.setattr( + warmup, "autotune_hisparse_flashinfer_attention", events.hisparse + ) + + if failure is None: + warmup.flashinfer_autotune(runner) + else: + with pytest.raises(RuntimeError, match="warmup failed"): + warmup.flashinfer_autotune(runner) + + kwargs = {"skip_ops": set(skip_ops)} if skip_ops else {} + expected = [] + if use_v2: + resolve_file.assert_not_called() + resolve_root.assert_called_once_with() + expected.append(call.autotune(mode="tune", cache_root=str(tmp_path), **kwargs)) + else: + resolve_root.assert_not_called() + resolve_file.assert_called_once_with(runner) + expected.extend( + [ + call.tuner.load_configs(str(rank_cache_path)), + call.autotune(tune_mode=True, **kwargs), + ] + ) + expected.extend( + [call.set_group(world.cpu_group if world_size > 1 else None), call.enter()] + ) + if hisparse: + expected.append(call.hisparse(runner)) + for name, arg, kwargs in [ + ("dummy", runner, {"skip_attn": hisparse}), + ("replayssm", runner, {}), + ("kimi", model, {}), + ]: + expected.append(getattr(call, name)(arg, **kwargs)) + if failure == name: + break + expected.append(call.exit()) + expected.append(call.set_group(None)) + if failure is None: + if world_size > 1: + expected.append(call.barrier()) + if use_v2 and world_size > 1: + expected.extend([call.reload(), call.barrier()]) + elif not use_v2: + expected.append(call.tuner.save_configs(str(rank_cache_path))) + assert events.mock_calls == expected + + +@pytest.mark.parametrize( + "node_count,elastic_ep,pp_size", + [(1, False, 1), (2, False, 1), (1, True, 1), (1, False, 2)], +) +@pytest.mark.parametrize("world_size,rank", [(1, 0), (2, 0), (2, 1)]) +@pytest.mark.parametrize("model_runner_v2,max_reqs", [(False, 2), (True, 1), (True, 2)]) +def test_sparse_autotune_respects_cache_selection( + monkeypatch, + tmp_path, + node_count, + elastic_ep, + pp_size, + world_size, + rank, + model_runner_v2, + max_reqs, +): + import vllm.model_executor.warmup.flashinfer_sparse_mla_warmup as sparse + + runner = _make_runner([]) + runner.max_num_reqs = max_reqs + runner.vllm_config.use_v2_model_runner = model_runner_v2 + runner.vllm_config.kernel_config = SimpleNamespace(enable_flashinfer_autotune=True) + runner.vllm_config.parallel_config.enable_elastic_ep = elastic_ep + runner.vllm_config.parallel_config.pipeline_parallel_size = pp_size + worker = SimpleNamespace( + model_runner=runner, + vllm_config=runner.vllm_config, + execute_model=Mock(), + sample_tokens=Mock(), + ) + world = SimpleNamespace( + rank_in_group=rank, + world_size=world_size, + broadcast_object=Mock(return_value=b"cached"), + barrier=Mock(), + ) + managed = Mock(return_value=nullcontext()) + legacy = Mock(return_value=nullcontext()) + reload = Mock() + tuner = Mock() + monkeypatch.setitem( + sys.modules, + "flashinfer", + SimpleNamespace(autotune_v2=managed, autotune_v2_reload=reload), + ) + monkeypatch.setitem( + sys.modules, + "flashinfer.autotuner", + SimpleNamespace(AutoTuner=SimpleNamespace(get=lambda: tuner)), + ) + monkeypatch.setattr( + "vllm.distributed.parallel_state.get_world_group", lambda: world + ) + monkeypatch.setattr( + "vllm.distributed.parallel_state.get_node_count", lambda: node_count + ) + monkeypatch.setattr(fi_utils, "has_flashinfer_autotune_v2", lambda: True) + monkeypatch.setattr(sparse, "has_flashinfer", lambda: True) + monkeypatch.setattr( + sparse.current_platform, "is_device_capability_family", lambda _: True + ) + monkeypatch.setattr( + sparse, "_flashinfer_sparse_mla_decode_label", lambda *_: "DSv3.2" + ) + monkeypatch.setattr(sparse, "resolve_flashinfer_autotune_v2_root", lambda: tmp_path) + cache_path = tmp_path / "autotune.json" + cache_path.write_bytes(b"cached") + monkeypatch.setattr( + sparse, "resolve_flashinfer_autotune_file", lambda _: cache_path + ) + monkeypatch.setattr(sparse, "flashinfer_autotune", legacy) + + def mixed_warmup(*args, mixed_step_context=None, **kwargs): + context = ( + mixed_step_context if mixed_step_context is not None else nullcontext() + ) + with context: + assert torch.is_inference_mode_enabled() + return True + + mixed = Mock(side_effect=mixed_warmup) + monkeypatch.setattr(sparse, "run_mixed_prefill_decode_warmup", mixed) + assert sparse._run_flashinfer_sparse_mla_decode_autotune(worker, 16, frozenset()) + + if node_count == 1 and not elastic_ep and pp_size == 1 and world_size == 1: + managed.assert_called_once_with(mode="tune", cache_root=tmp_path) + legacy.assert_not_called() + world.broadcast_object.assert_not_called() + tuner.load_configs.assert_not_called() + reload.assert_not_called() + world.barrier.assert_not_called() + else: + managed.assert_not_called() + reload.assert_not_called() + assert legacy.call_count == (1 if rank == 0 else 0) + world.broadcast_object.assert_called_once_with( + b"cached" if rank == 0 else None, src=0 + ) + tuner.load_configs.assert_called_once_with(str(cache_path)) + assert mixed.call_count == (1 if model_runner_v2 and max_reqs >= 2 else 0) + assert runner._dummy_run.call_count == (0 if mixed.called else 1) + + def test_flashinfer_autotune_buckets_cover_drafter_tokens(): """The drafter's M can exceed the pass size; buckets must include it.""" runner = _make_runner([]) @@ -266,7 +525,9 @@ class _AutotuneRun: for rank in range(self.pp * self.tp): self.rank = rank self.tuners[rank] = _AutotuneTuner(self) - flashinfer_autotune(_make_runner([])) + runner = _make_runner([]) + runner.vllm_config.parallel_config.pipeline_parallel_size = self.pp + flashinfer_autotune(runner) return self def assert_collectives_match(self): @@ -283,6 +544,8 @@ def autotune_run(monkeypatch, tmp_path): def make_run(*, pp=2, tp=4): run = _AutotuneRun(pp, tp, tmp_path) + monkeypatch.setattr(fi_utils, "has_flashinfer_autotune_v2", lambda: pp > 1) + monkeypatch.setattr(parallel_state, "get_node_count", lambda: 1) autotuner = ModuleType("flashinfer.autotuner") monkeypatch.setattr( autotuner, diff --git a/vllm/model_executor/warmup/flashinfer_autotune_cache.py b/vllm/model_executor/warmup/flashinfer_autotune_cache.py index ca2c4d1421..3c41743033 100644 --- a/vllm/model_executor/warmup/flashinfer_autotune_cache.py +++ b/vllm/model_executor/warmup/flashinfer_autotune_cache.py @@ -11,11 +11,14 @@ from typing import TYPE_CHECKING import vllm.envs as envs from vllm.config import CacheConfig, VllmConfig, replace +from vllm.logger import init_logger if TYPE_CHECKING: from vllm.distributed.parallel_state import GroupCoordinator from vllm.v1.worker.gpu_model_runner import GPUModelRunner +logger = init_logger(__name__) + def _normalize_cache_config_for_hash(cache_config: CacheConfig) -> CacheConfig: """Normalize cache config for hashing""" @@ -55,6 +58,53 @@ def resolve_flashinfer_autotune_file(runner: "GPUModelRunner") -> Path: return output_dir / "autotune_configs.json" +def use_flashinfer_autotune_v2(runner: "GPUModelRunner") -> bool: + from vllm.distributed.parallel_state import get_node_count, get_world_group + from vllm.utils.flashinfer import has_flashinfer_autotune_v2 + + # Elastic EP transfers tuning state to joining workers with v1 + # save_configs/load_configs, which do not export managed-store winners. + if runner.vllm_config.parallel_config.enable_elastic_ep: + return False + # PP stages profile independently. Keep their group-scoped v1 files until + # managed-cache reads and publishes can be isolated across tuning groups. + if runner.vllm_config.parallel_config.pipeline_parallel_size > 1: + return False + if not has_flashinfer_autotune_v2(): + return False + # Use the initialized topology: Ray may span nodes with config.nnodes=1. + # A cache hit skips profiling collectives, so node-local stores with + # different contents cannot safely participate in synchronized tuning. + if get_node_count() != 1: + logger.info_once( + "Using legacy FlashInfer autotune cache synchronization for " + "multi-node deployments; managed-cache synchronization is " + "currently supported only within one node." + ) + return False + # Lift this gate after pinning a release with the fix below and validating + # distributed managed-cache tuning: https://github.com/flashinfer-ai/flashinfer/pull/6185 + if get_world_group().world_size > 1: + logger.info_once( + "Using legacy FlashInfer autotune cache synchronization for " + "multi-rank deployments until FlashInfer coordinates cache hits." + ) + return False + return True + + +def resolve_flashinfer_autotune_v2_root() -> Path | None: + """Placement root for FlashInfer's managed autotune store. + + Unlike :func:`resolve_flashinfer_autotune_file`, no vLLM-side identity + (config hash, flashinfer version, arch) is folded into the path: the + managed store keys entries per op and hashes the software/hardware + environment itself. ``None`` selects FlashInfer's default root. + """ + override_dir = envs.VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR + return Path(override_dir).expanduser() if override_dir else None + + def sync_flashinfer_autotune_cache( runner: "GPUModelRunner", group: "GroupCoordinator", diff --git a/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py b/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py index 5c01f4ca2e..fee6a84762 100644 --- a/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py +++ b/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py @@ -9,6 +9,8 @@ import torch from vllm.logger import init_logger from vllm.model_executor.warmup.flashinfer_autotune_cache import ( resolve_flashinfer_autotune_file, + resolve_flashinfer_autotune_v2_root, + use_flashinfer_autotune_v2, write_flashinfer_autotune_cache, ) from vllm.platforms import current_platform @@ -145,7 +147,6 @@ def _run_flashinfer_sparse_mla_decode_autotune( world = get_world_group() is_leader = world.rank_in_group == 0 - cache_path = resolve_flashinfer_autotune_file(runner) dummy_run_kwargs = dict( num_tokens=num_tokens, @@ -155,6 +156,13 @@ def _run_flashinfer_sparse_mla_decode_autotune( create_mixed_batch=True, ) + if use_flashinfer_autotune_v2(runner): + return _run_sparse_mla_decode_autotune_v2( + worker, num_tokens, log_label, dummy_run_kwargs + ) + + cache_path = resolve_flashinfer_autotune_file(runner) + if is_leader: logger.info( "Autotuning FlashInfer SM120 sparse MLA %s decode with cache: %s", @@ -223,6 +231,58 @@ def _run_flashinfer_sparse_mla_decode_autotune( return True +def _run_sparse_mla_decode_autotune_v2( + worker: "Worker", + num_tokens: int, + log_label: str, + dummy_run_kwargs: dict, +) -> bool: + """Run every rank's sparse warmup in the managed autotune context. + + FlashInfer owns persistence; some sparse kernels use a separate calibration + cache instead of managed winner entries. + """ + from flashinfer import autotune_v2, autotune_v2_reload + + from vllm.distributed.parallel_state import get_world_group + + runner = worker.model_runner + world = get_world_group() + cache_root = resolve_flashinfer_autotune_v2_root() + if world.rank_in_group == 0: + logger.info( + "Autotuning FlashInfer SM120 sparse MLA %s decode with the managed " + "autotune context (root=%s)", + log_label, + cache_root if cache_root is not None else "flashinfer default", + ) + + with torch.inference_mode(): + if _uses_v2_model_runner(runner) and runner.max_num_reqs >= 2: + v2_runner = cast("V2GPUModelRunner", runner) + warmup_executed = run_mixed_prefill_decode_warmup( + v2_runner, + worker.execute_model, + worker.sample_tokens, + num_tokens, + mixed_step_context=autotune_v2(mode="tune", cache_root=cache_root), + req_id_prefix="_sparse_mla_v2_warmup", + ) + else: + warmup_executed = True + with autotune_v2(mode="tune", cache_root=cache_root): + runner._dummy_run(**dummy_run_kwargs) + + if not warmup_executed: + return False + + if world.world_size > 1: + world.barrier() + autotune_v2_reload() + world.barrier() + return True + + def _flashinfer_sparse_mla_decode_autotune( worker: "Worker", num_tokens: int, diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index b7bf50865f..ba97d2dfa3 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -23,6 +23,8 @@ from vllm.model_executor.warmup.cutedsl_warmup import cutedsl_warmup from vllm.model_executor.warmup.deep_gemm_warmup import deep_gemm_warmup from vllm.model_executor.warmup.flashinfer_autotune_cache import ( resolve_flashinfer_autotune_file, + resolve_flashinfer_autotune_v2_root, + use_flashinfer_autotune_v2, ) from vllm.model_executor.warmup.flashinfer_sparse_mla_warmup import ( autotune_hisparse_flashinfer_attention, @@ -419,7 +421,7 @@ def flashinfer_autotune(runner: "GPUModelRunner", *, skip_attn: bool = False) -> otherwise the world group tunes together. Per-tactic timings are averaged over the tuning group so all its ranks select the same tactic. - Results are persisted per rank: FlashInfer keys MoE entries by tp/ep rank + Legacy results are persisted per rank: FlashInfer keys MoE entries by tp/ep rank (``MoERunner.get_cache_key_extras``), so one rank's file only hits on that 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 @@ -450,38 +452,45 @@ def flashinfer_autotune(runner: "GPUModelRunner", *, skip_attn: bool = False) -> ) autotune_kwargs["skip_ops"] = skip_ops - cache_path = resolve_flashinfer_autotune_file(runner) - # The world group only spans DP ranks when vLLM folds DP into it. - dp_rank = runner.vllm_config.parallel_config.data_parallel_rank - cache_path = cache_path.with_name( - f"{cache_path.stem}_dp{dp_rank}_rank{world.rank_in_group}{cache_path.suffix}" - ) - if is_leader: - logger.info_once("Using FlashInfer autotune cache file: %s", cache_path) + use_autotune_v2 = use_flashinfer_autotune_v2(runner) + if use_autotune_v2: + from flashinfer import autotune_v2, autotune_v2_reload - # We skip EPLB here since we don't want to record dummy metrics. - # Randomize inputs to avoid every token pick the same experts, - # which lead to some EP ranks receiving no tokens and skipping their - # MoE kernel entirely, and cause hang due to all-reduce collective - # during synchronized autotuning. - if _all_ranks_have_matching_cache(cache_path, tune_group): - loaded = tuner.load_configs(str(cache_path)) - if tune_group.world_size > 1: - loaded_by_rank: list[bool | None] = [None] * tune_group.world_size - torch.distributed.all_gather_object( - loaded_by_rank, loaded, group=tune_group.cpu_group + cache_root = resolve_flashinfer_autotune_v2_root() + if is_leader: + logger.info_once( + "Using FlashInfer managed autotune cache (root=%s)", + cache_root if cache_root is not None else "flashinfer default", ) - loaded = all(loaded_by_rank) - if not loaded: - tuner.clear_cache() + autotune_context = autotune_v2( + mode="tune", cache_root=cache_root, **autotune_kwargs + ) + else: + cache_path = resolve_flashinfer_autotune_file(runner) + # The world group only spans DP ranks when vLLM folds DP into it. + dp_rank = runner.vllm_config.parallel_config.data_parallel_rank + cache_path = cache_path.with_name( + f"{cache_path.stem}_dp{dp_rank}_rank{world.rank_in_group}{cache_path.suffix}" + ) + if is_leader: + logger.info_once("Using FlashInfer autotune cache file: %s", cache_path) + + if _all_ranks_have_matching_cache(cache_path, tune_group): + loaded = tuner.load_configs(str(cache_path)) + if tune_group.world_size > 1: + loaded_by_rank: list[bool | None] = [None] * tune_group.world_size + torch.distributed.all_gather_object( + loaded_by_rank, loaded, group=tune_group.cpu_group + ) + loaded = all(loaded_by_rank) + if not loaded: + tuner.clear_cache() + autotune_context = fi_utils.autotune(tune_mode=True, **autotune_kwargs) group = tune_group.cpu_group if tune_group.world_size > 1 else None set_autotune_process_group(group) try: - with ( - torch.inference_mode(), - fi_utils.autotune(tune_mode=True, **autotune_kwargs), - ): + with torch.inference_mode(), autotune_context: hisparse_enabled = ( runner.vllm_config.attention_config.hisparse_config is not None ) @@ -498,7 +507,10 @@ def flashinfer_autotune(runner: "GPUModelRunner", *, skip_attn: bool = False) -> if world.world_size > 1: world.barrier() - # Skip the rewrite when nothing was tuned this start (every entry came from - # the file). FlashInfer gates its own autotune(cache=...) save the same way. - if tuner._dirty: + if use_autotune_v2: + if world.world_size > 1: + # All ranks serve the shared store's final state after publishing. + autotune_v2_reload() + world.barrier() + elif tuner._dirty: tuner.save_configs(str(cache_path)) diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index 69401febea..54fa01ed69 100644 --- a/vllm/utils/flashinfer.py +++ b/vllm/utils/flashinfer.py @@ -358,6 +358,19 @@ autotune = _lazy_import_wrapper( ) +@functools.cache +def has_flashinfer_autotune_v2() -> bool: + """Return ``True`` if FlashInfer ships the managed-cache autotuner + (``flashinfer.autotune_v2``, flashinfer-ai/flashinfer#3861).""" + if not has_flashinfer(): + return False + mod = _get_submodule("flashinfer") + return mod is not None and all( + callable(getattr(mod, name, None)) + for name in ("autotune_v2", "autotune_v2_reload") + ) + + @functools.cache def has_flashinfer_comm() -> bool: """Return `True` if FlashInfer comm module is available.""" @@ -1269,6 +1282,7 @@ __all__ = [ "flashinfer_fused_kda_decode", "autotune", "has_flashinfer_moe", + "has_flashinfer_autotune_v2", "has_flashinfer_comm", "has_flashinfer_nvlink_two_sided", "has_flashinfer_nvlink_one_sided",