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 run_mixed_prefill_decode_warmup(
        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."
        )