[FlashInfer] Migrate autotune warmup to flashinfer.autotune_v2 managed cache (#55671)

Signed-off-by: Yang Xu <yanxu@nvidia.com>
Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
Signed-off-by: YangXu <yanxu@nvidia.com>
Co-authored-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
Co-authored-by: Kimi Code <noreply@moonshot.cn>
Co-authored-by: Codex <noreply@openai.com>
Co-authored-by: Misha Goin <mgoin64@gmail.com>
Co-authored-by: vllm-agent <claw@inferact.ai>

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-version>/<arch>/<cache-hash>/` | FlashInfer autotune config cache. |
+| `VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR` | Managed: `~/.cache/flashinfer/autotune/`; legacy: `$VLLM_CACHE_ROOT/flashinfer_autotune_cache/<flashinfer-version>/<arch>/<cache-hash>/` | 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. |

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,

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",

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,

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,6 +452,20 @@ def flashinfer_autotune(runner: "GPUModelRunner", *, skip_attn: bool = False) ->
         )
         autotune_kwargs["skip_ops"] = skip_ops
 
+    use_autotune_v2 = use_flashinfer_autotune_v2(runner)
+    if use_autotune_v2:
+        from flashinfer import autotune_v2, autotune_v2_reload
+
+        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",
+            )
+        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
@@ -459,11 +475,6 @@ def flashinfer_autotune(runner: "GPUModelRunner", *, skip_attn: bool = False) ->
         if is_leader:
             logger.info_once("Using FlashInfer autotune cache file: %s", cache_path)
 
-    # 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:
@@ -474,14 +485,12 @@ def flashinfer_autotune(runner: "GPUModelRunner", *, skip_attn: bool = False) ->
                 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))

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",