vllm.v1.worker.gpu.warmup

Functions:

• run_context_prefill_warmup –

  Run a prefill step over about num_context_tokens of cached context.

• run_kv_cache_sizing_warmup –

  Run the steps an extensible KV cache measures its transient reserve on.

• run_mixed_prefill_decode_warmup –

  Run a V2 mixed prefill+decode step through normal scheduler inputs,

• warmup_kernels –

  Run scheduler-realistic prefill and decode steps to JIT compile kernels.

_reserved_block_count(num_tokens, kvcache_spec, *, num_lookahead_tokens, max_model_len, max_encoder_len, dcp_world_size=1)

Number of blocks the scheduler would hold for a request of num_tokens.

Warmup hand-builds its SchedulerOutputs, so it must reserve what KVCacheManager.allocate_slots reserves: the token range plus num_lookahead_tokens, where the speculator writes the KV of its drafts. A DCP-sharded block holds dcp_world_size times its tokens.

Source code in vllm/v1/worker/gpu/warmup.py

def _reserved_block_count(
    num_tokens: int,
    kvcache_spec: KVCacheSpec,
    *,
    num_lookahead_tokens: int,
    max_model_len: int,
    max_encoder_len: int,
    dcp_world_size: int = 1,
) -> int:
    """Number of blocks the scheduler would hold for a request of `num_tokens`.

    Warmup hand-builds its `SchedulerOutput`s, so it must reserve what
    `KVCacheManager.allocate_slots` reserves: the token range plus
    `num_lookahead_tokens`, where the speculator writes the KV of its drafts.
    A DCP-sharded block holds `dcp_world_size` times its tokens.
    """
    if isinstance(kvcache_spec, UniformTypeKVCacheSpecs):
        kvcache_spec = kvcache_spec.first_spec
    if isinstance(kvcache_spec, CircularBufferSpec):
        # Circular caches keep one physical ring block for the request lifetime.
        return 1
    if isinstance(kvcache_spec, CrossAttentionSpec):
        # Cross-attention blocks cover the encoder sequence only.
        return cdiv(max_encoder_len, kvcache_spec.block_size)
    num_speculative_blocks = 0
    if isinstance(kvcache_spec, MambaSpec):
        # MambaManager appends speculative running-state blocks in every cache
        # mode; align mode sizes from the uncapped, lookahead-free token range.
        num_speculative_blocks = kvcache_spec.num_speculative_blocks
        if kvcache_spec.mamba_cache_mode == "align":
            return cdiv(num_tokens, kvcache_spec.block_size) + num_speculative_blocks
    block_size = kvcache_spec.block_size
    if kvcache_spec.dcp_sharded:
        block_size *= dcp_world_size
    num_tokens = min(num_tokens + num_lookahead_tokens, max_model_len)
    return cdiv(num_tokens, block_size) + num_speculative_blocks

_warmup_block_counter(model_runner)

Bind _reserved_block_count to model_runner's reservation policy.

Source code in vllm/v1/worker/gpu/warmup.py

def _warmup_block_counter(
    model_runner: GPUModelRunner,
) -> Callable[[int, KVCacheSpec], int]:
    """Bind `_reserved_block_count` to `model_runner`'s reservation policy."""
    vllm_config = model_runner.vllm_config
    num_lookahead_tokens = vllm_config.num_lookahead_tokens
    max_model_len = model_runner.max_model_len
    max_encoder_len = getattr(model_runner.model_state, "max_encoder_len", 0)
    dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size

    def block_count(num_tokens: int, kvcache_spec: KVCacheSpec) -> int:
        return _reserved_block_count(
            num_tokens,
            kvcache_spec,
            num_lookahead_tokens=num_lookahead_tokens,
            max_model_len=max_model_len,
            max_encoder_len=max_encoder_len,
            dcp_world_size=dcp_world_size,
        )

    return block_count

_warmup_token_ids(model_runner, num_tokens)

Distinct-looking prompt token ids that stay within the vocabulary, which can be smaller than a full-budget or full-length prompt.

Source code in vllm/v1/worker/gpu/warmup.py

def _warmup_token_ids(model_runner: GPUModelRunner, num_tokens: int) -> list[int]:
    """Distinct-looking prompt token ids that stay within the vocabulary, which
    can be smaller than a full-budget or full-length prompt."""
    vocab_size = model_runner.model_config.get_vocab_size()
    return [i % vocab_size for i in range(num_tokens)]

