| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import enum |
| 5 | import time |
| 6 | from collections import deque |
| 7 | from collections.abc import Callable, Mapping |
| 8 | from dataclasses import dataclass |
| 9 | from typing import TYPE_CHECKING, Any |
| 10 | |
| 11 | import torch |
| 12 | |
| 13 | from vllm.multimodal.inputs import MultiModalFeatureSpec |
| 14 | from vllm.pooling_params import PoolingParams |
| 15 | from vllm.sampling_params import SamplingParams |
| 16 | from vllm.utils import length_from_prompt_token_ids_or_embeds |
| 17 | from vllm.v1.engine import ( |
| 18 | EngineCoreEvent, |
| 19 | EngineCoreEventType, |
| 20 | EngineCoreRequest, |
| 21 | FinishReason, |
| 22 | ) |
| 23 | from vllm.v1.kv_hints import KvHintsEnvelope |
| 24 | from vllm.v1.metrics.stats import PrefillStats, RequestSpecDecodeMetrics |
| 25 | from vllm.v1.structured_output.request import StructuredOutputRequest |
| 26 | from vllm.v1.utils import ConstantList |
| 27 | |
| 28 | if TYPE_CHECKING: |
| 29 | from vllm.lora.request import LoRARequest |
| 30 | from vllm.v1.core.kv_cache_utils import BlockHash |
| 31 | |
| 32 | |
| 33 | @dataclass |
| 34 | class StreamingUpdate: |
| 35 | """Lightweight data for streaming session continuation. |
| 36 | |
| 37 | Contains only the fields needed to update an existing streaming session |
| 38 | with new input data. |
| 39 | """ |
| 40 | |
| 41 | mm_features: list[MultiModalFeatureSpec] | None |
| 42 | prompt_token_ids: list[int] | None |
| 43 | max_tokens: int |
| 44 | arrival_time: float |
| 45 | sampling_params: SamplingParams | None |
| 46 | |
| 47 | @classmethod |
| 48 | def from_request(cls, request: "Request") -> "StreamingUpdate | None": |
| 49 | if not request.resumable: |
| 50 | return None |
| 51 | return cls( |
| 52 | mm_features=request.mm_features, |
| 53 | prompt_token_ids=request.prompt_token_ids, |
| 54 | max_tokens=request.max_tokens, |
| 55 | arrival_time=request.arrival_time, |
| 56 | sampling_params=request.sampling_params, |
| 57 | ) |
| 58 | |
| 59 | |
| 60 | class Request: |
| 61 | def __init__( |
| 62 | self, |
| 63 | request_id: str, |
| 64 | prompt_token_ids: list[int] | None, |
| 65 | sampling_params: SamplingParams | None, |
| 66 | pooling_params: PoolingParams | None, |
| 67 | client_index: int = 0, |
| 68 | arrival_time: float | None = None, |
| 69 | prompt_embeds: torch.Tensor | None = None, |
| 70 | prompt_is_token_ids: list[bool] | None = None, |
| 71 | mm_features: list[MultiModalFeatureSpec] | None = None, |
| 72 | lora_request: "LoRARequest | None" = None, |
| 73 | cache_salt: str | None = None, |
| 74 | priority: int = 0, |
| 75 | trace_headers: Mapping[str, str] | None = None, |
| 76 | block_hasher: Callable[["Request"], list["BlockHash"]] | None = None, |
| 77 | resumable: bool = False, |
| 78 | session_id: str | None = None, |
| 79 | reasoning_ended: bool | None = None, |
| 80 | reasoning_parser_kwargs: dict[str, Any] | None = None, |
| 81 | abort_immediately: bool = False, |
| 82 | kv_hints: KvHintsEnvelope | None = None, |
| 83 | ) -> None: |
| 84 | self.request_id = request_id |
| 85 | self.client_index = client_index |
| 86 | self.priority = priority |
| 87 | self.sampling_params = sampling_params |
| 88 | self.pooling_params = pooling_params |
| 89 | self.lora_request = lora_request |
| 90 | self.structured_output_request = StructuredOutputRequest.from_sampling_params( |
| 91 | sampling_params |
| 92 | ) |
| 93 | if self.structured_output_request is not None: |
| 94 | self.structured_output_request.reasoning_ended = reasoning_ended |
| 95 | self.structured_output_request.reasoning_parser_kwargs = ( |
| 96 | reasoning_parser_kwargs |
| 97 | ) |
| 98 | self.arrival_time = arrival_time if arrival_time is not None else time.time() |
| 99 | |
| 100 | self.status = RequestStatus.WAITING |
| 101 | self.events: list[EngineCoreEvent] = [] |
| 102 | self.stop_reason: int | str | None = None |
| 103 | |
| 104 | # P/D: Connector-specific KV transfer parameters. |
| 105 | self.kv_transfer_params: dict[str, Any] | None = None |
| 106 | # E/P/D: Connector-specific encoder-cache transfer parameters. |
| 107 | self.ec_transfer_params: dict[str, Any] | None = None |
| 108 | |
| 109 | if pooling_params is not None: |
| 110 | # Pooling models. |
| 111 | self.max_tokens = 1 |
| 112 | elif sampling_params is not None: |
| 113 | # Generative models. |
| 114 | assert sampling_params.max_tokens is not None |
| 115 | self.max_tokens = sampling_params.max_tokens |
| 116 | if self.structured_output_request is not None: |
| 117 | self.status = RequestStatus.WAITING_FOR_STRUCTURED_OUTPUT_GRAMMAR |
| 118 | |
| 119 | if sampling_params.extra_args is not None: |
| 120 | self.kv_transfer_params = sampling_params.extra_args.get( |
| 121 | "kv_transfer_params" |
| 122 | ) |
| 123 | self.ec_transfer_params = sampling_params.extra_args.get( |
| 124 | "ec_transfer_params" |
| 125 | ) |
| 126 | self.kv_cache_report_mode = sampling_params.extra_args.get( |
| 127 | "kv_cache_report_mode", "incremental" |
| 128 | ) |
| 129 | else: |
| 130 | self.kv_cache_report_mode = "incremental" |
| 131 | else: |
| 132 | raise ValueError("sampling_params and pooling_params can't both be unset") |
| 133 | |
| 134 | self.prompt_token_ids = prompt_token_ids |
| 135 | self.prompt_embeds = prompt_embeds |
| 136 | # Per-position mask used in mixed-mode (chat completion with |
| 137 | # prompt_embeds). `None` except when both `prompt_token_ids` and |
| 138 | # `prompt_embeds` are set and their positions are interleaved. |
| 139 | self.prompt_is_token_ids = prompt_is_token_ids |
| 140 | # Cache per-block prompt-embed hashes to avoid rehashing the same |
| 141 | # tensor slices when generating extra keys. |
| 142 | self._prompt_embeds_per_block_hashes: dict[tuple[int, int], bytes] = {} |
| 143 | self.num_prompt_tokens = length_from_prompt_token_ids_or_embeds( |
| 144 | prompt_token_ids, prompt_embeds |
| 145 | ) |
| 146 | self._output_token_ids: list[int] = [] |
| 147 | if self.prompt_token_ids is None: |
| 148 | self._all_token_ids: list[int] = [0] * self.num_prompt_tokens |
| 149 | elif self.prompt_is_token_ids is None: |
| 150 | self._all_token_ids = self.prompt_token_ids.copy() |
| 151 | else: |
| 152 | # Mixed-mode prompt: positions covered by prompt_embeds hold a sentinel |
| 153 | # special token id that may lie outside the embedding. Zero them, matching |
| 154 | # the no-token-ids case above, so embedding gathers over these placeholder |
| 155 | # ids stay in bounds; the actual inputs come from prompt_embeds. |
| 156 | self._all_token_ids = [ |
| 157 | t if is_tok else 0 |
| 158 | for t, is_tok in zip(self.prompt_token_ids, self.prompt_is_token_ids) |
| 159 | ] |
| 160 | |
| 161 | # Used in async scheduling. |
| 162 | self.num_output_placeholders = 0 |
| 163 | # Tokens of output in flight when the request was preempted: delivered |
| 164 | # on return, but must not mutate the reset counters. |
| 165 | self.num_stale_output_tokens = 0 |
| 166 | # Drop the stale output instead, for same-step preempt + resume |
| 167 | # (reset_prefix_cache). |
| 168 | self.drop_stale_output = False |
| 169 | |
| 170 | # Tokens of steps whose output is not yet processed (async scheduling |
| 171 | # and PP run ahead of the GPU); `num_computed_tokens` counts them |
| 172 | # optimistically. |
| 173 | self.num_in_flight_tokens = 0 |
| 174 | |
| 175 | # V2+PP+async: Enforces `pp_size` cadence between same-request decode steps |
| 176 | # so the worker's broadcast slot ring stays consistent. |
| 177 | self.next_decode_eligible_step = 0 |
| 178 | |
| 179 | # Seq of the most recent step this request was scheduled in; fences |
| 180 | # deferred block freeing (see Scheduler._free_request_blocks). |
| 181 | self.last_sched_seq = 0 |
| 182 | |
| 183 | self.spec_token_ids: list[int] = [] |
| 184 | self.num_computed_tokens = 0 |
| 185 | self.cache_salt: str | None = cache_salt |
| 186 | |
| 187 | # Multi-modal related |
| 188 | self.mm_features = mm_features or [] |
| 189 | |
| 190 | # Read-only views |
| 191 | # Prevent directly appending to these lists since |
| 192 | # they should also be updated simultaneously. |
| 193 | self.output_token_ids = ConstantList(self._output_token_ids) |
| 194 | self.all_token_ids = ConstantList(self._all_token_ids) |
| 195 | # trace_headers |
| 196 | self.trace_headers = trace_headers |
| 197 | self.session_id = session_id |
| 198 | self.kv_hints = kv_hints |
| 199 | |
| 200 | # True if this request is scheduled as a non-final prefill chunk. |
| 201 | self.is_prefill_chunk = False |
| 202 | |
| 203 | # Block-aligned token position of a proven shared prefix worth pinning |
| 204 | # in the (sparse) prefix cache; 0 means none. Set at admission for |
| 205 | # hybrid/Mamba models when a shared prefix is detected (Marconi-style). |
| 206 | self.shared_prefix_boundary = 0 |
| 207 | # DeepSeek-V4.1 only: SWA bounded replay. The request holds no |
| 208 | # sliding-window KV below this position; 0 when nothing replays. |
| 209 | self.replay_start = 0 |
| 210 | |
| 211 | # The number of NaNs in logits. A value greater than 0 |
| 212 | # indicates that the output is corrupted |
| 213 | self.num_nans_in_logits = 0 |
| 214 | |
| 215 | # The number of times this request has been preempted by the scheduler. |
| 216 | self.num_preemptions = 0 |
| 217 | |
| 218 | self.prefill_stats: PrefillStats | None = PrefillStats() |
| 219 | |
| 220 | # Per-request speculative-decoding acceptance accumulator. Populated by |
| 221 | # the scheduler when --per-request-spec-decode-metrics is set (eagerly on |
| 222 | # add_request, then observed each verify step); stays None otherwise. |
| 223 | self.spec_decode_metrics: RequestSpecDecodeMetrics | None = None |
| 224 | |
| 225 | self.block_hashes: list[BlockHash] = [] |
| 226 | # Store the block hasher without binding self to avoid creating a |
| 227 | # reference cycle (Request -> partial -> Request) that prevents |
| 228 | # immediate garbage collection via reference counting. |
| 229 | self._block_hasher: Callable[[Request], list[BlockHash]] | None = block_hasher |
| 230 | self.update_block_hashes() |
| 231 | |
| 232 | self.skip_reading_prefix_cache = self.get_skip_reading_prefix_cache() |
| 233 | |
| 234 | # Used for streaming |
| 235 | self.resumable = resumable |
| 236 | # None entry in the queue means finished. |
| 237 | self.streaming_queue: deque[StreamingUpdate | None] | None = None |
| 238 | |
| 239 | # If True, request should be aborted immediately after being added to |
| 240 | # the scheduler so the connector's request_finished hook runs. |
| 241 | self.abort_immediately = abort_immediately |
| 242 | |
| 243 | @classmethod |
| 244 | def from_engine_core_request( |
| 245 | cls, |
| 246 | request: EngineCoreRequest, |
| 247 | block_hasher: Callable[["Request"], list["BlockHash"]] | None, |
| 248 | ) -> "Request": |
| 249 | return cls( |
| 250 | request_id=request.request_id, |
| 251 | client_index=request.client_index, |
| 252 | prompt_token_ids=request.prompt_token_ids, |
| 253 | prompt_embeds=request.prompt_embeds, |
| 254 | prompt_is_token_ids=request.prompt_is_token_ids, |
| 255 | mm_features=request.mm_features, |
| 256 | sampling_params=request.sampling_params, |
| 257 | pooling_params=request.pooling_params, |
| 258 | arrival_time=request.arrival_time, |
| 259 | lora_request=request.lora_request, |
| 260 | cache_salt=request.cache_salt, |
| 261 | priority=request.priority, |
| 262 | trace_headers=request.trace_headers, |
| 263 | block_hasher=block_hasher, |
| 264 | resumable=request.resumable, |
| 265 | session_id=request.session_id, |
| 266 | kv_hints=request.kv_hints, |
| 267 | reasoning_ended=request.reasoning_ended, |
| 268 | reasoning_parser_kwargs=request.reasoning_parser_kwargs, |
| 269 | abort_immediately=request.abort_immediately, |
| 270 | ) |
| 271 | |
| 272 | def append_output_token_ids( |
| 273 | self, |
| 274 | token_ids: int | list[int], |
| 275 | ) -> None: |
| 276 | if isinstance(token_ids, int): |
| 277 | self._output_token_ids.append(token_ids) |
| 278 | self._all_token_ids.append(token_ids) |
| 279 | else: |
| 280 | self._output_token_ids.extend(token_ids) |
| 281 | self._all_token_ids.extend(token_ids) |
| 282 | |
| 283 | self.update_block_hashes() |
| 284 | |
| 285 | def update_block_hashes(self) -> None: |
| 286 | """Compute block hashes for any new full blocks and append them.""" |
| 287 | if self._block_hasher is not None: |
| 288 | self.block_hashes.extend(self._block_hasher(self)) |
| 289 | |
| 290 | @property |
| 291 | def use_structured_output(self) -> bool: |
| 292 | return self.structured_output_request is not None |
| 293 | |
| 294 | @property |
| 295 | def num_tokens(self) -> int: |
| 296 | return len(self._all_token_ids) |
| 297 | |
| 298 | @property |
| 299 | def num_tokens_with_spec(self) -> int: |
| 300 | return len(self._all_token_ids) + len(self.spec_token_ids) |
| 301 | |
| 302 | @property |
| 303 | def num_output_tokens(self) -> int: |
| 304 | return len(self._output_token_ids) |
| 305 | |
| 306 | @property |
| 307 | def num_encoder_inputs(self) -> int: |
| 308 | return len(self.mm_features) |
| 309 | |
| 310 | @property |
| 311 | def has_encoder_inputs(self) -> bool: |
| 312 | return self.num_encoder_inputs > 0 |
| 313 | |
| 314 | def get_skip_reading_prefix_cache(self) -> bool: |
| 315 | if ( |
| 316 | self.sampling_params is not None |
| 317 | and self.sampling_params.skip_reading_prefix_cache is not None |
| 318 | ): |
| 319 | return self.sampling_params.skip_reading_prefix_cache |
| 320 | elif ( |
| 321 | self.pooling_params is not None |
| 322 | and self.pooling_params.skip_reading_prefix_cache is not None |
| 323 | ): |
| 324 | return self.pooling_params.skip_reading_prefix_cache |
| 325 | return False |
| 326 | |
| 327 | def is_finished(self) -> bool: |
| 328 | return RequestStatus.is_finished(self.status) |
| 329 | |
| 330 | def get_finished_reason(self) -> FinishReason | None: |
| 331 | return RequestStatus.get_finished_reason(self.status) |
| 332 | |
| 333 | def get_num_encoder_embeds(self, input_id: int) -> int: |
| 334 | assert input_id < len(self.mm_features) |
| 335 | return self.mm_features[input_id].mm_position.get_num_embeds() |
| 336 | |
| 337 | def record_event( |
| 338 | self, |
| 339 | event_type: EngineCoreEventType, |
| 340 | timestamp: float | None = None, |
| 341 | ) -> None: |
| 342 | self.events.append(EngineCoreEvent.new_event(event_type, timestamp)) |
| 343 | |
| 344 | def take_events(self) -> list[EngineCoreEvent] | None: |
| 345 | if not self.events: |
| 346 | return None |
| 347 | events, self.events = self.events, [] |
| 348 | return events |
| 349 | |
| 350 | def take_prefill_stats(self) -> PrefillStats | None: |
| 351 | if self.prefill_stats is None: |
| 352 | return None |
| 353 | prefill_stats = self.prefill_stats |
| 354 | self.prefill_stats = None |
| 355 | return prefill_stats |
| 356 | |
| 357 | def __lt__(self, other: "Request") -> bool: |
| 358 | """Compare two requests based on priority, arrival time, and request ID. |
| 359 | Used in priority scheduling. |
| 360 | """ |
| 361 | if self.priority != other.priority: |
| 362 | return self.priority < other.priority |
| 363 | if self.arrival_time != other.arrival_time: |
| 364 | return self.arrival_time < other.arrival_time |
| 365 | if self.request_id != other.request_id: |
| 366 | return self.request_id < other.request_id |
| 367 | return id(self) < id(other) |
| 368 | |
| 369 | |
| 370 | class RequestStatus(enum.IntEnum): |
| 371 | """Status of a request.""" |
| 372 | |
| 373 | WAITING = enum.auto() |
| 374 | WAITING_FOR_STRUCTURED_OUTPUT_GRAMMAR = enum.auto() |
| 375 | WAITING_FOR_REMOTE_KVS = enum.auto() |
| 376 | WAITING_FOR_STREAMING_REQ = enum.auto() |
| 377 | RUNNING = enum.auto() |
| 378 | PREEMPTED = enum.auto() |
| 379 | # Note: anything after PREEMPTED will be considered |
| 380 | # as a finished status. |
| 381 | FINISHED_STOPPED = enum.auto() |
| 382 | FINISHED_LENGTH_CAPPED = enum.auto() |
| 383 | FINISHED_ABORTED = enum.auto() |
| 384 | FINISHED_IGNORED = enum.auto() |
| 385 | FINISHED_ERROR = enum.auto() |
| 386 | FINISHED_REPETITION = enum.auto() |
| 387 | |
| 388 | def __str__(self) -> str: |
| 389 | return self.name |
| 390 | |
| 391 | @staticmethod |
| 392 | def is_finished(status: "RequestStatus") -> bool: |
| 393 | return status > RequestStatus.PREEMPTED |
| 394 | |
| 395 | @staticmethod |
| 396 | def get_finished_reason(status: "RequestStatus") -> FinishReason | None: |
| 397 | return _FINISHED_REASON_MAP.get(status) |
| 398 | |
| 399 | |
| 400 | # Mapping of finished statuses to their finish reasons. |
| 401 | # NOTE: The ignored requests are the requests whose prompt lengths |
| 402 | # are longer than the model's length cap. Therefore, the stop |
| 403 | # reason should also be "length" as in OpenAI API. |
| 404 | _FINISHED_REASON_MAP = { |
| 405 | RequestStatus.FINISHED_STOPPED: FinishReason.STOP, |
| 406 | RequestStatus.FINISHED_LENGTH_CAPPED: FinishReason.LENGTH, |
| 407 | RequestStatus.FINISHED_ABORTED: FinishReason.ABORT, |
| 408 | RequestStatus.FINISHED_IGNORED: FinishReason.LENGTH, |
| 409 | RequestStatus.FINISHED_ERROR: FinishReason.ERROR, |
| 410 | RequestStatus.WAITING_FOR_STREAMING_REQ: FinishReason.STOP, |
| 411 | RequestStatus.FINISHED_REPETITION: FinishReason.REPETITION, |
| 412 | } |