Skip to content

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.

logger module-attribute 

logger = init_logger(__name__)

flashinfer_autotune 

flashinfer_autotune(runner: GPUModelRunner) -> 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.

Source code in vllm/model_executor/warmup/kernel_warmup.py 
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
def flashinfer_autotune(runner: "GPUModelRunner") -> 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.
    """
    from vllm.utils.flashinfer import autotune

    with torch.inference_mode(), autotune():
        # We skip EPLB here since we don't want to record dummy metrics
        # When autotuning with number of tokens m, flashinfer will autotune
        # operations for all number of tokens up to m.
        # So we only need to run with the max number of tokens.
        runner._dummy_run(
            runner.scheduler_config.max_num_batched_tokens,
            skip_eplb=True,
            is_profile=True,
        )

flashinfer_autotune_supported 

flashinfer_autotune_supported(
    vllm_config: VllmConfig,
) -> bool

Record known issues with vllm + flashinfer autotune here. Return True if and only if flashinfer autotune will run through without issues.

Source code in vllm/model_executor/warmup/kernel_warmup.py 
28
29
30
31
32
33
34
35
36
37
38
39
def flashinfer_autotune_supported(vllm_config: VllmConfig) -> bool:
    """
    Record known issues with vllm + flashinfer autotune here. Return True if
    and only if flashinfer autotune will run through without issues.
    """
    return not (
        vllm_config.parallel_config.data_parallel_size > 1
        and (
            envs.VLLM_USE_FLASHINFER_MOE_MXFP4_BF16
            or envs.VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8
        )
    )

kernel_warmup 

kernel_warmup(worker: Worker)
Source code in vllm/model_executor/warmup/kernel_warmup.py 
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
def kernel_warmup(worker: "Worker"):
    # Deep GEMM warmup
    do_deep_gemm_warmup = (
        envs.VLLM_USE_DEEP_GEMM
        and is_deep_gemm_supported()
        and envs.VLLM_DEEP_GEMM_WARMUP != "skip"
    )
    if do_deep_gemm_warmup:
        model = worker.get_model()
        max_tokens = worker.scheduler_config.max_num_batched_tokens
        deep_gemm_warmup(model, max_tokens)

    # FlashInfer autotune for Hopper (SM 9.0) and Blackwell (SM 10.0) GPUs
    if (
        has_flashinfer()
        and current_platform.has_device_capability(90)
        and flashinfer_autotune_supported(worker.vllm_config)
    ):
        flashinfer_autotune(worker.model_runner)

    # FlashInfer attention warmup
    # Only warmup if the model has FlashInfer attention groups
    # and is not a pooling model
    def _is_flashinfer_backend(backend):
        try:
            return backend.get_name() == "FLASHINFER"
        except NotImplementedError:
            return False

    if not worker.model_runner.is_pooling_model and all(
        _is_flashinfer_backend(group.backend)
        for groups in worker.model_runner.attn_groups
        for group in groups
    ):
        logger.info("Warming up FlashInfer attention.")
        # Warmup with mixed batch containing both prefill and decode tokens
        # This is to warm up both prefill and decode attention kernels
        worker.model_runner._dummy_run(
            num_tokens=16,
            skip_eplb=True,
            is_profile=True,
            force_attention=True,
            create_mixed_batch=True,
        )