vllm.model_executor.warmup.attention_warmup

Functions:

• mixed_batch_attention_warmup –

  Run a mixed prefill+decode batch to warm up both attention paths.

mixed_batch_attention_warmup(worker, num_tokens=16)

Run a mixed prefill+decode batch to warm up both attention paths.

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

def mixed_batch_attention_warmup(worker: "Worker", num_tokens: int = 16) -> None:
    """Run a mixed prefill+decode batch to warm up both attention paths."""
    num_tokens = min(
        num_tokens,
        worker.scheduler_config.max_num_batched_tokens,
        worker.model_config.max_model_len,
    )
    if not worker.use_v2_model_runner:
        worker.model_runner._dummy_run(
            num_tokens=num_tokens,
            skip_eplb=True,
            is_profile=True,
            force_attention=True,
            create_mixed_batch=True,
        )
        return

    if not run_mixed_prefill_decode_warmup(
        cast("GPUModelRunner", worker.model_runner),
        worker.execute_model,
        worker.sample_tokens,
        num_tokens,
    ):
        logger.info_once(
            "Skipping mixed prefill+decode attention warmup; attention kernels "
            "will be loaded on the first request."
        )