[Bugfix] Persist FlashInfer autotune cache per rank to fix the rank-0-only cache-hit deadlock (#57635)

Signed-off-by: Madhav Jivrajani <madhav.jiv@gmail.com>
Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
Co-authored-by: NickLucche <nicolo.lucchesi@mistral.ai>
Co-authored-by: Wei Zhao <51183510+wzhao18@users.noreply.github.com>

tests/model_executor/test_flashinfer_autotune_warmup.py

@@ -46,6 +46,7 @@ def _make_runner(modules, *, max_tokens: int = 8192, linear_backend: str = "auto
         vllm_config=SimpleNamespace(
             kernel_config=SimpleNamespace(linear_backend=linear_backend),
             attention_config=SimpleNamespace(hisparse_config=None),
+            parallel_config=SimpleNamespace(data_parallel_rank=0),
         ),
         get_model=Mock(
             return_value=SimpleNamespace(modules=Mock(return_value=modules))
@@ -146,12 +147,20 @@ class _AutotuneGroup:
     def record(self, operation):
         self.run.collectives[self.ranks][self.run.rank].append(operation)
 
-    def broadcast_object(self, obj, src=0):
-        assert src == 0
-        self.record(("broadcast", src))
-        if self.rank_in_group == src:
-            self.run.broadcasts[self.ranks] = obj
-        return self.run.broadcasts[self.ranks]
+    def all_gather_object(self, out, obj):
+        # Ranks run sequentially: snapshot every rank's file before any saves.
+        self.record(("all_gather",))
+        if isinstance(obj, bool):
+            out[:] = [rank not in self.run.rejected_load_ranks for rank in self.ranks]
+            assert out[self.rank_in_group] == obj
+            return
+        out[:] = self.run.gathered.setdefault(
+            self.ranks,
+            [
+                warmup._autotune_cache_fingerprint(self.run.cache_path(rank))
+                for rank in self.ranks
+            ],
+        )
 
     def barrier(self):
         self.record(("barrier",))
@@ -162,10 +171,18 @@ class _AutotuneTuner:
         self.run = run
         self.cache = {}
         self.loaded = None
+        self._dirty = False
 
     def load_configs(self, path):
+        if self.run.rank in self.run.rejected_load_ranks:
+            return False
         self.loaded = json.loads(Path(path).read_text())
         self.cache.update(self.loaded)
+        return True
+
+    def clear_cache(self):
+        self.cache.clear()
+        self._dirty = False
 
     def save_configs(self, path):
         self.run.saves.append((self.run.rank, Path(path), dict(self.cache)))
@@ -182,21 +199,27 @@ class _AutotuneTuner:
             if group is not None:
                 group.record(("all_reduce", operation, tactic))
         self.cache[operation] = self.run.rank // self.run.tp
+        self._dirty = True
 
 
 class _AutotuneRun:
-    def __init__(self, pp, tp):
+    def __init__(self, pp, tp, cache_dir):
         self.pp, self.tp = pp, tp
+        self.cache_dir = cache_dir
         self.rank = 0
         self.tuning_group = None
         self.collectives: dict[tuple[int, ...], dict[int, list[tuple[Any, ...]]]] = (
             defaultdict(lambda: defaultdict(list))
         )
-        self.broadcasts = {}
+        self.gathered = {}
+        self.rejected_load_ranks: set[int] = set()
         self.tuners = {}
         self.saves = []
         self.profile_groups = defaultdict(list)
 
+    def cache_path(self, rank):
+        return self.cache_dir / f"autotune_configs_dp0_rank{rank}.json"
+
     def world(self):
         return _AutotuneGroup(self, range(self.pp * self.tp))
 
@@ -235,7 +258,7 @@ def autotune_run(monkeypatch, tmp_path):
     from vllm.distributed import parallel_state
 
     def make_run(*, pp=2, tp=4):
-        run = _AutotuneRun(pp, tp)
+        run = _AutotuneRun(pp, tp, tmp_path)
         autotuner = ModuleType("flashinfer.autotuner")
         monkeypatch.setattr(
             autotuner,
@@ -251,6 +274,10 @@ def autotune_run(monkeypatch, tmp_path):
         monkeypatch.setattr(parallel_state, "get_tp_group", run.tensor_group)
         monkeypatch.setattr(parallel_state, "get_pp_group", run.pipeline_group)
         monkeypatch.setattr(fi_utils, "autotune", lambda **kwargs: nullcontext())
+        monkeypatch.setattr(
+            "torch.distributed.all_gather_object",
+            lambda out, obj, group: group.all_gather_object(out, obj),
+        )
         monkeypatch.setattr(
             warmup,
             "resolve_flashinfer_autotune_file",
@@ -286,19 +313,18 @@ def test_pp_stage_cache_roundtrip_isolated_and_asymmetric_hits_safe(
     legacy = tmp_path / "autotune_configs.json"
     legacy.write_text('{"legacy_world_cache": 99}')
     cold = autotune_run().execute()
-    assert [rank for rank, _, _ in cold.saves] == [0, 4]
-    paths = [path for _, path, _ in cold.saves]
-    assert len(set(paths)) == 2 and legacy not in paths
-    assert json.loads(legacy.read_text()) == {"legacy_world_cache": 99}
-    assert cold.saves[0][2] == {"shared_gemm": 0, "pp0_extra_gemm": 0}
-    assert cold.saves[1][2] == {"shared_gemm": 1}
     cold.assert_collectives_match()
+    assert [rank for rank, _, _ in cold.saves] == list(range(8))
+    assert [path for _, path, _ in cold.saves] == [cold.cache_path(r) for r in range(8)]
+    assert json.loads(legacy.read_text()) == {"legacy_world_cache": 99}
+    stage_caches = [{"shared_gemm": 0, "pp0_extra_gemm": 0}, {"shared_gemm": 1}]
+    assert all(cache == stage_caches[r // 4] for r, _, cache in cold.saves)
     warm = autotune_run().execute()
     warm.assert_collectives_match()
-    assert not warm.profile_groups
+    assert not warm.profile_groups and not warm.saves
     for rank, tuner in warm.tuners.items():
-        assert tuner.loaded == cold.saves[rank // 4][2]
-    paths[1].unlink()
+        assert tuner.loaded == stage_caches[rank // 4]
+    cold.cache_path(5).unlink()
     mixed = autotune_run().execute()
     mixed.assert_collectives_match()
     assert set(mixed.profile_groups) == {4, 5, 6, 7}
@@ -306,10 +332,33 @@ def test_pp_stage_cache_roundtrip_isolated_and_asymmetric_hits_safe(
     assert all(mixed.tuners[rank].loaded is None for rank in range(4, 8))
 
 
-def test_pp1_retains_world_synchronization_and_existing_cache_name(autotune_run):
+def test_pp1_tunes_world_group_and_saves_per_rank(autotune_run):
     run = autotune_run(pp=1, tp=4).execute()
     run.assert_collectives_match()
-    assert [(rank, path.name) for rank, path, _ in run.saves] == [
-        (0, "autotune_configs.json")
+    assert [(rank, path) for rank, path, _ in run.saves] == [
+        (rank, run.cache_path(rank)) for rank in range(4)
     ]
     assert all(groups == [(0, 1, 2, 3)] for groups in run.profile_groups.values())
+
+
+def test_mismatched_rank_caches_are_ignored_by_every_rank(autotune_run):
+    """One rank loading while others tune deadlocks the per-tactic reduce."""
+    cold = autotune_run(pp=1, tp=4).execute()
+    stale = cold.cache_path(2)
+    stale.write_text(json.dumps({**json.loads(stale.read_text()), "old_gemm": 0}))
+    rerun = autotune_run(pp=1, tp=4).execute()
+    rerun.assert_collectives_match()
+    assert all(tuner.loaded is None for tuner in rerun.tuners.values())
+    assert set(rerun.profile_groups) == {0, 1, 2, 3}
+
+
+@pytest.mark.parametrize("pp, tp", [(1, 1), (1, 4), (2, 4)])
+def test_rejected_cache_retunes_every_rank_in_its_tuning_group(autotune_run, pp, tp):
+    """Matching files can still be incompatible with one rank's runtime."""
+    autotune_run(pp=pp, tp=tp).execute()
+    rerun = autotune_run(pp=pp, tp=tp)
+    rerun.rejected_load_ranks = {tp - 1}
+    rerun.execute()
+    rerun.assert_collectives_match()
+    assert set(rerun.profile_groups) == set(range(tp))
+    assert {rank for rank, _, _ in rerun.saves} == set(range(tp))

vllm/model_executor/warmup/kernel_warmup.py

@@ -5,8 +5,10 @@ This is useful specifically for JIT'ed kernels as we don't want JIT'ing to
 happen during model execution.
 """
 
+import json
 import sys
 import time
+from pathlib import Path
 from typing import TYPE_CHECKING
 
 import torch
@@ -18,7 +20,6 @@ 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,
-    write_flashinfer_autotune_cache,
 )
 from vllm.model_executor.warmup.flashinfer_sparse_mla_warmup import (
     autotune_hisparse_flashinfer_attention,
@@ -408,6 +409,26 @@ def _run_flashinfer_bf16_autotune_dummy_run(
         )
 
 
+def _autotune_cache_fingerprint(path: Path) -> tuple[str, int] | None:
+    """Identify a saved autotune file by its FlashInfer metadata and size."""
+    try:
+        configs = json.loads(path.read_text())
+        metadata = configs.pop("_metadata", None)
+    except (OSError, ValueError, AttributeError):
+        return None
+    return json.dumps(metadata, sort_keys=True), len(configs)
+
+
+def _all_ranks_have_matching_cache(path: Path, group) -> bool:
+    """True iff every rank in ``group`` has its own, mutually consistent file."""
+    fingerprint = _autotune_cache_fingerprint(path)
+    if group.world_size == 1:
+        return fingerprint is not None
+    gathered: list[tuple[str, int] | None] = [None] * group.world_size
+    torch.distributed.all_gather_object(gathered, fingerprint, group=group.cpu_group)
+    return fingerprint is not None and all(f == fingerprint for f in gathered)
+
+
 def flashinfer_autotune(runner: "GPUModelRunner") -> None:
     """Autotune FlashInfer operations.
     FlashInfer have many implementations for the same operation,
@@ -421,6 +442,12 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:
     so each stage's TP group tunes separately with its own cache file;
     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
+    (``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
+    has a matching file and successfully loads it.
     """
     from flashinfer.autotuner import AutoTuner, set_autotune_process_group
 
@@ -447,11 +474,11 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:
         autotune_kwargs["skip_ops"] = skip_ops
 
     cache_path = resolve_flashinfer_autotune_file(runner)
-    if pp_size > 1:
-        ranks = "-".join(str(rank) for rank in tune_group.ranks)
-        cache_path = cache_path.with_name(
-            f"{cache_path.stem}_tp_{ranks}{cache_path.suffix}"
-        )
+    # 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)
 
@@ -460,16 +487,16 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:
     # 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.
-    # Read cached autotune results and broadcast within the tuning group.
-    cached_results: bytes | None = None
-    if is_leader and cache_path.exists():
-        with open(cache_path, "rb") as f:
-            cached_results = f.read()
-    cached_results = tune_group.broadcast_object(cached_results, src=0)
-    if cached_results is not None:
-        write_flashinfer_autotune_cache(cache_path, cached_results)
-        tune_group.barrier()
-        tuner.load_configs(str(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()
 
     group = tune_group.cpu_group if tune_group.world_size > 1 else None
     set_autotune_process_group(group)
@@ -497,5 +524,7 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:
 
     if world.world_size > 1:
         world.barrier()
-    if is_leader:
+    # 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:
         tuner.save_configs(str(cache_path))