| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from collections.abc import Callable |
| 5 | from dataclasses import dataclass |
| 6 | from functools import cached_property |
| 7 | from typing import TYPE_CHECKING |
| 8 | |
| 9 | from vllm.config.ec_manager_config import EncoderCacheManagerMetadata |
| 10 | from vllm.multimodal.utils import strip_covered_mm_data |
| 11 | |
| 12 | if TYPE_CHECKING: |
| 13 | import numpy as np |
| 14 | import numpy.typing as npt |
| 15 | import torch |
| 16 | |
| 17 | from vllm.distributed.aux_output_connector.connector import ( |
| 18 | AuxOutputConnectorMetadata, |
| 19 | ) |
| 20 | from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorMetadata |
| 21 | from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata |
| 22 | from vllm.lora.request import LoRARequest |
| 23 | from vllm.multimodal.inputs import MultiModalFeatureSpec |
| 24 | from vllm.pooling_params import PoolingParams |
| 25 | from vllm.sampling_params import SamplingParams |
| 26 | from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy |
| 27 | from vllm.v1.request import Request |
| 28 | else: |
| 29 | AuxOutputConnectorMetadata = object |
| 30 | ECConnectorMetadata = object |
| 31 | KVConnectorMetadata = object |
| 32 | KVCacheBlockCopy = object |
| 33 | LoRARequest = object |
| 34 | MultiModalFeatureSpec = object |
| 35 | PoolingParams = object |
| 36 | SamplingParams = object |
| 37 | Request = object |
| 38 | |
| 39 | |
| 40 | @dataclass |
| 41 | class NewRequestData: |
| 42 | req_id: str |
| 43 | prompt_token_ids: list[int] | None |
| 44 | mm_features: list[MultiModalFeatureSpec] |
| 45 | sampling_params: SamplingParams | None |
| 46 | pooling_params: PoolingParams | None |
| 47 | block_ids: tuple[list[int], ...] |
| 48 | num_computed_tokens: int |
| 49 | lora_request: LoRARequest | None |
| 50 | prompt_embeds: "torch.Tensor | None" = None |
| 51 | prompt_is_token_ids: list[bool] | None = None |
| 52 | |
| 53 | # Only used for v2 model runner. |
| 54 | prefill_token_ids: list[int] | None = None |
| 55 | # DeepSeek-V4.1 only: SWA bounded replay; see Request.replay_start. |
| 56 | replay_start: int = 0 |
| 57 | |
| 58 | @classmethod |
| 59 | def from_request( |
| 60 | cls, |
| 61 | request: Request, |
| 62 | block_ids: tuple[list[int], ...], |
| 63 | prefill_token_ids: list[int] | None = None, |
| 64 | uses_mrope: bool = False, |
| 65 | ) -> "NewRequestData": |
| 66 | return cls( |
| 67 | req_id=request.request_id, |
| 68 | prompt_token_ids=request.prompt_token_ids, |
| 69 | mm_features=strip_covered_mm_data( |
| 70 | request.mm_features, |
| 71 | request.num_computed_tokens, |
| 72 | uses_mrope=uses_mrope, |
| 73 | ), |
| 74 | sampling_params=request.sampling_params, |
| 75 | pooling_params=request.pooling_params, |
| 76 | block_ids=block_ids, |
| 77 | num_computed_tokens=request.num_computed_tokens, |
| 78 | lora_request=request.lora_request, |
| 79 | prompt_embeds=request.prompt_embeds, |
| 80 | prompt_is_token_ids=request.prompt_is_token_ids, |
| 81 | prefill_token_ids=prefill_token_ids, |
| 82 | replay_start=request.replay_start, |
| 83 | ) |
| 84 | |
| 85 | @property |
| 86 | def prompt_len(self) -> int: |
| 87 | if self.prompt_token_ids is not None: |
| 88 | return len(self.prompt_token_ids) |
| 89 | if self.prompt_embeds is not None: |
| 90 | return self.prompt_embeds.shape[0] |
| 91 | return 0 |
| 92 | |
| 93 | def __repr__(self) -> str: |
| 94 | prompt_embeds_shape = ( |
| 95 | self.prompt_embeds.shape if self.prompt_embeds is not None else None |
| 96 | ) |
| 97 | return ( |
| 98 | f"NewRequestData(" |
| 99 | f"req_id={self.req_id}," |
| 100 | f"prompt_token_ids={self.prompt_token_ids}," |
| 101 | f"prefill_token_ids={self.prefill_token_ids}," |
| 102 | f"mm_features={self.mm_features}," |
| 103 | f"sampling_params={self.sampling_params}," |
| 104 | f"block_ids={self.block_ids}," |
| 105 | f"num_computed_tokens={self.num_computed_tokens}," |
| 106 | f"lora_request={self.lora_request}," |
| 107 | f"prompt_embeds_shape={prompt_embeds_shape}" |
| 108 | ")" |
| 109 | ) |
| 110 | |
| 111 | # Version of __repr__ with the prompt data obfuscated |
| 112 | def anon_repr(self) -> str: |
| 113 | prompt_token_ids_len = ( |
| 114 | len(self.prompt_token_ids) if self.prompt_token_ids is not None else None |
| 115 | ) |
| 116 | prompt_embeds_shape = ( |
| 117 | self.prompt_embeds.shape if self.prompt_embeds is not None else None |
| 118 | ) |
| 119 | prefill_token_ids_len = ( |
| 120 | len(self.prefill_token_ids) if self.prefill_token_ids is not None else None |
| 121 | ) |
| 122 | return ( |
| 123 | f"NewRequestData(" |
| 124 | f"req_id={self.req_id}," |
| 125 | f"prompt_token_ids_len={prompt_token_ids_len}," |
| 126 | f"prefill_token_ids_len={prefill_token_ids_len}," |
| 127 | f"mm_features={self.mm_features}," |
| 128 | f"sampling_params={self.sampling_params}," |
| 129 | f"block_ids={self.block_ids}," |
| 130 | f"num_computed_tokens={self.num_computed_tokens}," |
| 131 | f"lora_request={self.lora_request}," |
| 132 | f"prompt_embeds_shape={prompt_embeds_shape}" |
| 133 | ")" |
| 134 | ) |
| 135 | |
| 136 | |
| 137 | @dataclass |
| 138 | class CachedRequestData: |
| 139 | req_ids: list[str] |
| 140 | # For request ids not in resumed_req_ids, new_block_ids will be appended to |
| 141 | # the request's block IDs. For those in the set, new_block_ids will be used as the |
| 142 | # request's block IDs instead of appending to the existing block IDs. |
| 143 | resumed_req_ids: set[str] |
| 144 | # NOTE(woosuk): new_token_ids is only used for pipeline parallelism. |
| 145 | # When PP is not used, new_token_ids will be empty. |
| 146 | new_token_ids: list[list[int]] |
| 147 | # MRV1-only: For requests not scheduled in the last step, propagate the token ids |
| 148 | # to the connector. Won't contain requests scheduled in the prior step. |
| 149 | all_token_ids: dict[str, list[int]] |
| 150 | new_block_ids: list[tuple[list[int], ...] | None] |
| 151 | num_computed_tokens: list[int] |
| 152 | num_output_tokens: list[int] |
| 153 | |
| 154 | # Version of dataclass repr with token IDs obfuscated. |
| 155 | def anon_repr(self) -> str: |
| 156 | new_token_ids_lens = [len(toks) for toks in self.new_token_ids] |
| 157 | all_token_ids_lens = { |
| 158 | req_id: len(toks) for req_id, toks in self.all_token_ids.items() |
| 159 | } |
| 160 | return ( |
| 161 | f"CachedRequestData(" |
| 162 | f"req_ids={self.req_ids}," |
| 163 | f"resumed_req_ids={self.resumed_req_ids}," |
| 164 | f"new_token_ids_lens={new_token_ids_lens}," |
| 165 | f"all_token_ids_lens={all_token_ids_lens}," |
| 166 | f"new_block_ids={self.new_block_ids}," |
| 167 | f"num_computed_tokens={self.num_computed_tokens}," |
| 168 | f"num_output_tokens={self.num_output_tokens}" |
| 169 | f")" |
| 170 | ) |
| 171 | |
| 172 | def __repr__(self) -> str: |
| 173 | return self.anon_repr() |
| 174 | |
| 175 | @property |
| 176 | def num_reqs(self) -> int: |
| 177 | return len(self.req_ids) |
| 178 | |
| 179 | @cached_property |
| 180 | def _req_id_to_num_output_tokens(self) -> dict[str, int]: |
| 181 | """Cache mapping of req_id to num_output_tokens for O(1) lookup. |
| 182 | |
| 183 | This cached property is safe because CachedRequestData instances |
| 184 | are created fresh each scheduling iteration and not mutated during |
| 185 | computation of iteration details. |
| 186 | """ |
| 187 | return dict(zip(self.req_ids, self.num_output_tokens)) |
| 188 | |
| 189 | def is_context_phase(self, req_id: str) -> bool: |
| 190 | num_output_tokens = self._req_id_to_num_output_tokens.get(req_id) |
| 191 | return num_output_tokens is not None and num_output_tokens == 0 |
| 192 | |
| 193 | @classmethod |
| 194 | def make_empty(cls) -> "CachedRequestData": |
| 195 | return cls( |
| 196 | req_ids=[], |
| 197 | resumed_req_ids=set(), |
| 198 | new_token_ids=[], |
| 199 | all_token_ids={}, |
| 200 | new_block_ids=[], |
| 201 | num_computed_tokens=[], |
| 202 | num_output_tokens=[], |
| 203 | ) |
| 204 | |
| 205 | |
| 206 | @dataclass |
| 207 | class ScheduledEncoderInputStats: |
| 208 | """Stats for encoder inputs scheduled in one iteration.""" |
| 209 | |
| 210 | num_inputs: int = 0 |
| 211 | output_tokens: int = 0 |
| 212 | |
| 213 | |
| 214 | @dataclass |
| 215 | class KVConnectorBlockState: |
| 216 | """Scheduler-local block state offered to a producer-side KV connector.""" |
| 217 | |
| 218 | # Requests scheduled this step and requests with a boundary-state hand-off. |
| 219 | req_ids: set[str] |
| 220 | # Resolve on access to avoid copying tables the connector never reads. |
| 221 | resolve_block_ids: Callable[[str], tuple[list[int], ...]] |
| 222 | # Exact Mamba "align" boundary-state hand-offs. |
| 223 | boundary_state_offloads: dict[str, list[tuple[int, int, int]]] |
| 224 | |
| 225 | def get_block_ids(self, req_id: str) -> tuple[list[int], ...] | None: |
| 226 | if req_id not in self.req_ids: |
| 227 | return None |
| 228 | return self.resolve_block_ids(req_id) |
| 229 | |
| 230 | |
| 231 | @dataclass |
| 232 | class SchedulerOutput: |
| 233 | # list of the requests that are scheduled for the first time. |
| 234 | # We cache the request's data in each worker process, so that we don't |
| 235 | # need to re-send it every scheduling step. |
| 236 | scheduled_new_reqs: list[NewRequestData] |
| 237 | # list of the requests that have been scheduled before. |
| 238 | # Since the request's data is already cached in the worker processes, |
| 239 | # we only send the diff to minimize the communication cost. |
| 240 | scheduled_cached_reqs: CachedRequestData |
| 241 | |
| 242 | # req_id -> num_scheduled_tokens |
| 243 | # Number of tokens scheduled for each request. |
| 244 | num_scheduled_tokens: dict[str, int] |
| 245 | # Total number of tokens scheduled for all requests. |
| 246 | # Equal to sum(num_scheduled_tokens.values()) |
| 247 | total_num_scheduled_tokens: int |
| 248 | # req_id -> spec_token_ids |
| 249 | # If a request does not have any spec decode tokens, it will not be |
| 250 | # included in the dictionary. |
| 251 | scheduled_spec_decode_tokens: dict[str, list[int]] |
| 252 | # req_id -> encoder input indices that need processing. |
| 253 | # E.g., if a request has [0, 1], it could mean the vision encoder needs |
| 254 | # to process that the request's 0-th and 1-th images in the current step. |
| 255 | scheduled_encoder_inputs: dict[str, list[int]] |
| 256 | # Number of common prefix blocks for all requests in each KV cache group. |
| 257 | # This can be used for cascade attention. |
| 258 | num_common_prefix_blocks: list[int] |
| 259 | |
| 260 | # Request IDs that are finished in between the previous and the current |
| 261 | # steps. This is used to notify the workers about the finished requests |
| 262 | # so that they can free the cached states for those requests. |
| 263 | finished_req_ids: set[str] |
| 264 | # list of mm_hash strings associated with the encoder outputs to be |
| 265 | # freed from the encoder cache. |
| 266 | free_encoder_mm_hashes: list[str] |
| 267 | |
| 268 | scheduled_encoder_input_stats: ScheduledEncoderInputStats | None = None |
| 269 | |
| 270 | # Request IDs that are preempted in this step. |
| 271 | # Only used for v2 model runner. |
| 272 | preempted_req_ids: set[str] | None = None |
| 273 | |
| 274 | # Whether any of the scheduled requests use structured output. |
| 275 | # Set only in async scheduling case. |
| 276 | has_structured_output_requests: bool = False |
| 277 | |
| 278 | # Whether the scheduled requests have all the output tokens they |
| 279 | # need to perform grammar bitmask computation. |
| 280 | pending_structured_output_tokens: bool = False |
| 281 | |
| 282 | # Used for adjusting acceptance rate calculation. |
| 283 | num_invalid_spec_tokens: dict[str, int] | None = None |
| 284 | |
| 285 | # KV Cache Connector metadata. |
| 286 | kv_connector_metadata: KVConnectorMetadata | None = None |
| 287 | |
| 288 | # Whether any scheduled request consumes KV that the connector loads |
| 289 | # synchronously during this step (load_async=False). |
| 290 | has_sync_kv_loads: bool = False |
| 291 | |
| 292 | # Execution-auxiliary output control metadata consumed by the worker connector. |
| 293 | aux_output_connector_metadata: AuxOutputConnectorMetadata | None = None |
| 294 | |
| 295 | # EC Cache Connector metadata |
| 296 | ec_connector_metadata: ECConnectorMetadata | None = None |
| 297 | # EC Cache Manager metadata |
| 298 | ec_manager_metadata: EncoderCacheManagerMetadata | None = None |
| 299 | # Block IDs freshly allocated from the pool during this scheduling step. |
| 300 | # The worker zeros the corresponding GPU memory before the blocks are used, |
| 301 | # preventing stale NaN/data from corrupting attention or SSM computation. |
| 302 | new_block_ids_to_zero: list[int] | None = None |
| 303 | |
| 304 | # CoW copies to apply after zeroing new blocks and before forward. |
| 305 | kv_cache_block_copies: list[KVCacheBlockCopy] | None = None |
| 306 | |
| 307 | # Complete block-table rows that replace incrementally appended block IDs. |
| 308 | block_table_updates: dict[str, tuple[list[int], ...]] | None = None |
| 309 | |
| 310 | # Scheduler-local; always None by the time this reaches a worker. |
| 311 | kv_connector_block_state: KVConnectorBlockState | None = None |
| 312 | |
| 313 | # Dynamic speculative decoding: optimal K chosen by scheduler. |
| 314 | # Number of spec tokens to schedule for the next step. |
| 315 | num_spec_tokens_to_schedule: int = 0 |
| 316 | |
| 317 | @classmethod |
| 318 | def make_empty(cls) -> "SchedulerOutput": |
| 319 | return cls( |
| 320 | scheduled_new_reqs=[], |
| 321 | scheduled_cached_reqs=CachedRequestData.make_empty(), |
| 322 | num_scheduled_tokens={}, |
| 323 | total_num_scheduled_tokens=0, |
| 324 | scheduled_spec_decode_tokens={}, |
| 325 | scheduled_encoder_inputs={}, |
| 326 | num_common_prefix_blocks=[], |
| 327 | finished_req_ids=set(), |
| 328 | free_encoder_mm_hashes=[], |
| 329 | ) |
| 330 | |
| 331 | |
| 332 | @dataclass |
| 333 | class GrammarOutput: |
| 334 | # ids of structured output requests. |
| 335 | structured_output_request_ids: list[str] |
| 336 | # Bitmask ordered as structured_output_request_ids. |
| 337 | grammar_bitmask: "npt.NDArray[np.int32]" |
| 338 | # Per request, ordered as structured_output_request_ids: how many leading |
| 339 | # drafts the bitmask constrained. `grammar_bitmask` fills every row after |
| 340 | # the first -1 placeholder with the all-permissive `_full_mask`, so drafts |
| 341 | # from this index on must not be accepted. None invalidates every draft. |
| 342 | num_acceptable_drafts: list[int] | None = None |