[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>

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

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