run_context_prefill_warmup(model_runner, worker_execute_model, worker_sample_tokens, num_context_tokens)

Run a prefill step over about num_context_tokens of cached context.

MLA up-projects a prefix cache hit's context chunk by chunk, a transient that no step starting from an empty prefix reaches. The queries are kept short: the full-budget step already measured theirs.

Source code in vllm/v1/worker/gpu/warmup.py

def run_context_prefill_warmup(
    model_runner: GPUModelRunner,
    worker_execute_model: Callable[[SchedulerOutput], Any],
    worker_sample_tokens: Callable[[GrammarOutput | None], Any],
    num_context_tokens: int,
) -> bool:
    """Run a prefill step over about `num_context_tokens` of cached context.

    MLA up-projects a prefix cache hit's context chunk by chunk, a transient
    that no step starting from an empty prefix reaches. The queries are kept
    short: the full-budget step already measured theirs.
    """
    query_len = model_runner.decode_query_len + 1
    # Leave room for a sampled token after the prompt; and no request needs
    # more context than the step does.
    context_len = min(model_runner.max_model_len - query_len - 1, num_context_tokens)
    if model_runner.is_pooling_model or context_len <= 0:
        return False

    kv_cache_specs = [
        g.kv_cache_spec for g in model_runner.kv_cache_config.kv_cache_groups
    ]
    block_count = _warmup_block_counter(model_runner)
    block_counts = [block_count(context_len + query_len, s) for s in kv_cache_specs]
    blocks_per_req = max(sum(block_counts), 1)
    num_reqs = min(
        cdiv(num_context_tokens, context_len),
        model_runner.max_num_reqs,
        model_runner.scheduler_config.max_num_batched_tokens // query_len,
        (num_committable_kv_blocks(model_runner) - 1) // blocks_per_req,
    )
    if num_reqs < 1:
        return False
    ensure_kv_cache_blocks(model_runner, 1 + num_reqs * blocks_per_req)

    next_block_id = 1

    def _alloc_blocks(num_blocks: int) -> list[int]:
        nonlocal next_block_id
        return list(range(next_block_id, next_block_id := next_block_id + num_blocks))

    req_ids = [f"_context_warmup_{i}_" for i in range(num_reqs)]
    token_ids = _warmup_token_ids(model_runner, context_len + query_len)
    sampling_params = SamplingParams(max_tokens=1, temperature=0.0)
    prefill_output = SchedulerOutput.make_empty()
    prefill_output.num_spec_tokens_to_schedule = model_runner.num_speculative_steps
    prefill_output.scheduled_new_reqs = [
        NewRequestData(
            req_id=req_id,
            prompt_token_ids=token_ids,
            mm_features=[],
            sampling_params=sampling_params,
            pooling_params=None,
            block_ids=tuple(_alloc_blocks(n) for n in block_counts),
            num_computed_tokens=context_len,
            lora_request=None,
            prefill_token_ids=token_ids,
        )
        for req_id in req_ids
    ]
    prefill_output.num_scheduled_tokens = dict.fromkeys(req_ids, query_len)
    prefill_output.total_num_scheduled_tokens = num_reqs * query_len

    cleanup_output = SchedulerOutput.make_empty()
    cleanup_output.finished_req_ids = set(req_ids)

    model_runner.kv_connector.set_disabled(True)
    try:
        worker_execute_model(prefill_output)
        worker_sample_tokens(None)
        worker_execute_model(cleanup_output)
    finally:
        model_runner.kv_connector.set_disabled(False)
    return True

run_kv_cache_sizing_warmup(model_runner, worker_execute_model, worker_sample_tokens)

Run the steps an extensible KV cache measures its transient reserve on. These size rather than compile, so they run without JIT warmup.

Source code in vllm/v1/worker/gpu/warmup.py

