iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source 
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>
This commit is contained in:
Yang Xu
2026-10-17 02:08:28 +00:00
committed by GitHub
co-authored by Zihua Wu Kimi Code Codex Misha Goin vllm-agent
parent b3a7e2ff6e
commit 71c20f73f2
6 changed files with 436 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/usage/security.md
+2 -2
View File
@@ -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
+267 -4
View File
@@ -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
+50
View File
@@ -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
+61 -1
View File
@@ -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
+42 -30
View File
@@ -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
+14
View File
@@ -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",