| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from collections.abc import Callable |
| 5 | from contextlib import AbstractContextManager, nullcontext |
| 6 | from typing import Any |
| 7 | |
| 8 | import numpy as np |
| 9 | import torch |
| 10 | |
| 11 | from vllm import PoolingParams, SamplingParams |
| 12 | from vllm.logger import init_logger |
| 13 | from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange |
| 14 | from vllm.utils.math_utils import cdiv |
| 15 | from vllm.v1.core.sched.output import ( |
| 16 | CachedRequestData, |
| 17 | GrammarOutput, |
| 18 | NewRequestData, |
| 19 | SchedulerOutput, |
| 20 | ) |
| 21 | from vllm.v1.kv_cache_interface import ( |
| 22 | CircularBufferSpec, |
| 23 | CrossAttentionSpec, |
| 24 | KVCacheSpec, |
| 25 | MambaSpec, |
| 26 | UniformTypeKVCacheSpecs, |
| 27 | ) |
| 28 | from vllm.v1.request import Request |
| 29 | from vllm.v1.worker.gpu.model_runner import GPUModelRunner |
| 30 | |
| 31 | logger = init_logger(__name__) |
| 32 | |
| 33 | |
| 34 | def _reserved_block_count( |
| 35 | num_tokens: int, |
| 36 | kvcache_spec: KVCacheSpec, |
| 37 | *, |
| 38 | num_lookahead_tokens: int, |
| 39 | max_model_len: int, |
| 40 | max_encoder_len: int, |
| 41 | ) -> int: |
| 42 | """Number of blocks the scheduler would hold for a request of `num_tokens`. |
| 43 | |
| 44 | Warmup hand-builds its `SchedulerOutput`s, so it must reserve what |
| 45 | `KVCacheManager.allocate_slots` reserves: the token range plus |
| 46 | `num_lookahead_tokens`, where the speculator writes the KV of its drafts. |
| 47 | """ |
| 48 | if isinstance(kvcache_spec, UniformTypeKVCacheSpecs): |
| 49 | kvcache_spec = kvcache_spec.first_spec |
| 50 | if isinstance(kvcache_spec, CircularBufferSpec): |
| 51 | # Circular caches keep one physical ring block for the request lifetime. |
| 52 | return 1 |
| 53 | if isinstance(kvcache_spec, CrossAttentionSpec): |
| 54 | # Cross-attention blocks cover the encoder sequence only. |
| 55 | return cdiv(max_encoder_len, kvcache_spec.block_size) |
| 56 | num_speculative_blocks = 0 |
| 57 | if isinstance(kvcache_spec, MambaSpec): |
| 58 | # MambaManager appends speculative running-state blocks in every cache |
| 59 | # mode; align mode sizes from the uncapped, lookahead-free token range. |
| 60 | num_speculative_blocks = kvcache_spec.num_speculative_blocks |
| 61 | if kvcache_spec.mamba_cache_mode == "align": |
| 62 | return cdiv(num_tokens, kvcache_spec.block_size) + num_speculative_blocks |
| 63 | num_tokens = min(num_tokens + num_lookahead_tokens, max_model_len) |
| 64 | return cdiv(num_tokens, kvcache_spec.block_size) + num_speculative_blocks |
| 65 | |
| 66 | |
| 67 | def _warmup_block_counter( |
| 68 | model_runner: GPUModelRunner, |
| 69 | ) -> Callable[[int, KVCacheSpec], int]: |
| 70 | """Bind `_reserved_block_count` to `model_runner`'s reservation policy.""" |
| 71 | num_lookahead_tokens = model_runner.vllm_config.num_lookahead_tokens |
| 72 | max_model_len = model_runner.max_model_len |
| 73 | max_encoder_len = getattr(model_runner.model_state, "max_encoder_len", 0) |
| 74 | |
| 75 | def block_count(num_tokens: int, kvcache_spec: KVCacheSpec) -> int: |
| 76 | return _reserved_block_count( |
| 77 | num_tokens, |
| 78 | kvcache_spec, |
| 79 | num_lookahead_tokens=num_lookahead_tokens, |
| 80 | max_model_len=max_model_len, |
| 81 | max_encoder_len=max_encoder_len, |
| 82 | ) |
| 83 | |
| 84 | return block_count |
| 85 | |
| 86 | |
| 87 | def run_mixed_prefill_decode_warmup( |
| 88 | model_runner: GPUModelRunner, |
| 89 | worker_execute_model: Callable[[SchedulerOutput], Any], |
| 90 | worker_sample_tokens: Callable[[GrammarOutput | None], Any], |
| 91 | num_tokens: int, |
| 92 | *, |
| 93 | mixed_step_context: AbstractContextManager[object] | None = None, |
| 94 | req_id_prefix: str = "_v2_mixed_warmup", |
| 95 | ) -> bool: |
| 96 | """Run a V2 mixed prefill+decode step through normal scheduler inputs.""" |
| 97 | if model_runner.is_pooling_model or model_runner.max_num_reqs < 2 or num_tokens < 3: |
| 98 | return False |
| 99 | |
| 100 | decode_req_id = f"{req_id_prefix}_decode_" |
| 101 | prefill_req_id = f"{req_id_prefix}_prefill_" |
| 102 | decode_prompt_len = 2 |
| 103 | decode_scheduled_tokens = 1 |
| 104 | prefill_len = num_tokens - decode_scheduled_tokens |
| 105 | decode_token_ids = list(range(decode_prompt_len)) |
| 106 | prefill_token_ids = list(range(prefill_len)) |
| 107 | |
| 108 | kv_cache_groups = model_runner.kv_cache_config.kv_cache_groups |
| 109 | num_kv_cache_groups = len(kv_cache_groups) |
| 110 | block_count = _warmup_block_counter(model_runner) |
| 111 | kv_cache_specs = [g.kv_cache_spec for g in kv_cache_groups] |
| 112 | decode_prefill_block_counts = [ |
| 113 | block_count(decode_prompt_len, s) for s in kv_cache_specs |
| 114 | ] |
| 115 | decode_block_counts = [ |
| 116 | block_count(decode_prompt_len + decode_scheduled_tokens, s) |
| 117 | for s in kv_cache_specs |
| 118 | ] |
| 119 | decode_block_deltas = [ |
| 120 | decode - prefill |
| 121 | for decode, prefill in zip(decode_block_counts, decode_prefill_block_counts) |
| 122 | ] |
| 123 | prefill_block_counts = [block_count(prefill_len, s) for s in kv_cache_specs] |
| 124 | required_blocks = sum(decode_block_counts) + sum(prefill_block_counts) |
| 125 | if model_runner.kv_cache_config.num_blocks <= required_blocks: |
| 126 | logger.warning( |
| 127 | "Skipping V2 mixed prefill+decode warmup because only %d KV blocks " |
| 128 | "are available for %d required warmup blocks.", |
| 129 | model_runner.kv_cache_config.num_blocks, |
| 130 | required_blocks, |
| 131 | ) |
| 132 | return False |
| 133 | |
| 134 | next_block_id = 1 |
| 135 | |
| 136 | def _alloc_blocks(num_blocks: int) -> list[int]: |
| 137 | nonlocal next_block_id |
| 138 | block_ids = list(range(next_block_id, next_block_id + num_blocks)) |
| 139 | next_block_id += num_blocks |
| 140 | return block_ids |
| 141 | |
| 142 | sampling_params = SamplingParams(max_tokens=2, temperature=0.0) |
| 143 | |
| 144 | decode_prefill_output = SchedulerOutput.make_empty() |
| 145 | decode_prefill_output.scheduled_new_reqs = [ |
| 146 | NewRequestData( |
| 147 | req_id=decode_req_id, |
| 148 | prompt_token_ids=decode_token_ids, |
| 149 | mm_features=[], |
| 150 | sampling_params=sampling_params, |
| 151 | pooling_params=None, |
| 152 | block_ids=tuple(_alloc_blocks(n) for n in decode_prefill_block_counts), |
| 153 | num_computed_tokens=0, |
| 154 | lora_request=None, |
| 155 | prefill_token_ids=decode_token_ids, |
| 156 | ), |
| 157 | ] |
| 158 | decode_prefill_output.num_scheduled_tokens = { |
| 159 | decode_req_id: decode_prompt_len, |
| 160 | } |
| 161 | decode_prefill_output.total_num_scheduled_tokens = decode_prompt_len |
| 162 | decode_prefill_output.num_common_prefix_blocks = [0] * num_kv_cache_groups |
| 163 | |
| 164 | decode_new_blocks = tuple(_alloc_blocks(n) for n in decode_block_deltas) |
| 165 | cached_decode_req = CachedRequestData.make_empty() |
| 166 | cached_decode_req.req_ids = [decode_req_id] |
| 167 | cached_decode_req.num_computed_tokens = [decode_prompt_len] |
| 168 | cached_decode_req.num_output_tokens = [1] |
| 169 | cached_decode_req.new_block_ids = [ |
| 170 | decode_new_blocks if any(decode_block_deltas) else None |
| 171 | ] |
| 172 | |
| 173 | mixed_output = SchedulerOutput.make_empty() |
| 174 | mixed_output.scheduled_cached_reqs = cached_decode_req |
| 175 | mixed_output.scheduled_new_reqs = [ |
| 176 | NewRequestData( |
| 177 | req_id=prefill_req_id, |
| 178 | prompt_token_ids=prefill_token_ids, |
| 179 | mm_features=[], |
| 180 | sampling_params=sampling_params, |
| 181 | pooling_params=None, |
| 182 | block_ids=tuple(_alloc_blocks(n) for n in prefill_block_counts), |
| 183 | num_computed_tokens=0, |
| 184 | lora_request=None, |
| 185 | prefill_token_ids=prefill_token_ids, |
| 186 | ), |
| 187 | ] |
| 188 | mixed_output.num_scheduled_tokens = { |
| 189 | decode_req_id: decode_scheduled_tokens, |
| 190 | prefill_req_id: prefill_len, |
| 191 | } |
| 192 | mixed_output.total_num_scheduled_tokens = num_tokens |
| 193 | mixed_output.num_common_prefix_blocks = [0] * num_kv_cache_groups |
| 194 | |
| 195 | cleanup_output = SchedulerOutput.make_empty() |
| 196 | cleanup_output.finished_req_ids = {decode_req_id, prefill_req_id} |
| 197 | |
| 198 | context = mixed_step_context or nullcontext() |
| 199 | model_runner.kv_connector.set_disabled(True) |
| 200 | try: |
| 201 | worker_execute_model(decode_prefill_output) |
| 202 | worker_sample_tokens(None) |
| 203 | with context: |
| 204 | worker_execute_model(mixed_output) |
| 205 | worker_sample_tokens(None) |
| 206 | worker_execute_model(cleanup_output) |
| 207 | finally: |
| 208 | model_runner.kv_connector.set_disabled(False) |
| 209 | return True |
| 210 | |
| 211 | |
| 212 | @torch.inference_mode() |
| 213 | def warmup_kernels( |
| 214 | model_runner: GPUModelRunner, |
| 215 | worker_execute_model: Callable[[SchedulerOutput], Any], |
| 216 | worker_sample_tokens: Callable[[GrammarOutput | None], Any], |
| 217 | ) -> None: |
| 218 | """Run scheduler-realistic prefill and decode steps to JIT compile kernels. |
| 219 | |
| 220 | We must call the provided worker's execute_model for pipeline parallel |
| 221 | coordination. |
| 222 | """ |
| 223 | # Adaptive costs are calibrated during capture, after this warmup. Exercise |
| 224 | # fixed draft counts here, then restore the manager for capture and serving. |
| 225 | adaptive_verification = model_runner.adaptive_verification |
| 226 | model_runner.adaptive_verification = None |
| 227 | rejection_sampler = model_runner.rejection_sampler |
| 228 | adaptive_sampling = ( |
| 229 | rejection_sampler is not None and rejection_sampler.enable_adaptive_verification |
| 230 | ) |
| 231 | if adaptive_sampling: |
| 232 | assert rejection_sampler is not None |
| 233 | rejection_sampler.enable_adaptive_verification = False |
| 234 | try: |
| 235 | _warmup_kernels(model_runner, worker_execute_model, worker_sample_tokens) |
| 236 | finally: |
| 237 | model_runner.adaptive_verification = adaptive_verification |
| 238 | if adaptive_sampling: |
| 239 | assert rejection_sampler is not None |
| 240 | rejection_sampler.enable_adaptive_verification = True |
| 241 | |
| 242 | |
| 243 | def _warmup_kernels( |
| 244 | model_runner: GPUModelRunner, |
| 245 | worker_execute_model: Callable[[SchedulerOutput], Any], |
| 246 | worker_sample_tokens: Callable[[GrammarOutput | None], Any], |
| 247 | ) -> None: |
| 248 | if model_runner.vllm_config.is_mm_encoder_only: |
| 249 | return |
| 250 | |
| 251 | num_spec_steps = model_runner.num_speculative_steps |
| 252 | decode_query_len = model_runner.decode_query_len |
| 253 | # Use decode_query_len + 1 tokens so the prefill batch's per-request query |
| 254 | # length exceeds decode_query_len, preventing it from being misclassified as |
| 255 | # a uniform decode batch. |
| 256 | prompt_len = decode_query_len + 1 |
| 257 | prompt_token_ids = list(range(prompt_len)) |
| 258 | # Upper bound on the decode steps built in `decode_steps` below. |
| 259 | num_decode_steps = 1 |
| 260 | if not model_runner.is_pooling_model: |
| 261 | num_decode_steps = 5 if num_spec_steps > 0 else 3 |
| 262 | # Size the block allocation for the worst case: every request advancing |
| 263 | # decode_query_len tokens on every decode step. |
| 264 | decode_len = prompt_len + num_decode_steps * decode_query_len |
| 265 | |
| 266 | kv_cache_groups = model_runner.kv_cache_config.kv_cache_groups |
| 267 | num_kv_cache_groups = len(kv_cache_groups) |
| 268 | |
| 269 | # Encoder-decoder models: give each warmup request a dummy encoder input so |
| 270 | # cross-attention warms up over a realistic, non-empty key sequence. |
| 271 | # The dummy mm_feature is registered in the encoder cache and only its encoder |
| 272 | # length is read (not the inputs themselves); the encoder itself is not scheduled. |
| 273 | max_encoder_len = getattr(model_runner.model_state, "max_encoder_len", 0) |
| 274 | warmup_mm_features: list[MultiModalFeatureSpec] = [] |
| 275 | if model_runner.is_encoder_decoder and max_encoder_len: |
| 276 | warmup_mm_features = [ |
| 277 | MultiModalFeatureSpec( |
| 278 | data=None, |
| 279 | modality="", |
| 280 | identifier="_warmup_encoder", |
| 281 | mm_position=PlaceholderRange(offset=0, length=max_encoder_len), |
| 282 | ) |
| 283 | ] |
| 284 | |
| 285 | # Compute per-request block counts for each KV cache group. |
| 286 | block_count = _warmup_block_counter(model_runner) |
| 287 | kv_cache_specs = [g.kv_cache_spec for g in kv_cache_groups] |
| 288 | prefill_block_counts = [block_count(prompt_len, s) for s in kv_cache_specs] |
| 289 | decode_block_counts = [block_count(decode_len, s) for s in kv_cache_specs] |
| 290 | max_blocks_per_req = sum(decode_block_counts) |
| 291 | |
| 292 | num_reqs = min( |
| 293 | model_runner.scheduler_config.max_num_seqs, |
| 294 | model_runner.scheduler_config.max_num_batched_tokens |
| 295 | // max(prompt_len, decode_query_len), |
| 296 | ) |
| 297 | block_tables = getattr(model_runner, "block_tables", None) |
| 298 | null_blocks = block_tables is not None and ( |
| 299 | block_tables.redirect_writes_to_null_block |
| 300 | ) |
| 301 | if max_blocks_per_req > 0 and not null_blocks: |
| 302 | # Reserve block 0 (null block) and ensure we have enough blocks. |
| 303 | # Encoder-only models allocate no KV blocks, so this cap doesn't apply. |
| 304 | num_reqs = min( |
| 305 | num_reqs, |
| 306 | max(1, (model_runner.kv_cache_config.num_blocks - 1) // max_blocks_per_req), |
| 307 | ) |
| 308 | |
| 309 | req_ids = [f"_warmup_{i}_" for i in range(num_reqs)] |
| 310 | |
| 311 | # SamplingParams exercising all sampling features. |
| 312 | if model_runner.is_pooling_model: |
| 313 | sampling_params = None |
| 314 | pooling_task = model_runner.model_config.get_pooling_task( |
| 315 | model_runner.get_supported_tasks() |
| 316 | ) |
| 317 | pooling_params = PoolingParams(task=pooling_task) |
| 318 | pooling_params.verify(model_runner.model_config) |
| 319 | else: |
| 320 | sampling_params = SamplingParams.for_sampler_warmup() |
| 321 | pooling_params = None |
| 322 | |
| 323 | # Assign distinct block IDs per request per group. 0 null block, start from 1. |
| 324 | next_block_id = 1 |
| 325 | |
| 326 | def _alloc_blocks(num_blocks: int) -> list[int]: |
| 327 | nonlocal next_block_id |
| 328 | return list(range(next_block_id, next_block_id := next_block_id + num_blocks)) |
| 329 | |
| 330 | # The KV-block zeroing kernel is driven by the scheduler's |
| 331 | # new_block_ids_to_zero, so none of the steps below reach it. |
| 332 | if model_runner.kv_block_zeroer is not None: |
| 333 | model_runner.kv_block_zeroer.warmup(model_runner.kv_cache_config.num_blocks) |
| 334 | |
| 335 | # Step 1: Prefill all requests with 1 + decode_query_len prompt tokens each. |
| 336 | new_reqs = [ |
| 337 | NewRequestData.from_request( |
| 338 | Request( |
| 339 | req_ids[i], |
| 340 | prompt_token_ids, |
| 341 | sampling_params, |
| 342 | pooling_params, |
| 343 | mm_features=warmup_mm_features, |
| 344 | ), |
| 345 | block_ids=tuple(_alloc_blocks(n) for n in prefill_block_counts), |
| 346 | prefill_token_ids=prompt_token_ids, |
| 347 | ) |
| 348 | for i in range(num_reqs) |
| 349 | ] |
| 350 | |
| 351 | prefill_output = SchedulerOutput.make_empty() |
| 352 | prefill_output.scheduled_new_reqs = new_reqs |
| 353 | prefill_output.num_scheduled_tokens = {rid: prompt_len for rid in req_ids} |
| 354 | prefill_output.total_num_scheduled_tokens = prompt_len * num_reqs |
| 355 | prefill_output.num_common_prefix_blocks = [0] * num_kv_cache_groups |
| 356 | |
| 357 | # Disable KV connector for warmup run. |
| 358 | model_runner.kv_connector.set_disabled(True) |
| 359 | worker_execute_model(prefill_output) |
| 360 | |
| 361 | if not model_runner.is_pooling_model: |
| 362 | # Warm up sampler and perform a decode step for non-pooling models. |
| 363 | |
| 364 | grammar_output = None |
| 365 | if model_runner.is_last_pp_rank: |
| 366 | # Build a GrammarOutput to exercise the structured output bitmask |
| 367 | # kernel during the prefill step. |
| 368 | vocab_size = model_runner.model_config.get_vocab_size() |
| 369 | bitmask_width = (vocab_size + 31) // 32 |
| 370 | grammar_bitmask = np.full( |
| 371 | (len(req_ids), bitmask_width), fill_value=-1, dtype=np.int32 |
| 372 | ) |
| 373 | grammar_output = GrammarOutput( |
| 374 | structured_output_request_ids=req_ids, grammar_bitmask=grammar_bitmask |
| 375 | ) |
| 376 | |
| 377 | worker_sample_tokens(grammar_output) |
| 378 | |
| 379 | # Per-request state carried across the decode steps. |
| 380 | req_computed = [prompt_len] * num_reqs |
| 381 | req_blocks = [list(prefill_block_counts) for _ in range(num_reqs)] |
| 382 | |
| 383 | def _run_decode_step(indices: list[int], spec_flags: list[bool]) -> None: |
| 384 | """Decode `indices`, spec-decoding the ones flagged in `spec_flags`.""" |
| 385 | cached_req_data = CachedRequestData.make_empty() |
| 386 | cached_req_data.req_ids = [req_ids[i] for i in indices] |
| 387 | cached_req_data.num_computed_tokens = [req_computed[i] for i in indices] |
| 388 | cached_req_data.num_output_tokens = [1] * len(indices) |
| 389 | cached_req_data.new_block_ids = [] |
| 390 | |
| 391 | step_num_scheduled_tokens: dict[str, int] = {} |
| 392 | step_spec_tokens: dict[str, list[int]] = {} |
| 393 | for i, use_spec in zip(indices, spec_flags): |
| 394 | num_tokens = decode_query_len if use_spec else 1 |
| 395 | after = req_computed[i] + num_tokens |
| 396 | deltas = [ |
| 397 | block_count(after, spec) - held |
| 398 | for spec, held in zip(kv_cache_specs, req_blocks[i]) |
| 399 | ] |
| 400 | cached_req_data.new_block_ids.append( |
| 401 | tuple(_alloc_blocks(n) for n in deltas) if any(deltas) else None |
| 402 | ) |
| 403 | req_blocks[i] = [ |
| 404 | held + delta for held, delta in zip(req_blocks[i], deltas) |
| 405 | ] |
| 406 | step_num_scheduled_tokens[req_ids[i]] = num_tokens |
| 407 | if use_spec: |
| 408 | step_spec_tokens[req_ids[i]] = [0] * num_spec_steps |
| 409 | |
| 410 | decode_output = SchedulerOutput.make_empty() |
| 411 | decode_output.scheduled_cached_reqs = cached_req_data |
| 412 | decode_output.num_scheduled_tokens = step_num_scheduled_tokens |
| 413 | decode_output.scheduled_spec_decode_tokens = step_spec_tokens |
| 414 | decode_output.total_num_scheduled_tokens = sum( |
| 415 | step_num_scheduled_tokens.values() |
| 416 | ) |
| 417 | decode_output.num_common_prefix_blocks = [0] * num_kv_cache_groups |
| 418 | |
| 419 | worker_execute_model(decode_output) |
| 420 | worker_sample_tokens(None) |
| 421 | |
| 422 | for i, use_spec in zip(indices, spec_flags): |
| 423 | req_computed[i] += decode_query_len if use_spec else 1 |
| 424 | |
| 425 | all_indices = list(range(num_reqs)) |
| 426 | use_spec_decode = num_spec_steps > 0 |
| 427 | |
| 428 | # Decode steps to warm, as (request indices, per-request spec flag). |
| 429 | # Under spec decoding the scheduler drops requests the drafter proposed |
| 430 | # nothing for, so warm each batch shape with and without draft tokens. |
| 431 | decode_steps: list[tuple[list[int], list[bool]]] = [ |
| 432 | (all_indices, [use_spec_decode] * num_reqs), |
| 433 | ] |
| 434 | if num_reqs >= 2: |
| 435 | # Mixed spec / non-spec: GDN and KDA reclassify the non-spec decode |
| 436 | # as a prefill and split the batch into spec/non-spec token indices. |
| 437 | decode_steps.append(([0, 1], [use_spec_decode, False])) |
| 438 | if use_spec_decode: |
| 439 | # Exercise the model paths that split a batch by whether each |
| 440 | # request received draft tokens. |
| 441 | decode_steps.append(([0, 1], [False, False])) |
| 442 | if num_reqs > 1: |
| 443 | decode_steps.append(([0], [use_spec_decode])) |
| 444 | if use_spec_decode: |
| 445 | decode_steps.append(([0], [False])) |
| 446 | elif use_spec_decode: |
| 447 | decode_steps.append(([0], [False])) |
| 448 | |
| 449 | for step_indices, step_spec_flags in decode_steps: |
| 450 | _run_decode_step(step_indices, step_spec_flags) |
| 451 | |
| 452 | # Clean up - process finish_req_ids. |
| 453 | cleanup_output = SchedulerOutput.make_empty() |
| 454 | cleanup_output.finished_req_ids = set(req_ids) |
| 455 | worker_execute_model(cleanup_output) |
| 456 | model_runner.kv_connector.set_disabled(False) |
| 457 | if model_runner.kv_block_zeroer is not None: |
| 458 | model_runner.kv_block_zeroer.zero_block_ids([0]) |
| 459 | torch.accelerator.synchronize() |