From 3403e0f176efb90b0d92cc0ff1d350f092ed9e9f Mon Sep 17 00:00:00 2001 From: Madhav Jivrajani <12381034+MadhavJivrajani@users.noreply.github.com> Date: Tue, 6 Oct 2026 13:51:01 -0400 Subject: [PATCH] [Bugfix] Persist FlashInfer autotune cache per rank to fix the rank-0-only cache-hit deadlock (#57635) Signed-off-by: Madhav Jivrajani Signed-off-by: NickLucche Co-authored-by: NickLucche Co-authored-by: Wei Zhao <51183510+wzhao18@users.noreply.github.com> --- .../test_flashinfer_autotune_warmup.py | 91 ++++++++++++++----- vllm/model_executor/warmup/kernel_warmup.py | 63 +++++++++---- 2 files changed, 116 insertions(+), 38 deletions(-) diff --git a/tests/model_executor/test_flashinfer_autotune_warmup.py b/tests/model_executor/test_flashinfer_autotune_warmup.py index 22bb54d71f..d65f2a4ad1 100644 --- a/tests/model_executor/test_flashinfer_autotune_warmup.py +++ b/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)) diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index 55d5ea7a7d..06f04b75d9 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/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))