| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """NOTE: Coding style guide for this file: |
| 4 | This model runner is shared by all models: text and multimodal, generative |
| 5 | and embedding, public and private. As a result, this file must only contain |
| 6 | code that is common to every model. Model-specific behavior belongs in the |
| 7 | appropriate model-specific files. |
| 8 | |
| 9 | In other words: |
| 10 | * Be paranoid about changing this file. It should remain stable. |
| 11 | * Be even more paranoid about adding new lines. It should remain minimal. |
| 12 | |
| 13 | Even for shared features (for example, different parallelism modes), keep the |
| 14 | complexity out of this path. The less common the feature, the more it should be |
| 15 | hidden. Prefer utility functions defined elsewhere and call them from here, |
| 16 | instead of embedding feature-specific logic directly. |
| 17 | """ |
| 18 | |
| 19 | import functools |
| 20 | import gc |
| 21 | import time |
| 22 | from contextlib import AbstractContextManager |
| 23 | from copy import deepcopy |
| 24 | from dataclasses import replace |
| 25 | from typing import Any, NamedTuple |
| 26 | |
| 27 | import numpy as np |
| 28 | import torch |
| 29 | import torch.nn as nn |
| 30 | |
| 31 | import vllm.envs as envs |
| 32 | from vllm.compilation.counter import compilation_counter |
| 33 | from vllm.compilation.cuda_graph import CUDAGraphStat |
| 34 | from vllm.compilation.wrapper import compile_model_with_stock_torch |
| 35 | from vllm.config import VllmConfig |
| 36 | from vllm.config.compilation import CompilationMode, CUDAGraphMode |
| 37 | from vllm.distributed.aux_output_connector.worker import ( |
| 38 | AuxOutputWorkerConnector, |
| 39 | get_aux_output_connector, |
| 40 | ) |
| 41 | from vllm.distributed.parallel_state import get_dcp_group, get_pp_group |
| 42 | from vllm.forward_context import BatchDescriptor, set_forward_context |
| 43 | from vllm.logger import init_logger |
| 44 | from vllm.model_executor.layers.fused_moe.all2all_utils import get_ep_all2all_manager |
| 45 | from vllm.model_executor.layers.mamba.ops.ssu_dispatch import ( |
| 46 | initialize_mamba_ssu_backend, |
| 47 | ) |
| 48 | from vllm.model_executor.model_loader import get_model_loader |
| 49 | from vllm.model_executor.models.interfaces import requires_raw_input_tokens |
| 50 | from vllm.model_executor.offloader import ( |
| 51 | create_offloader, |
| 52 | get_offloader, |
| 53 | set_offloader, |
| 54 | ) |
| 55 | from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry |
| 56 | from vllm.multimodal import MULTIMODAL_REGISTRY |
| 57 | from vllm.multimodal.encoder_budget import ( |
| 58 | MultiModalBudget, |
| 59 | ) |
| 60 | from vllm.platforms import current_platform |
| 61 | from vllm.sequence import IntermediateTensors |
| 62 | from vllm.tasks import SupportedTask |
| 63 | from vllm.utils.gc_utils import freeze_gc_for_cudagraph_capture |
| 64 | from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib |
| 65 | from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, async_tensor_h2d |
| 66 | from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput |
| 67 | from vllm.v1.kv_cache_interface import ( |
| 68 | KVCacheConfig, |
| 69 | MambaSpec, |
| 70 | UniformTypeKVCacheSpecs, |
| 71 | ) |
| 72 | from vllm.v1.outputs import ( |
| 73 | DraftTokenIds, |
| 74 | ECConnectorOutput, |
| 75 | ModelRunnerOutput, |
| 76 | ) |
| 77 | from vllm.v1.watermarking import create_watermarker |
| 78 | from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler |
| 79 | from vllm.v1.watermarking.spec_decode import ( |
| 80 | create_speculative_target_watermarker, |
| 81 | speculative_target_watermark_key, |
| 82 | ) |
| 83 | from vllm.v1.worker.block_table import get_block_table_width |
| 84 | from vllm.v1.worker.cp_utils import check_attention_cp_compatibility |
| 85 | from vllm.v1.worker.gpu import pcp_manager as pcp |
| 86 | from vllm.v1.worker.gpu.async_utils import ( |
| 87 | AsyncOutput, |
| 88 | AsyncPoolingOutput, |
| 89 | StepTimingCollector, |
| 90 | ) |
| 91 | from vllm.v1.worker.gpu.attn_utils import ( |
| 92 | FastPrefillHelper, |
| 93 | build_slot_mappings_by_layer, |
| 94 | get_kv_cache_spec, |
| 95 | init_attn_backend, |
| 96 | init_kv_cache, |
| 97 | ) |
| 98 | from vllm.v1.worker.gpu.block_table import BlockTables |
| 99 | from vllm.v1.worker.gpu.buffer_utils import set_default_max_concurrency |
| 100 | from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens |
| 101 | from vllm.v1.worker.gpu.cudagraph_utils import ( |
| 102 | BatchExecutionDescriptor, |
| 103 | ModelCudaGraphManager, |
| 104 | has_compiled_submodule, |
| 105 | make_cudagraph_stats, |
| 106 | ) |
| 107 | from vllm.v1.worker.gpu.cudagraph_utils import ( |
| 108 | profile_cudagraph_memory as _profile_cudagraph_memory, |
| 109 | ) |
| 110 | from vllm.v1.worker.gpu.dp_utils import DPSyncState, dispatch_cg_and_sync_dp |
| 111 | from vllm.v1.worker.gpu.ec_connector import get_ec_connector |
| 112 | from vllm.v1.worker.gpu.eplb_utils import EPLBController, step_eplb_after |
| 113 | from vllm.v1.worker.gpu.input_batch import ( |
| 114 | InputBatch, |
| 115 | InputBuffers, |
| 116 | combine_sampled_and_draft_tokens, |
| 117 | expand_idx_mapping, |
| 118 | post_update, |
| 119 | post_update_num_computed_tokens, |
| 120 | prepare_pos_seq_lens, |
| 121 | prepare_prefill_inputs, |
| 122 | set_dummy_context, |
| 123 | ) |
| 124 | from vllm.v1.worker.gpu.kv_connector import ( |
| 125 | NO_OP_KV_CONNECTOR, |
| 126 | KVConnector, |
| 127 | get_kv_connector, |
| 128 | ) |
| 129 | from vllm.v1.worker.gpu.lora_utils import ( |
| 130 | LoraState, |
| 131 | create_lora_capture_hook, |
| 132 | get_lora_capture_cases, |
| 133 | get_num_active_loras_for_dispatch, |
| 134 | ) |
| 135 | from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache |
| 136 | from vllm.v1.worker.gpu.mm.lora import set_active_mm_loras |
| 137 | from vllm.v1.worker.gpu.model_states import init_model_state |
| 138 | from vllm.v1.worker.gpu.pool.pooling_runner import PoolingRunner |
| 139 | from vllm.v1.worker.gpu.pp_utils import PPHandler |
| 140 | from vllm.v1.worker.gpu.sample.batch_shard import ( |
| 141 | BatchSharder, |
| 142 | all_to_all_logits, |
| 143 | gather_sampler_output, |
| 144 | ) |
| 145 | from vllm.v1.worker.gpu.sample.logits_processor import build_custom_logits_processors |
| 146 | from vllm.v1.worker.gpu.sample.output import SamplerOutput |
| 147 | from vllm.v1.worker.gpu.sample.prompt_logprob import PromptLogprobsWorker |
| 148 | from vllm.v1.worker.gpu.sample.sampler import Sampler |
| 149 | from vllm.v1.worker.gpu.shutdown import free_before_shutdown |
| 150 | from vllm.v1.worker.gpu.spec_decode import init_speculator |
| 151 | from vllm.v1.worker.gpu.spec_decode.adaptive_verification import ( |
| 152 | AdaptiveVerificationManager, |
| 153 | maybe_create_adaptive_verification_manager, |
| 154 | resolve_adaptive_cudagraph_mode, |
| 155 | ) |
| 156 | from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import ( |
| 157 | set_eagle3_aux_hidden_state_layers, |
| 158 | verify_supports_aux_hidden_states_over_pp, |
| 159 | ) |
| 160 | from vllm.v1.worker.gpu.spec_decode.rejection_sampler import ( |
| 161 | RejectionSampler, |
| 162 | get_max_chunk_logits, |
| 163 | ) |
| 164 | from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator |
| 165 | from vllm.v1.worker.gpu.spec_decode.utils import DraftTokensHandler |
| 166 | from vllm.v1.worker.gpu.states import RequestState |
| 167 | from vllm.v1.worker.gpu.structured_outputs import ( |
| 168 | StructuredOutputsWorker, |
| 169 | grammar_invalid_drafts, |
| 170 | ) |
| 171 | from vllm.v1.worker.gpu.ubatch_utils import ( |
| 172 | UBatchRunner, |
| 173 | UBatchState, |
| 174 | maybe_build_ubatch_runner, |
| 175 | ) |
| 176 | from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin |
| 177 | from vllm.v1.worker.utils import ( |
| 178 | KVBlockZeroer, |
| 179 | clear_layer_kv_caches, |
| 180 | copy_kv_cache_blocks_inplace, |
| 181 | get_uniform_decode_token_count, |
| 182 | ) |
| 183 | from vllm.v1.worker.workspace import lock_workspace, use_workspace_lane |
| 184 | |
| 185 | logger = init_logger(__name__) |
| 186 | |
| 187 | |
| 188 | class GPUModelRunner(LoRAModelRunnerMixin): |
| 189 | def __init__(self, vllm_config: VllmConfig, device: torch.device): |
| 190 | self.vllm_config = vllm_config |
| 191 | self.model_config = vllm_config.model_config |
| 192 | self.cache_config = vllm_config.cache_config |
| 193 | self.compilation_config = vllm_config.compilation_config |
| 194 | self.lora_config = vllm_config.lora_config |
| 195 | self.load_config = vllm_config.load_config |
| 196 | self.parallel_config = vllm_config.parallel_config |
| 197 | self.scheduler_config = vllm_config.scheduler_config |
| 198 | self.speculative_config = vllm_config.speculative_config |
| 199 | self._draft_workspace_lane = int( |
| 200 | self.speculative_config is not None and self.speculative_config.use_dspark() |
| 201 | ) |
| 202 | self.observability_config = vllm_config.observability_config |
| 203 | self.jit_warmup_registry = JitWarmupRegistry(vllm_config) |
| 204 | |
| 205 | self.device = device |
| 206 | self.dtype = self.model_config.dtype |
| 207 | self.kv_cache_dtype = self.dtype |
| 208 | if self.cache_config.cache_dtype != "auto": |
| 209 | # Quantized KV cache. |
| 210 | self.kv_cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[ |
| 211 | self.cache_config.cache_dtype |
| 212 | ] |
| 213 | |
| 214 | # Lazily initialized in _init_kv_zero_meta() when the KV cache needs |
| 215 | # zeroing (e.g. hybrid models with fp8 KV cache). |
| 216 | self.kv_block_zeroer: KVBlockZeroer | None = None |
| 217 | |
| 218 | self.vocab_size = self.model_config.get_vocab_size() |
| 219 | self.max_model_len = self.model_config.max_model_len |
| 220 | self.max_num_tokens = self.scheduler_config.max_num_batched_tokens |
| 221 | self.max_num_reqs = self.scheduler_config.max_num_seqs |
| 222 | self.is_encoder_decoder = self.model_config.is_encoder_decoder |
| 223 | |
| 224 | self.output_copy_stream = torch.cuda.Stream(self.device) |
| 225 | |
| 226 | # Pipeline parallelism. |
| 227 | self.use_pp = self.parallel_config.pipeline_parallel_size > 1 |
| 228 | self.is_first_pp_rank = get_pp_group().is_first_rank |
| 229 | self.is_last_pp_rank = get_pp_group().is_last_rank |
| 230 | |
| 231 | # Size the UVA buffer pools to the max number of concurrent in-flight |
| 232 | # steps. Must run before any pooled buffer is constructed |
| 233 | set_default_max_concurrency(vllm_config.max_concurrent_batches) |
| 234 | |
| 235 | # PP broadcast/recv helper. Runs the collective on a side stream. |
| 236 | self.pp_handler: PPHandler | None = None |
| 237 | |
| 238 | # Persistent buffer for intermediate tensors (non-first PP ranks). |
| 239 | self.intermediate_tensors: IntermediateTensors | None = None |
| 240 | |
| 241 | # Dual batch overlap. Created in initialize_kv_cache(), once everything |
| 242 | # it runs the microbatched forward with exists. |
| 243 | self.ubatch_runner: UBatchRunner | None = None |
| 244 | |
| 245 | # Detect EP all2all peer faults to prevent emitting corrupted output. |
| 246 | # Only meaningful for MoE + DP with an FT-capable all2all backend. |
| 247 | self.check_ep_fault = False |
| 248 | if self.dp_size > 1 and self.model_config.is_moe: |
| 249 | self.check_ep_fault = get_ep_all2all_manager().support_fault_tolerance |
| 250 | |
| 251 | # Decode context parallelism. |
| 252 | self.dcp_size = self.parallel_config.decode_context_parallel_size |
| 253 | self.use_dcp = self.dcp_size > 1 |
| 254 | self.dcp_rank = get_dcp_group().rank_in_group if self.use_dcp else 0 |
| 255 | self.cp_interleave = self.parallel_config.cp_kv_cache_interleave_size |
| 256 | |
| 257 | # Multimodal |
| 258 | self.mm_registry = MULTIMODAL_REGISTRY |
| 259 | self.supports_mm_inputs = self.model_config.supports_multimodal_inputs |
| 260 | self.uses_inputs_embeds = ( |
| 261 | self.supports_mm_inputs or self.model_config.enable_prompt_embeds |
| 262 | ) |
| 263 | self.encoder_cache = None |
| 264 | if self.supports_mm_inputs and self.is_first_pp_rank: |
| 265 | self.encoder_cache = EncoderCache() |
| 266 | self.ec_connector = get_ec_connector(vllm_config, self.encoder_cache) |
| 267 | |
| 268 | # Speculative decoding. |
| 269 | self.speculator = None |
| 270 | self.use_aux_hidden_state_outputs = False |
| 271 | self.num_speculative_steps = vllm_config.num_speculative_tokens |
| 272 | if self.speculative_config is not None: |
| 273 | if self.is_last_pp_rank: |
| 274 | self.speculator = init_speculator(self.vllm_config, self.device) |
| 275 | |
| 276 | if self.speculative_config.method in ( |
| 277 | "eagle3", |
| 278 | "dflash", |
| 279 | "dspark", |
| 280 | "extract_hidden_states", |
| 281 | ): |
| 282 | # Drafting may require auxiliary hidden states from target model outputs |
| 283 | self.use_aux_hidden_state_outputs = True |
| 284 | |
| 285 | # Draft tokens propagation - for spec-dec + struct outputs. |
| 286 | self.draft_tokens_handler = DraftTokensHandler(self.device) |
| 287 | |
| 288 | self.pcp_manager: pcp.PCPManager | None = None |
| 289 | |
| 290 | # Pooling models. |
| 291 | self.is_pooling_model = self.model_config.runner_type == "pooling" |
| 292 | self.pooling_runner: PoolingRunner | None = None |
| 293 | |
| 294 | num_prefill_lookahead = max(1, vllm_config.num_prefill_lookahead_tokens) |
| 295 | |
| 296 | self.step_timing = StepTimingCollector() |
| 297 | |
| 298 | # General request states. |
| 299 | self.req_states = RequestState( |
| 300 | max_num_reqs=self.max_num_reqs, |
| 301 | max_model_len=self.max_model_len, |
| 302 | max_num_batched_tokens=self.max_num_tokens, |
| 303 | num_speculative_steps=self.num_speculative_steps, |
| 304 | vocab_size=self.vocab_size, |
| 305 | device=self.device, |
| 306 | num_prefill_lookahead=num_prefill_lookahead, |
| 307 | ) |
| 308 | self.adaptive_verification: AdaptiveVerificationManager | None = None |
| 309 | self.input_buffers = InputBuffers( |
| 310 | max_num_reqs=self.max_num_reqs, |
| 311 | max_num_tokens=self.max_num_tokens, |
| 312 | device=self.device, |
| 313 | ) |
| 314 | self.fast_prefill: FastPrefillHelper | None = None |
| 315 | if self.use_pp: |
| 316 | self.pp_handler = PPHandler( |
| 317 | max_num_reqs=self.max_num_reqs, |
| 318 | num_speculative_steps=self.num_speculative_steps, |
| 319 | device=self.device, |
| 320 | ) |
| 321 | |
| 322 | # Samplers and decode_query_len created in load_model() after |
| 323 | # model_state exists (num_new_sampled_tokens_per_step from ModelState). |
| 324 | self.sampler: Sampler | None = None |
| 325 | self.rejection_sampler: RejectionSampler | None = None |
| 326 | self.batch_sharder: BatchSharder | None = None |
| 327 | self.prompt_logprobs_worker: PromptLogprobsWorker | None = None |
| 328 | self.structured_outputs_worker: StructuredOutputsWorker | None = None |
| 329 | self.cudagraph_manager: ModelCudaGraphManager | None = None |
| 330 | |
| 331 | # LoRA-related workers. |
| 332 | self.lora_state = LoraState(max_num_reqs=self.max_num_reqs) |
| 333 | self.lora_capture_cases = [0] |
| 334 | if self.lora_config: |
| 335 | self.lora_capture_cases = get_lora_capture_cases( |
| 336 | self.lora_config, self.compilation_config |
| 337 | ) |
| 338 | |
| 339 | # KV Connector if configured. |
| 340 | self.kv_connector: KVConnector = NO_OP_KV_CONNECTOR |
| 341 | |
| 342 | # For transferring state from execute_model to subsequent sample_tokens call. |
| 343 | self.execute_model_state: ExecuteModelState | None = None |
| 344 | |
| 345 | # Expert parallelism load balancer. |
| 346 | self.eplb = EPLBController(self.parallel_config, self.device) |
| 347 | # The AuxOutput Connector owns R3 capture, copying, and storage. |
| 348 | self.aux_output_connector: AuxOutputWorkerConnector | None = None |
| 349 | |
| 350 | set_offloader(create_offloader(self.vllm_config.offload_config)) |
| 351 | |
| 352 | def update_max_model_len(self, max_model_len: int) -> None: |
| 353 | self.max_model_len = max_model_len |
| 354 | self.req_states.max_model_len = max_model_len |
| 355 | |
| 356 | def get_supported_tasks(self) -> tuple[SupportedTask, ...]: |
| 357 | tasks: list[SupportedTask] = [] |
| 358 | if self.model_config.runner_type == "generate": |
| 359 | tasks.extend(self.model_state.get_supported_generation_tasks()) |
| 360 | if self.is_pooling_model: |
| 361 | # Do not rely on pooling_runner here, since this information is needed |
| 362 | # on the first PP rank, while pooling_runner is only initialized |
| 363 | # on the last PP rank. |
| 364 | tasks.extend(PoolingRunner.get_supported_tasks(self.model)) |
| 365 | return tuple(tasks) |
| 366 | |
| 367 | def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None: |
| 368 | time_before_load = time.perf_counter() |
| 369 | if load_dummy_weights: |
| 370 | self.load_config.load_format = "dummy" |
| 371 | self.eplb.prepare_load() |
| 372 | eplb_models_added = False |
| 373 | with DeviceMemoryProfiler() as m: |
| 374 | model_loader = get_model_loader(self.vllm_config.load_config) |
| 375 | logger.info_once("Loading model from scratch...") |
| 376 | |
| 377 | # Capture warmup providers selected while constructing the model. |
| 378 | with self.jit_warmup_registry.activate(): |
| 379 | self.model = model_loader.load_model( |
| 380 | vllm_config=self.vllm_config, |
| 381 | model_config=self.vllm_config.model_config, |
| 382 | ) |
| 383 | if self.lora_config: |
| 384 | self.model = self.load_lora_model( |
| 385 | self.model, self.vllm_config, self.device |
| 386 | ) |
| 387 | |
| 388 | if self.use_aux_hidden_state_outputs: |
| 389 | assert self.speculative_config is not None |
| 390 | set_eagle3_aux_hidden_state_layers(self.model, self.speculative_config) |
| 391 | if self.use_pp: |
| 392 | assert self.speculative_config.method is not None |
| 393 | verify_supports_aux_hidden_states_over_pp( |
| 394 | self.model, self.speculative_config.method |
| 395 | ) |
| 396 | assert self.pp_handler is not None |
| 397 | self.pp_handler.configure_aux_hidden_state_relay(self.model) |
| 398 | if isinstance(self.speculator, DraftModelSpeculator): |
| 399 | with use_workspace_lane(self._draft_workspace_lane): |
| 400 | self.speculator.load_model(self.model) |
| 401 | eplb_models_added = self.eplb.maybe_register_speculator( |
| 402 | self.speculator, self.speculative_config, load_dummy_weights |
| 403 | ) |
| 404 | time_after_load = time.perf_counter() |
| 405 | |
| 406 | self.model_memory_usage = m.consumed_memory |
| 407 | logger.info( |
| 408 | "Model loading took %s GiB memory and %.6f seconds", |
| 409 | format_gib(m.consumed_memory), |
| 410 | time_after_load - time_before_load, |
| 411 | ) |
| 412 | |
| 413 | # Initialize the components that require the model. |
| 414 | self.model_state = init_model_state( |
| 415 | self.vllm_config, self.model, self.encoder_cache, self.device |
| 416 | ) |
| 417 | |
| 418 | self.decode_query_len = ( |
| 419 | self.num_speculative_steps |
| 420 | + self.model_state.num_new_sampled_tokens_per_step |
| 421 | ) |
| 422 | |
| 423 | if self.parallel_config.enable_batch_sharded_sampling: |
| 424 | if hasattr(self.model, "compute_logits_local"): |
| 425 | self.batch_sharder = BatchSharder( |
| 426 | max_num_reqs=self.max_num_reqs, |
| 427 | max_num_logits_per_req=self.decode_query_len, |
| 428 | device=self.device, |
| 429 | ) |
| 430 | logger.info("Batch-sharded sampling enabled.") |
| 431 | else: |
| 432 | logger.warning_once( |
| 433 | "Disabling batch-sharded sampling: %s does not implement " |
| 434 | "compute_logits_local", |
| 435 | type(self.model).__name__, |
| 436 | ) |
| 437 | |
| 438 | # Initialize samplers. Model states may override via custom_sampler(). |
| 439 | config_processors = self.model_config.logits_processors or () |
| 440 | custom_logits_processors = build_custom_logits_processors( |
| 441 | self.vllm_config, self.req_states, self.is_pooling_model, config_processors |
| 442 | ) |
| 443 | if self.is_last_pp_rank and not self.is_pooling_model: |
| 444 | # V2 bypasses TopKTopPSampler, which registers native warmups. |
| 445 | # ROCm only: on CUDA these warmups add ~2 min to every engine start. |
| 446 | if current_platform.is_rocm(): |
| 447 | from vllm.v1.sample.ops.topk_topp_sampler import ( |
| 448 | register_top_k_top_p_warmups, |
| 449 | ) |
| 450 | |
| 451 | with self.jit_warmup_registry.activate(): |
| 452 | register_top_k_top_p_warmups() |
| 453 | |
| 454 | sampler_kwargs: dict[str, Any] = { |
| 455 | "vllm_config": self.vllm_config, |
| 456 | "max_num_reqs": self.max_num_reqs, |
| 457 | "vocab_size": self.vocab_size, |
| 458 | "device": self.device, |
| 459 | "req_states": self.req_states, |
| 460 | "logprobs_mode": self.model_config.logprobs_mode, |
| 461 | "num_speculative_tokens": self.decode_query_len, |
| 462 | "use_fp64_gumbel": self.model_config.use_fp64_gumbel, |
| 463 | "enable_trace_replay": self.model_config.enable_trace_replay, |
| 464 | "return_sampling_mask": self.model_config.return_sampling_mask, |
| 465 | "custom_logits_processors": custom_logits_processors, |
| 466 | } |
| 467 | if self.vllm_config.watermark_config is None: |
| 468 | self.sampler = Sampler(**sampler_kwargs) |
| 469 | else: |
| 470 | wm_config = self.vllm_config.watermark_config |
| 471 | watermarker = create_watermarker(wm_config) |
| 472 | if self.speculative_config is not None: |
| 473 | watermarker = create_speculative_target_watermarker(watermarker) |
| 474 | self.sampler = GPUWatermarkSampler( |
| 475 | watermarker, |
| 476 | deduplicate_contexts=wm_config.deduplicate_contexts, |
| 477 | deduplicate_contexts_max_history=( |
| 478 | wm_config.deduplicate_contexts_max_history |
| 479 | ), |
| 480 | **sampler_kwargs, |
| 481 | ) |
| 482 | custom = self.model_state.custom_sampler(self.sampler) |
| 483 | |
| 484 | if custom: |
| 485 | self.vllm_config._check_watermarking_unsupported(custom_sampler=True) |
| 486 | self.sampler, self.rejection_sampler = custom |
| 487 | elif self.speculative_config is not None: |
| 488 | self.rejection_sampler = RejectionSampler( |
| 489 | self.sampler, |
| 490 | self.speculative_config, |
| 491 | self.device, |
| 492 | watermark_key=speculative_target_watermark_key( |
| 493 | self.vllm_config.watermark_config |
| 494 | ), |
| 495 | ) |
| 496 | self.prompt_logprobs_worker = PromptLogprobsWorker( |
| 497 | self.max_num_reqs, |
| 498 | self.device, |
| 499 | logprobs_mode=self.model_config.logprobs_mode, |
| 500 | ) |
| 501 | self.structured_outputs_worker = StructuredOutputsWorker( |
| 502 | max_num_logits=self.max_num_reqs * self.decode_query_len, |
| 503 | vocab_size=self.vocab_size, |
| 504 | device=self.device, |
| 505 | mask_stride=self.decode_query_len, |
| 506 | num_bonus_tokens=self.model_state.num_new_sampled_tokens_per_step, |
| 507 | ) |
| 508 | |
| 509 | if self.is_pooling_model and self.is_last_pp_rank: |
| 510 | self.pooling_runner = PoolingRunner(self.model, self.vllm_config) |
| 511 | eplb_models_added |= self.eplb.maybe_register_model( |
| 512 | self.model, self.model_config |
| 513 | ) |
| 514 | self.eplb.maybe_start_async_loop(eplb_models_added, load_dummy_weights) |
| 515 | |
| 516 | if not self.is_first_pp_rank: |
| 517 | # For non-first PP ranks, create intermediate tensors sized |
| 518 | # for the max capture size so they can be sliced per batch. |
| 519 | # Save as persistent member so runtime can copy received data |
| 520 | # into the same addresses that the CUDA graphs captured. |
| 521 | self.intermediate_tensors = self.model.make_empty_intermediate_tensors( |
| 522 | batch_size=self.max_num_tokens, |
| 523 | dtype=self.model_config.dtype, |
| 524 | device=self.device, |
| 525 | ) |
| 526 | |
| 527 | get_offloader().post_init() |
| 528 | |
| 529 | if self.compilation_config.mode == CompilationMode.STOCK_TORCH_COMPILE: |
| 530 | compile_model_with_stock_torch(self.model, self.vllm_config) |
| 531 | |
| 532 | @property |
| 533 | def dp_size(self) -> int: |
| 534 | # Not cached: elastic EP rewrites parallel_config in place on reconfigure. |
| 535 | return self.parallel_config.data_parallel_size |
| 536 | |
| 537 | @property |
| 538 | def dp_rank(self) -> int: |
| 539 | # Not cached: elastic EP rewrites parallel_config in place on reconfigure. |
| 540 | return self.parallel_config.data_parallel_rank |
| 541 | |
| 542 | def get_model(self) -> nn.Module: |
| 543 | return self.model |
| 544 | |
| 545 | def get_draft_model(self) -> nn.Module | None: |
| 546 | speculator = self.speculator |
| 547 | if not isinstance(speculator, DraftModelSpeculator): |
| 548 | return None |
| 549 | return speculator.model |
| 550 | |
| 551 | def reload_weights(self, *args, **kwargs) -> None: |
| 552 | # TODO(Wentao): Use full version instead of import when fully migrated to v2 |
| 553 | from vllm.v1.worker.gpu_model_runner import GPUModelRunner as GPUModelRunnerV1 |
| 554 | |
| 555 | GPUModelRunnerV1.reload_weights(self, *args, **kwargs) # type: ignore[arg-type] |
| 556 | |
| 557 | def update_config(self, *args, **kwargs) -> None: |
| 558 | # TODO(Wentao): Use full version instead of import when fully migrated to v2 |
| 559 | from vllm.v1.worker.gpu_model_runner import GPUModelRunner as GPUModelRunnerV1 |
| 560 | |
| 561 | GPUModelRunnerV1.update_config(self, *args, **kwargs) # type: ignore[arg-type] |
| 562 | |
| 563 | # v2 reads config via self.vllm_config (e.g. in load_model), so keep it |
| 564 | # in sync with the attributes the v1 helper just replaced. |
| 565 | self.vllm_config.model_config = self.model_config |
| 566 | self.vllm_config.load_config = self.load_config |
| 567 | |
| 568 | @functools.cached_property |
| 569 | def main_stream(self) -> torch.cuda.Stream: |
| 570 | # Cache the default CUDA stream to avoid lookup overhead. |
| 571 | return torch.cuda.current_stream(self.device) |
| 572 | |
| 573 | def get_encoder_timing_stats(self) -> dict[str, dict[str, float | int]]: |
| 574 | encoder_runner = getattr(self.model_state, "encoder_runner", None) |
| 575 | if encoder_runner is None: |
| 576 | return {} |
| 577 | return encoder_runner.get_encoder_timing_stats() |
| 578 | |
| 579 | def get_kv_cache_spec(self): |
| 580 | kv_cache_spec = get_kv_cache_spec(self.vllm_config) |
| 581 | # A draft running at DCP=1 under a DCP target keeps a replicated cache. |
| 582 | if ( |
| 583 | isinstance(self.speculator, DraftModelSpeculator) |
| 584 | and self.speculator.dcp_size == 1 |
| 585 | and self.dcp_size > 1 |
| 586 | ): |
| 587 | for name in self.speculator.draft_attn_layer_names & kv_cache_spec.keys(): |
| 588 | kv_cache_spec[name] = replace(kv_cache_spec[name], dcp_sharded=False) |
| 589 | return kv_cache_spec |
| 590 | |
| 591 | def initialize_kv_cache( |
| 592 | self, |
| 593 | kv_cache_config: KVCacheConfig, |
| 594 | is_profiling: bool = False, |
| 595 | kv_cache_allocation_context: AbstractContextManager | None = None, |
| 596 | ) -> None: |
| 597 | # GPUWorker finalizes the PD interleave before KV cache initialization. |
| 598 | self.cp_interleave = self.parallel_config.cp_kv_cache_interleave_size |
| 599 | kv_cache_config = deepcopy(kv_cache_config) |
| 600 | self.kv_cache_config = kv_cache_config |
| 601 | |
| 602 | block_table_max_model_len = self.max_model_len |
| 603 | if self.is_encoder_decoder: |
| 604 | # Cross-attention block tables need to index encoder tokens, which |
| 605 | # can exceed the decoder's max_model_len. |
| 606 | block_table_max_model_len = max( |
| 607 | block_table_max_model_len, |
| 608 | self.scheduler_config.max_num_encoder_input_tokens, |
| 609 | getattr(self.model_config.hf_config, "max_source_positions", 0), |
| 610 | ) |
| 611 | |
| 612 | block_sizes = [] |
| 613 | max_num_blocks_per_group = [] |
| 614 | slot_mapping_enabled = [] |
| 615 | dcp_sharded = [] |
| 616 | for kv_cache_group in kv_cache_config.kv_cache_groups: |
| 617 | spec = kv_cache_group.kv_cache_spec |
| 618 | block_sizes.append(spec.block_size) |
| 619 | layer_spec = ( |
| 620 | spec.first_spec if isinstance(spec, UniformTypeKVCacheSpecs) else spec |
| 621 | ) |
| 622 | slot_mapping_enabled.append(layer_spec.uses_slot_mapping) |
| 623 | dcp_sharded.append(spec.dcp_sharded) |
| 624 | # Let each cache type account for CP. Attention KV is DCP-sharded, |
| 625 | # while Mamba/GDN recurrent state is replicated across DCP ranks. |
| 626 | max_num_blocks = spec.max_num_blocks_per_req( |
| 627 | self.vllm_config, block_table_max_model_len |
| 628 | ) |
| 629 | # Preserve each cache type's alignment requirements after applying |
| 630 | # its topology-aware block-table width. |
| 631 | max_num_blocks = get_block_table_width( |
| 632 | max_num_blocks, |
| 633 | spec.block_size, |
| 634 | token_alignment=spec.block_table_token_alignment, |
| 635 | ) |
| 636 | max_num_blocks_per_group.append(max_num_blocks) |
| 637 | |
| 638 | target_attn_layer_names = None |
| 639 | if isinstance(self.speculator, DraftModelSpeculator): |
| 640 | # Adaptive verification validates target attention separately. |
| 641 | target_attn_layer_names = { |
| 642 | layer_name |
| 643 | for group in self.kv_cache_config.kv_cache_groups |
| 644 | for layer_name in group.layer_names |
| 645 | } - self.speculator.draft_attn_layer_names |
| 646 | draft_attn_layer_names = None |
| 647 | if isinstance(self.speculator, DraftModelSpeculator): |
| 648 | draft_attn_layer_names = self.speculator.draft_attn_layer_names |
| 649 | # Metadata builders select attention kernels that need JIT warmup. |
| 650 | with self.jit_warmup_registry.activate(): |
| 651 | ( |
| 652 | self.attn_groups, |
| 653 | attn_cg_support, |
| 654 | self.kernel_block_sizes, |
| 655 | ) = init_attn_backend( |
| 656 | self.kv_cache_config, |
| 657 | self.vllm_config, |
| 658 | self.device, |
| 659 | draft_layer_names=draft_attn_layer_names, |
| 660 | ) |
| 661 | additional_attn_cg_support = self.model_state.get_additional_cg_support() |
| 662 | attn_cg_support = attn_cg_support.narrow(*additional_attn_cg_support) |
| 663 | # The speculator clears the flag at load time when the checkpoint has |
| 664 | # no confidence head, so it holds the effective value. |
| 665 | self.adaptive_verification = maybe_create_adaptive_verification_manager( |
| 666 | enable_adaptive_verification=getattr( |
| 667 | self.speculator, "enable_adaptive_verification", False |
| 668 | ), |
| 669 | attn_groups=self.attn_groups, |
| 670 | req_states=self.req_states, |
| 671 | query_start_loc=self.input_buffers.query_start_loc, |
| 672 | num_bonus_tokens=self.model_state.num_new_sampled_tokens_per_step, |
| 673 | max_total_logits=get_max_chunk_logits(self.vocab_size), |
| 674 | vllm_config=self.vllm_config, |
| 675 | target_layer_names=target_attn_layer_names, |
| 676 | additional_attn_cg_support=additional_attn_cg_support, |
| 677 | ) |
| 678 | |
| 679 | self.block_tables = BlockTables( |
| 680 | block_sizes=block_sizes, |
| 681 | max_num_reqs=self.max_num_reqs, |
| 682 | max_num_batched_tokens=self.max_num_tokens, |
| 683 | max_num_blocks_per_group=max_num_blocks_per_group, |
| 684 | device=self.device, |
| 685 | kernel_block_sizes=self.kernel_block_sizes, |
| 686 | slot_mapping_enabled=slot_mapping_enabled, |
| 687 | dcp_sharded=dcp_sharded, |
| 688 | cp_size=self.dcp_size, |
| 689 | cp_rank=self.dcp_rank, |
| 690 | cp_interleave=self.cp_interleave, |
| 691 | ) |
| 692 | self.pcp_manager = pcp.maybe_build_pcp_manager( |
| 693 | self.vllm_config, |
| 694 | self.device, |
| 695 | self.supports_mm_inputs, |
| 696 | self.block_tables, |
| 697 | cls=self.pcp_manager_cls, |
| 698 | ) |
| 699 | self.ubatch_runner = maybe_build_ubatch_runner( |
| 700 | self.vllm_config, |
| 701 | self.device, |
| 702 | self.model_state, |
| 703 | self.attn_groups, |
| 704 | self.kv_cache_config, |
| 705 | self.max_num_reqs, |
| 706 | ) |
| 707 | if self.speculator is not None: |
| 708 | self.speculator.pcp_manager = self.pcp_manager |
| 709 | initialize_mamba_ssu_backend( |
| 710 | self.vllm_config.mamba_config, |
| 711 | self.kv_cache_config, |
| 712 | use_replayssm=self.cache_config.use_replayssm, |
| 713 | ) |
| 714 | piecewise_capture_available = bool( |
| 715 | envs.VLLM_USE_BREAKABLE_CUDAGRAPH or has_compiled_submodule(self.model) |
| 716 | ) |
| 717 | if self.adaptive_verification is not None: |
| 718 | self.compilation_config.cudagraph_mode = resolve_adaptive_cudagraph_mode( |
| 719 | self.compilation_config.cudagraph_mode, |
| 720 | piecewise_capture_available=piecewise_capture_available, |
| 721 | ) |
| 722 | cudagraph_mode = self.compilation_config.resolve_cudagraph_mode_and_sizes( |
| 723 | attn_cg_support.min_cg_support, |
| 724 | attn_cg_support.min_cg_attn_backend, |
| 725 | self.decode_query_len, |
| 726 | use_v2_model_runner=True, |
| 727 | tensor_parallel_size=self.parallel_config.tensor_parallel_size, |
| 728 | kv_cache_config=self.kv_cache_config, |
| 729 | max_num_reqs=self.max_num_reqs, |
| 730 | is_profiling=is_profiling, |
| 731 | piecewise_capture_available=piecewise_capture_available, |
| 732 | ) |
| 733 | self.cudagraph_manager = ModelCudaGraphManager( |
| 734 | self.vllm_config, |
| 735 | self.device, |
| 736 | cudagraph_mode, |
| 737 | decode_query_len=self.decode_query_len, |
| 738 | lora_capture_cases=self.lora_capture_cases, |
| 739 | varlen_decode=self.adaptive_verification is not None, |
| 740 | ubatch_runner=self.ubatch_runner, |
| 741 | ) |
| 742 | if self.cache_config.kv_sharing_fast_prefill and self.pcp_manager is None: |
| 743 | self.fast_prefill = FastPrefillHelper( |
| 744 | self.cudagraph_manager, self.max_num_tokens |
| 745 | ) |
| 746 | check_attention_cp_compatibility(self.vllm_config, target_attn_layer_names) |
| 747 | if isinstance(self.speculator, DraftModelSpeculator): |
| 748 | # HACK(woosuk) |
| 749 | self.speculator.set_attn( |
| 750 | self.model_state, |
| 751 | self.kv_cache_config, |
| 752 | self.block_tables, |
| 753 | self.input_buffers, |
| 754 | self.attn_groups, |
| 755 | ) |
| 756 | if self.speculator is not None: |
| 757 | # After set_attn, so the speculator can size its cudagraph mode |
| 758 | # to its own attention support. |
| 759 | self.speculator.init_cudagraph_manager(cudagraph_mode) |
| 760 | |
| 761 | # Capture warmup providers that depend on allocated KV-cache strides. |
| 762 | with self.jit_warmup_registry.activate(): |
| 763 | kv_caches_dict = init_kv_cache( |
| 764 | self.compilation_config.static_forward_context, |
| 765 | self.kv_cache_config, |
| 766 | self.device, |
| 767 | self.kernel_block_sizes, |
| 768 | self.vllm_config, |
| 769 | kv_cache_allocation_context=kv_cache_allocation_context, |
| 770 | block_tables=self.block_tables, |
| 771 | ) |
| 772 | self.kv_caches = [ |
| 773 | cache for cache in kv_caches_dict.values() if cache.device == self.device |
| 774 | ] |
| 775 | if is_profiling: |
| 776 | self.kv_connector = NO_OP_KV_CONNECTOR |
| 777 | else: |
| 778 | self.kv_connector = get_kv_connector(self.vllm_config, kv_caches_dict) |
| 779 | |
| 780 | # AuxOutput connector requires resolved kv_cache_config. |
| 781 | if self.vllm_config.aux_output_config.enabled: |
| 782 | self.aux_output_connector = get_aux_output_connector( |
| 783 | self.model, self.vllm_config, kv_cache_config |
| 784 | ) |
| 785 | |
| 786 | def _init_kv_zero_meta(self) -> None: |
| 787 | """Build KV-block zeroing metadata; invoked from gpu_worker.""" |
| 788 | self.kv_block_zeroer = KVBlockZeroer( |
| 789 | self.device, |
| 790 | attn_groups_iter=(g for groups in self.attn_groups for g in groups), |
| 791 | kernel_block_sizes=self.kernel_block_sizes, |
| 792 | static_forward_context=self.compilation_config.static_forward_context, |
| 793 | num_blocks=self.kv_cache_config.num_blocks, |
| 794 | ) |
| 795 | |
| 796 | @torch.inference_mode() |
| 797 | @step_eplb_after(is_dummy=True) |
| 798 | def _dummy_run( |
| 799 | self, |
| 800 | num_tokens: int, |
| 801 | *args, |
| 802 | skip_attn: bool = False, |
| 803 | uniform_decode: bool = False, |
| 804 | context_len: int = 0, |
| 805 | skip_eplb: bool = False, |
| 806 | is_profile: bool = False, |
| 807 | valid_dummy_state_slots: bool = False, |
| 808 | randomize_inputs: bool = False, |
| 809 | **kwargs, |
| 810 | ) -> tuple[torch.Tensor | None, torch.Tensor | None]: |
| 811 | if skip_attn and not is_profile: |
| 812 | raise ValueError( |
| 813 | "skip_attn must only be True for initial memory profiling." |
| 814 | ) |
| 815 | |
| 816 | # Create a dummy scheduler output. Plain draft-model speculation adds |
| 817 | # one correction slot per request during prefill. The scheduler |
| 818 | # accounts for these slots, while dummy runs bypass the scheduler. |
| 819 | # Adjust the dummy token count only when the expanded draft batch |
| 820 | # would exceed the token budget. |
| 821 | num_reqs = min(num_tokens, self.max_num_reqs) |
| 822 | if uniform_decode: |
| 823 | # HACK(lucas): for now since the worker is shared between MRV1 and MRV2, |
| 824 | # and for spec-decode with MTP we want to make sure the dummy runs use |
| 825 | # 1+num_speculative_tokens we use max here, this will likely be eventually |
| 826 | # changed in the worker: https://github.com/vllm-project/vllm/pull/35243 |
| 827 | num_tokens = max(num_tokens, self.decode_query_len) |
| 828 | num_reqs = num_tokens // self.decode_query_len |
| 829 | assert num_tokens % self.decode_query_len == 0 |
| 830 | elif self.speculator is not None: |
| 831 | num_reqs = min( |
| 832 | num_reqs, self.max_num_tokens // self.speculator.num_query_per_req |
| 833 | ) |
| 834 | num_tokens = min( |
| 835 | num_tokens, |
| 836 | self.max_num_tokens |
| 837 | - num_reqs * self.speculator.num_extra_query_per_req, |
| 838 | ) |
| 839 | |
| 840 | # Distribute the remainder evenly so no dummy request exceeds |
| 841 | # ceil(num_tokens / num_reqs) <= max_model_len tokens. |
| 842 | |
| 843 | num_tokens_per_request = [ |
| 844 | num_tokens // num_reqs + (i >= num_reqs - num_tokens % num_reqs) |
| 845 | for i in range(num_reqs) |
| 846 | ] |
| 847 | |
| 848 | assert sum(num_tokens_per_request) == num_tokens |
| 849 | num_scheduled_tokens = { |
| 850 | f"_dummy_req_{i}": n for i, n in enumerate(num_tokens_per_request) |
| 851 | } |
| 852 | dummy_scheduler_output = SchedulerOutput.make_empty() |
| 853 | dummy_scheduler_output.total_num_scheduled_tokens = num_tokens |
| 854 | dummy_scheduler_output.num_scheduled_tokens = num_scheduled_tokens |
| 855 | |
| 856 | # Disable any use of KVConnector for dummy runs. |
| 857 | self.kv_connector.set_disabled(True) |
| 858 | |
| 859 | # Get the intermediate tensors for the dummy run. |
| 860 | intermediate_tensors = None |
| 861 | if not self.is_first_pp_rank: |
| 862 | assert self.intermediate_tensors is not None |
| 863 | intermediate_tensors = self.intermediate_tensors[:num_tokens] |
| 864 | |
| 865 | max_loras = self.lora_config.max_loras if self.lora_config is not None else 0 |
| 866 | with self.maybe_dummy_run_with_lora( |
| 867 | self.lora_config, |
| 868 | num_scheduled_tokens=np.array(num_tokens_per_request, dtype=np.int32), |
| 869 | num_sampled_tokens=None, |
| 870 | remove_lora=True, |
| 871 | num_active_loras=max_loras, |
| 872 | ): |
| 873 | # Execute the model. |
| 874 | self.execute_model( |
| 875 | dummy_scheduler_output, |
| 876 | intermediate_tensors=intermediate_tensors, |
| 877 | dummy_run=True, |
| 878 | skip_attn_for_dummy_run=skip_attn, |
| 879 | is_profile=is_profile, |
| 880 | context_len=context_len, |
| 881 | valid_dummy_state_slots=valid_dummy_state_slots, |
| 882 | randomize_inputs=randomize_inputs, |
| 883 | ) |
| 884 | self.kv_connector.set_disabled(False) |
| 885 | |
| 886 | # Non-last PP ranks don't produce output for sampling. |
| 887 | if not self.is_last_pp_rank: |
| 888 | return None, None |
| 889 | |
| 890 | assert self.execute_model_state is not None |
| 891 | input_batch = self.execute_model_state.input_batch |
| 892 | attn_metadata = self.execute_model_state.attn_metadata |
| 893 | slot_mappings_by_layer = self.execute_model_state.slot_mappings_by_layer |
| 894 | hidden_states = self.execute_model_state.hidden_states |
| 895 | aux_hidden_states = self.execute_model_state.aux_hidden_states |
| 896 | dp_sync = self.execute_model_state.dp_sync |
| 897 | self.execute_model_state = None |
| 898 | |
| 899 | self.step_timing.forward_end() |
| 900 | |
| 901 | # dummy run the eagle speculator's propose to ensure DP/EP sync. |
| 902 | if self.speculator is not None: |
| 903 | assert self.sampler is not None |
| 904 | self.step_timing.drafter_start() |
| 905 | mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None |
| 906 | if self.speculator.supports_mm_inputs: |
| 907 | all_false = torch.zeros( |
| 908 | input_batch.num_tokens, dtype=torch.bool, device="cpu" |
| 909 | ) |
| 910 | mm_inputs = [], all_false |
| 911 | |
| 912 | # Let the target override the hidden state fed to the drafter |
| 913 | # (e.g. DeepSeek V4 MTP needs the pre-hc_head residual). The |
| 914 | # target returns a persistent buffer sized at max_num_batched_tokens; |
| 915 | # slice to the active token count that propose() expects. |
| 916 | spec_hidden_states = hidden_states |
| 917 | if hasattr(self.model, "get_mtp_target_hidden_states"): |
| 918 | pre_hc_hidden_states = self.model.get_mtp_target_hidden_states() |
| 919 | spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]] # type: ignore[union-attr] |
| 920 | if isinstance(self.sampler, GPUWatermarkSampler): |
| 921 | self.speculator.prepare_watermarking( |
| 922 | self.sampler, input_batch.idx_mapping |
| 923 | ) |
| 924 | with use_workspace_lane(self._draft_workspace_lane): |
| 925 | self.speculator.propose( |
| 926 | input_batch=input_batch, |
| 927 | attn_metadata=attn_metadata, |
| 928 | slot_mappings=slot_mappings_by_layer, |
| 929 | last_hidden_states=spec_hidden_states, |
| 930 | aux_hidden_states=aux_hidden_states, |
| 931 | num_sampled=torch.ones( |
| 932 | input_batch.num_reqs, dtype=torch.int32, device=self.device |
| 933 | ), |
| 934 | num_rejected=torch.zeros( |
| 935 | input_batch.num_reqs, dtype=torch.int32, device=self.device |
| 936 | ), |
| 937 | last_sampled=self.req_states.last_sampled_tokens, |
| 938 | next_prefill_tokens=self.req_states.next_prefill_tokens, |
| 939 | temperature=self.sampler.sampling_states.temperature.gpu, |
| 940 | seeds=self.sampler.sampling_states.seeds.gpu, |
| 941 | dp_sync=dp_sync, |
| 942 | dummy_run=True, |
| 943 | skip_attn_for_dummy_run=skip_attn, |
| 944 | mm_inputs=mm_inputs, |
| 945 | is_profile=is_profile, |
| 946 | ) |
| 947 | self.step_timing.drafter_end() |
| 948 | |
| 949 | assert hidden_states is not None # Last PP rank always has hidden_states |
| 950 | sample_hidden_states = hidden_states[input_batch.logits_indices] |
| 951 | return hidden_states, sample_hidden_states |
| 952 | |
| 953 | @torch.inference_mode() |
| 954 | def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> None: |
| 955 | num_reqs = hidden_states.shape[0] |
| 956 | logits = self.model.compute_logits(hidden_states) |
| 957 | dummy_input_batch = InputBatch.make_dummy( |
| 958 | num_reqs, num_reqs, self.input_buffers |
| 959 | ) |
| 960 | |
| 961 | # NOTE(woosuk): During the initial memory profiling, the sampler may skip |
| 962 | # top_k, top_p, and logprobs, using less GPU memory than what is possible |
| 963 | # during actual execution. |
| 964 | assert self.sampler is not None |
| 965 | self.sampler(logits, dummy_input_batch) |
| 966 | |
| 967 | @torch.inference_mode() |
| 968 | def _dummy_pooler_run(self, hidden_states: torch.Tensor) -> None: |
| 969 | assert self.pooling_runner is not None |
| 970 | self.pooling_runner.dummy_pooler_run(hidden_states) |
| 971 | |
| 972 | @torch.inference_mode() |
| 973 | def profile_run(self, randomize_inputs: bool = False) -> None: |
| 974 | if self.supports_mm_inputs and self.is_first_pp_rank: |
| 975 | mm_config = self.model_config.multimodal_config |
| 976 | if mm_config is not None and not mm_config.skip_mm_profiling: |
| 977 | mm_budget = MultiModalBudget( |
| 978 | self.vllm_config, self.mm_registry, enable_cache=False |
| 979 | ) |
| 980 | dummy_mm_inputs = mm_budget.get_dummy_encoder_profile_inputs() |
| 981 | self.model_state.encoder_runner.profile_encoder_cache( |
| 982 | dummy_mm_inputs, mm_budget |
| 983 | ) |
| 984 | |
| 985 | hidden_states, sample_hidden_states = self._dummy_run( |
| 986 | self.max_num_tokens, |
| 987 | skip_attn=True, |
| 988 | is_profile=True, |
| 989 | randomize_inputs=randomize_inputs, |
| 990 | ) |
| 991 | |
| 992 | # Only run sampler/pooler on last PP rank (non-last ranks return None). |
| 993 | if self.is_last_pp_rank: |
| 994 | assert sample_hidden_states is not None |
| 995 | if self.pooling_runner is None: |
| 996 | self._dummy_sampler_run(sample_hidden_states) |
| 997 | else: |
| 998 | self._dummy_pooler_run(hidden_states) |
| 999 | |
| 1000 | torch.accelerator.synchronize() |
| 1001 | del hidden_states, sample_hidden_states |
| 1002 | self.reset_encoder_cache() |
| 1003 | gc.collect() |
| 1004 | |
| 1005 | def reset_mm_cache(self) -> None: |
| 1006 | if self.encoder_cache is not None: |
| 1007 | self.encoder_cache.reset_mm_cache() |
| 1008 | |
| 1009 | def reset_encoder_cache(self) -> None: |
| 1010 | if self.encoder_cache is not None: |
| 1011 | self.encoder_cache.reset_encoder_cache() |
| 1012 | if self.pooling_runner is not None: |
| 1013 | self.pooling_runner.clear() |
| 1014 | |
| 1015 | def release_late_interaction_query_cache(self, query_keys: list[str]) -> None: |
| 1016 | if self.pooling_runner is not None: |
| 1017 | self.pooling_runner.release_late_interaction_query_cache(query_keys) |
| 1018 | |
| 1019 | @torch.inference_mode() |
| 1020 | def profile_cudagraph_memory(self) -> int: |
| 1021 | """Estimate the GPU memory required to capture CUDA graphs.""" |
| 1022 | return _profile_cudagraph_memory(self) |
| 1023 | |
| 1024 | def needs_cudagraph_capture(self) -> bool: |
| 1025 | """Whether capture_model() has any CUDA graphs left to capture.""" |
| 1026 | return ( |
| 1027 | self.model_state.supports_mm_inputs |
| 1028 | and self.model_state.encoder_runner.has_cudagraph() |
| 1029 | ) or ( |
| 1030 | self.cudagraph_manager is not None |
| 1031 | and self.cudagraph_manager.needs_capture() |
| 1032 | ) |
| 1033 | |
| 1034 | @torch.inference_mode() |
| 1035 | def capture_model(self, *, profile_only: bool = False) -> int: |
| 1036 | assert self.cudagraph_manager is not None |
| 1037 | capture_encoder = ( |
| 1038 | self.model_state.supports_mm_inputs |
| 1039 | and self.model_state.encoder_runner.has_cudagraph() |
| 1040 | ) |
| 1041 | capture_decoder = self.cudagraph_manager.needs_capture() |
| 1042 | if not capture_encoder and not capture_decoder: |
| 1043 | logger.warning( |
| 1044 | "Skipping encoder and decoder CUDA graph capture. To enable " |
| 1045 | "encoder capture, ensure `cudagraph_mm_encoder` is enabled; " |
| 1046 | "to enable decoder capture, ensure `cudagraph_mode` is not `NONE`." |
| 1047 | ) |
| 1048 | return 0 |
| 1049 | |
| 1050 | compilation_counter.num_gpu_runner_capture_triggers += 1 |
| 1051 | |
| 1052 | start_time = time.perf_counter() |
| 1053 | with freeze_gc_for_cudagraph_capture(): |
| 1054 | torch.accelerator.empty_cache() |
| 1055 | start_free_gpu_memory = torch.accelerator.get_memory_info()[0] |
| 1056 | |
| 1057 | with self.maybe_setup_dummy_loras(self.lora_config): |
| 1058 | if capture_encoder: |
| 1059 | self.model_state.encoder_runner.capture() |
| 1060 | |
| 1061 | if capture_decoder: |
| 1062 | input_buffers = self.input_buffers |
| 1063 | if self.pcp_manager is not None: |
| 1064 | input_buffers = self.pcp_manager.input_buffers |
| 1065 | self.cudagraph_manager.capture( |
| 1066 | self.model, |
| 1067 | self.model_state, |
| 1068 | input_buffers, |
| 1069 | self.intermediate_tensors, |
| 1070 | self.block_tables, |
| 1071 | self.attn_groups, |
| 1072 | self.kv_cache_config, |
| 1073 | pcp_manager=self.pcp_manager, |
| 1074 | has_lora=self.lora_config is not None, |
| 1075 | use_aux_hidden_state_outputs=self.use_aux_hidden_state_outputs, |
| 1076 | lora_capture_hook=create_lora_capture_hook( |
| 1077 | self.lora_config, self |
| 1078 | ), |
| 1079 | ) |
| 1080 | if self.speculator is not None: |
| 1081 | with use_workspace_lane(self._draft_workspace_lane): |
| 1082 | self.speculator.capture() |
| 1083 | if self.adaptive_verification is not None: |
| 1084 | with self.step_timing.collect() as timings: |
| 1085 | for batch in self.adaptive_verification.batches_to_profile( |
| 1086 | self.cudagraph_manager.captured_token_counts() |
| 1087 | ): |
| 1088 | self._dummy_run(**batch) |
| 1089 | self.adaptive_verification.set_initial_cost_curves(timings) |
| 1090 | self.kv_connector.reset_capture_state() |
| 1091 | |
| 1092 | end_free_gpu_memory = torch.accelerator.get_memory_info()[0] |
| 1093 | |
| 1094 | if not profile_only: |
| 1095 | # Lock workspace to prevent resizing during execution. A resize after |
| 1096 | # capture frees the static cuda graph buffer. |
| 1097 | lock_workspace() |
| 1098 | |
| 1099 | end_time = time.perf_counter() |
| 1100 | elapsed_time = end_time - start_time |
| 1101 | cuda_graph_size = start_free_gpu_memory - end_free_gpu_memory |
| 1102 | # This usually takes 5~20 seconds. |
| 1103 | logger.info( |
| 1104 | "Graph capturing finished in %.0f secs, took %.2f GiB", |
| 1105 | elapsed_time, |
| 1106 | cuda_graph_size / (1 << 30), |
| 1107 | ) |
| 1108 | return cuda_graph_size |
| 1109 | |
| 1110 | def _remove_request(self, req_id: str) -> bool: |
| 1111 | # Call model_state.remove_request *before* req_states.remove_request |
| 1112 | # so the model_state can still look up the slot index. |
| 1113 | self.model_state.remove_request(req_id) |
| 1114 | req_idx = self.req_states.remove_request(req_id) |
| 1115 | if req_idx is None: |
| 1116 | return False |
| 1117 | if self.pooling_runner is not None: |
| 1118 | self.pooling_runner.remove_request(req_idx) |
| 1119 | if self.pp_handler is not None: |
| 1120 | self.pp_handler.on_req_idx_freed(req_idx) |
| 1121 | if self.encoder_cache is not None: |
| 1122 | self.encoder_cache.remove_request(req_id) |
| 1123 | if self.prompt_logprobs_worker is not None: |
| 1124 | self.prompt_logprobs_worker.remove_request(req_id) |
| 1125 | self.lora_state.remove_request(req_id) |
| 1126 | return True |
| 1127 | |
| 1128 | def finish_requests(self, scheduler_output: SchedulerOutput) -> None: |
| 1129 | finished_req_ids = scheduler_output.finished_req_ids |
| 1130 | if self.pooling_runner is not None: |
| 1131 | # Preempted docs keep their query-use reservation until rescheduled. |
| 1132 | self.pooling_runner.on_requests_finished(finished_req_ids) |
| 1133 | preempted_req_ids = scheduler_output.preempted_req_ids |
| 1134 | if preempted_req_ids: |
| 1135 | finished_req_ids = finished_req_ids.union(preempted_req_ids) |
| 1136 | # Sorted so every TP rank frees request slots in the same order. |
| 1137 | # Features like batch-sharded sampling derive rank request ownership |
| 1138 | # from the slot index. |
| 1139 | for req_id in sorted(finished_req_ids): |
| 1140 | self._remove_request(req_id) |
| 1141 | |
| 1142 | def free_states(self, scheduler_output: SchedulerOutput) -> None: |
| 1143 | if self.encoder_cache is not None: |
| 1144 | for mm_hash in scheduler_output.free_encoder_mm_hashes: |
| 1145 | self.encoder_cache.free_encoder_cache(mm_hash) |
| 1146 | |
| 1147 | def update_pp_decode_requests(self): |
| 1148 | # For non-last PP ranks, update decode requests with sampler output from |
| 1149 | # the prior step in which they were scheduled (pp_size steps ago). |
| 1150 | if self.pp_handler is not None: |
| 1151 | outputs = self.pp_handler.get_prev_sampled_outputs( |
| 1152 | self.req_states.draft_tokens |
| 1153 | ) |
| 1154 | if outputs is not None: |
| 1155 | self.postprocess_sampled(**outputs) |
| 1156 | |
| 1157 | def add_requests(self, scheduler_output: SchedulerOutput) -> None: |
| 1158 | for new_req_data in scheduler_output.scheduled_new_reqs: |
| 1159 | assert new_req_data.prefill_token_ids is not None |
| 1160 | req_id = new_req_data.req_id |
| 1161 | |
| 1162 | # Streaming input update: request already exists from a prior |
| 1163 | # chunk. Remove old state so it can be cleanly re-added below |
| 1164 | # with the updated prompt_token_ids and mm_features. |
| 1165 | self._remove_request(req_id) |
| 1166 | |
| 1167 | prompt_len = new_req_data.prompt_len |
| 1168 | sampling_params = new_req_data.sampling_params |
| 1169 | self.req_states.add_request( |
| 1170 | req_id=req_id, |
| 1171 | prompt_len=prompt_len, |
| 1172 | all_token_ids=new_req_data.prefill_token_ids, |
| 1173 | num_computed_tokens=new_req_data.num_computed_tokens, |
| 1174 | max_tokens=sampling_params.max_tokens if sampling_params else 1, # type: ignore[arg-type] |
| 1175 | ) |
| 1176 | req_index = self.req_states.req_id_to_index[req_id] |
| 1177 | if self.adaptive_verification is not None: |
| 1178 | self.adaptive_verification.add_request(req_index) |
| 1179 | |
| 1180 | if self.pooling_runner is not None: |
| 1181 | assert new_req_data.pooling_params is not None |
| 1182 | self.pooling_runner.add_request( |
| 1183 | req_id, |
| 1184 | req_index, |
| 1185 | new_req_data.pooling_params, |
| 1186 | new_req_data.prompt_token_ids, |
| 1187 | ) |
| 1188 | |
| 1189 | if self.encoder_cache is not None: |
| 1190 | self.encoder_cache.add_request(req_id, new_req_data.mm_features) |
| 1191 | |
| 1192 | self.model_state.add_request(req_index, new_req_data) |
| 1193 | self.block_tables.append_block_ids( |
| 1194 | req_index, new_req_data.block_ids, overwrite=True |
| 1195 | ) |
| 1196 | self.lora_state.add_request(req_id, req_index, new_req_data.lora_request) |
| 1197 | |
| 1198 | if self.is_last_pp_rank and new_req_data.sampling_params is not None: |
| 1199 | assert self.sampler is not None |
| 1200 | self.sampler.add_request(req_index, new_req_data.sampling_params) |
| 1201 | assert self.prompt_logprobs_worker is not None |
| 1202 | self.prompt_logprobs_worker.add_request( |
| 1203 | req_id, req_index, new_req_data.sampling_params |
| 1204 | ) |
| 1205 | |
| 1206 | if scheduler_output.scheduled_new_reqs: |
| 1207 | self.req_states.apply_staged_writes() |
| 1208 | self.model_state.apply_staged_writes() |
| 1209 | if self.sampler is not None: |
| 1210 | self.sampler.apply_staged_writes() |
| 1211 | |
| 1212 | def update_requests(self, scheduler_output: SchedulerOutput) -> None: |
| 1213 | # Add new blocks and update num_computed_tokens for the existing requests. |
| 1214 | reqs = scheduler_output.scheduled_cached_reqs |
| 1215 | table_updates = scheduler_output.block_table_updates or {} |
| 1216 | for req_id, block_ids in table_updates.items(): |
| 1217 | if (idx := self.req_states.req_id_to_index.get(req_id)) is not None: |
| 1218 | self.block_tables.append_block_ids(idx, block_ids, overwrite=True) |
| 1219 | num_computed_tokens_np = self.req_states.num_computed_tokens_np |
| 1220 | for req_id, num_computed_tokens, req_new_block_ids in zip( |
| 1221 | reqs.req_ids, reqs.num_computed_tokens, reqs.new_block_ids |
| 1222 | ): |
| 1223 | req_index = self.req_states.req_id_to_index[req_id] |
| 1224 | num_computed_tokens_np[req_index] = num_computed_tokens |
| 1225 | if req_new_block_ids is not None and req_id not in table_updates: |
| 1226 | self.block_tables.append_block_ids( |
| 1227 | req_index, req_new_block_ids, overwrite=False |
| 1228 | ) |
| 1229 | |
| 1230 | # Update CPU num_computed_prefill_tokens. |
| 1231 | np.minimum( |
| 1232 | self.req_states.num_computed_tokens_np, |
| 1233 | self.req_states.prefill_len.np, |
| 1234 | out=self.req_states.num_computed_prefill_tokens, |
| 1235 | ) |
| 1236 | |
| 1237 | # Zero GPU memory for freshly allocated cache blocks to prevent |
| 1238 | # stale NaN/data from corrupting attention or SSM computation. |
| 1239 | if scheduler_output.new_block_ids_to_zero: |
| 1240 | assert self.kv_block_zeroer is not None |
| 1241 | self.kv_block_zeroer.zero_block_ids(scheduler_output.new_block_ids_to_zero) |
| 1242 | |
| 1243 | # Apply copy-on-write block copies for partial prefix-cache hits, after |
| 1244 | # zeroing new blocks and before the forward pass reads them. |
| 1245 | if scheduler_output.kv_cache_block_copies: |
| 1246 | copy_kv_cache_blocks_inplace( |
| 1247 | self.kv_caches, |
| 1248 | self.kv_cache_config.num_blocks, |
| 1249 | scheduler_output.kv_cache_block_copies, |
| 1250 | ) |
| 1251 | |
| 1252 | def gather_batch_req_state( |
| 1253 | self, scheduler_output: SchedulerOutput, dummy_run: bool |
| 1254 | ) -> tuple["BatchReqState | None", int | None]: |
| 1255 | """Gather CPU request state for the scheduled batch, in batch order. |
| 1256 | Returns (batch_state, uniform_decode_token_count) |
| 1257 | """ |
| 1258 | num_tokens_per_req = scheduler_output.num_scheduled_tokens |
| 1259 | num_reqs = len(num_tokens_per_req) |
| 1260 | num_toks = scheduler_output.total_num_scheduled_tokens |
| 1261 | max_query_len = max(scheduler_output.num_scheduled_tokens.values()) |
| 1262 | |
| 1263 | if dummy_run: |
| 1264 | # Dummy batches are uniform by construction. |
| 1265 | return None, get_uniform_decode_token_count( |
| 1266 | num_reqs, num_toks, max_query_len, decode_graph_eligible=True |
| 1267 | ) |
| 1268 | |
| 1269 | draft_tokens = scheduler_output.scheduled_spec_decode_tokens |
| 1270 | # batch_idx -> req_id |
| 1271 | req_ids = sort_batch_req_ids( |
| 1272 | num_tokens_per_req, draft_tokens, self.decode_query_len |
| 1273 | ) |
| 1274 | |
| 1275 | numtoks_iter = map(num_tokens_per_req.__getitem__, req_ids) |
| 1276 | num_scheduled_tokens = np.fromiter(numtoks_iter, dtype=np.int32, count=num_reqs) |
| 1277 | |
| 1278 | idx_mapping_iter = map(self.req_states.req_id_to_index.__getitem__, req_ids) |
| 1279 | idx_mapping_np = np.fromiter(idx_mapping_iter, dtype=np.intp, count=num_reqs) |
| 1280 | prefill_len_np = self.req_states.prefill_len.np[idx_mapping_np] |
| 1281 | num_computed_prefill_tokens_np = self.req_states.num_computed_prefill_tokens[ |
| 1282 | idx_mapping_np |
| 1283 | ] |
| 1284 | is_prefilling_np = num_computed_prefill_tokens_np < prefill_len_np |
| 1285 | has_prefill = bool(is_prefilling_np.any()) |
| 1286 | max_seq_len_np = None |
| 1287 | if self.use_pp and has_prefill: |
| 1288 | # consumed by the PP `compute_need_sampled_mask` for batches with prefill. |
| 1289 | max_seq_len_np = self.req_states.max_seq_len[idx_mapping_np] |
| 1290 | |
| 1291 | num_draft_tokens_np = None |
| 1292 | if draft_tokens: |
| 1293 | num_drafts_iter = (len(draft_tokens.get(req_id, ())) for req_id in req_ids) |
| 1294 | num_draft_tokens_np = np.fromiter( |
| 1295 | num_drafts_iter, dtype=np.int32, count=num_reqs |
| 1296 | ) |
| 1297 | if self.adaptive_verification is not None: |
| 1298 | num_toks = self.adaptive_verification.get_num_tokens( |
| 1299 | num_tokens_per_req, draft_tokens |
| 1300 | ) |
| 1301 | |
| 1302 | prefill_runs_as_decode_np = None |
| 1303 | decode_graph_eligible = not has_prefill |
| 1304 | if has_prefill and self.pcp_manager is None: |
| 1305 | # One new prompt token over existing context (possibly padded with |
| 1306 | # placeholder drafts) runs as a decode. PCP partitions prefill rows. |
| 1307 | num_new_tokens = num_scheduled_tokens |
| 1308 | if num_draft_tokens_np is not None: |
| 1309 | num_new_tokens = num_new_tokens - num_draft_tokens_np |
| 1310 | one_token_ext = (num_new_tokens == 1) & (num_computed_prefill_tokens_np > 0) |
| 1311 | prefill_runs_as_decode_np = is_prefilling_np & one_token_ext |
| 1312 | decode_graph_eligible = bool( |
| 1313 | (prefill_runs_as_decode_np | ~is_prefilling_np).all() |
| 1314 | ) |
| 1315 | |
| 1316 | batch_state = BatchReqState( |
| 1317 | req_ids=req_ids, |
| 1318 | num_scheduled_tokens=num_scheduled_tokens, |
| 1319 | num_tokens=num_toks, |
| 1320 | num_draft_tokens_np=num_draft_tokens_np, |
| 1321 | idx_mapping_np=idx_mapping_np, |
| 1322 | prefill_len_np=prefill_len_np, |
| 1323 | num_computed_prefill_tokens_np=num_computed_prefill_tokens_np, |
| 1324 | is_prefilling_np=is_prefilling_np, |
| 1325 | has_prefill=has_prefill, |
| 1326 | prefill_runs_as_decode_np=prefill_runs_as_decode_np, |
| 1327 | decode_graph_eligible=decode_graph_eligible, |
| 1328 | max_seq_len_np=max_seq_len_np, |
| 1329 | ) |
| 1330 | return batch_state, get_uniform_decode_token_count( |
| 1331 | num_reqs, num_toks, max_query_len, decode_graph_eligible |
| 1332 | ) |
| 1333 | |
| 1334 | def prepare_inputs( |
| 1335 | self, |
| 1336 | scheduler_output: SchedulerOutput, |
| 1337 | batch_req_state: "BatchReqState", |
| 1338 | batch_desc: BatchExecutionDescriptor, |
| 1339 | num_active_loras: int, |
| 1340 | ) -> InputBatch: |
| 1341 | num_tokens = batch_req_state.num_tokens |
| 1342 | num_tokens_after_padding = max(num_tokens, batch_desc.num_tokens) |
| 1343 | assert num_tokens > 0 |
| 1344 | is_padding = self.input_buffers.is_padding[:num_tokens_after_padding] |
| 1345 | if envs.VLLM_MOE_SKIP_PADDING: |
| 1346 | is_padding[:num_tokens].fill_(False) |
| 1347 | is_padding[num_tokens:].fill_(True) |
| 1348 | |
| 1349 | req_ids = batch_req_state.req_ids |
| 1350 | num_scheduled_tokens_np = batch_req_state.num_scheduled_tokens |
| 1351 | idx_mapping_np = batch_req_state.idx_mapping_np |
| 1352 | idx_mapping = async_tensor_h2d( |
| 1353 | idx_mapping_np, device=self.device, dtype=torch.int32 |
| 1354 | ) |
| 1355 | num_reqs = len(req_ids) |
| 1356 | |
| 1357 | # Get the number of draft tokens for each request. |
| 1358 | draft_tokens = scheduler_output.scheduled_spec_decode_tokens |
| 1359 | num_draft_tokens_per_req = None |
| 1360 | if not draft_tokens and self.model_state.num_new_sampled_tokens_per_step == 1: |
| 1361 | # No draft token scheduled (common case). |
| 1362 | total_num_draft_tokens = 0 |
| 1363 | total_num_logits = num_reqs |
| 1364 | cu_num_logits_np = np.arange(num_reqs + 1, dtype=np.int32) |
| 1365 | cu_num_logits = torch.arange( |
| 1366 | num_reqs + 1, device=self.device, dtype=torch.int32 |
| 1367 | ) |
| 1368 | expanded_idx_mapping = idx_mapping |
| 1369 | expanded_local_pos = torch.zeros( |
| 1370 | num_reqs, dtype=torch.int32, device=self.device |
| 1371 | ) |
| 1372 | else: |
| 1373 | num_draft_tokens_per_req = batch_req_state.num_draft_tokens_np |
| 1374 | if num_draft_tokens_per_req is None: |
| 1375 | num_draft_tokens_per_req = np.zeros(num_reqs, dtype=np.int32) |
| 1376 | num_bonus_tokens = self.model_state.num_new_sampled_tokens_per_step |
| 1377 | total_num_draft_tokens = int(num_draft_tokens_per_req.sum()) |
| 1378 | total_num_logits = num_reqs * num_bonus_tokens + total_num_draft_tokens |
| 1379 | num_logits = num_draft_tokens_per_req + num_bonus_tokens |
| 1380 | # combine_sampled_and_draft_tokens places a request's logits rows |
| 1381 | # at [query_end - num_logits, query_end). Fewer query rows than |
| 1382 | # that would silently select the preceding request's hidden states. |
| 1383 | assert (num_scheduled_tokens_np >= num_logits).all() |
| 1384 | cu_num_logits_np = np.empty(num_reqs + 1, dtype=np.int32) |
| 1385 | cu_num_logits_np[0] = 0 |
| 1386 | np.cumsum(num_logits, out=cu_num_logits_np[1:]) |
| 1387 | cu_num_logits = async_tensor_h2d(cu_num_logits_np, device=self.device) |
| 1388 | |
| 1389 | # The general branch also serves a no-draft batch with k != 1, and |
| 1390 | # compact_batch's budget is only primed when drafts are scheduled. |
| 1391 | adaptive_verification = self.adaptive_verification if draft_tokens else None |
| 1392 | num_scheduled_tokens_upper_bound = num_scheduled_tokens_np |
| 1393 | if adaptive_verification is not None: |
| 1394 | # num_scheduled_tokens represents the draft budget evenly distributed across |
| 1395 | # all verification requests, `reallocate_drafts` will unevenly assign the |
| 1396 | # draft budget to requests on the GPU side only. |
| 1397 | num_scheduled_tokens_np, cu_num_logits_np = ( |
| 1398 | adaptive_verification.compact_batch( |
| 1399 | num_draft_tokens_per_req, |
| 1400 | num_scheduled_tokens_np, |
| 1401 | cu_num_logits_np, |
| 1402 | ) |
| 1403 | ) |
| 1404 | |
| 1405 | # Get query_start_loc. |
| 1406 | # num_reqs_padded is None for PIECEWISE graphs (no request padding needed) |
| 1407 | num_reqs_padded = batch_desc.num_reqs or num_reqs |
| 1408 | query_start_loc_np = np.empty(self.max_num_reqs + 1, dtype=np.int32) |
| 1409 | query_start_loc_np[0] = 0 |
| 1410 | np.cumsum(num_scheduled_tokens_np, out=query_start_loc_np[1 : num_reqs + 1]) |
| 1411 | # Pad for full CUDA graph mode. |
| 1412 | # Some attention backends like FA3 require query_start_loc to be non-decreasing. |
| 1413 | query_start_loc_np[num_reqs + 1 :] = num_tokens |
| 1414 | query_start_loc = self.input_buffers.query_start_loc |
| 1415 | async_tensor_h2d(query_start_loc_np, out=query_start_loc) |
| 1416 | if adaptive_verification is not None: |
| 1417 | cu_num_logits, query_start_loc, total_num_draft_tokens = ( |
| 1418 | adaptive_verification.reallocate_drafts(req_ids, idx_mapping) |
| 1419 | ) |
| 1420 | total_num_logits = num_reqs * num_bonus_tokens + total_num_draft_tokens |
| 1421 | if num_draft_tokens_per_req is not None: |
| 1422 | expanded_idx_mapping, expanded_local_pos = expand_idx_mapping( |
| 1423 | idx_mapping, total_num_logits, cu_num_logits, self.decode_query_len |
| 1424 | ) |
| 1425 | query_start_loc_np = query_start_loc_np[: num_reqs_padded + 1] |
| 1426 | query_start_loc = query_start_loc[: num_reqs_padded + 1] |
| 1427 | |
| 1428 | # Get prefill tokens if any. |
| 1429 | if batch_req_state.has_prefill: |
| 1430 | prepare_prefill_inputs( |
| 1431 | self.input_buffers.input_ids, |
| 1432 | self.req_states.next_prefill_tokens, |
| 1433 | idx_mapping, |
| 1434 | query_start_loc, |
| 1435 | self.req_states.all_token_ids.gpu, |
| 1436 | self.req_states.prefill_len.gpu, |
| 1437 | self.req_states.num_computed_tokens.gpu, |
| 1438 | ) |
| 1439 | |
| 1440 | # Prepare positions and seq_lens. |
| 1441 | prepare_pos_seq_lens( |
| 1442 | idx_mapping, |
| 1443 | query_start_loc, |
| 1444 | self.req_states.num_computed_tokens.gpu, |
| 1445 | self.input_buffers.positions, |
| 1446 | self.input_buffers.seq_lens, |
| 1447 | ) |
| 1448 | seq_lens = self.input_buffers.seq_lens[:num_reqs_padded] |
| 1449 | |
| 1450 | # Some input token ids are directly read from the last sampled tokens |
| 1451 | # and draft tokens. Also, get the logits indices to sample tokens from. |
| 1452 | logits_indices = combine_sampled_and_draft_tokens( |
| 1453 | self.input_buffers.input_ids, |
| 1454 | idx_mapping, |
| 1455 | self.req_states.last_sampled_tokens, |
| 1456 | query_start_loc, |
| 1457 | seq_lens, |
| 1458 | self.req_states.prefill_len.gpu, |
| 1459 | self.req_states.draft_tokens, |
| 1460 | cu_num_logits, |
| 1461 | total_num_logits, |
| 1462 | self.model_state.num_new_sampled_tokens_per_step, |
| 1463 | ) |
| 1464 | |
| 1465 | fast_prefill = None |
| 1466 | if self.fast_prefill is not None: |
| 1467 | fast_prefill = self.fast_prefill.prepare( |
| 1468 | logits_indices, |
| 1469 | num_reqs, |
| 1470 | cu_num_logits_np, |
| 1471 | batch_req_state.has_prefill, |
| 1472 | batch_desc, |
| 1473 | num_active_loras=num_active_loras, |
| 1474 | ) |
| 1475 | |
| 1476 | # CPU upper bound on seq_lens; padded entries left at zero. |
| 1477 | num_computed_tokens_np = self.req_states.num_computed_tokens_np[idx_mapping_np] |
| 1478 | seq_lens_cpu_upper_bound_np = np.zeros(num_reqs_padded, dtype=np.int32) |
| 1479 | np.add( |
| 1480 | num_computed_tokens_np, |
| 1481 | num_scheduled_tokens_upper_bound, |
| 1482 | out=seq_lens_cpu_upper_bound_np[:num_reqs], |
| 1483 | ) |
| 1484 | seq_lens_cpu_upper_bound = torch.from_numpy(seq_lens_cpu_upper_bound_np) |
| 1485 | |
| 1486 | prompt_lens = None |
| 1487 | if self.model_config.rswa_window is not None: |
| 1488 | # prompt_lens is only used in R-SWA case. |
| 1489 | prompt_lens = self.req_states.prompt_len.gpu[idx_mapping] |
| 1490 | |
| 1491 | input_batch = InputBatch( |
| 1492 | req_ids=req_ids, |
| 1493 | num_reqs=num_reqs, |
| 1494 | num_reqs_after_padding=num_reqs_padded, |
| 1495 | idx_mapping=idx_mapping, |
| 1496 | idx_mapping_np=idx_mapping_np, |
| 1497 | expanded_idx_mapping=expanded_idx_mapping, |
| 1498 | expanded_local_pos=expanded_local_pos, |
| 1499 | num_scheduled_tokens=num_scheduled_tokens_upper_bound, |
| 1500 | num_tokens=num_tokens, |
| 1501 | num_tokens_after_padding=num_tokens_after_padding, |
| 1502 | num_draft_tokens=total_num_draft_tokens, |
| 1503 | num_draft_tokens_per_req=num_draft_tokens_per_req, |
| 1504 | query_start_loc=query_start_loc, |
| 1505 | query_start_loc_np=query_start_loc_np, |
| 1506 | seq_lens=seq_lens, |
| 1507 | seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, |
| 1508 | dcp_local_seq_lens=None, |
| 1509 | num_computed_tokens_np=num_computed_tokens_np, |
| 1510 | prefill_len_np=batch_req_state.prefill_len_np, |
| 1511 | num_computed_prefill_tokens_np=batch_req_state.num_computed_prefill_tokens_np, |
| 1512 | is_prefilling_np=batch_req_state.is_prefilling_np, |
| 1513 | has_prefill=batch_req_state.has_prefill, |
| 1514 | prefill_runs_as_decode_np=batch_req_state.prefill_runs_as_decode_np, |
| 1515 | decode_graph_eligible=batch_req_state.decode_graph_eligible, |
| 1516 | max_seq_len_np=batch_req_state.max_seq_len_np, |
| 1517 | input_ids=self.input_buffers.input_ids[:num_tokens_after_padding], |
| 1518 | positions=self.input_buffers.positions[:num_tokens_after_padding], |
| 1519 | is_padding=is_padding, |
| 1520 | logits_indices=logits_indices, |
| 1521 | cu_num_logits=cu_num_logits, |
| 1522 | cu_num_logits_np=cu_num_logits_np, |
| 1523 | has_structured_output_reqs=scheduler_output.has_structured_output_requests, |
| 1524 | prompt_lens=prompt_lens, |
| 1525 | fast_prefill=fast_prefill, |
| 1526 | max_query_len=( |
| 1527 | int(num_scheduled_tokens_upper_bound.max()) |
| 1528 | if adaptive_verification is not None |
| 1529 | else None |
| 1530 | ), |
| 1531 | ) |
| 1532 | if self.pcp_manager is not None: |
| 1533 | input_batch = self.pcp_manager.partition_batch(input_batch, batch_desc) |
| 1534 | return input_batch |
| 1535 | |
| 1536 | def prepare_attn( |
| 1537 | self, input_batch: InputBatch |
| 1538 | ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]: |
| 1539 | if self.pcp_manager is not None: |
| 1540 | return self.pcp_manager.prepare_attn(input_batch) |
| 1541 | |
| 1542 | # Block tables: num_kv_cache_groups x [num_reqs_padded, max_num_blocks]. |
| 1543 | block_tables = self.block_tables.gather_block_tables( |
| 1544 | input_batch.idx_mapping, num_reqs_padded=input_batch.num_reqs_after_padding |
| 1545 | ) |
| 1546 | # Slot mappings: [num_kv_cache_groups, num_tokens_padded]. |
| 1547 | # Kernel pads beyond num_tokens with PAD_SLOT_ID. |
| 1548 | slot_mappings = self.block_tables.compute_slot_mappings( |
| 1549 | input_batch.idx_mapping, |
| 1550 | input_batch.query_start_loc, |
| 1551 | input_batch.positions, |
| 1552 | num_tokens_padded=input_batch.num_tokens_after_padding, |
| 1553 | ) |
| 1554 | return block_tables, slot_mappings |
| 1555 | |
| 1556 | def prepare_dummy_attn( |
| 1557 | self, input_batch: InputBatch, valid_state_slots: bool = False |
| 1558 | ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]: |
| 1559 | block_table_provider = self.pcp_manager or self.block_tables |
| 1560 | block_tables = block_table_provider.get_dummy_block_tables(input_batch.num_reqs) |
| 1561 | if valid_state_slots: |
| 1562 | state_slots = torch.arange( |
| 1563 | 1, input_batch.num_reqs + 1, dtype=torch.int32, device=self.device |
| 1564 | ) |
| 1565 | for block_table in block_tables: |
| 1566 | block_table[:, 0].copy_(state_slots) |
| 1567 | slot_mappings = block_table_provider.get_dummy_slot_mappings( |
| 1568 | input_batch.num_tokens |
| 1569 | ) |
| 1570 | return block_tables, slot_mappings |
| 1571 | |
| 1572 | def sample( |
| 1573 | self, |
| 1574 | hidden_states: torch.Tensor, |
| 1575 | input_batch: InputBatch, |
| 1576 | grammar_output: GrammarOutput | None, |
| 1577 | ) -> tuple[SamplerOutput, torch.Tensor, torch.Tensor]: |
| 1578 | shard_metadata = None |
| 1579 | global_input_batch = input_batch |
| 1580 | if self.batch_sharder is not None: |
| 1581 | # Shard the inputs along the batch dimension to sample in parallel |
| 1582 | # across TP ranks. |
| 1583 | input_batch, sorted_logits_indices, grammar_output, shard_metadata = ( |
| 1584 | self.batch_sharder.shard_sampler_inputs(input_batch, grammar_output) |
| 1585 | ) |
| 1586 | # The hidden states must be gathered in rank-owner-sorted order |
| 1587 | # before computing the partial-vocab logits, so that the all-to-all |
| 1588 | # produces full-vocab logits for just the locally-owned requests. |
| 1589 | sample_hidden_states = hidden_states[sorted_logits_indices] |
| 1590 | local_logits = self.model.compute_logits_local(sample_hidden_states) |
| 1591 | logits = all_to_all_logits(local_logits, shard_metadata) |
| 1592 | logits = logits[:, : self.vocab_size] |
| 1593 | else: |
| 1594 | sample_hidden_states = hidden_states[input_batch.logits_indices] |
| 1595 | logits = self.model.compute_logits(sample_hidden_states) |
| 1596 | |
| 1597 | invalid_drafts = None |
| 1598 | # A diffusion prefill has no logit rows even when a bitmask row |
| 1599 | # arrived for it. |
| 1600 | if grammar_output is not None and logits.shape[0] > 0: |
| 1601 | # Apply grammar bitmask to the logits in-place. |
| 1602 | assert self.structured_outputs_worker is not None |
| 1603 | self.structured_outputs_worker.apply_grammar_bitmask( |
| 1604 | logits, |
| 1605 | input_batch, |
| 1606 | grammar_output.structured_output_request_ids, |
| 1607 | grammar_output.grammar_bitmask, |
| 1608 | ) |
| 1609 | invalid_drafts = grammar_invalid_drafts( |
| 1610 | input_batch, |
| 1611 | grammar_output.structured_output_request_ids, |
| 1612 | grammar_output.num_acceptable_drafts, |
| 1613 | ) |
| 1614 | |
| 1615 | sampler_output: SamplerOutput | None |
| 1616 | if input_batch.num_reqs == 0: |
| 1617 | # This rank owns no requests this step. It contributes an |
| 1618 | # all-padding block to the gather below. |
| 1619 | sampler_output = None |
| 1620 | elif input_batch.num_draft_tokens == 0 or self.rejection_sampler is None: |
| 1621 | assert self.sampler is not None |
| 1622 | sampler_output = self.sampler(logits, input_batch) |
| 1623 | else: |
| 1624 | # Rejection sampling for spec decoding. |
| 1625 | assert self.rejection_sampler is not None |
| 1626 | assert self.speculator is not None |
| 1627 | sampler_output = self.rejection_sampler( |
| 1628 | logits, |
| 1629 | input_batch, |
| 1630 | # Draft logits are needed for probabilistic rejection sampling. |
| 1631 | self.speculator.draft_logits, |
| 1632 | invalid_drafts, |
| 1633 | ) |
| 1634 | |
| 1635 | if shard_metadata is not None: |
| 1636 | # Gather the sharded sampler outputs from the TP ranks into a single |
| 1637 | # sampler output. |
| 1638 | assert self.sampler is not None |
| 1639 | sampler_output = gather_sampler_output( |
| 1640 | sampler_output, |
| 1641 | shard_metadata, |
| 1642 | device=self.device, |
| 1643 | global_batch=global_input_batch, |
| 1644 | local_batch=input_batch, |
| 1645 | gather_num_nans=self.sampler.compute_nans, |
| 1646 | logprobs_dims=self.sampler.get_logprobs_dims( |
| 1647 | global_input_batch.idx_mapping_np, |
| 1648 | # Rejection sampler does not return logprob token ids. |
| 1649 | include_token_ids=( |
| 1650 | global_input_batch.num_draft_tokens == 0 |
| 1651 | or self.rejection_sampler is None |
| 1652 | ), |
| 1653 | ), |
| 1654 | ) |
| 1655 | |
| 1656 | assert sampler_output is not None |
| 1657 | return sampler_output, sampler_output.num_sampled, sampler_output.num_rejected |
| 1658 | |
| 1659 | def postprocess_sampled( |
| 1660 | self, |
| 1661 | idx_mapping: torch.Tensor, # May include -1 for masked entries |
| 1662 | sampled_tokens: torch.Tensor, |
| 1663 | num_sampled: torch.Tensor, |
| 1664 | num_rejected: torch.Tensor, |
| 1665 | query_start_loc: torch.Tensor | None = None, |
| 1666 | ) -> None: |
| 1667 | # Update the number of computed tokens. |
| 1668 | output_bin_counts = None |
| 1669 | if self.is_last_pp_rank: |
| 1670 | assert self.sampler is not None |
| 1671 | output_bin_counts = self.sampler.penalties_state.output_bin_counts |
| 1672 | post_update( |
| 1673 | idx_mapping, |
| 1674 | self.req_states.num_computed_tokens.gpu, |
| 1675 | self.req_states.last_sampled_tokens, |
| 1676 | output_bin_counts, |
| 1677 | sampled_tokens, |
| 1678 | num_sampled, |
| 1679 | num_rejected, |
| 1680 | query_start_loc, |
| 1681 | self.req_states.all_token_ids.gpu, |
| 1682 | self.req_states.total_len.gpu, |
| 1683 | ) |
| 1684 | |
| 1685 | self.model_state.postprocess_state( |
| 1686 | idx_mapping, num_sampled, self.req_states.num_computed_tokens.gpu |
| 1687 | ) |
| 1688 | |
| 1689 | def _merge_ec_connector_no_forward( |
| 1690 | self, scheduler_output: SchedulerOutput, output: ModelRunnerOutput |
| 1691 | ) -> ModelRunnerOutput: |
| 1692 | """Let the EC connector send/recv on a step with no work to run.""" |
| 1693 | return ModelRunnerOutput.with_ec_conn_output( |
| 1694 | output, |
| 1695 | self.ec_connector.no_forward(scheduler_output).ec_connector_output, |
| 1696 | ) |
| 1697 | |
| 1698 | @torch.inference_mode() |
| 1699 | def execute_model( |
| 1700 | self, |
| 1701 | scheduler_output: SchedulerOutput, |
| 1702 | intermediate_tensors: IntermediateTensors | None = None, |
| 1703 | dummy_run: bool = False, |
| 1704 | skip_attn_for_dummy_run: bool = False, |
| 1705 | is_profile: bool = False, |
| 1706 | context_len: int = 0, |
| 1707 | valid_dummy_state_slots: bool = False, |
| 1708 | randomize_inputs: bool = False, |
| 1709 | ) -> ModelRunnerOutput | IntermediateTensors | None: |
| 1710 | if not dummy_run: |
| 1711 | # Update the request states. |
| 1712 | self.update_pp_decode_requests() |
| 1713 | self.finish_requests(scheduler_output) |
| 1714 | self.free_states(scheduler_output) |
| 1715 | self.add_requests(scheduler_output) |
| 1716 | self.update_requests(scheduler_output) |
| 1717 | self.block_tables.apply_staged_writes() |
| 1718 | if self.aux_output_connector is not None: |
| 1719 | # Register this step before the GPU forward. |
| 1720 | self.aux_output_connector.begin_step( |
| 1721 | scheduler_output.aux_output_connector_metadata |
| 1722 | ) |
| 1723 | if scheduler_output.total_num_scheduled_tokens == 0: |
| 1724 | # No need to run the model. |
| 1725 | empty_output = self.kv_connector.no_forward(scheduler_output) |
| 1726 | return self._merge_ec_connector_no_forward( |
| 1727 | scheduler_output, empty_output |
| 1728 | ) |
| 1729 | |
| 1730 | # Get batch descriptor and sync across DP ranks. |
| 1731 | num_reqs = len(scheduler_output.num_scheduled_tokens) |
| 1732 | num_toks = scheduler_output.total_num_scheduled_tokens |
| 1733 | max_query_len: int | None = max(scheduler_output.num_scheduled_tokens.values()) |
| 1734 | batch_req_state, uniform_tok_count = self.gather_batch_req_state( |
| 1735 | scheduler_output, dummy_run |
| 1736 | ) |
| 1737 | if batch_req_state is not None: |
| 1738 | num_toks = batch_req_state.num_tokens |
| 1739 | if batch_req_state.has_prefill: |
| 1740 | # Varlen decode graphs replay decodes only, and their bound |
| 1741 | # alone would admit a short prefill. |
| 1742 | max_query_len = None |
| 1743 | if self.pcp_manager is not None: |
| 1744 | num_toks = self.pcp_manager.get_num_tokens_for_dispatch( |
| 1745 | batch_req_state.num_scheduled_tokens, |
| 1746 | batch_req_state.is_prefilling_np, |
| 1747 | ) |
| 1748 | |
| 1749 | num_active_loras = 0 |
| 1750 | if self.lora_config: |
| 1751 | req_ids = list(scheduler_output.num_scheduled_tokens.keys()) |
| 1752 | num_active_loras = get_num_active_loras_for_dispatch( |
| 1753 | self.lora_config, self.lora_state, req_ids, dummy_run |
| 1754 | ) |
| 1755 | |
| 1756 | skip_compiled = False |
| 1757 | if self.is_encoder_decoder and scheduler_output.scheduled_encoder_inputs: |
| 1758 | # Encoder-decoder models such as Whisper should run eager/non-compiled |
| 1759 | # when encoder inputs are scheduled, because this step updates |
| 1760 | # cross-attention cache with dynamic encoder outputs. |
| 1761 | skip_compiled = True |
| 1762 | |
| 1763 | batch_desc, dp_sync = dispatch_cg_and_sync_dp( |
| 1764 | self.cudagraph_manager, |
| 1765 | num_reqs, |
| 1766 | num_toks, |
| 1767 | uniform_tok_count, |
| 1768 | self.dp_size, |
| 1769 | self.dp_rank, |
| 1770 | max_query_len=max_query_len, |
| 1771 | need_eager=is_profile or skip_compiled, |
| 1772 | num_active_loras=num_active_loras, |
| 1773 | parallel_config=self.parallel_config, |
| 1774 | allow_ubatching=( |
| 1775 | self.ubatch_runner is not None and not skip_attn_for_dummy_run |
| 1776 | ), |
| 1777 | uniform_decode=uniform_tok_count == self.decode_query_len, |
| 1778 | ) |
| 1779 | |
| 1780 | if batch_desc.num_tokens == 0: |
| 1781 | # All DP ranks have zero tokens to run. |
| 1782 | empty_output = self.kv_connector.no_forward(scheduler_output) |
| 1783 | return self._merge_ec_connector_no_forward(scheduler_output, empty_output) |
| 1784 | |
| 1785 | cudagraph_stats = None |
| 1786 | if not dummy_run: |
| 1787 | # Common case. |
| 1788 | # Prepare all the inputs and copy to the input buffers. |
| 1789 | if self.observability_config.cudagraph_metrics: |
| 1790 | cudagraph_stats = make_cudagraph_stats(batch_desc, num_toks) |
| 1791 | assert batch_req_state is not None |
| 1792 | input_batch = self.prepare_inputs( |
| 1793 | scheduler_output, batch_req_state, batch_desc, num_active_loras |
| 1794 | ) |
| 1795 | block_tables, slot_mappings = self.prepare_attn(input_batch) |
| 1796 | # Mamba "align" pre-copy: migrate recurrent state across block |
| 1797 | # boundaries before the forward. Runs only on real batches, and |
| 1798 | # before model_state.prepare_attn gathers num_accepted_tokens so the |
| 1799 | # boundary reset is visible to the attention metadata. |
| 1800 | self.model_state.preprocess_state( |
| 1801 | input_batch, |
| 1802 | block_tables, |
| 1803 | self.kv_cache_config, |
| 1804 | self.req_states.num_computed_tokens.gpu, |
| 1805 | ) |
| 1806 | |
| 1807 | if self.lora_config: |
| 1808 | # Activate LoRA adapters. |
| 1809 | lora_inputs = self.lora_state.make_lora_inputs( |
| 1810 | input_batch.req_ids, |
| 1811 | input_batch.idx_mapping_np, |
| 1812 | input_batch.num_scheduled_tokens, |
| 1813 | ) |
| 1814 | self._set_active_loras(*lora_inputs) |
| 1815 | else: |
| 1816 | # No actual tokens to run. A dummy run for DP or memory profiling. |
| 1817 | dummy_num_reqs = batch_desc.num_reqs or num_reqs |
| 1818 | input_batch = InputBatch.make_dummy( |
| 1819 | dummy_num_reqs, |
| 1820 | batch_desc.num_tokens, |
| 1821 | self.input_buffers, |
| 1822 | max_query_len=batch_desc.max_query_len, |
| 1823 | # Profiling and warmup must route the dummy tokens to experts |
| 1824 | # so MoE memory is measured and MoE kernels are exercised. |
| 1825 | is_padding=not is_profile, |
| 1826 | ) |
| 1827 | if randomize_inputs: |
| 1828 | # All-zero input_ids route every token to the same experts. |
| 1829 | input_batch.input_ids.random_(0, self.vocab_size) |
| 1830 | if self.pcp_manager is not None: |
| 1831 | input_batch = self.pcp_manager.prepare_inputs_to_capture(input_batch) |
| 1832 | if skip_attn_for_dummy_run: |
| 1833 | assert batch_desc.cg_mode != CUDAGraphMode.FULL, ( |
| 1834 | "Attention metadata must be prepared for dummy runs when using " |
| 1835 | "FULL cudagraph mode." |
| 1836 | ) |
| 1837 | block_tables, slot_mappings = None, None |
| 1838 | else: |
| 1839 | block_tables, slot_mappings = self.prepare_dummy_attn( |
| 1840 | input_batch, valid_dummy_state_slots |
| 1841 | ) |
| 1842 | if context_len: |
| 1843 | set_dummy_context( |
| 1844 | input_batch, |
| 1845 | self.block_tables, |
| 1846 | context_len, |
| 1847 | self.kv_cache_config.num_blocks, |
| 1848 | self.max_model_len, |
| 1849 | input_block_tables=block_tables, |
| 1850 | ) |
| 1851 | |
| 1852 | if self.dcp_size > 1 and not (dummy_run and skip_attn_for_dummy_run): |
| 1853 | input_buffers = self.input_buffers |
| 1854 | if self.pcp_manager is not None: |
| 1855 | input_buffers = self.pcp_manager.input_buffers |
| 1856 | input_batch.dcp_local_seq_lens = prepare_dcp_local_seq_lens( |
| 1857 | input_buffers.dcp_local_seq_lens, |
| 1858 | input_batch.seq_lens, |
| 1859 | input_batch.num_reqs, |
| 1860 | self.dcp_size, |
| 1861 | self.dcp_rank, |
| 1862 | self.cp_interleave, |
| 1863 | num_reqs_padded=input_batch.num_reqs_after_padding, |
| 1864 | ) |
| 1865 | |
| 1866 | attn_metadata = None |
| 1867 | slot_mappings_by_layer = None |
| 1868 | ubatch_state: UBatchState | None = None |
| 1869 | if batch_desc.num_ubatches > 1: |
| 1870 | assert self.ubatch_runner is not None |
| 1871 | assert block_tables is not None and slot_mappings is not None |
| 1872 | ubatch_state = self.ubatch_runner.prepare( |
| 1873 | input_batch, |
| 1874 | block_tables, |
| 1875 | slot_mappings, |
| 1876 | cg_mode=batch_desc.cg_mode, |
| 1877 | for_capture=dummy_run and batch_desc.cg_mode == CUDAGraphMode.FULL, |
| 1878 | ) |
| 1879 | elif not (dummy_run and skip_attn_for_dummy_run): |
| 1880 | assert slot_mappings is not None |
| 1881 | slot_mappings_by_layer = build_slot_mappings_by_layer( |
| 1882 | slot_mappings, self.kv_cache_config |
| 1883 | ) |
| 1884 | assert block_tables is not None |
| 1885 | attn_groups = self.attn_groups |
| 1886 | if dummy_run and is_profile: |
| 1887 | # Mamba layers take a cheap warmup path with no metadata; |
| 1888 | # attention metadata is still built so those kernels tune. |
| 1889 | attn_groups = [ |
| 1890 | [g for g in groups if not isinstance(g.kv_cache_spec, MambaSpec)] |
| 1891 | for groups in attn_groups |
| 1892 | ] |
| 1893 | attn_metadata = self.model_state.prepare_attn( |
| 1894 | input_batch, |
| 1895 | batch_desc.cg_mode, |
| 1896 | block_tables, |
| 1897 | slot_mappings, |
| 1898 | attn_groups, |
| 1899 | self.kv_cache_config, |
| 1900 | # FULL replay reads capture-time metadata buffers. Re-stage them |
| 1901 | # from the zeroed dummy block tables instead of retaining state |
| 1902 | # indices from the previous real batch. |
| 1903 | for_capture=dummy_run and batch_desc.cg_mode == CUDAGraphMode.FULL, |
| 1904 | ) |
| 1905 | |
| 1906 | input_ids = input_batch.input_ids |
| 1907 | inputs_embeds = None |
| 1908 | ec_connector_output = None |
| 1909 | if self.uses_inputs_embeds and self.is_first_pp_rank: |
| 1910 | # Prepare inputs_embeds (MM encoder outputs and/or prompt_embeds |
| 1911 | # overlay). Only first PP rank prepares them. |
| 1912 | if dummy_run: |
| 1913 | # Obtain embeddings of correct shape for compiled model. |
| 1914 | inputs_embeds = self.model_state.dummy_inputs_embeds( |
| 1915 | input_batch.num_tokens_after_padding |
| 1916 | ) |
| 1917 | else: |
| 1918 | scheduled_encoder_inputs = scheduler_output.scheduled_encoder_inputs |
| 1919 | if self.supports_mm_inputs and self.lora_config is not None: |
| 1920 | set_active_mm_loras( |
| 1921 | model=self.model, |
| 1922 | lora_manager=self.lora_manager, |
| 1923 | encoder_cache=self.encoder_cache, |
| 1924 | req_id_to_index=self.req_states.req_id_to_index, |
| 1925 | lora_state=self.lora_state, |
| 1926 | scheduled_encoder_inputs=scheduled_encoder_inputs, |
| 1927 | ) |
| 1928 | with self.ec_connector.maybe_get_output( |
| 1929 | scheduler_output |
| 1930 | ) as ec_connector_output: |
| 1931 | inputs_embeds = self.model_state.prepare_inputs_embeds( |
| 1932 | scheduled_encoder_inputs, input_batch, self.req_states |
| 1933 | ) |
| 1934 | if inputs_embeds is not None and not requires_raw_input_tokens(self.model): |
| 1935 | input_ids = None |
| 1936 | |
| 1937 | model_inputs = { |
| 1938 | "input_ids": input_ids, |
| 1939 | "positions": input_batch.positions, |
| 1940 | "inputs_embeds": inputs_embeds, |
| 1941 | "intermediate_tensors": None, |
| 1942 | # NOTE: Values returned by `prepare_inputs` will override the default |
| 1943 | # values above. |
| 1944 | **self.model_state.prepare_inputs(input_batch, self.req_states), |
| 1945 | } |
| 1946 | if not self.is_first_pp_rank: |
| 1947 | # Update for non-first PP ranks. |
| 1948 | model_inputs["input_ids"] = None |
| 1949 | model_inputs["inputs_embeds"] = None |
| 1950 | |
| 1951 | # Prepare the intermediate tensors. |
| 1952 | assert intermediate_tensors is not None |
| 1953 | assert self.intermediate_tensors is not None |
| 1954 | n = input_batch.num_tokens_after_padding |
| 1955 | new_tensors = { |
| 1956 | k: v[:n] |
| 1957 | if dummy_run |
| 1958 | else v[:n].copy_(intermediate_tensors.tensors[k][:n]) |
| 1959 | for k, v in self.intermediate_tensors.tensors.items() |
| 1960 | } |
| 1961 | model_inputs["intermediate_tensors"] = IntermediateTensors(new_tensors) |
| 1962 | del intermediate_tensors |
| 1963 | |
| 1964 | # Update the EPLB meta. |
| 1965 | ubatch_slices = ubatch_state.slices if ubatch_state is not None else None |
| 1966 | self.eplb.prepare_forward( |
| 1967 | self.model_config, input_batch.num_tokens, ubatch_slices |
| 1968 | ) |
| 1969 | |
| 1970 | self.step_timing.record_batch( |
| 1971 | input_batch, batch_desc.cg_mode == CUDAGraphMode.FULL |
| 1972 | ) |
| 1973 | self.step_timing.forward_start() |
| 1974 | |
| 1975 | connector_kwargs = dict( |
| 1976 | scheduler_output=scheduler_output, |
| 1977 | request_state_indices=input_batch.idx_mapping, |
| 1978 | request_ids=input_batch.req_ids, |
| 1979 | num_tokens=input_batch.num_tokens, |
| 1980 | ) |
| 1981 | |
| 1982 | # Run model. |
| 1983 | if batch_desc.cg_mode == CUDAGraphMode.FULL: |
| 1984 | # Use explicit cudagraph replay for FULL mode. |
| 1985 | # NOTE(woosuk): Here, we don't need to pass the input tensors, |
| 1986 | # because they are already copied to the CUDA graph input buffers. |
| 1987 | assert self.cudagraph_manager is not None |
| 1988 | self.kv_connector.pre_forward( |
| 1989 | **connector_kwargs, attn_metadata=attn_metadata |
| 1990 | ) |
| 1991 | model_output = self.cudagraph_manager.run_fullgraph(batch_desc) |
| 1992 | else: |
| 1993 | # For piecewise and eager mode, just call model(). |
| 1994 | batch_descriptor = BatchDescriptor( |
| 1995 | num_tokens=input_batch.num_tokens_after_padding, |
| 1996 | has_lora=self.lora_config is not None, |
| 1997 | num_active_loras=batch_desc.num_active_loras, |
| 1998 | ) |
| 1999 | |
| 2000 | with set_forward_context( |
| 2001 | attn_metadata, |
| 2002 | self.vllm_config, |
| 2003 | num_tokens=input_batch.num_tokens_after_padding, |
| 2004 | cudagraph_runtime_mode=batch_desc.cg_mode, |
| 2005 | num_tokens_across_dp=( |
| 2006 | dp_sync.num_tokens_across_dp if dp_sync is not None else None |
| 2007 | ), |
| 2008 | batch_descriptor=batch_descriptor, |
| 2009 | ubatch_slices=ubatch_slices, |
| 2010 | slot_mapping=slot_mappings_by_layer, |
| 2011 | skip_compiled=skip_compiled, |
| 2012 | is_padding=input_batch.is_padding, |
| 2013 | ): |
| 2014 | self.kv_connector.pre_forward(**connector_kwargs) |
| 2015 | if ubatch_state is not None: |
| 2016 | assert self.ubatch_runner is not None |
| 2017 | model_output = self.ubatch_runner.run( |
| 2018 | self.model, model_inputs, ubatch_state |
| 2019 | ) |
| 2020 | elif batch_desc.cg_mode == CUDAGraphMode.PIECEWISE: |
| 2021 | # Run the PIECEWISE graph (compiled PW cudagraph or breakable |
| 2022 | # cudagraph, chosen inside run_pw_graph). cg_mode is only |
| 2023 | # PIECEWISE after the cudagraph manager exists. |
| 2024 | assert self.cudagraph_manager is not None |
| 2025 | model_output = self.cudagraph_manager.run_pw_graph( |
| 2026 | self.model, model_inputs |
| 2027 | ) |
| 2028 | else: |
| 2029 | # Eager (NONE): call the raw model directly. |
| 2030 | model_output = self.model(**model_inputs) |
| 2031 | |
| 2032 | self.kv_connector.finish_forward() |
| 2033 | |
| 2034 | if self.is_last_pp_rank: |
| 2035 | if self.use_aux_hidden_state_outputs: |
| 2036 | assert isinstance(model_output, tuple) |
| 2037 | hidden_states, aux_hidden_states = model_output |
| 2038 | else: |
| 2039 | assert isinstance(model_output, torch.Tensor) |
| 2040 | hidden_states = model_output |
| 2041 | aux_hidden_states = None |
| 2042 | output_intermediate_tensors = None |
| 2043 | else: |
| 2044 | assert isinstance(model_output, IntermediateTensors) |
| 2045 | hidden_states = None |
| 2046 | aux_hidden_states = None |
| 2047 | output_intermediate_tensors = model_output |
| 2048 | |
| 2049 | finished_req_ids = scheduler_output.finished_req_ids |
| 2050 | self.execute_model_state = ExecuteModelState( |
| 2051 | input_batch=input_batch, |
| 2052 | attn_metadata=attn_metadata, |
| 2053 | slot_mappings_by_layer=slot_mappings_by_layer, |
| 2054 | hidden_states=hidden_states, |
| 2055 | aux_hidden_states=aux_hidden_states, |
| 2056 | dp_sync=dp_sync, |
| 2057 | finished_req_ids=finished_req_ids, |
| 2058 | ec_connector_output=ec_connector_output, |
| 2059 | cudagraph_stats=cudagraph_stats, |
| 2060 | ) |
| 2061 | |
| 2062 | if not self.is_last_pp_rank: |
| 2063 | # Non-last PP rank: return IntermediateTensors for sending. |
| 2064 | assert output_intermediate_tensors is not None |
| 2065 | assert self.pp_handler is not None |
| 2066 | return self.pp_handler.relay_aux_hidden_states( |
| 2067 | model_inputs["intermediate_tensors"], output_intermediate_tensors |
| 2068 | ) |
| 2069 | return None |
| 2070 | |
| 2071 | @torch.inference_mode() |
| 2072 | @step_eplb_after() |
| 2073 | def sample_tokens( |
| 2074 | self, grammar_output: GrammarOutput | None |
| 2075 | ) -> AsyncOutput | ModelRunnerOutput | None: |
| 2076 | if self.execute_model_state is None: |
| 2077 | # The prior execute_model call must have failed. |
| 2078 | return None |
| 2079 | |
| 2080 | input_batch = self.execute_model_state.input_batch |
| 2081 | attn_metadata = self.execute_model_state.attn_metadata |
| 2082 | slot_mappings_by_layer = self.execute_model_state.slot_mappings_by_layer |
| 2083 | hidden_states = self.execute_model_state.hidden_states |
| 2084 | aux_hidden_states = self.execute_model_state.aux_hidden_states |
| 2085 | dp_sync = self.execute_model_state.dp_sync |
| 2086 | finished_req_ids = self.execute_model_state.finished_req_ids |
| 2087 | ec_connector_output = self.execute_model_state.ec_connector_output |
| 2088 | cudagraph_stats = self.execute_model_state.cudagraph_stats |
| 2089 | self.execute_model_state = None |
| 2090 | |
| 2091 | if not self.is_last_pp_rank: |
| 2092 | # Non-last PP rank: hidden_states is None because this rank produced |
| 2093 | # IntermediateTensors instead of final hidden states. Receive the |
| 2094 | # sampled tokens broadcast from the last rank and update local state. |
| 2095 | if self.pcp_manager is not None: |
| 2096 | input_batch = self.pcp_manager.global_batch |
| 2097 | assert self.pp_handler is not None |
| 2098 | all_decode_next = self.pp_handler.receive(input_batch) |
| 2099 | # Optimistically update num_computed_tokens for entire batch here. |
| 2100 | # Will be adjusted for rejections if necessary in update_requests. |
| 2101 | self.postprocess_num_computed_tokens(input_batch) |
| 2102 | if not all_decode_next: |
| 2103 | # Might contain non-final prefill chunks, which will be scheduled |
| 2104 | # in the immediate next step (rather than in pp_size steps). |
| 2105 | self.model_state.postprocess_state(input_batch.idx_mapping, 0) |
| 2106 | |
| 2107 | # Post-step KV connector related operations. |
| 2108 | kv_connector_output = self.kv_connector.post_forward(finished_req_ids) |
| 2109 | # The first PP rank holds the encoder cache, so pass its EC output on. |
| 2110 | output = ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) |
| 2111 | return ModelRunnerOutput.with_ec_conn_output(output, ec_connector_output) |
| 2112 | |
| 2113 | # Last rank: sample tokens |
| 2114 | assert hidden_states is not None |
| 2115 | draft_hidden_states = hidden_states |
| 2116 | if self.pcp_manager is not None: |
| 2117 | hidden_states, aux_hidden_states, input_batch = ( |
| 2118 | self.pcp_manager.restore_for_sampling(hidden_states, aux_hidden_states) |
| 2119 | ) |
| 2120 | |
| 2121 | sampler_output, num_sampled, num_rejected = self.sample( |
| 2122 | hidden_states, input_batch, grammar_output |
| 2123 | ) |
| 2124 | |
| 2125 | if self.pp_handler is not None: |
| 2126 | # Broadcast to non-last PP ranks (handles spec decode multi-token). |
| 2127 | self.pp_handler.broadcast( |
| 2128 | sampler_output.sampled_token_ids, |
| 2129 | num_sampled, |
| 2130 | num_rejected, |
| 2131 | input_batch, |
| 2132 | ) |
| 2133 | |
| 2134 | assert self.prompt_logprobs_worker is not None |
| 2135 | prompt_logprobs_dict = self.prompt_logprobs_worker.compute_prompt_logprobs( |
| 2136 | self.model.compute_logits, |
| 2137 | hidden_states, |
| 2138 | input_batch, |
| 2139 | self.req_states.all_token_ids.gpu, |
| 2140 | self.req_states.num_computed_tokens.gpu, |
| 2141 | self.req_states.prompt_len.np, |
| 2142 | ) |
| 2143 | prompt_token_id_logprobs_dict = ( |
| 2144 | self.prompt_logprobs_worker.compute_prompt_token_id_logprobs( |
| 2145 | self.model.compute_logits, |
| 2146 | hidden_states, |
| 2147 | input_batch, |
| 2148 | self.req_states.prompt_len.np, |
| 2149 | ) |
| 2150 | ) |
| 2151 | |
| 2152 | # Prepare the model runner output. |
| 2153 | model_runner_output = ModelRunnerOutput( |
| 2154 | req_ids=input_batch.req_ids, |
| 2155 | # NOTE(woosuk): req_id_to_index is unused in this model runner. |
| 2156 | # Only for compatibility with the existing model runner and scheduler. |
| 2157 | req_id_to_index={req_id: i for i, req_id in enumerate(input_batch.req_ids)}, |
| 2158 | sampled_token_ids=None, # type: ignore |
| 2159 | prompt_logprobs_dict=prompt_logprobs_dict, # type: ignore[arg-type] |
| 2160 | prompt_token_id_logprobs_dict=prompt_token_id_logprobs_dict, |
| 2161 | cudagraph_stats=cudagraph_stats, |
| 2162 | ) |
| 2163 | pending_aux_output = None |
| 2164 | if self.aux_output_connector is not None: |
| 2165 | pending_aux_output = self.aux_output_connector.prepare_output(input_batch) |
| 2166 | |
| 2167 | # Start async output copy here so that it can overlap with speculator proposal. |
| 2168 | async_output = AsyncOutput( |
| 2169 | model_runner_output=model_runner_output, |
| 2170 | sampler_output=sampler_output, |
| 2171 | num_sampled_tokens=num_sampled, |
| 2172 | main_stream=self.main_stream, |
| 2173 | copy_stream=self.output_copy_stream, |
| 2174 | check_ep_fault=self.check_ep_fault, |
| 2175 | pending_aux_output=pending_aux_output, |
| 2176 | ) |
| 2177 | |
| 2178 | mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None |
| 2179 | # The encoder runner exists only on the first PP rank, so later ranks |
| 2180 | # have no cached embeddings to gather. |
| 2181 | if ( |
| 2182 | self.speculator is not None |
| 2183 | and self.speculator.supports_mm_inputs |
| 2184 | and self.model_state.supports_mm_inputs |
| 2185 | ): |
| 2186 | # Get cached multimodal embeddings for draft forward. |
| 2187 | # NOTE: This is done here because postprocess updates |
| 2188 | # num_computed_prefill_tokens. |
| 2189 | # The EAGLE/MTP drafter reads one position ahead of the target. |
| 2190 | # TODO(TheEpicDolphin): Gather MM embeddings for all speculative |
| 2191 | # steps during multi-module MTP. |
| 2192 | mm_inputs = self.model_state.gather_mm_embeddings( |
| 2193 | input_batch, draft_lookahead=1 |
| 2194 | ) |
| 2195 | |
| 2196 | # Postprocess results and update request states. |
| 2197 | # NOTE: This is intentionally done after creating the AsyncOutput, |
| 2198 | # ensuring that `copy_event` is recorded before calling postprocess. |
| 2199 | # This sequencing may slightly reduce latency as async D2H copy does not |
| 2200 | # need to wait for the postprocess to finish. |
| 2201 | self.postprocess_sampled( |
| 2202 | input_batch.idx_mapping, |
| 2203 | sampler_output.sampled_token_ids, |
| 2204 | num_sampled, |
| 2205 | num_rejected, |
| 2206 | input_batch.query_start_loc, |
| 2207 | ) |
| 2208 | |
| 2209 | if self.speculator is not None: |
| 2210 | assert self.sampler is not None |
| 2211 | if isinstance(self.speculator, DraftModelSpeculator): |
| 2212 | self.speculator.observe_verification( |
| 2213 | input_batch.idx_mapping, num_sampled, num_rejected |
| 2214 | ) |
| 2215 | # Let the target override the hidden state fed to the drafter |
| 2216 | # (e.g. DeepSeek V4 MTP needs the pre-hc_head residual). The |
| 2217 | # target returns a persistent buffer sized at max_num_batched_tokens; |
| 2218 | # slice to the active token count that propose() expects. |
| 2219 | spec_hidden_states = draft_hidden_states |
| 2220 | if hasattr(self.model, "get_mtp_target_hidden_states"): |
| 2221 | pre_hc_hidden_states = self.model.get_mtp_target_hidden_states() |
| 2222 | spec_hidden_states = pre_hc_hidden_states[: draft_hidden_states.size(0)] |
| 2223 | if isinstance(self.sampler, GPUWatermarkSampler): |
| 2224 | self.speculator.prepare_watermarking( |
| 2225 | self.sampler, input_batch.idx_mapping |
| 2226 | ) |
| 2227 | with use_workspace_lane(self._draft_workspace_lane): |
| 2228 | draft_tokens = self.speculator.propose( |
| 2229 | input_batch, |
| 2230 | attn_metadata, |
| 2231 | slot_mappings_by_layer, |
| 2232 | spec_hidden_states, |
| 2233 | aux_hidden_states, |
| 2234 | num_sampled, |
| 2235 | num_rejected, |
| 2236 | self.req_states.last_sampled_tokens, |
| 2237 | self.req_states.next_prefill_tokens, |
| 2238 | self.sampler.sampling_states.temperature.gpu, |
| 2239 | self.sampler.sampling_states.seeds.gpu, |
| 2240 | dp_sync=dp_sync, |
| 2241 | mm_inputs=mm_inputs, |
| 2242 | ) |
| 2243 | self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens |
| 2244 | if self.adaptive_verification is not None: |
| 2245 | self.adaptive_verification.record_confidences( |
| 2246 | self.speculator.draft_token_confidence_probs, input_batch |
| 2247 | ) |
| 2248 | |
| 2249 | if self.num_speculative_steps > 0: |
| 2250 | # Spec-decode and diffusion LLMs both use draft tokens but the latter does |
| 2251 | # not have a speculator (i.e. self.speculator is None) |
| 2252 | self.draft_tokens_handler.set_draft_tokens( |
| 2253 | input_batch, |
| 2254 | self.req_states.draft_tokens[input_batch.idx_mapping], |
| 2255 | ) |
| 2256 | if self.pp_handler is not None: |
| 2257 | self.pp_handler.broadcast_drafts( |
| 2258 | self.req_states.draft_tokens, input_batch |
| 2259 | ) |
| 2260 | |
| 2261 | # Post-step KV connector related operations. |
| 2262 | kv_connector_output = self.kv_connector.post_forward(finished_req_ids) |
| 2263 | model_runner_output.kv_connector_output = kv_connector_output |
| 2264 | model_runner_output.ec_connector_output = ec_connector_output |
| 2265 | |
| 2266 | return async_output |
| 2267 | |
| 2268 | def take_draft_token_ids(self) -> DraftTokenIds | None: |
| 2269 | return self.draft_tokens_handler.get_draft_tokens() |
| 2270 | |
| 2271 | @torch.inference_mode() |
| 2272 | @step_eplb_after() |
| 2273 | def pool(self) -> AsyncPoolingOutput | ModelRunnerOutput | None: |
| 2274 | if self.execute_model_state is None: |
| 2275 | # The prior execute_model call must have failed. |
| 2276 | return None |
| 2277 | |
| 2278 | input_batch = self.execute_model_state.input_batch |
| 2279 | hidden_states = self.execute_model_state.hidden_states |
| 2280 | finished_req_ids = self.execute_model_state.finished_req_ids |
| 2281 | ec_connector_output = self.execute_model_state.ec_connector_output |
| 2282 | self.execute_model_state = None |
| 2283 | |
| 2284 | # Post-step KV connector related operations. |
| 2285 | kv_connector_output = self.kv_connector.post_forward(finished_req_ids) |
| 2286 | |
| 2287 | if not self.is_last_pp_rank: |
| 2288 | self.postprocess_num_computed_tokens(input_batch) |
| 2289 | output = ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) |
| 2290 | return ModelRunnerOutput.with_ec_conn_output(output, ec_connector_output) |
| 2291 | |
| 2292 | assert self.pooling_runner is not None |
| 2293 | pooler_output, finished_mask = self.pooling_runner.pool( |
| 2294 | hidden_states, input_batch, self.req_states |
| 2295 | ) |
| 2296 | |
| 2297 | # Build the model runner output. |
| 2298 | model_runner_output = ModelRunnerOutput( |
| 2299 | req_ids=input_batch.req_ids, |
| 2300 | req_id_to_index={req_id: i for i, req_id in enumerate(input_batch.req_ids)}, |
| 2301 | kv_connector_output=kv_connector_output, |
| 2302 | ec_connector_output=ec_connector_output, |
| 2303 | ) |
| 2304 | async_output = AsyncPoolingOutput( |
| 2305 | model_runner_output=model_runner_output, |
| 2306 | pooler_output=pooler_output, |
| 2307 | finished_mask=finished_mask, |
| 2308 | main_stream=self.main_stream, |
| 2309 | copy_stream=self.output_copy_stream, |
| 2310 | ) |
| 2311 | |
| 2312 | self.postprocess_num_computed_tokens(input_batch) |
| 2313 | return async_output |
| 2314 | |
| 2315 | def postprocess_num_computed_tokens(self, input_batch: InputBatch) -> None: |
| 2316 | # Update the number of computed tokens. |
| 2317 | post_update_num_computed_tokens( |
| 2318 | input_batch.idx_mapping, |
| 2319 | self.req_states.num_computed_tokens.gpu, |
| 2320 | input_batch.query_start_loc, |
| 2321 | ) |
| 2322 | |
| 2323 | def shutdown(self) -> None: |
| 2324 | """Release GPU tensors (model weights, KV caches, workspace) so that |
| 2325 | memory is reclaimable when running in the same process.""" |
| 2326 | torch.accelerator.synchronize() |
| 2327 | if self.aux_output_connector is not None: |
| 2328 | self.aux_output_connector.close() |
| 2329 | set_offloader(None) |
| 2330 | self.cudagraph_manager = None |
| 2331 | self.fast_prefill = None |
| 2332 | self.pooling_runner = None |
| 2333 | if hasattr(self, "kv_caches"): |
| 2334 | self.kv_caches.clear() |
| 2335 | if hasattr(self, "attn_groups"): |
| 2336 | self.attn_groups.clear() |
| 2337 | if hasattr(self, "kv_cache_config"): |
| 2338 | del self.kv_cache_config |
| 2339 | if hasattr(self, "model_state") and self.model_state.supports_mm_inputs: |
| 2340 | self.model_state.encoder_runner.clear() |
| 2341 | free_before_shutdown(self.vllm_config) |
| 2342 | if hasattr(self, "model_state"): |
| 2343 | del self.model_state |
| 2344 | # Detach the layer-level KV/state cache tensors before dropping the |
| 2345 | # models; the model objects can outlive this runner. |
| 2346 | speculator = getattr(self, "speculator", None) |
| 2347 | if speculator is not None: |
| 2348 | if draft_model := getattr(speculator, "model", None): |
| 2349 | clear_layer_kv_caches(draft_model.modules()) |
| 2350 | self.speculator = None |
| 2351 | if hasattr(self, "model"): |
| 2352 | clear_layer_kv_caches(self.model.modules()) |
| 2353 | del self.model |
| 2354 | |
| 2355 | gc.collect() |
| 2356 | torch.accelerator.empty_cache() |
| 2357 | logger.debug("Cleaned up model weights, KV caches, and workspace") |
| 2358 | |
| 2359 | ########### EPLB methods start ########### |
| 2360 | @property |
| 2361 | def eplb_state(self): |
| 2362 | return self.eplb.state |
| 2363 | |
| 2364 | @eplb_state.setter |
| 2365 | def eplb_state(self, state) -> None: |
| 2366 | self.eplb.state = state |
| 2367 | |
| 2368 | @property |
| 2369 | def eep_eplb_suppressed(self) -> bool: |
| 2370 | return self.eplb.suppressed |
| 2371 | |
| 2372 | @eep_eplb_suppressed.setter |
| 2373 | def eep_eplb_suppressed(self, suppressed: bool) -> None: |
| 2374 | self.eplb.suppressed = suppressed |
| 2375 | |
| 2376 | def setup_eplb_from_mapping( |
| 2377 | self, expanded_physical_to_logical: torch.Tensor |
| 2378 | ) -> None: |
| 2379 | self.eplb.setup_from_mapping(self.model_config, expanded_physical_to_logical) |
| 2380 | |
| 2381 | ########### EPLB methods end ########### |
| 2382 | |
| 2383 | # Out-of-tree hardware runners can select a PCP manager class. |
| 2384 | @property |
| 2385 | def pcp_manager_cls(self) -> type[pcp.PCPManager]: |
| 2386 | return pcp.PCPManager |
| 2387 | |
| 2388 | |
| 2389 | class ExecuteModelState(NamedTuple): |
| 2390 | input_batch: InputBatch |
| 2391 | attn_metadata: dict[str, Any] | None |
| 2392 | slot_mappings_by_layer: dict[str, torch.Tensor] | None |
| 2393 | hidden_states: torch.Tensor | None |
| 2394 | aux_hidden_states: list[torch.Tensor] | None |
| 2395 | dp_sync: DPSyncState | None |
| 2396 | finished_req_ids: set[str] |
| 2397 | ec_connector_output: ECConnectorOutput | None |
| 2398 | cudagraph_stats: CUDAGraphStat | None |
| 2399 | |
| 2400 | |
| 2401 | class BatchReqState(NamedTuple): |
| 2402 | """CPU request state for a scheduled batch, in batch (sorted) order.""" |
| 2403 | |
| 2404 | req_ids: list[str] |
| 2405 | num_scheduled_tokens: np.ndarray # [num_reqs] |
| 2406 | # May be less than scheduler_output.total_num_scheduled_tokens: |
| 2407 | # adaptive verification trims the draft budget before running. |
| 2408 | num_tokens: int |
| 2409 | num_draft_tokens_np: np.ndarray | None # [num_reqs], None if no drafts |
| 2410 | idx_mapping_np: np.ndarray # [num_reqs] |
| 2411 | prefill_len_np: np.ndarray # [num_reqs] |
| 2412 | num_computed_prefill_tokens_np: np.ndarray # [num_reqs] |
| 2413 | is_prefilling_np: np.ndarray # [num_reqs] |
| 2414 | max_seq_len_np: np.ndarray | None # [num_reqs], None in non-pp case |
| 2415 | has_prefill: bool |
| 2416 | prefill_runs_as_decode_np: np.ndarray | None # [num_reqs] |
| 2417 | decode_graph_eligible: bool |
| 2418 | |
| 2419 | |
| 2420 | def sort_batch_req_ids( |
| 2421 | num_tokens_per_req: dict[str, int], |
| 2422 | draft_tokens: dict[str, list[int]], |
| 2423 | decode_query_len: int, |
| 2424 | ) -> list[str]: |
| 2425 | # Order verification/decode -> short_extend -> prefill; |
| 2426 | # split_decodes_and_prefills relies on decode-like requests leading. |
| 2427 | key = lambda r: ( |
| 2428 | not draft_tokens.get(r), |
| 2429 | (num := num_tokens_per_req[r]) != decode_query_len, |
| 2430 | num, |
| 2431 | ) |
| 2432 | return sorted(num_tokens_per_req, key=key) |