vllm.model_executor.warmup.kernel_warmup

Warmup kernels used during model execution. This is useful specifically for JIT'ed kernels as we don't want JIT'ing to happen during model execution.

Functions:

• flashinfer_autotune –

  Autotune FlashInfer operations.

_aiter_gemm_m_buckets(max_m)

One token count per AITER tuned-GEMM bucket up to max_m.

AITER pads M to 1, 2, 4 or 8 below 16, then to multiples of 16 / 32 / 64 / 128 up to 256 / 1024 / 4096 / beyond, and builds each bucket's kernel on first use.

Source code in vllm/model_executor/warmup/kernel_warmup.py

def _aiter_gemm_m_buckets(max_m: int) -> list[int]:
    """One token count per AITER tuned-GEMM bucket up to ``max_m``.

    AITER pads M to 1, 2, 4 or 8 below 16, then to multiples of 16 / 32 / 64 /
    128 up to 256 / 1024 / 4096 / beyond, and builds each bucket's kernel on
    first use.
    """
    buckets = [m for m in (1, 2, 4, 8) if m <= max_m]
    for step, upper in ((16, 256), (32, 1024), (64, 4096), (128, max_m)):
        start = (buckets[-1] // step + 1) * step if buckets else step
        buckets.extend(range(start, min(upper, max_m) + 1, step))
    if max_m > 0 and buckets[-1] != max_m:
        buckets.append(max_m)
    return buckets

_all_ranks_have_matching_cache(path, group)

True iff every rank in group has its own, mutually consistent file.

Source code in vllm/model_executor/warmup/kernel_warmup.py

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)

_autotune_cache_fingerprint(path)

Identify a saved autotune file by its FlashInfer metadata and size.

Source code in vllm/model_executor/warmup/kernel_warmup.py

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)

_flashinfer_deferred_moe_token_counts(runner)

Return bounded token counts that exercise deferred MoE finalization.

Source code in vllm/model_executor/warmup/kernel_warmup.py

def _flashinfer_deferred_moe_token_counts(
    runner: "GPUModelRunner",
) -> tuple[int, ...]:
    """Return bounded token counts that exercise deferred MoE finalization."""
    from vllm.model_executor.layers.fused_moe import MoERunner

    max_tokens = runner.scheduler_config.max_num_batched_tokens
    token_counts: list[int] = []
    for module in runner.get_model().modules():
        if not isinstance(module, MoERunner):
            continue

        moe_config = module.moe_config
        max_deferred_tokens = moe_config.defer_moe_finalize_max_num_tokens
        if moe_config.use_deferred_moe_finalize and max_deferred_tokens > 0:
            token_counts.append(min(max_tokens, max_deferred_tokens))

    return tuple(dict.fromkeys(token_counts))

_load_shared_autotune_cache(cache_path, tuner, tune_group)

Load the per-rank autotune cache into the tuner, group-wide.

Returns True only when every rank of the tune group had a matching file and loaded it: the dummy-run sweep contains collectives, so skipping it is safe only if the decision is unanimous.

Source code in vllm/model_executor/warmup/kernel_warmup.py

def _load_shared_autotune_cache(cache_path: Path, tuner, tune_group) -> bool:
    """Load the per-rank autotune cache into the tuner, group-wide.

    Returns True only when every rank of the tune group had a matching file
    and loaded it: the dummy-run sweep contains collectives, so skipping it
    is safe only if the decision is unanimous.
    """
    if not _all_ranks_have_matching_cache(cache_path, tune_group):
        return False
    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()
    return loaded

flashinfer_autotune(runner, *, skip_attn=False)

Autotune FlashInfer operations. FlashInfer have many implementations for the same operation, autotuning runs benchmarks for each implementation and stores the results. The results are cached transparently and future calls to FlashInfer will use the best implementation. Without autotuning, FlashInfer will rely on heuristics, which may be significantly slower.

Parameters:

• runner

  (GPUModelRunner) –

  The model runner to tune.

• skip_attn

  (bool, default: False ) –

  Additionally skip attention in the tuning dummy runs.

With PP > 1, stages run different layers and may profile different ops, 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.

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 has a matching file and successfully loads it.

Source code in vllm/model_executor/warmup/kernel_warmup.py

def flashinfer_autotune(runner: "GPUModelRunner", *, skip_attn: bool = False) -> None:
    """Autotune FlashInfer operations.
    FlashInfer have many implementations for the same operation,
    autotuning runs benchmarks for each implementation and stores
    the results. The results are cached transparently and
    future calls to FlashInfer will use the best implementation.
    Without autotuning, FlashInfer will rely on heuristics, which may
    be significantly slower.

    Args:
        runner: The model runner to tune.
        skip_attn: Additionally skip attention in the tuning dummy runs.

    With PP > 1, stages run different layers and may profile different ops,
    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.

    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
    has a matching file and successfully loads it.

    """
    from flashinfer.autotuner import AutoTuner, set_autotune_process_group

    import vllm.utils.flashinfer as fi_utils
    from vllm.distributed.parallel_state import (
        get_pp_group,
        get_tp_group,
        get_world_group,
    )

    world = get_world_group()
    pp_size = get_pp_group().world_size
    tune_group = get_tp_group() if pp_size > 1 else world
    is_leader = tune_group.rank_in_group == 0
    tuner = AutoTuner.get()

    autotune_kwargs: dict = {}
    skip_ops = _flashinfer_autotune_skip_ops(runner)
    if skip_ops:
        logger.info_once(
            "Skipping FlashInfer autotuning for ops %s",
            tuple(sorted(skip_ops)),
        )
        autotune_kwargs["skip_ops"] = skip_ops

    use_autotune_v2 = use_flashinfer_autotune_v2(runner)
    cache_hit = False
    if use_autotune_v2:
        from flashinfer import autotune_v2, autotune_v2_reload

        cache_root = resolve_flashinfer_autotune_v2_root()
        if is_leader:
            logger.info_once(
                "Using FlashInfer managed autotune cache (root=%s). Pass "
                "--disable-flashinfer-autotune to skip autotuning.",
                cache_root if cache_root is not None else "flashinfer default",
            )
        autotune_context = autotune_v2(
            mode="tune", cache_root=cache_root, **autotune_kwargs
        )
    else:
        cache_path = resolve_rank_flashinfer_autotune_file(runner)
        if is_leader:
            logger.info_once(
                "Using FlashInfer autotune cache file: %s. Pass "
                "--disable-flashinfer-autotune to skip autotuning.",
                cache_path,
            )

        cache_hit = _load_shared_autotune_cache(cache_path, tuner, tune_group)
        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(), autotune_context:
            hisparse_enabled = (
                runner.vllm_config.attention_config.hisparse_config is not None
            )
            skip_attn = skip_attn or hisparse_enabled
            if hisparse_enabled:
                # HiSparse hot-buffer attention is bounded by decode batch
                # size, not the prefill-sized batch used for the full model.
                autotune_hisparse_flashinfer_attention(runner)
            if not cache_hit:
                _run_flashinfer_autotune_dummy_runs(runner, skip_attn=skip_attn)
            replayssm_autotune_warmup(runner)
            _autotune_kimi_k3_kda_qkvg(runner.get_model())
    finally:
        set_autotune_process_group(None)

    if world.world_size > 1:
        world.barrier()
    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))