Skip to content

vllm.model_executor.warmup.flashinfer_autotune_cache

FlashInfer autotune cache helpers.

Functions:

_normalize_cache_config_for_hash(cache_config)

Normalize cache config for hashing

Source code in vllm/model_executor/warmup/flashinfer_autotune_cache.py 
23
24
25
26
27
28
29
def _normalize_cache_config_for_hash(cache_config: CacheConfig) -> CacheConfig:
    """Normalize cache config for hashing"""
    if cache_config.mamba_block_size is None and cache_config.kv_cache_layout is None:
        return cache_config
    normalized = replace(cache_config, mamba_block_size=None)
    assert normalized.kv_cache_layout is None
    return normalized

flashinfer_autotune_filename(dp_rank, rank)

Per-rank cache filename, shared by the engine reader and D2D exchange.

Source code in vllm/model_executor/warmup/flashinfer_autotune_cache.py 
57
58
59
def flashinfer_autotune_filename(dp_rank: int, rank: int) -> str:
    """Per-rank cache filename, shared by the engine reader and D2D exchange."""
    return f"autotune_configs_dp{dp_rank}_rank{rank}.json"

flashinfer_autotune_root()

Root directory of the FlashInfer autotune cache (no runner needed).

Source code in vllm/model_executor/warmup/flashinfer_autotune_cache.py 
41
42
43
44
45
46
47
48
49
50
51
52
53
54
def flashinfer_autotune_root() -> Path:
    """Root directory of the FlashInfer autotune cache (no runner needed)."""
    override_dir = envs.VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR
    if override_dir:
        return Path(override_dir).expanduser()
    from flashinfer.jit import env as flashinfer_jit_env

    flashinfer_workspace = flashinfer_jit_env.FLASHINFER_WORKSPACE_DIR
    return (
        Path(envs.VLLM_CACHE_ROOT)
        / "flashinfer_autotune_cache"
        / flashinfer_workspace.parent.name
        / flashinfer_workspace.name
    )

resolve_flashinfer_autotune_v2_root()

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.

Source code in vllm/model_executor/warmup/flashinfer_autotune_cache.py 
114
115
116
117
118
119
120
121
122
123
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

resolve_rank_flashinfer_autotune_file(runner)

The per-rank variant of :func:resolve_flashinfer_autotune_file.

Source code in vllm/model_executor/warmup/flashinfer_autotune_cache.py 
68
69
70
71
72
73
74
75
76
def resolve_rank_flashinfer_autotune_file(runner: "GPUModelRunner") -> Path:
    """The per-rank variant of :func:`resolve_flashinfer_autotune_file`."""
    from vllm.distributed.parallel_state import get_world_group

    path = resolve_flashinfer_autotune_file(runner)
    dp_rank = runner.vllm_config.parallel_config.data_parallel_rank
    return path.with_name(
        flashinfer_autotune_filename(dp_rank, get_world_group().rank_in_group)
    )