Sync FlashInfer autotune tactic choice across TP ranks (#35343)

Co-authored-by: Mohammad Angkad <mohammad.angkad@radixark.ai>

python/sglang/srt/model_executor/runner/flashinfer_autotune.py

@@ -17,6 +17,7 @@ import contextlib
 import datetime
 import functools
 import hashlib
+import json
 import logging
 from pathlib import Path
 from typing import TYPE_CHECKING, Callable, Optional
@@ -34,6 +35,7 @@ from sglang.srt.runtime_context import (
 from sglang.srt.utils import empty_context, log_info_on_rank0
 
 if TYPE_CHECKING:
+    from sglang.srt.distributed.parallel_state import GroupCoordinator
     from sglang.srt.model_executor.model_runner import ModelRunner
     from sglang.srt.model_executor.runner.base_runner import BaseRunner
 
@@ -170,14 +172,90 @@ def flashinfer_autotune_cache_path(model_runner: ModelRunner) -> Path:
     )
 
 
+def _autotune_tactic_sync_group(
+    tp_group: GroupCoordinator,
+) -> Optional[torch.distributed.ProcessGroup]:
+    """CPU group over the ranks that must agree on the tuned tactics.
+
+    Per-rank timing noise alone makes each rank's ``argmin`` pick a different
+    tactic for the same shape. FlashInfer all-reduces the timings over this
+    group so every rank minimizes over the same numbers. TP is the scope: those
+    ranks run the same dummy forward, and PP stages are already separate groups.
+    """
+    if tp_group.world_size <= 1:
+        return None
+    # The CPU group keeps the reduction of these scalars off the profiled stream.
+    return tp_group.cpu_group
+
+
+@contextlib.contextmanager
+def _autotune_process_group(group: Optional[torch.distributed.ProcessGroup]):
+    """Set FlashInfer's timing-reduction group, restoring the previous one after."""
+    from flashinfer.autotuner import (
+        get_autotune_process_group,
+        set_autotune_process_group,
+    )
+
+    previous = get_autotune_process_group()
+    set_autotune_process_group(group)
+    try:
+        yield
+    finally:
+        set_autotune_process_group(previous)
+
+
+def _autotune_cache_digest(cache_path: Path, env: dict[str, str]) -> str:
+    """Hash of what this rank would load from ``cache_path`` ("" for nothing).
+
+    Includes the environment: ``load_configs`` ignores the whole file when its
+    ``_metadata`` stamp disagrees with the environment reading it, so equal
+    tactics alone do not mean two ranks load the same thing.
+    """
+    if not cache_path.is_file():
+        return ""
+    try:
+        configs = json.loads(cache_path.read_text())
+    except (OSError, ValueError):
+        return ""
+    if not isinstance(configs, dict):
+        return ""
+    payload = {"file": configs, "env": env}
+    return hashlib.sha256(json.dumps(payload, sort_keys=True).encode()).hexdigest()
+
+
+def _drop_diverged_autotune_cache(
+    cache_path: Path, group: torch.distributed.ProcessGroup, env: dict[str, str]
+) -> None:
+    """Enter tuning with the same cache on every rank, or with none at all.
+
+    A cache hit skips a profile, so caches that disagree desync the reduction.
+    """
+    digests: list[str] = [""] * torch.distributed.get_world_size(group)
+    torch.distributed.all_gather_object(
+        digests, _autotune_cache_digest(cache_path, env), group=group
+    )
+    if len(set(digests)) == 1:
+        return
+    log_info_on_rank0(
+        logger,
+        "FlashInfer autotune: per-rank caches disagree, discarding them and "
+        "tuning from scratch so all ranks agree on the tactics.",
+    )
+    cache_path.unlink(missing_ok=True)
+
+
 @contextlib.contextmanager
 def flashinfer_autotune_context(model_runner: ModelRunner, *, run_lm_head: bool):
-    from flashinfer.autotuner import autotune
+    # The gate below decides on the same inputs load_configs does.
+    from flashinfer.autotuner import _collect_metadata, autotune
 
     mr = model_runner
     cache_path = flashinfer_autotune_cache_path(mr)
+    sync_group = _autotune_tactic_sync_group(mr.tp_group)
     if envs.SGLANG_FLASHINFER_AUTOTUNE_CACHE.get():
         autotune_cache = cache_path
+        if sync_group is not None:
+            _drop_diverged_autotune_cache(cache_path, sync_group, _collect_metadata())
         logger.info("Running FlashInfer autotune with cache: %s", autotune_cache)
     else:
         timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
@@ -197,7 +275,7 @@ def flashinfer_autotune_context(model_runner: ModelRunner, *, run_lm_head: bool)
         from sglang.srt.layers.logits_processor import autotune_dummy_run_mode
 
         skip_ops = get_flashinfer_autotune_skip_ops(mr)
-        with autotune(
+        with _autotune_process_group(sync_group), autotune(
             True,
             cache=str(autotune_cache),
             skip_ops=skip_ops,
@@ -321,6 +399,11 @@ def maybe_flashinfer_autotune_extend(
     try:
         run_flashinfer_autotune_forward(mr, forward_fn, run_lm_head=False)
     except torch.OutOfMemoryError:
+        if _autotune_tactic_sync_group(mr.tp_group) is not None:
+            # Tuning is collective: this rank has stopped reducing while its
+            # peers wait on the next tactic, so skipping the pass would hang
+            # them. Fail instead of degrading alone.
+            raise
         # The pass is an optimization; without headroom for the extend-shaped
         # forward, fall back to untuned extend buckets instead of failing.
         log_info_on_rank0(