@torch.inference_mode()
def run_kv_cache_sizing_warmup(
    model_runner: GPUModelRunner,
    worker_execute_model: Callable[[SchedulerOutput], Any],
    worker_sample_tokens: Callable[[GrammarOutput | None], Any],
) -> None:
    """Run the steps an extensible KV cache measures its transient reserve on.
    These size rather than compile, so they run without JIT warmup.
    """
    if model_runner.is_encoder_decoder:
        return
    # Profiling skipped attention; measure it on the largest possible prefill.
    # Its prompt logprobs materialize full-vocab logits a chunk at a time.
    scheduler_config = model_runner.scheduler_config
    max_logprobs = model_runner.model_config.max_logprobs
    run_mixed_prefill_decode_warmup(
        model_runner,
        worker_execute_model,
        worker_sample_tokens,
        min(scheduler_config.max_num_batched_tokens, model_runner.max_model_len),
        prompt_logprobs=max_logprobs if max_logprobs != 0 else None,
    )
    if model_runner.model_config.use_mla:
        from vllm.model_executor.layers.attention.mla_attention import (
            MLACommonMetadataBuilder,
        )

        # Fill one chunked-context workspace, so the measured peak covers
        # up-projecting a prefix cache hit's context.
        workspace_size = (
            MLACommonMetadataBuilder.determine_chunked_prefill_workspace_size(
                model_runner.vllm_config
            )
        )
        run_context_prefill_warmup(
            model_runner, worker_execute_model, worker_sample_tokens, workspace_size
        )

run_mixed_prefill_decode_warmup(model_runner, worker_execute_model, worker_sample_tokens, num_tokens, *, mixed_step_context=None, req_id_prefix='_v2_mixed_warmup', prompt_logprobs=None)

Run a V2 mixed prefill+decode step through normal scheduler inputs, requesting prompt_logprobs on the prefill request.

Source code in vllm/v1/worker/gpu/warmup.py

def run_mixed_prefill_decode_warmup(
    model_runner: GPUModelRunner,
    worker_execute_model: Callable[[SchedulerOutput], Any],
    worker_sample_tokens: Callable[[GrammarOutput | None], Any],
    num_tokens: int,
    *,
    mixed_step_context: AbstractContextManager[object] | None = None,
    req_id_prefix: str = "_v2_mixed_warmup",
    prompt_logprobs: int | None = None,
) -> bool:
    """Run a V2 mixed prefill+decode step through normal scheduler inputs,
    requesting prompt_logprobs on the prefill request."""
    if model_runner.is_pooling_model or model_runner.max_num_reqs < 2 or num_tokens < 3:
        return False

    decode_req_id = f"{req_id_prefix}_decode_"
    prefill_req_id = f"{req_id_prefix}_prefill_"
    decode_prompt_len = 2
    decode_scheduled_tokens = 1
    prefill_len = num_tokens - decode_scheduled_tokens
    decode_token_ids = _warmup_token_ids(model_runner, decode_prompt_len)
    prefill_token_ids = _warmup_token_ids(model_runner, prefill_len)

    kv_cache_groups = model_runner.kv_cache_config.kv_cache_groups
    block_count = _warmup_block_counter(model_runner)
    kv_cache_specs = [g.kv_cache_spec for g in kv_cache_groups]
    decode_prefill_block_counts = [
        block_count(decode_prompt_len, s) for s in kv_cache_specs
    ]
    decode_block_counts = [
        block_count(decode_prompt_len + decode_scheduled_tokens, s)
        for s in kv_cache_specs
    ]
    decode_block_deltas = [
        decode - prefill
        for decode, prefill in zip(decode_block_counts, decode_prefill_block_counts)
    ]
    prefill_block_counts = [block_count(prefill_len, s) for s in kv_cache_specs]
    required_blocks = sum(decode_block_counts) + sum(prefill_block_counts)
    num_blocks = num_committable_kv_blocks(model_runner)
    if num_blocks <= required_blocks:
        logger.warning(
            "Skipping V2 mixed prefill+decode warmup because only %d KV blocks "
            "are available for %d required warmup blocks.",
            num_blocks,
            required_blocks,
        )
        return False
    ensure_kv_cache_blocks(model_runner, 1 + required_blocks)

    next_block_id = 1

    def _alloc_blocks(num_blocks: int) -> list[int]:
        nonlocal next_block_id
        block_ids = list(range(next_block_id, next_block_id + num_blocks))
        next_block_id += num_blocks
        return block_ids

    sampling_params = SamplingParams(
        max_tokens=2,
        temperature=0.0,
        watermarking=False,
    )

    decode_prefill_output = SchedulerOutput.make_empty()
    decode_prefill_output.num_spec_tokens_to_schedule = (
        model_runner.num_speculative_steps
    )
    decode_prefill_output.scheduled_new_reqs = [
        NewRequestData(
            req_id=decode_req_id,
            prompt_token_ids=decode_token_ids,
            mm_features=[],
            sampling_params=sampling_params,
            pooling_params=None,
            block_ids=tuple(_alloc_blocks(n) for n in decode_prefill_block_counts),
            num_computed_tokens=0,
            lora_request=None,
            prefill_token_ids=decode_token_ids,
        ),
    ]
    decode_prefill_output.num_scheduled_tokens = {
        decode_req_id: decode_prompt_len,
    }
    decode_prefill_output.total_num_scheduled_tokens = decode_prompt_len

    decode_new_blocks = tuple(_alloc_blocks(n) for n in decode_block_deltas)
    cached_decode_req = CachedRequestData.make_empty()
    cached_decode_req.req_ids = [decode_req_id]
    cached_decode_req.num_computed_tokens = [decode_prompt_len]
    cached_decode_req.num_output_tokens = [1]
    cached_decode_req.new_block_ids = [
        decode_new_blocks if any(decode_block_deltas) else None
    ]

    mixed_output = SchedulerOutput.make_empty()
    mixed_output.num_spec_tokens_to_schedule = model_runner.num_speculative_steps
    mixed_output.scheduled_cached_reqs = cached_decode_req
    mixed_output.scheduled_new_reqs = [
        NewRequestData(
            req_id=prefill_req_id,
            prompt_token_ids=prefill_token_ids,
            mm_features=[],
            sampling_params=SamplingParams(
                max_tokens=2,
                temperature=0.0,
                watermarking=False,
                prompt_logprobs=prompt_logprobs,
            ),
            pooling_params=None,
            block_ids=tuple(_alloc_blocks(n) for n in prefill_block_counts),
            num_computed_tokens=0,
            lora_request=None,
            prefill_token_ids=prefill_token_ids,
        ),
    ]
    mixed_output.num_scheduled_tokens = {
        decode_req_id: decode_scheduled_tokens,
        prefill_req_id: prefill_len,
    }
    mixed_output.total_num_scheduled_tokens = num_tokens

    cleanup_output = SchedulerOutput.make_empty()
    cleanup_output.finished_req_ids = {decode_req_id, prefill_req_id}

    context = mixed_step_context or nullcontext()
    model_runner.kv_connector.set_disabled(True)
    try:
        worker_execute_model(decode_prefill_output)
        worker_sample_tokens(None)
        with context:
            worker_execute_model(mixed_output)
            worker_sample_tokens(None)
        worker_execute_model(cleanup_output)
    finally:
        model_runner.kv_connector.set_disabled(False)
    return True

