# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Callable from contextlib import AbstractContextManager, nullcontext from typing import Any import numpy as np import torch from vllm import PoolingParams, SamplingParams from vllm.logger import init_logger from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange from vllm.utils.math_utils import cdiv from vllm.v1.core.sched.output import ( CachedRequestData, GrammarOutput, NewRequestData, SchedulerOutput, ) from vllm.v1.kv_cache_interface import ( CircularBufferSpec, CrossAttentionSpec, KVCacheSpec, MambaSpec, UniformTypeKVCacheSpecs, ) from vllm.v1.request import Request from vllm.v1.worker.gpu.model_runner import GPUModelRunner logger = init_logger(__name__) def _draft_ids(model_runner, group_id: int, ids: list[int]) -> list[int]: """GLM53-DRAFT-POOL: warmup block ids for a group on the drafter's own, smaller pool, wrapped into it (warmup K/V contents do not matter).""" cfg = model_runner.kv_cache_config n = getattr(cfg, "draft_num_blocks", None) if not n or not set(cfg.draft_layer_names).intersection(cfg.kv_cache_groups[group_id].layer_names): return ids return [1 + (i - 1) % (n - 1) for i in ids] def _reserved_block_count( num_tokens: int, kvcache_spec: KVCacheSpec, *, num_lookahead_tokens: int, max_model_len: int, max_encoder_len: int, ) -> 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. """ 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 num_tokens = min(num_tokens + num_lookahead_tokens, max_model_len) return cdiv(num_tokens, kvcache_spec.block_size) + num_speculative_blocks def _warmup_block_counter( model_runner: GPUModelRunner, ) -> Callable[[int, KVCacheSpec], int]: """Bind `_reserved_block_count` to `model_runner`'s reservation policy.""" num_lookahead_tokens = model_runner.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) 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, ) return block_count 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", ) -> bool: """Run a V2 mixed prefill+decode step through normal scheduler inputs.""" 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 = list(range(decode_prompt_len)) prefill_token_ids = list(range(prefill_len)) kv_cache_groups = model_runner.kv_cache_config.kv_cache_groups num_kv_cache_groups = len(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) if model_runner.kv_cache_config.num_blocks <= required_blocks: logger.warning( "Skipping V2 mixed prefill+decode warmup because only %d KV blocks " "are available for %d required warmup blocks.", model_runner.kv_cache_config.num_blocks, required_blocks, ) return False 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) decode_prefill_output = SchedulerOutput.make_empty() 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(_draft_ids(model_runner, g, _alloc_blocks(n)) for g, n in enumerate(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_prefill_output.num_common_prefix_blocks = [0] * num_kv_cache_groups decode_new_blocks = tuple(_draft_ids(model_runner, g, _alloc_blocks(n)) for g, n in enumerate(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.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=sampling_params, pooling_params=None, block_ids=tuple(_draft_ids(model_runner, g, _alloc_blocks(n)) for g, n in enumerate(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 mixed_output.num_common_prefix_blocks = [0] * num_kv_cache_groups 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 @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) finally: model_runner.adaptive_verification = adaptive_verification if adaptive_sampling: assert rejection_sampler is not None rejection_sampler.enable_adaptive_verification = True def _warmup_kernels( model_runner: GPUModelRunner, worker_execute_model: Callable[[SchedulerOutput], Any], worker_sample_tokens: Callable[[GrammarOutput | None], Any], ) -> None: if model_runner.vllm_config.is_mm_encoder_only: return num_spec_steps = model_runner.num_speculative_steps decode_query_len = model_runner.decode_query_len # Use decode_query_len + 1 tokens so the prefill batch's per-request query # length exceeds decode_query_len, preventing it from being misclassified as # a uniform decode batch. prompt_len = decode_query_len + 1 prompt_token_ids = list(range(prompt_len)) # Upper bound on the decode steps built in `decode_steps` below. num_decode_steps = 1 if not model_runner.is_pooling_model: num_decode_steps = 5 if num_spec_steps > 0 else 3 # Size the block allocation for the worst case: every request advancing # decode_query_len tokens on every decode step. decode_len = prompt_len + num_decode_steps * decode_query_len kv_cache_groups = model_runner.kv_cache_config.kv_cache_groups num_kv_cache_groups = len(kv_cache_groups) # Encoder-decoder models: give each warmup request a dummy encoder input so # cross-attention warms up over a realistic, non-empty key sequence. # The dummy mm_feature is registered in the encoder cache and only its encoder # length is read (not the inputs themselves); the encoder itself is not scheduled. max_encoder_len = getattr(model_runner.model_state, "max_encoder_len", 0) warmup_mm_features: list[MultiModalFeatureSpec] = [] if model_runner.is_encoder_decoder and max_encoder_len: warmup_mm_features = [ MultiModalFeatureSpec( data=None, modality="", identifier="_warmup_encoder", mm_position=PlaceholderRange(offset=0, length=max_encoder_len), ) ] # Compute per-request block counts for each KV cache group. block_count = _warmup_block_counter(model_runner) kv_cache_specs = [g.kv_cache_spec for g in kv_cache_groups] prefill_block_counts = [block_count(prompt_len, s) for s in kv_cache_specs] decode_block_counts = [block_count(decode_len, s) for s in kv_cache_specs] max_blocks_per_req = sum(decode_block_counts) num_reqs = min( model_runner.scheduler_config.max_num_seqs, model_runner.scheduler_config.max_num_batched_tokens // max(prompt_len, decode_query_len), ) if max_blocks_per_req > 0: # Reserve block 0 (null block) and ensure we have enough blocks. # Encoder-only models allocate no KV blocks, so this cap doesn't apply. num_reqs = min( num_reqs, max(1, (model_runner.kv_cache_config.num_blocks - 1) // max_blocks_per_req), ) req_ids = [f"_warmup_{i}_" for i in range(num_reqs)] # SamplingParams exercising all sampling features. if model_runner.is_pooling_model: sampling_params = None pooling_task = model_runner.model_config.get_pooling_task( model_runner.get_supported_tasks() ) pooling_params = PoolingParams(task=pooling_task) pooling_params.verify(model_runner.model_config) else: sampling_params = SamplingParams.for_sampler_warmup() pooling_params = None # Assign distinct block IDs per request per group. 0 null block, start from 1. 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)) # The KV-block zeroing kernel is driven by the scheduler's # new_block_ids_to_zero, so none of the steps below reach it. if model_runner.kv_block_zeroer is not None: model_runner.kv_block_zeroer.warmup(model_runner.kv_cache_config.num_blocks) # Step 1: Prefill all requests with 1 + decode_query_len prompt tokens each. new_reqs = [ NewRequestData.from_request( Request( req_ids[i], prompt_token_ids, sampling_params, pooling_params, mm_features=warmup_mm_features, ), block_ids=tuple(_draft_ids(model_runner, g, _alloc_blocks(n)) for g, n in enumerate(prefill_block_counts)), prefill_token_ids=prompt_token_ids, ) for i in range(num_reqs) ] prefill_output = SchedulerOutput.make_empty() prefill_output.scheduled_new_reqs = new_reqs prefill_output.num_scheduled_tokens = {rid: prompt_len for rid in req_ids} prefill_output.total_num_scheduled_tokens = prompt_len * num_reqs prefill_output.num_common_prefix_blocks = [0] * num_kv_cache_groups # Disable KV connector for warmup run. model_runner.kv_connector.set_disabled(True) worker_execute_model(prefill_output) if not model_runner.is_pooling_model: # Warm up sampler and perform a decode step for non-pooling models. grammar_output = None if model_runner.is_last_pp_rank: # Build a GrammarOutput to exercise the structured output bitmask # kernel during the prefill step. vocab_size = model_runner.model_config.get_vocab_size() bitmask_width = (vocab_size + 31) // 32 grammar_bitmask = np.full( (len(req_ids), bitmask_width), fill_value=-1, dtype=np.int32 ) grammar_output = GrammarOutput( structured_output_request_ids=req_ids, grammar_bitmask=grammar_bitmask ) worker_sample_tokens(grammar_output) # Per-request state carried across the decode steps. req_computed = [prompt_len] * num_reqs req_blocks = [list(prefill_block_counts) for _ in range(num_reqs)] def _run_decode_step(indices: list[int], spec_flags: list[bool]) -> None: """Decode `indices`, spec-decoding the ones flagged in `spec_flags`.""" cached_req_data = CachedRequestData.make_empty() cached_req_data.req_ids = [req_ids[i] for i in indices] cached_req_data.num_computed_tokens = [req_computed[i] for i in indices] cached_req_data.num_output_tokens = [1] * len(indices) cached_req_data.new_block_ids = [] step_num_scheduled_tokens: dict[str, int] = {} step_spec_tokens: dict[str, list[int]] = {} for i, use_spec in zip(indices, spec_flags): num_tokens = decode_query_len if use_spec else 1 after = req_computed[i] + num_tokens deltas = [ block_count(after, spec) - held for spec, held in zip(kv_cache_specs, req_blocks[i]) ] cached_req_data.new_block_ids.append( tuple(_draft_ids(model_runner, g, _alloc_blocks(n)) for g, n in enumerate(deltas)) if any(deltas) else None ) req_blocks[i] = [ held + delta for held, delta in zip(req_blocks[i], deltas) ] step_num_scheduled_tokens[req_ids[i]] = num_tokens if use_spec: step_spec_tokens[req_ids[i]] = [0] * num_spec_steps decode_output = SchedulerOutput.make_empty() decode_output.scheduled_cached_reqs = cached_req_data decode_output.num_scheduled_tokens = step_num_scheduled_tokens decode_output.scheduled_spec_decode_tokens = step_spec_tokens decode_output.total_num_scheduled_tokens = sum( step_num_scheduled_tokens.values() ) decode_output.num_common_prefix_blocks = [0] * num_kv_cache_groups worker_execute_model(decode_output) worker_sample_tokens(None) for i, use_spec in zip(indices, spec_flags): req_computed[i] += decode_query_len if use_spec else 1 all_indices = list(range(num_reqs)) use_spec_decode = num_spec_steps > 0 # Decode steps to warm, as (request indices, per-request spec flag). # Under spec decoding the scheduler drops requests the drafter proposed # nothing for, so warm each batch shape with and without draft tokens. decode_steps: list[tuple[list[int], list[bool]]] = [ (all_indices, [use_spec_decode] * num_reqs), ] if num_reqs >= 2: # Mixed spec / non-spec: GDN and KDA reclassify the non-spec decode # as a prefill and split the batch into spec/non-spec token indices. decode_steps.append(([0, 1], [use_spec_decode, False])) if use_spec_decode: # Exercise the model paths that split a batch by whether each # request received draft tokens. decode_steps.append(([0, 1], [False, False])) if num_reqs > 1: decode_steps.append(([0], [use_spec_decode])) if use_spec_decode: decode_steps.append(([0], [False])) elif use_spec_decode: decode_steps.append(([0], [False])) for step_indices, step_spec_flags in decode_steps: _run_decode_step(step_indices, step_spec_flags) # Clean up - process finish_req_ids. cleanup_output = SchedulerOutput.make_empty() cleanup_output.finished_req_ids = set(req_ids) worker_execute_model(cleanup_output) model_runner.kv_connector.set_disabled(False) torch.accelerator.synchronize()