| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import asyncio |
| 5 | from collections import defaultdict, deque |
| 6 | from collections.abc import Iterable |
| 7 | from dataclasses import dataclass |
| 8 | from typing import TYPE_CHECKING, Any, cast |
| 9 | |
| 10 | import numpy as np |
| 11 | import torch |
| 12 | |
| 13 | from vllm.lora.request import LoRARequest |
| 14 | from vllm.outputs import ( |
| 15 | STREAM_FINISHED, |
| 16 | CompletionOutput, |
| 17 | PoolingOutput, |
| 18 | PoolingRequestOutput, |
| 19 | RequestError, |
| 20 | RequestOutput, |
| 21 | SamplingMask, |
| 22 | ) |
| 23 | from vllm.sampling_params import RequestOutputKind |
| 24 | from vllm.tokenizers import TokenizerLike |
| 25 | from vllm.tracing import ( |
| 26 | SpanAttributes, |
| 27 | SpanKind, |
| 28 | extract_trace_context, |
| 29 | instrument_manual, |
| 30 | ) |
| 31 | from vllm.utils import length_from_prompt_token_ids_or_embeds |
| 32 | from vllm.v1.engine import EngineCoreOutput, EngineCoreRequest, FinishReason |
| 33 | from vllm.v1.engine.detokenizer import IncrementalDetokenizer |
| 34 | from vllm.v1.engine.logprobs import LogprobsProcessor |
| 35 | from vllm.v1.engine.parallel_sampling import ParentRequest |
| 36 | from vllm.v1.metrics.stats import ( |
| 37 | IterationStats, |
| 38 | LoRARequestStates, |
| 39 | RequestSpecDecodeMetrics, |
| 40 | RequestStateStats, |
| 41 | SchedulerStats, |
| 42 | ) |
| 43 | from vllm.v1.outputs import SamplingMaskLists |
| 44 | |
| 45 | if TYPE_CHECKING: |
| 46 | from vllm.v1.engine.admission_control import SharedAdmissionStats |
| 47 | |
| 48 | # shared empty CPU tensor used as a placeholder pooling output |
| 49 | EMPTY_CPU_TENSOR = torch.empty(0, device="cpu") |
| 50 | |
| 51 | |
| 52 | class RequestOutputCollector: |
| 53 | """Collects streamed RequestOutputs per individual request, |
| 54 | for hand-off to the consuming asyncio generate task. |
| 55 | |
| 56 | When streaming deltas, RequestOutputs are merged if the |
| 57 | producer gets ahead of the consumer. |
| 58 | """ |
| 59 | |
| 60 | def __init__(self, output_kind: RequestOutputKind, request_id: str): |
| 61 | self.aggregate = output_kind == RequestOutputKind.DELTA |
| 62 | self.request_id = request_id |
| 63 | self.output: RequestOutput | PoolingRequestOutput | Exception | None = None |
| 64 | self.ready = asyncio.Event() |
| 65 | |
| 66 | self._input_stream_task: asyncio.Task | None = None |
| 67 | |
| 68 | def put(self, output: RequestOutput | PoolingRequestOutput | Exception) -> None: |
| 69 | """Non-blocking put operation.""" |
| 70 | if self.output is None or isinstance(output, Exception): |
| 71 | self.output = output |
| 72 | self.ready.set() |
| 73 | elif isinstance(self.output, RequestOutput) and isinstance( |
| 74 | output, RequestOutput |
| 75 | ): |
| 76 | # This ensures that request outputs with different request indexes |
| 77 | # (if n > 1) do not override each other. |
| 78 | self.output.add(output, aggregate=self.aggregate) |
| 79 | elif isinstance(self.output, PoolingRequestOutput) and isinstance( |
| 80 | output, PoolingRequestOutput |
| 81 | ): |
| 82 | self.output = output |
| 83 | |
| 84 | async def get(self) -> RequestOutput | PoolingRequestOutput: |
| 85 | """Get operation blocks on put event.""" |
| 86 | while (output := self.output) is None: |
| 87 | await self.ready.wait() |
| 88 | self.output = None |
| 89 | self.ready.clear() |
| 90 | if isinstance(output, Exception): |
| 91 | raise output |
| 92 | return output |
| 93 | |
| 94 | def get_nowait(self) -> RequestOutput | PoolingRequestOutput | None: |
| 95 | """Non-blocking get operation.""" |
| 96 | output = self.output |
| 97 | if output is not None: |
| 98 | self.output = None |
| 99 | self.ready.clear() |
| 100 | if isinstance(output, Exception): |
| 101 | raise output |
| 102 | return output |
| 103 | |
| 104 | def close(self): |
| 105 | if self._input_stream_task is not None: |
| 106 | self._input_stream_task.cancel() |
| 107 | self._input_stream_task = None |
| 108 | |
| 109 | def __del__(self): |
| 110 | if (task := self._input_stream_task) is not None: |
| 111 | task.get_loop().call_soon_threadsafe(task.cancel) |
| 112 | self._input_stream_task = None |
| 113 | |
| 114 | |
| 115 | @dataclass |
| 116 | class OutputProcessorOutput: |
| 117 | request_outputs: list[RequestOutput | PoolingRequestOutput] |
| 118 | reqs_to_abort: list[str] |
| 119 | |
| 120 | |
| 121 | @dataclass |
| 122 | class StreamingUpdate: |
| 123 | """Streaming input update data for output processor. |
| 124 | |
| 125 | Contains the incremental prompt data to be applied to a request state |
| 126 | when the current sub-request completes. |
| 127 | """ |
| 128 | |
| 129 | prompt: str | None |
| 130 | prompt_token_ids: list[int] | None |
| 131 | arrival_time: float |
| 132 | final: bool = False |
| 133 | |
| 134 | |
| 135 | class RequestState: |
| 136 | def __init__( |
| 137 | self, |
| 138 | request_id: str, |
| 139 | external_req_id: str, |
| 140 | parent_req: ParentRequest | None, |
| 141 | request_index: int, |
| 142 | lora_request: LoRARequest | None, |
| 143 | output_kind: RequestOutputKind, |
| 144 | prompt: str | None, |
| 145 | prompt_token_ids: list[int] | None, |
| 146 | prompt_embeds: torch.Tensor | None, |
| 147 | logprobs_processor: LogprobsProcessor | None, |
| 148 | detokenizer: IncrementalDetokenizer | None, |
| 149 | max_tokens_param: int | None, |
| 150 | arrival_time: float, |
| 151 | queue: RequestOutputCollector | None, |
| 152 | log_stats: bool, |
| 153 | stream_interval: int, |
| 154 | top_p: float | None = None, |
| 155 | n: int | None = None, |
| 156 | temperature: float | None = None, |
| 157 | stream_input: bool = False, |
| 158 | remote_prefill_cached_tokens: int | None = None, |
| 159 | ): |
| 160 | self.request_id = request_id |
| 161 | self.external_req_id = external_req_id |
| 162 | self.parent_req = parent_req |
| 163 | self.request_index = request_index |
| 164 | self.lora_request = lora_request |
| 165 | self.lora_name = lora_request.lora_name if lora_request is not None else None |
| 166 | self.output_kind = output_kind |
| 167 | self.prompt = prompt |
| 168 | self.prompt_token_ids = prompt_token_ids |
| 169 | self.prompt_embeds = prompt_embeds |
| 170 | self.prompt_len = length_from_prompt_token_ids_or_embeds( |
| 171 | self.prompt_token_ids, self.prompt_embeds |
| 172 | ) |
| 173 | self.logprobs_processor = logprobs_processor |
| 174 | self.detokenizer = detokenizer |
| 175 | self.max_tokens_param = max_tokens_param |
| 176 | self.top_p = top_p |
| 177 | self.n = n |
| 178 | self.temperature = temperature |
| 179 | self.is_prefilling = True |
| 180 | self.queue = queue |
| 181 | self.num_cached_tokens = 0 |
| 182 | self.num_cache_creation_tokens = 0 |
| 183 | self.remote_prefill_cached_tokens = remote_prefill_cached_tokens |
| 184 | # Per-sequence spec-decode accumulator; arrives once (on finish) via |
| 185 | # EngineCoreOutput, then attached to this sequence's CompletionOutput. |
| 186 | self.spec_decode_metrics: RequestSpecDecodeMetrics | None = None |
| 187 | |
| 188 | self.stats = RequestStateStats(arrival_time=arrival_time) if log_stats else None |
| 189 | |
| 190 | # Routed experts accumulation (prompt + sample chunks) |
| 191 | self.routed_experts_chunks: list[np.ndarray] = [] |
| 192 | self.sampling_mask_chunks: list[SamplingMaskLists] = [] |
| 193 | |
| 194 | # Stream Interval |
| 195 | self.stream_interval = stream_interval |
| 196 | self.sent_tokens_offset = 0 # Offset of sent tokens |
| 197 | |
| 198 | # Streaming input queue |
| 199 | self.streaming_input = stream_input |
| 200 | self.input_chunk_queue: deque[StreamingUpdate] | None = ( |
| 201 | deque() if stream_input else None |
| 202 | ) |
| 203 | |
| 204 | def apply_streaming_update(self, update: StreamingUpdate) -> None: |
| 205 | # Apply the update to the request state. |
| 206 | self.streaming_input = not update.final |
| 207 | # TODO also include relevant output tokens in new prompt here |
| 208 | # (match scheduler behavior). |
| 209 | if update.prompt: |
| 210 | self.prompt = ( |
| 211 | (self.prompt + update.prompt) if self.prompt else update.prompt |
| 212 | ) |
| 213 | if self.prompt_token_ids: |
| 214 | self.prompt_token_ids.extend(update.prompt_token_ids or ()) |
| 215 | else: |
| 216 | self.prompt_token_ids = update.prompt_token_ids or [] |
| 217 | assert self.prompt_token_ids is not None |
| 218 | self.prompt_len = len(self.prompt_token_ids) |
| 219 | if self.stats is not None: |
| 220 | self.stats.arrival_time = update.arrival_time |
| 221 | self.is_prefilling = True |
| 222 | |
| 223 | @classmethod |
| 224 | def from_new_request( |
| 225 | cls, |
| 226 | tokenizer: TokenizerLike | None, |
| 227 | request: EngineCoreRequest, |
| 228 | prompt: str | None, |
| 229 | parent_req: ParentRequest | None, |
| 230 | request_index: int, |
| 231 | queue: RequestOutputCollector | None, |
| 232 | log_stats: bool, |
| 233 | stream_interval: int, |
| 234 | ) -> "RequestState": |
| 235 | remote_prefill_cached_tokens = None |
| 236 | if sampling_params := request.sampling_params: |
| 237 | # In a remote prefill scenario, report cached tokens |
| 238 | # as cache hit rate on the remote P worker. |
| 239 | if sampling_params.extra_args: |
| 240 | kv_transfer_params = sampling_params.extra_args.get( |
| 241 | "kv_transfer_params" |
| 242 | ) |
| 243 | if kv_transfer_params and kv_transfer_params.get("do_remote_prefill"): |
| 244 | cached = kv_transfer_params.get("remote_prefill_cached_tokens") |
| 245 | if isinstance(cached, int): |
| 246 | remote_prefill_cached_tokens = cached |
| 247 | if not sampling_params.detokenize: |
| 248 | tokenizer = None |
| 249 | output_kind = sampling_params.output_kind |
| 250 | if sampling_params.stream_interval is not None: |
| 251 | # clamp to the engine-level stream interval. |
| 252 | stream_interval = max(sampling_params.stream_interval, stream_interval) |
| 253 | logprobs_processor = LogprobsProcessor.from_new_request( |
| 254 | tokenizer=tokenizer, |
| 255 | request=request, |
| 256 | ) |
| 257 | detokenizer = IncrementalDetokenizer.from_new_request( |
| 258 | tokenizer=tokenizer, |
| 259 | request=request, |
| 260 | ) |
| 261 | max_tokens_param = sampling_params.max_tokens |
| 262 | top_p = sampling_params.top_p |
| 263 | n = sampling_params.n |
| 264 | temperature = sampling_params.temperature |
| 265 | else: |
| 266 | logprobs_processor = None |
| 267 | detokenizer = None |
| 268 | max_tokens_param = None |
| 269 | top_p = None |
| 270 | n = None |
| 271 | temperature = None |
| 272 | assert request.pooling_params is not None |
| 273 | output_kind = request.pooling_params.output_kind |
| 274 | |
| 275 | assert request.external_req_id is not None |
| 276 | return cls( |
| 277 | request_id=request.request_id, |
| 278 | external_req_id=request.external_req_id, |
| 279 | parent_req=parent_req, |
| 280 | request_index=request_index, |
| 281 | lora_request=request.lora_request, |
| 282 | output_kind=output_kind, |
| 283 | prompt=prompt, |
| 284 | prompt_token_ids=request.prompt_token_ids, |
| 285 | prompt_embeds=request.prompt_embeds, |
| 286 | logprobs_processor=logprobs_processor, |
| 287 | detokenizer=detokenizer, |
| 288 | max_tokens_param=max_tokens_param, |
| 289 | top_p=top_p, |
| 290 | n=n, |
| 291 | temperature=temperature, |
| 292 | arrival_time=request.arrival_time, |
| 293 | queue=queue, |
| 294 | log_stats=log_stats, |
| 295 | stream_interval=stream_interval, |
| 296 | stream_input=request.resumable, |
| 297 | remote_prefill_cached_tokens=remote_prefill_cached_tokens, |
| 298 | ) |
| 299 | |
| 300 | def make_request_output( |
| 301 | self, |
| 302 | new_token_ids: list[int], |
| 303 | pooling_output: torch.Tensor | None, |
| 304 | finish_reason: FinishReason | None, |
| 305 | stop_reason: int | str | None, |
| 306 | kv_transfer_params: dict[str, Any] | None = None, |
| 307 | ec_transfer_params: dict[str, Any] | None = None, |
| 308 | error: RequestError | None = None, |
| 309 | ) -> RequestOutput | PoolingRequestOutput | None: |
| 310 | finished = finish_reason is not None |
| 311 | final_only = self.output_kind == RequestOutputKind.FINAL_ONLY |
| 312 | |
| 313 | if not finished and final_only: |
| 314 | # Only the final output is required in FINAL_ONLY mode. |
| 315 | return None |
| 316 | |
| 317 | if self.stream_interval > 1: |
| 318 | assert self.detokenizer is not None |
| 319 | |
| 320 | # Send output request only when |
| 321 | # 1. It has finished, or |
| 322 | # 2. It is the first token, or |
| 323 | # 3. It has reached the stream interval number of tokens |
| 324 | if not ( |
| 325 | finished |
| 326 | or self.sent_tokens_offset == 0 |
| 327 | or self.detokenizer.num_output_tokens() - self.sent_tokens_offset |
| 328 | >= self.stream_interval |
| 329 | ): |
| 330 | return None |
| 331 | |
| 332 | if self.output_kind == RequestOutputKind.DELTA: |
| 333 | # Send tokens from the offset in DELTA mode, otherwise all |
| 334 | # tokens are sent. |
| 335 | new_token_ids = self.detokenizer.output_token_ids[ |
| 336 | self.sent_tokens_offset : |
| 337 | ] |
| 338 | self.sent_tokens_offset = self.detokenizer.num_output_tokens() |
| 339 | |
| 340 | external_req_id = self.external_req_id |
| 341 | |
| 342 | if pooling_output is not None: |
| 343 | return self._new_request_output( |
| 344 | external_req_id, |
| 345 | [self._new_pooling_output(pooling_output)], |
| 346 | finished, |
| 347 | error=error, |
| 348 | ) |
| 349 | |
| 350 | output = self._new_completion_output(new_token_ids, finish_reason, stop_reason) |
| 351 | |
| 352 | if self.parent_req is None: |
| 353 | outputs = [output] |
| 354 | else: |
| 355 | outputs, finished = self.parent_req.get_outputs(self.request_id, output) |
| 356 | if not outputs: |
| 357 | return None |
| 358 | external_req_id = self.parent_req.external_req_id |
| 359 | |
| 360 | return self._new_request_output( |
| 361 | external_req_id, |
| 362 | outputs, |
| 363 | finished, |
| 364 | kv_transfer_params, |
| 365 | ec_transfer_params, |
| 366 | ) |
| 367 | |
| 368 | def _new_request_output( |
| 369 | self, |
| 370 | external_req_id: str, |
| 371 | outputs: list[CompletionOutput] | list[PoolingOutput], |
| 372 | finished: bool, |
| 373 | kv_transfer_params: dict[str, Any] | None = None, |
| 374 | ec_transfer_params: dict[str, Any] | None = None, |
| 375 | error: RequestError | None = None, |
| 376 | ) -> RequestOutput | PoolingRequestOutput: |
| 377 | # If prompt embeds were used, put placeholder prompt token ids |
| 378 | prompt_token_ids = self.prompt_token_ids |
| 379 | if prompt_token_ids is None and self.prompt_embeds is not None: |
| 380 | prompt_token_ids = [0] * len(self.prompt_embeds) |
| 381 | assert prompt_token_ids is not None |
| 382 | |
| 383 | first_output = outputs[0] |
| 384 | if isinstance(first_output, PoolingOutput): |
| 385 | assert len(outputs) == 1 |
| 386 | return PoolingRequestOutput( |
| 387 | request_id=external_req_id, |
| 388 | outputs=first_output, |
| 389 | num_cached_tokens=self.num_cached_tokens, |
| 390 | prompt_token_ids=prompt_token_ids, |
| 391 | finished=finished, |
| 392 | error=error, |
| 393 | ) |
| 394 | assert self.logprobs_processor is not None |
| 395 | if self.output_kind == RequestOutputKind.DELTA: |
| 396 | # Side effect: logprobs processor forgets prompt logprobs |
| 397 | prompt_logprobs = self.logprobs_processor.pop_prompt_logprobs() |
| 398 | prompt_token_id_logprobs = ( |
| 399 | self.logprobs_processor.pop_prompt_token_id_logprobs() |
| 400 | ) |
| 401 | else: |
| 402 | prompt_logprobs = self.logprobs_processor.prompt_logprobs |
| 403 | prompt_token_id_logprobs = self.logprobs_processor.prompt_token_id_logprobs |
| 404 | |
| 405 | return RequestOutput( |
| 406 | request_id=external_req_id, # request_id is what was provided externally |
| 407 | lora_request=self.lora_request, |
| 408 | prompt=self.prompt, |
| 409 | prompt_token_ids=prompt_token_ids, |
| 410 | prompt_logprobs=prompt_logprobs, |
| 411 | prompt_token_id_logprobs=prompt_token_id_logprobs, |
| 412 | outputs=cast(list[CompletionOutput], outputs), |
| 413 | finished=finished, |
| 414 | kv_transfer_params=kv_transfer_params, |
| 415 | ec_transfer_params=ec_transfer_params, |
| 416 | num_cached_tokens=self.num_cached_tokens, |
| 417 | num_cache_creation_tokens=self.num_cache_creation_tokens, |
| 418 | metrics=self.stats, |
| 419 | ) |
| 420 | |
| 421 | def _new_completion_output( |
| 422 | self, |
| 423 | token_ids: list[int], |
| 424 | finish_reason: FinishReason | None, |
| 425 | stop_reason: int | str | None, |
| 426 | ) -> CompletionOutput: |
| 427 | assert self.detokenizer is not None |
| 428 | assert self.logprobs_processor is not None |
| 429 | finished = finish_reason is not None |
| 430 | delta = self.output_kind == RequestOutputKind.DELTA |
| 431 | |
| 432 | # Prepare text and token_ids, based on delta mode |
| 433 | text = self.detokenizer.get_next_output_text(finished, delta) |
| 434 | if not delta: |
| 435 | token_ids = self.detokenizer.output_token_ids |
| 436 | |
| 437 | # Prepare logprobs, based on delta mode |
| 438 | logprobs = self.logprobs_processor.logprobs |
| 439 | if delta and logprobs: |
| 440 | num_new_tokens = len(token_ids) |
| 441 | # Avoid [-0:], which returns the full accumulated history when a |
| 442 | # delta contains no new token IDs. [:0] preserves the concrete |
| 443 | # list or FlatLogprobs representation while returning no entries. |
| 444 | logprobs = logprobs[-num_new_tokens:] if num_new_tokens else logprobs[:0] |
| 445 | |
| 446 | sampling_mask = None |
| 447 | if (delta or finished) and self.sampling_mask_chunks: |
| 448 | sampling_mask = SamplingMask( |
| 449 | [ |
| 450 | position |
| 451 | for chunk in self.sampling_mask_chunks |
| 452 | for position in chunk.to_nested_list() |
| 453 | ] |
| 454 | ) |
| 455 | if delta: |
| 456 | self.sampling_mask_chunks.clear() |
| 457 | |
| 458 | # Concatenate routed experts on finish |
| 459 | routed_experts = None |
| 460 | if finished and self.routed_experts_chunks: |
| 461 | routed_experts = np.concatenate(self.routed_experts_chunks, axis=0) |
| 462 | |
| 463 | return CompletionOutput( |
| 464 | index=self.request_index, |
| 465 | text=text, |
| 466 | token_ids=token_ids, |
| 467 | routed_experts=routed_experts, |
| 468 | sampling_mask=sampling_mask, |
| 469 | logprobs=logprobs, |
| 470 | cumulative_logprob=self.logprobs_processor.cumulative_logprob, |
| 471 | finish_reason=str(finish_reason) if finished else None, |
| 472 | stop_reason=stop_reason if finished else None, |
| 473 | spec_decode_metrics=self.spec_decode_metrics if finished else None, |
| 474 | ) |
| 475 | |
| 476 | def _new_pooling_output(self, pooling_output: torch.Tensor) -> PoolingOutput: |
| 477 | return PoolingOutput(data=pooling_output) |
| 478 | |
| 479 | |
| 480 | class OutputProcessor: |
| 481 | """Process EngineCoreOutputs into RequestOutputs.""" |
| 482 | |
| 483 | def __init__( |
| 484 | self, |
| 485 | tokenizer: TokenizerLike | None, |
| 486 | *, |
| 487 | log_stats: bool, |
| 488 | stream_interval: int = 1, |
| 489 | tracing_enabled: bool = False, |
| 490 | admission_stats: "SharedAdmissionStats | None" = None, |
| 491 | ): |
| 492 | self.log_stats = log_stats |
| 493 | self.tokenizer = tokenizer |
| 494 | self.stream_interval = stream_interval |
| 495 | self.request_states: dict[str, RequestState] = {} |
| 496 | self.parent_requests: dict[str, ParentRequest] = {} |
| 497 | self.external_req_ids: defaultdict[str, list[str]] = defaultdict(list) |
| 498 | self.lora_states = LoRARequestStates(log_stats) |
| 499 | self.tracing_enabled = tracing_enabled |
| 500 | self.admission_stats = admission_stats |
| 501 | |
| 502 | def get_num_unfinished_requests(self): |
| 503 | return len(self.request_states) |
| 504 | |
| 505 | def has_request(self, request_id: str) -> bool: |
| 506 | return request_id in self.request_states |
| 507 | |
| 508 | def get_num_queued_tokens(self) -> int: |
| 509 | """Total prompt tokens of requests currently in the prefill phase. |
| 510 | |
| 511 | Uses ``prompt_len`` rather than remaining prefill work because the |
| 512 | scheduler's ``num_computed_tokens`` is not propagated to the API |
| 513 | server until prefill completes. See ``SchedulerConfig`` docs. |
| 514 | """ |
| 515 | return sum( |
| 516 | req.prompt_len for req in self.request_states.values() if req.is_prefilling |
| 517 | ) |
| 518 | |
| 519 | def has_unfinished_requests(self) -> bool: |
| 520 | return len(self.request_states) > 0 |
| 521 | |
| 522 | def propagate_error(self, e: Exception): |
| 523 | """Propagate error to all generate() tasks.""" |
| 524 | for _, state in self.request_states.items(): |
| 525 | assert state.queue is not None |
| 526 | state.queue.put(e) |
| 527 | |
| 528 | def abort_requests(self, request_ids: Iterable[str], internal: bool) -> list[str]: |
| 529 | """Abort a list of requests. |
| 530 | |
| 531 | The request_ids may be either external request IDs (those passed to |
| 532 | InputProcessor.process_inputs()) or internal request IDs (those randomly |
| 533 | generated when creating the EngineCoreRequest). |
| 534 | |
| 535 | If an external request ID is provided, and that external request ID |
| 536 | was used for multiple requests, all requests associated with that external |
| 537 | request ID are aborted. |
| 538 | |
| 539 | In the case of parallel sampling, a request ID may be used to identify |
| 540 | a parent request, in which case the associated child requests are aborted |
| 541 | also. |
| 542 | """ |
| 543 | internal_req_ids = [] |
| 544 | for request_id in request_ids: |
| 545 | if internal: |
| 546 | # Internal ID - this may be a parent request |
| 547 | internal_req_ids.append(request_id) |
| 548 | |
| 549 | # Remove internal ID from the external->internal mapping |
| 550 | if req_state := self.request_states.get(request_id): |
| 551 | external_req_id = req_state.external_req_id |
| 552 | internal_ids = self.external_req_ids[external_req_id] |
| 553 | internal_ids.remove(request_id) |
| 554 | if not internal_ids: |
| 555 | del self.external_req_ids[external_req_id] |
| 556 | elif internal_ids := self.external_req_ids.pop(request_id, []): |
| 557 | # External ID - abort all requests in the external->internal mapping |
| 558 | internal_req_ids.extend(internal_ids) |
| 559 | |
| 560 | request_ids_to_abort = [] |
| 561 | for request_id in internal_req_ids: |
| 562 | req_state = self.request_states.pop(request_id, None) |
| 563 | if req_state is not None: |
| 564 | self.lora_states.request_finished(request_id, req_state.lora_name) |
| 565 | request_ids_to_abort.append(request_id) |
| 566 | # Produce final abort output. |
| 567 | if req_state.queue is not None and ( |
| 568 | request_output := req_state.make_request_output( |
| 569 | new_token_ids=[], |
| 570 | # Set pooling_output is not None to |
| 571 | # correctly enter the abort pooling branch |
| 572 | pooling_output=EMPTY_CPU_TENSOR |
| 573 | if req_state.detokenizer is None |
| 574 | else None, |
| 575 | finish_reason=FinishReason.ABORT, |
| 576 | stop_reason=None, |
| 577 | kv_transfer_params=None, |
| 578 | ec_transfer_params=None, |
| 579 | ) |
| 580 | ): |
| 581 | req_state.queue.put(request_output) |
| 582 | elif parent := self.parent_requests.get(request_id): |
| 583 | # Abort children prior to removing the parent. |
| 584 | if parent.child_requests: |
| 585 | child_reqs = list(parent.child_requests) |
| 586 | child_reqs = self.abort_requests(child_reqs, internal=True) |
| 587 | request_ids_to_abort.extend(child_reqs) |
| 588 | self.parent_requests.pop(request_id, None) |
| 589 | self._update_admission_stats() |
| 590 | return request_ids_to_abort |
| 591 | |
| 592 | def add_request( |
| 593 | self, |
| 594 | request: EngineCoreRequest, |
| 595 | prompt: str | None, |
| 596 | parent_req: ParentRequest | None = None, |
| 597 | request_index: int = 0, |
| 598 | queue: RequestOutputCollector | None = None, |
| 599 | ) -> None: |
| 600 | request_id = request.request_id |
| 601 | req_state = self.request_states.get(request_id) |
| 602 | if req_state is not None: |
| 603 | self._update_streaming_request_state(req_state, request, prompt) |
| 604 | return |
| 605 | |
| 606 | req_state = RequestState.from_new_request( |
| 607 | tokenizer=self.tokenizer, |
| 608 | request=request, |
| 609 | prompt=prompt, |
| 610 | parent_req=parent_req, |
| 611 | request_index=request_index, |
| 612 | queue=queue, |
| 613 | log_stats=self.log_stats, |
| 614 | stream_interval=self.stream_interval, |
| 615 | ) |
| 616 | self.request_states[request_id] = req_state |
| 617 | if parent_req: |
| 618 | self.parent_requests[parent_req.request_id] = parent_req |
| 619 | |
| 620 | # Track the external_req_id -> [internal_req_id, ...] mapping |
| 621 | self.external_req_ids[req_state.external_req_id].append(request_id) |
| 622 | self._update_admission_stats() |
| 623 | |
| 624 | def _update_streaming_request_state( |
| 625 | self, req_state: RequestState, request: EngineCoreRequest, prompt: str | None |
| 626 | ) -> None: |
| 627 | """Queue a streaming update instead of immediately applying it.""" |
| 628 | if not request.resumable: |
| 629 | # Final request - just mark completion, don't add its dummy tokens. |
| 630 | if req_state.input_chunk_queue is None: |
| 631 | # Engine already finished - emit final output and clean up. |
| 632 | self._finish_request(req_state) |
| 633 | if req_state.queue is not None: |
| 634 | # Emit a final output with finished=True |
| 635 | # to unblock the generate() loop. |
| 636 | req_state.queue.put(STREAM_FINISHED) |
| 637 | elif req_state.input_chunk_queue: |
| 638 | req_state.input_chunk_queue[-1].final = True |
| 639 | else: |
| 640 | req_state.streaming_input = False |
| 641 | return |
| 642 | |
| 643 | update = StreamingUpdate( |
| 644 | prompt=prompt, |
| 645 | prompt_token_ids=request.prompt_token_ids, |
| 646 | arrival_time=request.arrival_time, |
| 647 | ) |
| 648 | |
| 649 | # Apply request updates now if the last input already completed. |
| 650 | if req_state.input_chunk_queue is None: |
| 651 | req_state.apply_streaming_update(update) |
| 652 | req_state.input_chunk_queue = deque() |
| 653 | else: |
| 654 | # Queue the streaming update otherwise. |
| 655 | req_state.input_chunk_queue.append(update) |
| 656 | |
| 657 | def process_outputs( |
| 658 | self, |
| 659 | engine_core_outputs: list[EngineCoreOutput], |
| 660 | engine_core_timestamp: float | None = None, |
| 661 | iteration_stats: IterationStats | None = None, |
| 662 | ) -> OutputProcessorOutput: |
| 663 | """Process the EngineCoreOutputs: |
| 664 | 1) Compute stats for logging |
| 665 | 2) Detokenize |
| 666 | 3) Create and handle RequestOutput objects: |
| 667 | * If there is a queue (for usage with AsyncLLM), |
| 668 | put the RequestOutput objects into the queue for |
| 669 | handling by the per-request generate() tasks. |
| 670 | |
| 671 | * If there is no queue (for usage with LLMEngine), |
| 672 | return a list of RequestOutput objects. |
| 673 | |
| 674 | NOTE FOR DEVELOPERS |
| 675 | |
| 676 | vLLM V1 minimizes the number of python loops over the full |
| 677 | batch to ensure system overheads are minimized. This is the |
| 678 | only function that should loop over EngineCoreOutputs. |
| 679 | |
| 680 | If you need to touch every element of the batch, do it from |
| 681 | within the loop below. |
| 682 | """ |
| 683 | request_outputs: list[RequestOutput | PoolingRequestOutput] = [] |
| 684 | reqs_to_abort: list[str] = [] |
| 685 | for engine_core_output in engine_core_outputs: |
| 686 | req_id = engine_core_output.request_id |
| 687 | req_state = self.request_states.get(req_id) |
| 688 | if req_state is None: |
| 689 | # Ignore output for already-aborted request. |
| 690 | continue |
| 691 | |
| 692 | # 1) Compute stats for this iteration. |
| 693 | self._update_stats_from_output( |
| 694 | req_state, engine_core_output, engine_core_timestamp, iteration_stats |
| 695 | ) |
| 696 | |
| 697 | new_token_ids = engine_core_output.new_token_ids |
| 698 | pooling_output = engine_core_output.pooling_output |
| 699 | finish_reason = engine_core_output.finish_reason |
| 700 | stop_reason = engine_core_output.stop_reason |
| 701 | kv_transfer_params = engine_core_output.kv_transfer_params |
| 702 | ec_transfer_params = engine_core_output.ec_transfer_params |
| 703 | if engine_core_output.routed_experts is not None: |
| 704 | req_state.routed_experts_chunks.append( |
| 705 | engine_core_output.routed_experts |
| 706 | ) |
| 707 | |
| 708 | if req_state.is_prefilling: |
| 709 | if engine_core_output.prefill_stats is not None: |
| 710 | req_state.num_cached_tokens = ( |
| 711 | engine_core_output.prefill_stats.num_cached_tokens |
| 712 | ) |
| 713 | req_state.num_cache_creation_tokens = ( |
| 714 | engine_core_output.prefill_stats.num_cache_creation_tokens |
| 715 | ) |
| 716 | if req_state.remote_prefill_cached_tokens is not None: |
| 717 | req_state.num_cached_tokens = req_state.remote_prefill_cached_tokens |
| 718 | req_state.is_prefilling = False |
| 719 | |
| 720 | if engine_core_output.spec_decode_metrics is not None: |
| 721 | req_state.spec_decode_metrics = engine_core_output.spec_decode_metrics |
| 722 | |
| 723 | request_error = None |
| 724 | if ( |
| 725 | req_state.detokenizer is None |
| 726 | and engine_core_output.mm_cache_miss_hashes |
| 727 | and finish_reason == FinishReason.ERROR |
| 728 | ): |
| 729 | request_error = RequestError( |
| 730 | code="multimodal_cache_miss", |
| 731 | message="Multi-modal processor cache miss.", |
| 732 | retryable=True, |
| 733 | ) |
| 734 | pooling_output = EMPTY_CPU_TENSOR |
| 735 | |
| 736 | if pooling_output is None: |
| 737 | assert req_state.detokenizer is not None |
| 738 | assert req_state.logprobs_processor is not None |
| 739 | if engine_core_output.new_sampling_mask is not None: |
| 740 | req_state.sampling_mask_chunks.append( |
| 741 | engine_core_output.new_sampling_mask |
| 742 | ) |
| 743 | # 2) Detokenize the token ids into text and perform stop checks. |
| 744 | stop_string = req_state.detokenizer.update( |
| 745 | new_token_ids, finish_reason == FinishReason.STOP |
| 746 | ) |
| 747 | if stop_string: |
| 748 | finish_reason = FinishReason.STOP |
| 749 | stop_reason = stop_string |
| 750 | |
| 751 | # 3) Compute sample and prompt logprobs for request, |
| 752 | # if required. |
| 753 | req_state.logprobs_processor.update_from_output(engine_core_output) |
| 754 | |
| 755 | # 4) Create and handle RequestOutput objects. |
| 756 | if request_output := req_state.make_request_output( |
| 757 | new_token_ids, |
| 758 | pooling_output, |
| 759 | finish_reason, |
| 760 | stop_reason, |
| 761 | kv_transfer_params, |
| 762 | ec_transfer_params, |
| 763 | error=request_error, |
| 764 | ): |
| 765 | if req_state.streaming_input: |
| 766 | request_output.finished = False |
| 767 | |
| 768 | if req_state.queue is not None: |
| 769 | # AsyncLLM: put into queue for handling by generate(). |
| 770 | req_state.queue.put(request_output) |
| 771 | else: |
| 772 | # LLMEngine: return list of RequestOutputs. |
| 773 | request_outputs.append(request_output) |
| 774 | |
| 775 | # Free completed requests. |
| 776 | if finish_reason is not None: |
| 777 | if req_state.streaming_input: |
| 778 | if req_state.input_chunk_queue: |
| 779 | update = req_state.input_chunk_queue.popleft() |
| 780 | req_state.apply_streaming_update(update) |
| 781 | else: |
| 782 | req_state.input_chunk_queue = None |
| 783 | else: |
| 784 | self._finish_request(req_state) |
| 785 | if not engine_core_output.finished: |
| 786 | # If req not finished in EngineCore, but Detokenizer |
| 787 | # detected stop string, abort needed in EngineCore. |
| 788 | reqs_to_abort.append(req_id) |
| 789 | |
| 790 | # Track per-request stats |
| 791 | self._update_stats_from_finished( |
| 792 | req_state, finish_reason, iteration_stats |
| 793 | ) |
| 794 | if self.tracing_enabled: |
| 795 | self.do_tracing(engine_core_output, req_state, iteration_stats) |
| 796 | |
| 797 | return OutputProcessorOutput( |
| 798 | request_outputs=request_outputs, |
| 799 | reqs_to_abort=reqs_to_abort, |
| 800 | ) |
| 801 | |
| 802 | def _finish_request(self, req_state: RequestState) -> None: |
| 803 | req_id = req_state.request_id |
| 804 | self.request_states.pop(req_id) |
| 805 | |
| 806 | internal_ids = self.external_req_ids[req_state.external_req_id] |
| 807 | internal_ids.remove(req_id) |
| 808 | if not internal_ids: |
| 809 | del self.external_req_ids[req_state.external_req_id] |
| 810 | |
| 811 | # Remove parent request if applicable. |
| 812 | parent_req = req_state.parent_req |
| 813 | if parent_req and not parent_req.child_requests: |
| 814 | self.parent_requests.pop(parent_req.request_id, None) |
| 815 | self._update_admission_stats() |
| 816 | |
| 817 | def _update_admission_stats(self) -> None: |
| 818 | if self.admission_stats is not None: |
| 819 | self.admission_stats.set_num_requests(self.get_num_unfinished_requests()) |
| 820 | |
| 821 | def update_scheduler_stats(self, scheduler_stats: SchedulerStats | None): |
| 822 | self.lora_states.update_scheduler_stats(scheduler_stats) |
| 823 | |
| 824 | def do_tracing( |
| 825 | self, |
| 826 | engine_core_output: EngineCoreOutput, |
| 827 | req_state: RequestState, |
| 828 | iteration_stats: IterationStats | None, |
| 829 | ) -> None: |
| 830 | assert req_state.stats is not None |
| 831 | assert iteration_stats is not None |
| 832 | |
| 833 | metrics = req_state.stats |
| 834 | arrival_time_ns = int(metrics.arrival_time * 1e9) |
| 835 | trace_context = extract_trace_context(engine_core_output.trace_headers) |
| 836 | prompt_length = length_from_prompt_token_ids_or_embeds( |
| 837 | req_state.prompt_token_ids, req_state.prompt_embeds |
| 838 | ) |
| 839 | |
| 840 | # Calculate timing metrics |
| 841 | e2e_time = iteration_stats.iteration_timestamp - metrics.arrival_time |
| 842 | queued_time = metrics.scheduled_ts - metrics.queued_ts |
| 843 | prefill_time = metrics.first_token_ts - metrics.scheduled_ts |
| 844 | decode_time = metrics.last_token_ts - metrics.first_token_ts |
| 845 | inference_time = metrics.last_token_ts - metrics.scheduled_ts |
| 846 | |
| 847 | # Build attributes dict |
| 848 | attributes: dict[str, Any] = { |
| 849 | SpanAttributes.GEN_AI_LATENCY_TIME_TO_FIRST_TOKEN: ( |
| 850 | metrics.first_token_latency |
| 851 | ), |
| 852 | SpanAttributes.GEN_AI_LATENCY_E2E: e2e_time, |
| 853 | SpanAttributes.GEN_AI_LATENCY_TIME_IN_QUEUE: queued_time, |
| 854 | SpanAttributes.GEN_AI_USAGE_PROMPT_TOKENS: prompt_length, |
| 855 | SpanAttributes.GEN_AI_USAGE_COMPLETION_TOKENS: ( |
| 856 | metrics.num_generation_tokens |
| 857 | ), |
| 858 | SpanAttributes.GEN_AI_LATENCY_TIME_IN_MODEL_PREFILL: prefill_time, |
| 859 | SpanAttributes.GEN_AI_LATENCY_TIME_IN_MODEL_DECODE: decode_time, |
| 860 | SpanAttributes.GEN_AI_LATENCY_TIME_IN_MODEL_INFERENCE: inference_time, |
| 861 | SpanAttributes.GEN_AI_REQUEST_ID: req_state.external_req_id, |
| 862 | } |
| 863 | |
| 864 | # Add optional request parameters |
| 865 | if req_state.top_p: |
| 866 | attributes[SpanAttributes.GEN_AI_REQUEST_TOP_P] = req_state.top_p |
| 867 | if req_state.max_tokens_param: |
| 868 | attributes[SpanAttributes.GEN_AI_REQUEST_MAX_TOKENS] = ( |
| 869 | req_state.max_tokens_param |
| 870 | ) |
| 871 | if req_state.temperature: |
| 872 | attributes[SpanAttributes.GEN_AI_REQUEST_TEMPERATURE] = ( |
| 873 | req_state.temperature |
| 874 | ) |
| 875 | if req_state.n: |
| 876 | attributes[SpanAttributes.GEN_AI_REQUEST_N] = req_state.n |
| 877 | |
| 878 | instrument_manual( |
| 879 | span_name="llm_request", |
| 880 | start_time=arrival_time_ns, |
| 881 | attributes=attributes, |
| 882 | context=trace_context, |
| 883 | kind=SpanKind.SERVER, |
| 884 | ) |
| 885 | |
| 886 | def _update_stats_from_output( |
| 887 | self, |
| 888 | req_state: RequestState, |
| 889 | engine_core_output: EngineCoreOutput, |
| 890 | engine_core_timestamp: float | None, |
| 891 | iteration_stats: IterationStats | None, |
| 892 | ): |
| 893 | if iteration_stats is None: |
| 894 | return |
| 895 | |
| 896 | assert engine_core_timestamp is not None |
| 897 | assert req_state.stats is not None |
| 898 | iteration_stats.update_from_output( |
| 899 | engine_core_output, |
| 900 | engine_core_timestamp, |
| 901 | req_state.is_prefilling, |
| 902 | req_state.stats, |
| 903 | self.lora_states, |
| 904 | req_state.lora_name, |
| 905 | ) |
| 906 | |
| 907 | def _update_stats_from_finished( |
| 908 | self, |
| 909 | req_state: RequestState, |
| 910 | finish_reason: FinishReason | None, |
| 911 | iteration_stats: IterationStats | None, |
| 912 | ): |
| 913 | if iteration_stats is None: |
| 914 | return |
| 915 | |
| 916 | assert finish_reason is not None |
| 917 | assert req_state.stats is not None |
| 918 | iteration_stats.update_from_finished_request( |
| 919 | finish_reason=finish_reason, |
| 920 | request_id=req_state.external_req_id, |
| 921 | num_prompt_tokens=req_state.prompt_len, |
| 922 | max_tokens_param=req_state.max_tokens_param, |
| 923 | req_stats=req_state.stats, |
| 924 | num_cached_tokens=req_state.num_cached_tokens, |
| 925 | ) |
| 926 | self.lora_states.request_finished(req_state.request_id, req_state.lora_name) |
| 927 | |
| 928 | ParentRequest.observe_finished_request( |
| 929 | req_state.parent_req, iteration_stats, req_state.stats.num_generation_tokens |
| 930 | ) |