warmup_kernels(model_runner, worker_execute_model, worker_sample_tokens)

Run scheduler-realistic prefill and decode steps to JIT compile kernels.

We must call the provided worker's execute_model for pipeline parallel coordination.

Source code in vllm/v1/worker/gpu/warmup.py

@torch.inference_mode()
def warmup_kernels(
    model_runner: GPUModelRunner,
    worker_execute_model: Callable[[SchedulerOutput], Any],
    worker_sample_tokens: Callable[[GrammarOutput | None], Any],
) -> None:
    """Run scheduler-realistic prefill and decode steps to JIT compile kernels.

    We must call the provided worker's execute_model for pipeline parallel
    coordination.
    """
    # Adaptive costs are calibrated during capture, after this warmup. Exercise
    # fixed draft counts here, then restore the manager for capture and serving.
    adaptive_verification = model_runner.adaptive_verification
    model_runner.adaptive_verification = None
    rejection_sampler = model_runner.rejection_sampler
    adaptive_sampling = (
        rejection_sampler is not None and rejection_sampler.enable_adaptive_verification
    )
    if adaptive_sampling:
        assert rejection_sampler is not None
        rejection_sampler.enable_adaptive_verification = False
    try:
        _warmup_kernels(model_runner, worker_execute_model, worker_sample_tokens)
        # All PP ranks must run the same forwards to match sends and receives.
        if model_runner.speculative_config is not None:
            # Greedy batches skip logits processing, so rejection kernels
            # also see model-dtype logits instead of FP32.
            _warmup_kernels(
                model_runner, worker_execute_model, worker_sample_tokens, greedy=True
            )
    finally:
        model_runner.adaptive_verification = adaptive_verification
        if adaptive_sampling:
            assert rejection_sampler is not None
            rejection_sampler.enable_adaptive_verification = True
    if adaptive_verification is not None:
        adaptive_verification.warmup_budget_kernel()
    acceptance_estimator = getattr(
        model_runner.speculator, "acceptance_estimator", None
    )
    if acceptance_estimator is not None:
        acceptance_estimator.warmup()