| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from abc import ABC, abstractmethod |
| 5 | from collections.abc import Sequence |
| 6 | from copy import copy |
| 7 | from dataclasses import dataclass, field |
| 8 | from typing import TYPE_CHECKING, NamedTuple, TypeAlias |
| 9 | |
| 10 | import numpy as np |
| 11 | import torch |
| 12 | |
| 13 | from vllm.compilation.cuda_graph import CUDAGraphStat |
| 14 | from vllm.utils.torch_utils import PIN_MEMORY |
| 15 | from vllm.v1.core.sched.output import SchedulerOutput |
| 16 | |
| 17 | if TYPE_CHECKING: |
| 18 | from vllm.distributed.aux_output_connector.connector import AuxRequestOutput |
| 19 | from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorWorkerMetadata |
| 20 | from vllm.distributed.ec_transfer.ec_connector.metrics import ECConnectorStats |
| 21 | from vllm.distributed.kv_events import KVConnectorKVEvents |
| 22 | from vllm.distributed.kv_transfer.kv_connector.v1.base import ( |
| 23 | KVConnectorWorkerMetadata, |
| 24 | ) |
| 25 | from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats |
| 26 | else: |
| 27 | AuxRequestOutput = object |
| 28 | KVConnectorStats = object |
| 29 | KVConnectorWorkerMetadata = object |
| 30 | KVConnectorKVEvents = object |
| 31 | ECConnectorWorkerMetadata = object |
| 32 | |
| 33 | |
| 34 | class LogprobsLists(NamedTuple): |
| 35 | # [num_reqs x num_generated_tokens, max_num_logprobs + 1] |
| 36 | logprob_token_ids: np.ndarray |
| 37 | # [num_reqs x num_generated_tokens, max_num_logprobs + 1] |
| 38 | logprobs: np.ndarray |
| 39 | # [num_reqs x num_generated_tokens] |
| 40 | sampled_token_ranks: np.ndarray |
| 41 | # [num_reqs] |
| 42 | # Used for slicing the logprobs in cases like speculative |
| 43 | # decoding where the number of generated tokens may be |
| 44 | # different for each request. |
| 45 | cu_num_generated_tokens: list[int] | None = None |
| 46 | |
| 47 | def slice_request(self, req_idx: int, num_positions: int): |
| 48 | if self.cu_num_generated_tokens is not None: |
| 49 | req_idx = self.cu_num_generated_tokens[req_idx] |
| 50 | end_idx = req_idx + num_positions |
| 51 | return LogprobsLists( |
| 52 | self.logprob_token_ids[req_idx:end_idx], |
| 53 | self.logprobs[req_idx:end_idx], |
| 54 | self.sampled_token_ranks[req_idx:end_idx], |
| 55 | None, |
| 56 | ) |
| 57 | |
| 58 | |
| 59 | class SamplingMaskLists(NamedTuple): |
| 60 | """CSR sampling masks; a request slice may contain multiple positions.""" |
| 61 | |
| 62 | # [num_kept_tokens] |
| 63 | token_ids: np.ndarray |
| 64 | # [num_positions + 1], or None for a single position |
| 65 | offsets: np.ndarray | None = None |
| 66 | # [num_requests + 1] for multi-position request batches. |
| 67 | cu_num_generated_tokens: list[int] | None = None |
| 68 | |
| 69 | def slice_request(self, req_idx: int, num_positions: int) -> "SamplingMaskLists": |
| 70 | assert self.offsets is not None |
| 71 | cu = self.cu_num_generated_tokens |
| 72 | start = req_idx if cu is None else cu[req_idx] |
| 73 | end = start + num_positions |
| 74 | lo, hi = self.offsets[start], self.offsets[end] |
| 75 | if num_positions == 1: |
| 76 | return SamplingMaskLists(self.token_ids[lo:hi]) |
| 77 | return SamplingMaskLists( |
| 78 | self.token_ids[lo:hi], self.offsets[start : end + 1] - lo |
| 79 | ) |
| 80 | |
| 81 | def to_nested_list(self) -> list[list[int]]: |
| 82 | token_ids = self.token_ids.tolist() |
| 83 | if self.offsets is None: |
| 84 | return [token_ids] |
| 85 | offsets = self.offsets.tolist() |
| 86 | return [token_ids[offsets[i] : offsets[i + 1]] for i in range(len(offsets) - 1)] |
| 87 | |
| 88 | |
| 89 | class LogprobsTensors(NamedTuple): |
| 90 | # [num_reqs x num_generated_tokens, max_num_logprobs + 1] |
| 91 | logprob_token_ids: torch.Tensor |
| 92 | # [num_reqs x num_generated_tokens, max_num_logprobs + 1] |
| 93 | logprobs: torch.Tensor |
| 94 | # [num_reqs x num_generated_tokens] |
| 95 | selected_token_ranks: torch.Tensor |
| 96 | # [num_reqs + 1] |
| 97 | cu_num_generated_tokens: list[int] | None = None |
| 98 | # [num_reqs + 1]. Set instead of cu_num_generated_tokens when the |
| 99 | # boundaries only exist on device (adaptive verification); rides along on |
| 100 | # the async D2H copy. |
| 101 | cu_num_generated_tokens_tensor: torch.Tensor | None = None |
| 102 | |
| 103 | def tolists(self, cu_num_generated_tokens: list[int] | None = None): |
| 104 | if cu_num_generated_tokens is None: |
| 105 | if self.cu_num_generated_tokens_tensor is not None: |
| 106 | cu_num_generated_tokens = self.cu_num_generated_tokens_tensor.tolist() |
| 107 | else: |
| 108 | cu_num_generated_tokens = self.cu_num_generated_tokens |
| 109 | return LogprobsLists( |
| 110 | self.logprob_token_ids.cpu().numpy(), |
| 111 | self.logprobs.cpu().numpy(), |
| 112 | self.selected_token_ranks.cpu().numpy(), |
| 113 | cu_num_generated_tokens, |
| 114 | ) |
| 115 | |
| 116 | def to_cpu_nonblocking(self) -> "LogprobsTensors": |
| 117 | if self.logprob_token_ids.device.type == "cpu": |
| 118 | return self |
| 119 | cu_tensor = self.cu_num_generated_tokens_tensor |
| 120 | if cu_tensor is not None: |
| 121 | cu_tensor = cu_tensor.to("cpu", non_blocking=True) |
| 122 | return LogprobsTensors( |
| 123 | self.logprob_token_ids.to("cpu", non_blocking=True), |
| 124 | self.logprobs.to("cpu", non_blocking=True), |
| 125 | self.selected_token_ranks.to("cpu", non_blocking=True), |
| 126 | self.cu_num_generated_tokens, |
| 127 | cu_tensor, |
| 128 | ) |
| 129 | |
| 130 | def filter(self, mask: torch.Tensor) -> "LogprobsTensors": |
| 131 | """Filter the logprobs tensors with the given bool mask.""" |
| 132 | assert self.cu_num_generated_tokens is None, ( |
| 133 | "filter can't be used with cu_num_generated_tokens" |
| 134 | ) |
| 135 | assert self.cu_num_generated_tokens_tensor is None, ( |
| 136 | "filter can't be used with cu_num_generated_tokens_tensor" |
| 137 | ) |
| 138 | return LogprobsTensors( |
| 139 | self.logprob_token_ids[mask], |
| 140 | self.logprobs[mask], |
| 141 | self.selected_token_ranks[mask], |
| 142 | ) |
| 143 | |
| 144 | @staticmethod |
| 145 | def cat( |
| 146 | tensors: Sequence["LogprobsTensors"], |
| 147 | cu_num_generated_tokens: list[int] | None = None, |
| 148 | ) -> "LogprobsTensors": |
| 149 | """Concatenate flattened logprob tensors.""" |
| 150 | assert tensors |
| 151 | assert cu_num_generated_tokens is not None or all( |
| 152 | tensor.cu_num_generated_tokens is None for tensor in tensors |
| 153 | ) |
| 154 | if len(tensors) == 1: |
| 155 | tensor = tensors[0] |
| 156 | if cu_num_generated_tokens is None: |
| 157 | return tensor |
| 158 | return tensor._replace(cu_num_generated_tokens=cu_num_generated_tokens) |
| 159 | # The multi-chunk path rebuilds boundaries from the CPU layout, which |
| 160 | # the device-only boundaries of adaptive verification never use. |
| 161 | assert all(tensor.cu_num_generated_tokens_tensor is None for tensor in tensors) |
| 162 | return LogprobsTensors( |
| 163 | logprob_token_ids=torch.cat( |
| 164 | [tensor.logprob_token_ids for tensor in tensors] |
| 165 | ), |
| 166 | logprobs=torch.cat([tensor.logprobs for tensor in tensors]), |
| 167 | selected_token_ranks=torch.cat( |
| 168 | [tensor.selected_token_ranks for tensor in tensors] |
| 169 | ), |
| 170 | cu_num_generated_tokens=cu_num_generated_tokens, |
| 171 | ) |
| 172 | |
| 173 | @staticmethod |
| 174 | def empty_cpu( |
| 175 | num_positions: int, num_tokens_per_position: int |
| 176 | ) -> "LogprobsTensors": |
| 177 | """Create empty LogprobsTensors on CPU.""" |
| 178 | logprob_token_ids = torch.empty( |
| 179 | (num_positions, num_tokens_per_position), |
| 180 | dtype=torch.int32, |
| 181 | device="cpu", |
| 182 | pin_memory=PIN_MEMORY, |
| 183 | ) |
| 184 | logprobs = logprob_token_ids.new_empty( |
| 185 | (num_positions, num_tokens_per_position), |
| 186 | dtype=torch.float32, |
| 187 | pin_memory=PIN_MEMORY, |
| 188 | ) |
| 189 | selected_token_ranks = logprob_token_ids.new_empty( |
| 190 | num_positions, pin_memory=PIN_MEMORY |
| 191 | ) |
| 192 | return LogprobsTensors( |
| 193 | logprob_token_ids=logprob_token_ids, |
| 194 | logprobs=logprobs, |
| 195 | selected_token_ranks=selected_token_ranks, |
| 196 | ) |
| 197 | |
| 198 | |
| 199 | # [num_reqs, <dynamic>] |
| 200 | # The shape of each element depends on the pooler used |
| 201 | PoolerOutput: TypeAlias = torch.Tensor | list[torch.Tensor] | list[torch.Tensor | None] |
| 202 | |
| 203 | |
| 204 | @dataclass |
| 205 | class SamplerOutput: |
| 206 | # [num_reqs, max_num_generated_tokens] |
| 207 | # Different requests can have different number of generated tokens. |
| 208 | # All requests are padded to max_num_generated_tokens. |
| 209 | # PLACEHOLDER_TOKEN_ID (-1 by default) is used for padding. |
| 210 | sampled_token_ids: torch.Tensor |
| 211 | logprobs_tensors: LogprobsTensors | None |
| 212 | |
| 213 | |
| 214 | @dataclass |
| 215 | class KVConnectorOutput: |
| 216 | # [req_ids] |
| 217 | finished_sending: set[str] | None = None |
| 218 | finished_recving: set[str] | None = None |
| 219 | kv_connector_stats: KVConnectorStats | None = None |
| 220 | kv_cache_events: KVConnectorKVEvents | None = None |
| 221 | kv_connector_worker_meta: KVConnectorWorkerMetadata | None = None |
| 222 | # IDs of externally computed KV blocks that failed to load. |
| 223 | # Requests referencing these blocks should be rescheduled to recompute them |
| 224 | invalid_block_ids: set[int] = field(default_factory=set) |
| 225 | # Receive failures keyed by request identity. This remains unambiguous for |
| 226 | # hybrid/multi-pool cache layouts where numeric block IDs overlap. |
| 227 | failed_recving: set[str] = field(default_factory=set) |
| 228 | # Configuration describing how many finished sending/receiving |
| 229 | # notifications should be expected for each request. This allows |
| 230 | # handshake-based connectors like Nixl to update the KVOutputAggregator. |
| 231 | # It captures a static setup info and should almost always remain constant |
| 232 | # for a given connector after discovery. Default value entails no change. |
| 233 | expected_finished_count: int = 0 |
| 234 | |
| 235 | def is_empty(self): |
| 236 | return ( |
| 237 | not self.finished_sending |
| 238 | and not self.finished_recving |
| 239 | and not self.kv_connector_stats |
| 240 | and not self.kv_cache_events |
| 241 | and not self.invalid_block_ids |
| 242 | and not self.failed_recving |
| 243 | and not self.kv_connector_worker_meta |
| 244 | ) |
| 245 | |
| 246 | |
| 247 | @dataclass |
| 248 | class ECConnectorOutput: |
| 249 | # [mm_hash] |
| 250 | finished_sending: set[str] | None = None |
| 251 | finished_recving: set[str] | None = None |
| 252 | ec_connector_stats: "ECConnectorStats | None" = None |
| 253 | ec_connector_worker_meta: ECConnectorWorkerMetadata | None = None |
| 254 | |
| 255 | def is_empty(self): |
| 256 | return ( |
| 257 | not self.finished_sending |
| 258 | and not self.finished_recving |
| 259 | and not self.ec_connector_stats |
| 260 | and not self.ec_connector_worker_meta |
| 261 | ) |
| 262 | |
| 263 | |
| 264 | # ModelRunnerOutput is serialized and sent to the scheduler process. |
| 265 | # This is expensive for torch.Tensor so prefer to use list instead. |
| 266 | @dataclass |
| 267 | class ModelRunnerOutput: |
| 268 | # [num_reqs] |
| 269 | req_ids: list[str] |
| 270 | # req_id -> index |
| 271 | req_id_to_index: dict[str, int] |
| 272 | |
| 273 | # num_reqs x num_generated_tokens |
| 274 | # num_generated_tokens is the number of tokens |
| 275 | # generated in the current step. It can be different for |
| 276 | # each request due to speculative/jump decoding. |
| 277 | sampled_token_ids: list[list[int]] = field(default_factory=list) |
| 278 | |
| 279 | # [num_reqs, max_num_logprobs + 1] |
| 280 | # [num_reqs, max_num_logprobs + 1] |
| 281 | # [num_reqs] |
| 282 | logprobs: LogprobsLists | None = None |
| 283 | |
| 284 | # req_id -> (token_ids, logprobs, ranks) |
| 285 | # [prompt_len, num_prompt_logprobs] |
| 286 | # [prompt_len, num_prompt_logprobs] |
| 287 | # [prompt_len] |
| 288 | prompt_logprobs_dict: dict[str, LogprobsTensors | None] = field( |
| 289 | default_factory=dict |
| 290 | ) |
| 291 | |
| 292 | # req_id -> [num_scored_rows, num_token_ids] prompt_logprob_token_ids scores. |
| 293 | prompt_token_id_logprobs_dict: dict[str, torch.Tensor] = field(default_factory=dict) |
| 294 | |
| 295 | # [num_reqs, hidden_size] |
| 296 | pooler_output: list[torch.Tensor | None] | None = None |
| 297 | |
| 298 | kv_connector_output: KVConnectorOutput | None = None |
| 299 | |
| 300 | ec_connector_output: ECConnectorOutput | None = None |
| 301 | |
| 302 | # req_id -> num_nans_in_logits |
| 303 | num_nans_in_logits: dict[str, int] | None = None |
| 304 | |
| 305 | # information related to cudagraph execution |
| 306 | cudagraph_stats: CUDAGraphStat | None = None |
| 307 | |
| 308 | aux_output_connector_output: dict[str, AuxRequestOutput] | None = None |
| 309 | |
| 310 | # ``None`` when ``return_sampling_mask`` is off. |
| 311 | sampling_masks: SamplingMaskLists | None = None |
| 312 | |
| 313 | @staticmethod |
| 314 | def with_kv_conn_output_only( |
| 315 | kv_connector_output: KVConnectorOutput | None, |
| 316 | ) -> "ModelRunnerOutput": |
| 317 | """Return ModelRunnerOutput containing the provided KVConnectorOutput, |
| 318 | otherwise empty. Returns None if kv_connector_output is passed as None. |
| 319 | """ |
| 320 | if kv_connector_output is None or kv_connector_output.is_empty(): |
| 321 | return EMPTY_MODEL_RUNNER_OUTPUT |
| 322 | output = copy(EMPTY_MODEL_RUNNER_OUTPUT) |
| 323 | output.kv_connector_output = kv_connector_output |
| 324 | return output |
| 325 | |
| 326 | @staticmethod |
| 327 | def with_ec_conn_output_only( |
| 328 | ec_connector_output: ECConnectorOutput | None, |
| 329 | ) -> "ModelRunnerOutput": |
| 330 | """Return an otherwise-empty output carrying `ec_connector_output`.""" |
| 331 | return ModelRunnerOutput.with_ec_conn_output( |
| 332 | EMPTY_MODEL_RUNNER_OUTPUT, ec_connector_output |
| 333 | ) |
| 334 | |
| 335 | @staticmethod |
| 336 | def with_ec_conn_output( |
| 337 | output: "ModelRunnerOutput", |
| 338 | ec_connector_output: ECConnectorOutput | None, |
| 339 | ) -> "ModelRunnerOutput": |
| 340 | """Return `output` carrying `ec_connector_output`. |
| 341 | |
| 342 | The shared empty output is copied rather than written to, so callers |
| 343 | must use the return value. |
| 344 | """ |
| 345 | if ec_connector_output is None or ec_connector_output.is_empty(): |
| 346 | return output |
| 347 | if output is EMPTY_MODEL_RUNNER_OUTPUT: |
| 348 | output = copy(EMPTY_MODEL_RUNNER_OUTPUT) |
| 349 | output.ec_connector_output = ec_connector_output |
| 350 | return output |
| 351 | |
| 352 | |
| 353 | # ModelRunnerOutput wrapper for async scheduling. |
| 354 | class AsyncModelRunnerOutput(ABC): |
| 355 | @abstractmethod |
| 356 | def get_output(self) -> ModelRunnerOutput: |
| 357 | """Get the ModelRunnerOutput for this async output. |
| 358 | |
| 359 | This is a blocking call that waits until the results are ready, which |
| 360 | might involve copying device tensors to the host. |
| 361 | This method should only be called once per AsyncModelRunnerOutput. |
| 362 | """ |
| 363 | pass |
| 364 | |
| 365 | |
| 366 | @dataclass |
| 367 | class DraftTokenIds: |
| 368 | # [num_reqs] |
| 369 | req_ids: list[str] |
| 370 | # num_reqs x num_draft_tokens |
| 371 | draft_token_ids: list[list[int]] |
| 372 | |
| 373 | |
| 374 | def make_empty_encoder_model_runner_output( |
| 375 | scheduler_output: "SchedulerOutput", |
| 376 | ) -> ModelRunnerOutput: |
| 377 | """Create a ModelRunnerOutput stub that contains the correct |
| 378 | per-request bookkeeping but no generated data yet. |
| 379 | """ |
| 380 | if not scheduler_output.num_scheduled_tokens: |
| 381 | return EMPTY_MODEL_RUNNER_OUTPUT |
| 382 | |
| 383 | # Convert to list so we get a deterministic, indexable sequence |
| 384 | req_ids: list[str] = list(scheduler_output.num_scheduled_tokens.keys()) |
| 385 | |
| 386 | # Give every request its own contiguous index |
| 387 | req_id_to_index: dict[str, int] = {rid: idx for idx, rid in enumerate(req_ids)} |
| 388 | |
| 389 | # An encoder instance never samples, so it emits no tokens at all. The |
| 390 | # scheduler finishes these requests once their prompt is fully encoded |
| 391 | # (see `Scheduler.update_from_output`). |
| 392 | sampled_token_ids: list[list[int]] = [[] for _ in req_ids] |
| 393 | |
| 394 | # Pooler outputs are not available yet ⇒ use None placeholders |
| 395 | pooler_output: list[torch.Tensor | None] = [None for _ in req_ids] |
| 396 | |
| 397 | return ModelRunnerOutput( |
| 398 | req_ids=req_ids, |
| 399 | req_id_to_index=req_id_to_index, |
| 400 | sampled_token_ids=sampled_token_ids, |
| 401 | pooler_output=pooler_output, |
| 402 | ) |
| 403 | |
| 404 | |
| 405 | EMPTY_MODEL_RUNNER_OUTPUT = ModelRunnerOutput(req_ids=[], req_id_to_index={}) |