| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import itertools |
| 4 | import time |
| 5 | from collections import defaultdict, deque |
| 6 | from collections.abc import Iterable |
| 7 | from dataclasses import replace |
| 8 | from typing import Any |
| 9 | |
| 10 | from vllm.compilation.cuda_graph import CUDAGraphStat |
| 11 | from vllm.config import KVEventsConfig, VllmConfig |
| 12 | from vllm.distributed.aux_output_connector.connector import AuxOutputSchedulerConnector |
| 13 | from vllm.distributed.ec_transfer.ec_connector.base import ( |
| 14 | ECConnectorBase, |
| 15 | ECConnectorMetadata, |
| 16 | ECConnectorRole, |
| 17 | ) |
| 18 | from vllm.distributed.ec_transfer.ec_connector.factory import ECConnectorFactory |
| 19 | from vllm.distributed.ec_transfer.ec_connector.metrics import ECConnectorStats |
| 20 | from vllm.distributed.kv_events import EventPublisherFactory, KVEventBatch |
| 21 | from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory |
| 22 | from vllm.distributed.kv_transfer.kv_connector.v1 import ( |
| 23 | KVConnectorBase_V1, |
| 24 | KVConnectorRole, |
| 25 | SupportsHMA, |
| 26 | ) |
| 27 | from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata |
| 28 | from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats |
| 29 | from vllm.logger import init_logger |
| 30 | from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry |
| 31 | from vllm.multimodal.encoder_budget import MultiModalBudget |
| 32 | from vllm.multimodal.utils import get_mm_features_in_window |
| 33 | from vllm.utils.math_utils import cdiv |
| 34 | from vllm.v1.core.encoder_cache_manager import ( |
| 35 | EncoderCacheManager, |
| 36 | EncoderDecoderCacheManager, |
| 37 | ) |
| 38 | from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager |
| 39 | from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector |
| 40 | from vllm.v1.core.kv_cache_utils import KVCacheBlock |
| 41 | from vllm.v1.core.sched.interface import PauseState, SchedulerInterface |
| 42 | from vllm.v1.core.sched.output import ( |
| 43 | CachedRequestData, |
| 44 | GrammarOutput, |
| 45 | KVConnectorBlockState, |
| 46 | NewRequestData, |
| 47 | ScheduledEncoderInputStats, |
| 48 | SchedulerOutput, |
| 49 | ) |
| 50 | from vllm.v1.core.sched.request_queue import ( |
| 51 | RequestQueue, |
| 52 | SchedulingPolicy, |
| 53 | create_request_queue, |
| 54 | ) |
| 55 | from vllm.v1.core.sched.utils import check_stop, remove_all |
| 56 | from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs |
| 57 | from vllm.v1.kv_cache_interface import ( |
| 58 | KVCacheConfig, |
| 59 | MambaSpec, |
| 60 | get_mamba_prefill_checkpoint_position, |
| 61 | is_mamba_prefill_checkpoint_valid, |
| 62 | ) |
| 63 | from vllm.v1.metrics.perf import ModelMetrics, PerfStats |
| 64 | from vllm.v1.metrics.stats import ( |
| 65 | KV_FETCH_COMPLETED_WAITING, |
| 66 | KV_FETCH_IN_PROGRESS, |
| 67 | KV_FETCH_STAGES, |
| 68 | KV_FETCH_WAITING_TO_START, |
| 69 | PrefixCacheStats, |
| 70 | RequestSpecDecodeMetrics, |
| 71 | SchedulerStats, |
| 72 | ) |
| 73 | from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput |
| 74 | from vllm.v1.request import Request, RequestStatus, StreamingUpdate |
| 75 | from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup |
| 76 | from vllm.v1.spec_decode.metrics import SpecDecodingStats |
| 77 | from vllm.v1.structured_output import StructuredOutputManager |
| 78 | from vllm.v1.structured_output.utils import strip_speculative_padding |
| 79 | from vllm.v1.utils import record_function_or_nullcontext |
| 80 | |
| 81 | logger = init_logger(__name__) |
| 82 | |
| 83 | |
| 84 | class Scheduler(SchedulerInterface): |
| 85 | def __init__( |
| 86 | self, |
| 87 | vllm_config: VllmConfig, |
| 88 | kv_cache_config: KVCacheConfig, |
| 89 | structured_output_manager: StructuredOutputManager, |
| 90 | block_size: int, |
| 91 | hash_block_size: int | None = None, |
| 92 | mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY, |
| 93 | include_finished_set: bool = False, |
| 94 | log_stats: bool = False, |
| 95 | ) -> None: |
| 96 | self.vllm_config = vllm_config |
| 97 | self.scheduler_config = vllm_config.scheduler_config |
| 98 | self.cache_config = vllm_config.cache_config |
| 99 | self.lora_config = vllm_config.lora_config |
| 100 | self.model_uses_mrope = vllm_config.model_config.uses_mrope |
| 101 | self.kv_cache_config = kv_cache_config |
| 102 | self.kv_events_config = vllm_config.kv_events_config |
| 103 | self.parallel_config = vllm_config.parallel_config |
| 104 | self.log_stats = log_stats |
| 105 | self.observability_config = vllm_config.observability_config |
| 106 | self.spec_decode_metrics_level = ( |
| 107 | self.observability_config.per_request_spec_decode_metrics |
| 108 | ) |
| 109 | self.kv_metrics_collector: KVCacheMetricsCollector | None = None |
| 110 | if self.observability_config.kv_cache_metrics: |
| 111 | self.kv_metrics_collector = KVCacheMetricsCollector( |
| 112 | self.observability_config.kv_cache_metrics_sample, |
| 113 | ) |
| 114 | self.structured_output_manager = structured_output_manager |
| 115 | self.is_encoder_decoder = vllm_config.model_config.is_encoder_decoder |
| 116 | self.is_mm_encoder_only = vllm_config.is_mm_encoder_only |
| 117 | |
| 118 | # include_finished_set controls whether a separate set of finished |
| 119 | # request ids should be included in the EngineCoreOutputs returned |
| 120 | # by update_from_outputs(). This is currently used in the multi-engine |
| 121 | # case to track request lifetimes efficiently. |
| 122 | self.finished_req_ids_dict: dict[int, set[str]] | None = ( |
| 123 | defaultdict(set) if include_finished_set else None |
| 124 | ) |
| 125 | # Track requests scheduled in prior step (MRV1-only). |
| 126 | self.prev_step_scheduled_req_ids: set[str] = set() |
| 127 | |
| 128 | # Scheduling constraints. |
| 129 | self.max_num_running_reqs = self.scheduler_config.max_num_seqs |
| 130 | # Admission into RUNNING. Separate from max_num_running_reqs, which |
| 131 | # is the model-runner slot count. |
| 132 | self.max_num_active_reqs = ( |
| 133 | self.scheduler_config.max_num_active_seqs |
| 134 | if self.scheduler_config.max_num_active_seqs is not None |
| 135 | else self.max_num_running_reqs |
| 136 | ) |
| 137 | self.max_num_scheduled_tokens = ( |
| 138 | self.scheduler_config.max_num_scheduled_tokens |
| 139 | if self.scheduler_config.max_num_scheduled_tokens is not None |
| 140 | else self.scheduler_config.max_num_batched_tokens |
| 141 | ) |
| 142 | self.adaptive_long_prefill_threshold = ( |
| 143 | self.scheduler_config.long_prefill_token_threshold_adaptive |
| 144 | ) |
| 145 | self.max_model_len = vllm_config.model_config.max_model_len |
| 146 | self.enable_kv_cache_events = ( |
| 147 | self.kv_events_config is not None |
| 148 | and self.kv_events_config.enable_kv_cache_events |
| 149 | ) |
| 150 | # Diffusion models may not sample any tokens for a denoising step. |
| 151 | self.num_sampled_tokens_per_step = ( |
| 152 | 1 if not vllm_config.model_config.is_diffusion else 0 |
| 153 | ) |
| 154 | |
| 155 | # Create KVConnector for the Scheduler. Note that each Worker |
| 156 | # will have a corresponding KVConnector with Role=WORKER. |
| 157 | # KV Connector pushes/pull of remote KVs for P/D and offloading. |
| 158 | self.connector = None |
| 159 | self.connector_prefix_cache_stats: PrefixCacheStats | None = None |
| 160 | self.recompute_kv_load_failures = True |
| 161 | self.defer_block_free = False |
| 162 | # Whether a preempted request's in-flight output must be dropped; see |
| 163 | # KVConnectorBase_V1.requires_kv_delivery. |
| 164 | self.requires_kv_delivery = False |
| 165 | kv_transfer_config = self.vllm_config.kv_transfer_config |
| 166 | if kv_transfer_config is not None: |
| 167 | assert not self.is_encoder_decoder, ( |
| 168 | "Encoder-decoder models are not currently supported with KV connectors" |
| 169 | ) |
| 170 | self.connector = KVConnectorFactory.create_connector( |
| 171 | config=self.vllm_config, |
| 172 | role=KVConnectorRole.SCHEDULER, |
| 173 | kv_cache_config=self.kv_cache_config, |
| 174 | ) |
| 175 | if self.log_stats: |
| 176 | self.connector_prefix_cache_stats = PrefixCacheStats() |
| 177 | kv_load_failure_policy = kv_transfer_config.kv_load_failure_policy |
| 178 | self.recompute_kv_load_failures = kv_load_failure_policy == "recompute" |
| 179 | |
| 180 | # With overlapping batches (async scheduling or PP), a step may |
| 181 | # still be writing a freed request's KV blocks. A consumer KV |
| 182 | # Connector can reallocate and fill those blocks via a load that |
| 183 | # isn't ordered against that write, so defer freeing them. |
| 184 | multiple_inflight_batches = self.vllm_config.max_concurrent_batches > 1 |
| 185 | if multiple_inflight_batches and kv_transfer_config.is_kv_consumer: |
| 186 | self.defer_block_free = True |
| 187 | |
| 188 | self.requires_kv_delivery = self.connector.requires_kv_delivery |
| 189 | |
| 190 | self.kv_event_publisher = EventPublisherFactory.create( |
| 191 | self.kv_events_config, |
| 192 | self.parallel_config.data_parallel_index, |
| 193 | ) |
| 194 | self.ec_connector = None |
| 195 | if self.vllm_config.ec_transfer_config is not None: |
| 196 | self.ec_connector = ECConnectorFactory.create_connector( |
| 197 | config=self.vllm_config, role=ECConnectorRole.SCHEDULER |
| 198 | ) |
| 199 | |
| 200 | num_gpu_blocks = self.cache_config.num_gpu_blocks |
| 201 | assert num_gpu_blocks is not None and num_gpu_blocks > 0 |
| 202 | |
| 203 | self.block_size = block_size |
| 204 | self.dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size |
| 205 | self.pcp_world_size = vllm_config.parallel_config.prefill_context_parallel_size |
| 206 | |
| 207 | # req_id -> Request |
| 208 | self.requests: dict[str, Request] = {} |
| 209 | # Scheduling policy |
| 210 | try: |
| 211 | self.policy = SchedulingPolicy(self.scheduler_config.policy) |
| 212 | except ValueError as e: |
| 213 | raise ValueError( |
| 214 | f"Unknown scheduling policy: {self.scheduler_config.policy}" |
| 215 | ) from e |
| 216 | # Priority queues for requests. |
| 217 | self.waiting = create_request_queue(self.policy) |
| 218 | # Waiting requests that hold KV blocks are always drained before |
| 219 | # self.waiting: a block-holder must never sit behind a request whose |
| 220 | # failed allocation would stop the scheduling scan. |
| 221 | self.kv_holding_waiting = create_request_queue(self.policy) |
| 222 | # Waiting requests that are deferred, i.e. were skipped by the scheduling |
| 223 | # scan or enqueued in a blocked status. |
| 224 | self.deferred_waiting: set[Request] = set() |
| 225 | self.running: list[Request] = [] |
| 226 | |
| 227 | # The request IDs that are finished in between the previous and the |
| 228 | # current steps. This is used to notify the workers about the finished |
| 229 | # requests so that they can free the cached states for those requests. |
| 230 | # This is flushed at the end of each scheduling step. |
| 231 | self.finished_req_ids: set[str] = set() |
| 232 | |
| 233 | # IDs of requests preempted since the last call to schedule(). |
| 234 | self.reset_preempted_req_ids: set[str] = set() |
| 235 | |
| 236 | # Counter for requests waiting for streaming input. Used to calculate |
| 237 | # number of unfinished requests |
| 238 | self.num_waiting_for_streaming_input: int = 0 |
| 239 | |
| 240 | # KV Connector: requests in process of async KV loading or recving |
| 241 | self.finished_recving_kv_req_ids: set[str] = set() |
| 242 | self.failed_recving_kv_req_ids: set[str] = set() |
| 243 | |
| 244 | # Grammar compilation failures to finish as per-request errors in |
| 245 | # update_from_output. |
| 246 | self.grammar_compile_error_reqs: set[str] = set() |
| 247 | self.encoder_cache_mismatch_reqs: set[str] = set() |
| 248 | |
| 249 | # Encoder-related. |
| 250 | # Calculate encoder cache size if applicable |
| 251 | supports_mm_inputs = vllm_config.model_config.supports_multimodal_inputs |
| 252 | mm_budget = ( |
| 253 | MultiModalBudget(vllm_config, mm_registry) if supports_mm_inputs else None |
| 254 | ) |
| 255 | |
| 256 | # NOTE: Text-only encoder-decoder models are implemented as |
| 257 | # multi-modal models for convenience |
| 258 | # Example: https://github.com/vllm-project/bart-plugin |
| 259 | if self.is_encoder_decoder: |
| 260 | assert mm_budget and len(mm_budget.mm_max_toks_per_item) <= 1, ( |
| 261 | "Encoder-decoder models are expected to implement the " |
| 262 | "multimodal interface with at most one modality." |
| 263 | ) |
| 264 | |
| 265 | self.max_num_encoder_input_tokens = ( |
| 266 | mm_budget.encoder_compute_budget if mm_budget else 0 |
| 267 | ) |
| 268 | encoder_cache_size = mm_budget.encoder_cache_size if mm_budget else 0 |
| 269 | manager_cls_obj = vllm_config.ec_manager_config.get_encoder_cache_manager_obj() |
| 270 | if manager_cls_obj is None: |
| 271 | manager_cls_obj = ( |
| 272 | EncoderDecoderCacheManager |
| 273 | if self.is_encoder_decoder |
| 274 | else EncoderCacheManager |
| 275 | ) |
| 276 | self.encoder_cache_manager = manager_cls_obj.create_manager( |
| 277 | cache_size=encoder_cache_size, vllm_config=vllm_config |
| 278 | ) |
| 279 | speculative_config = vllm_config.speculative_config |
| 280 | self.use_eagle = False |
| 281 | self.use_eagle_block_drop = False |
| 282 | self.num_spec_tokens = vllm_config.num_speculative_tokens |
| 283 | self.num_lookahead_tokens = vllm_config.num_lookahead_tokens |
| 284 | # DSV41 SWA bounded replay: groups that declare a replay window are rebuilt |
| 285 | # after a prefix hit by recomputing its trailing tokens. One window |
| 286 | # for all such groups, so the rewind matches every group's allocation. |
| 287 | replay_windows = { |
| 288 | group.kv_cache_spec.prefix_replay_tokens |
| 289 | for group in kv_cache_config.kv_cache_groups |
| 290 | if group.kv_cache_spec.prefix_replay_tokens > 0 |
| 291 | } |
| 292 | assert len(replay_windows) <= 1, ( |
| 293 | f"Prefix replay windows should agree: {sorted(replay_windows)}" |
| 294 | ) |
| 295 | self.prefix_replay_tokens = replay_windows.pop() if replay_windows else 0 |
| 296 | self.num_prefill_lookahead = vllm_config.num_prefill_lookahead_tokens |
| 297 | self.dynamic_sd_lookup: list[int] | None = None |
| 298 | if speculative_config is not None: |
| 299 | if speculative_config.num_speculative_tokens_per_batch_size: |
| 300 | self.dynamic_sd_lookup = build_dynamic_sd_schedule_lookup( |
| 301 | speculative_config.num_speculative_tokens_per_batch_size, |
| 302 | vllm_max_batch_size=self.scheduler_config.max_num_seqs, |
| 303 | vllm_num_speculative_tokens=self.num_spec_tokens, |
| 304 | ) |
| 305 | self.use_eagle = speculative_config.use_eagle() |
| 306 | self.use_eagle_block_drop = speculative_config.use_eagle_block_drop() |
| 307 | if self.use_eagle and not self.use_eagle_block_drop: |
| 308 | logger.warning( |
| 309 | "EAGLE trailing prefix-cache block dropping is disabled. " |
| 310 | "This is experimental and may affect speculative-token " |
| 311 | "acceptance rates." |
| 312 | ) |
| 313 | |
| 314 | # Create the KV cache manager. |
| 315 | if hash_block_size is None: |
| 316 | hash_block_size = block_size |
| 317 | self.hash_block_size = hash_block_size |
| 318 | self.kv_cache_manager = KVCacheManager( |
| 319 | kv_cache_config=kv_cache_config, |
| 320 | max_model_len=self.max_model_len, |
| 321 | max_in_flight_tokens=vllm_config.max_in_flight_tokens, |
| 322 | enable_caching=self.cache_config.enable_prefix_caching, |
| 323 | use_eagle=self.use_eagle_block_drop, |
| 324 | num_prefill_lookahead=self.num_prefill_lookahead, |
| 325 | log_stats=self.log_stats, |
| 326 | enable_kv_cache_events=self.enable_kv_cache_events, |
| 327 | dcp_world_size=self.dcp_world_size, |
| 328 | pcp_world_size=1, |
| 329 | scheduler_block_size=self.block_size, |
| 330 | hash_block_size=hash_block_size, |
| 331 | metrics_collector=self.kv_metrics_collector, |
| 332 | watermark=self.scheduler_config.watermark, |
| 333 | enable_mamba_shared_prefix_checkpoint=( |
| 334 | self.cache_config.enable_mamba_shared_prefix_checkpoint |
| 335 | ), |
| 336 | ) |
| 337 | # Bind after construction so connectors can access the cache manager. |
| 338 | if self.connector is not None: |
| 339 | self.connector.bind_kv_cache_manager(self.kv_cache_manager) |
| 340 | |
| 341 | self.use_pp = self.parallel_config.pipeline_parallel_size > 1 |
| 342 | self.use_v2_model_runner = vllm_config.use_v2_model_runner |
| 343 | # Scheduler iteration counter. Drives the V2+PP+async decode-throttle |
| 344 | # cadence (`next_decode_eligible_step`). |
| 345 | self.current_step = 0 |
| 346 | # DP prefill balancing: Flag to track whether the last cadence-aligned |
| 347 | # prefill batch fully drained the waiting queue. Prefill throttling |
| 348 | # is disabled in this case. |
| 349 | self.prefill_capacity_bound = False |
| 350 | self.scheduler_reserve_full_isl = ( |
| 351 | self.scheduler_config.scheduler_reserve_full_isl |
| 352 | ) |
| 353 | |
| 354 | self.has_mamba_layers = kv_cache_config.has_mamba_layers |
| 355 | self.needs_kv_cache_zeroing = kv_cache_config.needs_kv_cache_zeroing |
| 356 | # Blocks that async KV loads will overwrite this step, skipped from |
| 357 | # zeroing since the zeroing could race the out-of-band write. |
| 358 | self._skip_zero_block_ids: set[int] = set() |
| 359 | self.need_mamba_block_aligned_split = ( |
| 360 | self.has_mamba_layers and self.cache_config.mamba_cache_mode == "align" |
| 361 | ) |
| 362 | # TODO: Support models with multiple Mamba specs that require different |
| 363 | # prefill checkpoint alignments instead of selecting the first one. |
| 364 | self.mamba_prefill_checkpoint_alignment = next( |
| 365 | ( |
| 366 | group.kv_cache_spec.prefill_checkpoint_alignment |
| 367 | for group in kv_cache_config.kv_cache_groups |
| 368 | if isinstance(group.kv_cache_spec, MambaSpec) |
| 369 | ), |
| 370 | None, |
| 371 | ) |
| 372 | self.mamba_has_prefill_checkpoint_blocks = self.has_mamba_layers and all( |
| 373 | not isinstance(group.kv_cache_spec, MambaSpec) |
| 374 | or group.kv_cache_spec.num_prefill_checkpoint_blocks > 0 |
| 375 | for group in kv_cache_config.kv_cache_groups |
| 376 | ) |
| 377 | # A finer prefix_match_unit is configured: a mamba partial tail entry |
| 378 | # can only be registered by a step ending exactly at the prompt's last |
| 379 | # hash boundary, so the split adds that stop. |
| 380 | self.mamba_partial_cache_hit = ( |
| 381 | self.need_mamba_block_aligned_split |
| 382 | and self.hash_block_size < self.block_size |
| 383 | and self.kv_cache_manager.coordinator.enable_partial_hash_hits |
| 384 | ) |
| 385 | # Opt-in: also stop at the junction, where an eagle sibling resumes. The |
| 386 | # manager decides whether it can check-point there (per-group eagle bit, |
| 387 | # no MTP re-prefill tail); splitting for a stop it would refuse costs a |
| 388 | # forward pass and displaces the block-boundary stop. |
| 389 | self.mamba_shared_prefix_checkpoint = ( |
| 390 | self.mamba_partial_cache_hit |
| 391 | and self.kv_cache_manager.mamba_shared_prefix_checkpoint |
| 392 | ) |
| 393 | |
| 394 | # Counts of non-empty steps scheduled / processed. update_from_output |
| 395 | # is called once per scheduled step in FIFO order, so these stay in sync. |
| 396 | self.sched_step_seq = 0 |
| 397 | self.processed_step_seq = 0 |
| 398 | # FIFO of (fence_seq, blocks): blocks become safe to free once |
| 399 | # processed_step_seq >= fence_seq. |
| 400 | self.deferred_frees: deque[tuple[int, list[KVCacheBlock]]] = deque() |
| 401 | |
| 402 | self.perf_metrics: ModelMetrics | None = None |
| 403 | if self.log_stats and vllm_config.observability_config.enable_mfu_metrics: |
| 404 | self.perf_metrics = ModelMetrics(vllm_config) |
| 405 | |
| 406 | self.aux_output_connector = ( |
| 407 | AuxOutputSchedulerConnector() |
| 408 | if vllm_config.aux_output_config.enabled |
| 409 | else None |
| 410 | ) |
| 411 | self.return_sampling_mask = vllm_config.model_config.return_sampling_mask |
| 412 | |
| 413 | self._pause_state: PauseState = PauseState.UNPAUSED |
| 414 | |
| 415 | # In-flight requests still prefilling (prefill chunks + in-progress |
| 416 | # async KV loads). Their remaining-block reservation gates async loads. |
| 417 | self._inflight_prefills: set[Request] = set() |
| 418 | |
| 419 | # Requests with an async KV load, by stage (KV_FETCH_STAGES). |
| 420 | # None when not tracked (stats disabled or no KV connector). |
| 421 | self._kv_fetch_stages: dict[Request, str] | None = ( |
| 422 | {} if self.log_stats and self.connector is not None else None |
| 423 | ) |
| 424 | self._kv_fetch_counts: dict[str, int] = dict.fromkeys(KV_FETCH_STAGES, 0) |
| 425 | |
| 426 | def _mamba_block_aligned_split( |
| 427 | self, |
| 428 | request: Request, |
| 429 | num_new_tokens: int, |
| 430 | num_new_local_computed_tokens: int = 0, |
| 431 | num_external_computed_tokens: int = 0, |
| 432 | ) -> int: |
| 433 | """Clip a prefill chunk so it ends where Mamba state must be cached. |
| 434 | |
| 435 | In "align" cache mode reusable SSM states are materialized at block |
| 436 | boundaries, plus mandatory early stops (the prompt's partial-tail hash |
| 437 | boundary, a detected shared-prefix junction). If a block is larger |
| 438 | than the configured prefill chunk limit, intermediate chunks keep |
| 439 | private running state until they reach the next cacheable position. |
| 440 | """ |
| 441 | start = ( |
| 442 | request.num_computed_tokens |
| 443 | + num_new_local_computed_tokens |
| 444 | + num_external_computed_tokens |
| 445 | ) |
| 446 | # Split only during prefill: `request.num_tokens - 1` extends this to |
| 447 | # resumed requests replaying their output tokens. |
| 448 | prefill_end = max(request.num_prompt_tokens, request.num_tokens - 1) |
| 449 | if start >= prefill_end: |
| 450 | return num_new_tokens |
| 451 | |
| 452 | block_size = self.cache_config.block_size |
| 453 | # The last block-aligned position whose state can be cached. With |
| 454 | # Eagle, FullAttn prunes the last matching block, so back off one |
| 455 | # block to avoid a Mamba cache miss. |
| 456 | last_cache_position = request.num_tokens - request.num_tokens % block_size |
| 457 | if self.use_eagle_block_drop: |
| 458 | last_cache_position = max(last_cache_position - block_size, 0) |
| 459 | |
| 460 | end = start + num_new_tokens |
| 461 | checkpoint_position = get_mamba_prefill_checkpoint_position( |
| 462 | prefill_end, |
| 463 | self.hash_block_size, |
| 464 | drop_eagle_block=self.use_eagle_block_drop, |
| 465 | ) |
| 466 | use_internal_checkpoint = ( |
| 467 | self.mamba_has_prefill_checkpoint_blocks |
| 468 | and end >= prefill_end |
| 469 | and is_mamba_prefill_checkpoint_valid( |
| 470 | query_start=start, |
| 471 | query_end=end, |
| 472 | checkpoint_position=checkpoint_position, |
| 473 | hash_block_size=self.hash_block_size, |
| 474 | mamba_block_size=block_size, |
| 475 | checkpoint_alignment=self.mamba_prefill_checkpoint_alignment, |
| 476 | ) |
| 477 | ) |
| 478 | if use_internal_checkpoint: |
| 479 | last_cache_position = 0 |
| 480 | # Invariant: slot p holds the state after exactly (p + 1) * block_size |
| 481 | # tokens. State is written at chunk ends, so chunk ends must be block |
| 482 | # aligned. Exempt: the prompt's last chunk, whose slot decode advances |
| 483 | # to the boundary. A block too wide for one chunk advances sub-block |
| 484 | # and re-aligns at the next boundary. |
| 485 | if end < prefill_end: |
| 486 | max_prefill_tokens = self.max_num_scheduled_tokens |
| 487 | long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold |
| 488 | if long_prefill_threshold > 0: |
| 489 | max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) |
| 490 | aligned_end = end // block_size * block_size |
| 491 | if aligned_end > start or block_size <= max_prefill_tokens: |
| 492 | end = aligned_end |
| 493 | |
| 494 | next_block_boundary = (start // block_size + 1) * block_size |
| 495 | tail_boundary = ( |
| 496 | request.num_prompt_tokens // self.hash_block_size * self.hash_block_size |
| 497 | if self.mamba_partial_cache_hit and not use_internal_checkpoint |
| 498 | else 0 |
| 499 | ) |
| 500 | if tail_boundary and self.use_eagle_block_drop: |
| 501 | # Eagle matches one hash unit past the candidate and drops it, so |
| 502 | # nothing proves the prompt's own last hash boundary. Materialize |
| 503 | # the state one unit lower, where the hit can actually land. Keyed on |
| 504 | # the block-drop bit, not plain use_eagle: this shift exists only to |
| 505 | # compensate for the drop, and the Mamba manager's matching gate |
| 506 | # reads the same bit (the coordinator is handed use_eagle_block_drop). |
| 507 | tail_boundary = max(tail_boundary - self.hash_block_size, 0) |
| 508 | junction = request.shared_prefix_boundary |
| 509 | # Block-floored: a sub-block junction's state is not separately cacheable. |
| 510 | block_floored = start + (junction - start) // block_size * block_size |
| 511 | # Past the prompt the manager writes nothing, so fall back to the |
| 512 | # block-floored stop rather than dropping it: a resumed request replaying |
| 513 | # output tokens can still observe a junction there. |
| 514 | junction_stop = ( |
| 515 | junction |
| 516 | if self.mamba_shared_prefix_checkpoint |
| 517 | and junction <= request.num_prompt_tokens |
| 518 | else block_floored |
| 519 | ) |
| 520 | stops = ( |
| 521 | # Same invariant: a chunk starting mid-block stops at the boundary |
| 522 | # rather than running past it. |
| 523 | next_block_boundary |
| 524 | if start % block_size != 0 and not use_internal_checkpoint |
| 525 | else 0, |
| 526 | # Never run past the last cacheable block boundary mid-chunk. |
| 527 | last_cache_position, |
| 528 | # Fine-grained hits: the prompt's partial-tail entry can only be |
| 529 | # registered by a chunk ending exactly at its last hash boundary. |
| 530 | tail_boundary |
| 531 | if last_cache_position < tail_boundary < request.num_prompt_tokens |
| 532 | else 0, |
| 533 | # Marconi shared-prefix junction: cache its state so sibling |
| 534 | # requests sharing the prefix can reuse it. |
| 535 | junction_stop if start < junction < end else 0, |
| 536 | ) |
| 537 | # Stop at the earliest mandatory position strictly inside the chunk. |
| 538 | end = min((s for s in stops if start < s < end), default=end) |
| 539 | return max(end - start, 0) |
| 540 | |
| 541 | def _get_local_prefix_cache_hit( |
| 542 | self, request: Request |
| 543 | ) -> tuple[KVCacheBlocks, int, int, bool]: |
| 544 | connector = self.connector |
| 545 | if connector is not None and connector.supports_divergent_local_hybrid_hits: |
| 546 | return self.kv_cache_manager.get_computed_blocks_for_connector(request) |
| 547 | |
| 548 | blocks, num_local, shared_prefix_boundary = ( |
| 549 | self.kv_cache_manager.get_computed_blocks(request) |
| 550 | ) |
| 551 | return blocks, num_local, shared_prefix_boundary, False |
| 552 | |
| 553 | def _reserve_prefill_lookahead( |
| 554 | self, |
| 555 | request: Request, |
| 556 | num_computed_tokens: int, |
| 557 | num_new_tokens: int, |
| 558 | ) -> int: |
| 559 | """Never end a prefill chunk within num_prefill_lookahead of the |
| 560 | prefill end. |
| 561 | |
| 562 | At a chunked-prefill boundary, the multi-module MTP drafter consumes |
| 563 | the next num_prefill_lookahead known prefill tokens as draft inputs. A |
| 564 | boundary closer to the end than that would make it fall back to |
| 565 | sampled drafts, permanently polluting the trailing modules' KV caches. |
| 566 | Either finish the prefill or leave at least num_prefill_lookahead for |
| 567 | the next chunk. No-op for eagle-family drafters (lookahead 1). |
| 568 | """ |
| 569 | remaining = request.num_tokens - num_computed_tokens - num_new_tokens |
| 570 | if 0 < remaining < self.num_prefill_lookahead: |
| 571 | num_new_tokens -= self.num_prefill_lookahead - remaining |
| 572 | return max(num_new_tokens, 0) |
| 573 | |
| 574 | def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: |
| 575 | self.current_step += 1 |
| 576 | # NOTE(woosuk) on the scheduling algorithm: |
| 577 | # There's no "decoding phase" nor "prefill phase" in the scheduler. |
| 578 | # Each request just has the num_computed_tokens and |
| 579 | # num_tokens_with_spec. num_tokens_with_spec = |
| 580 | # len(prompt_token_ids) + len(output_token_ids) + len(spec_token_ids). |
| 581 | # At each step, the scheduler tries to assign tokens to the requests |
| 582 | # so that each request's num_computed_tokens can catch up its |
| 583 | # num_tokens_with_spec. This is general enough to cover |
| 584 | # chunked prefills, prefix caching, speculative decoding, |
| 585 | # and the "jump decoding" optimization in the future. |
| 586 | |
| 587 | scheduled_new_reqs: list[Request] = [] |
| 588 | scheduled_resumed_reqs: list[Request] = [] |
| 589 | scheduled_running_reqs: list[Request] = [] |
| 590 | preempted_reqs: list[Request] = [] |
| 591 | |
| 592 | req_to_new_blocks: dict[str, KVCacheBlocks] = {} |
| 593 | num_scheduled_tokens: dict[str, int] = {} |
| 594 | token_budget = self.max_num_scheduled_tokens |
| 595 | spec = self.vllm_config.speculative_config |
| 596 | draft_slots = spec.max_num_new_slots_for_drafting if spec is not None else 0 |
| 597 | input_budget = self.scheduler_config.max_num_batched_tokens |
| 598 | if self._pause_state == PauseState.PAUSED_ALL: |
| 599 | # Do not schedule any requests when paused. |
| 600 | token_budget = 0 |
| 601 | |
| 602 | # Encoder-related. |
| 603 | scheduled_encoder_inputs: dict[str, list[int]] = {} |
| 604 | encoder_compute_budget = self.max_num_encoder_input_tokens |
| 605 | # Spec decode-related. |
| 606 | scheduled_spec_decode_tokens: dict[str, list[int]] = {} |
| 607 | # Whether the running batch contains any prefill requests. |
| 608 | prefill_scheduled = False |
| 609 | # Whether any scheduled request has a synchronous connector KV load. |
| 610 | has_sync_kv_loads = False |
| 611 | |
| 612 | # For logging. |
| 613 | scheduled_timestamp = time.monotonic() |
| 614 | |
| 615 | self.kv_cache_manager.new_step_starts() |
| 616 | |
| 617 | # DP prefill balancing: on a throttled (non-cadence-aligned) step, defer |
| 618 | # all prefill compute unless saturated. |
| 619 | defer_prefills = ( |
| 620 | throttle_prefills and not self.prefill_capacity_bound |
| 621 | ) and any(not r.is_prefill_chunk for r in self.running) |
| 622 | |
| 623 | # `long_prefill_token_threshold` exists to stop a long prefill from |
| 624 | # starving other requests of the token budget. When it is the only |
| 625 | # request there is nobody to starve, so let it use the whole budget. |
| 626 | num_running, num_waiting = self.get_request_counts() |
| 627 | num_eligible_reqs = num_running + num_waiting |
| 628 | long_prefill_token_threshold = ( |
| 629 | self.scheduler_config.long_prefill_token_threshold |
| 630 | if num_eligible_reqs > 1 |
| 631 | else 0 |
| 632 | ) |
| 633 | if long_prefill_token_threshold > 0 and self.adaptive_long_prefill_threshold: |
| 634 | # Floor the cap at a fair share of the input budget so it never |
| 635 | # cuts a request below max_num_batched_tokens / num requests. |
| 636 | long_prefill_token_threshold = max( |
| 637 | long_prefill_token_threshold, input_budget // num_eligible_reqs |
| 638 | ) |
| 639 | |
| 640 | # First, schedule the RUNNING requests. |
| 641 | req_index = 0 |
| 642 | while req_index < len(self.running) and token_budget > 0: |
| 643 | request = self.running[req_index] |
| 644 | if input_budget <= draft_slots: |
| 645 | break |
| 646 | |
| 647 | if ( |
| 648 | request.num_output_placeholders > 0 |
| 649 | # This is (num_computed_tokens + 1) - (num_output_placeholders - 1). |
| 650 | # Since output placeholders are also included in the computed tokens |
| 651 | # count, we subtract (num_output_placeholders - 1) to remove any draft |
| 652 | # tokens, so that we can be sure no further steps are needed even if |
| 653 | # they are all rejected. |
| 654 | and request.num_computed_tokens + 2 - request.num_output_placeholders |
| 655 | >= request.num_prompt_tokens + request.max_tokens |
| 656 | ): |
| 657 | # Async scheduling: Avoid scheduling an extra step when we are sure that |
| 658 | # the previous step has reached request.max_tokens. We don't schedule |
| 659 | # partial draft tokens since this prevents uniform decode optimizations. |
| 660 | req_index += 1 |
| 661 | continue |
| 662 | |
| 663 | if self.current_step < request.next_decode_eligible_step: |
| 664 | # V2+PP+async: enforce `pp_size` steps between same-req decodes |
| 665 | # to match worker-side sampled-tokens broadcast slot ring cadence. |
| 666 | req_index += 1 |
| 667 | continue |
| 668 | |
| 669 | if defer_prefills and request.is_prefill_chunk: |
| 670 | # DP prefill balancing: defer this in-progress prefill chunk to a |
| 671 | # cadence-aligned step; decodes still run to fill this step. |
| 672 | req_index += 1 |
| 673 | continue |
| 674 | |
| 675 | if ( |
| 676 | self.ec_connector is not None |
| 677 | and request.mm_features |
| 678 | and not self.ec_connector.ensure_cache_available( |
| 679 | request, |
| 680 | request.num_computed_tokens - request.num_output_placeholders, |
| 681 | ) |
| 682 | ): |
| 683 | req_index += 1 |
| 684 | continue |
| 685 | |
| 686 | num_new_tokens = ( |
| 687 | request.num_tokens_with_spec |
| 688 | + request.num_output_placeholders |
| 689 | - request.num_computed_tokens |
| 690 | ) |
| 691 | if 0 < long_prefill_token_threshold < num_new_tokens: |
| 692 | num_new_tokens = long_prefill_token_threshold |
| 693 | num_new_tokens = min( |
| 694 | num_new_tokens, token_budget, input_budget - draft_slots |
| 695 | ) |
| 696 | |
| 697 | # Make sure the input position does not exceed the max model len. |
| 698 | # This is necessary when using spec decoding. |
| 699 | num_new_tokens = min( |
| 700 | num_new_tokens, |
| 701 | self.max_model_len |
| 702 | - request.num_computed_tokens |
| 703 | - self.num_sampled_tokens_per_step, |
| 704 | ) |
| 705 | |
| 706 | # Apply Mamba alignment before encoder caps. |
| 707 | if self.need_mamba_block_aligned_split: |
| 708 | num_new_tokens = self._mamba_block_aligned_split( |
| 709 | request, num_new_tokens |
| 710 | ) |
| 711 | |
| 712 | # Schedule encoder inputs. |
| 713 | encoder_inputs_to_schedule = None |
| 714 | external_load_encoder_input: list[int] = [] |
| 715 | new_encoder_compute_budget = encoder_compute_budget |
| 716 | if request.has_encoder_inputs: |
| 717 | ( |
| 718 | encoder_inputs_to_schedule, |
| 719 | num_new_tokens, |
| 720 | new_encoder_compute_budget, |
| 721 | external_load_encoder_input, |
| 722 | ) = self._try_schedule_encoder_inputs( |
| 723 | request, |
| 724 | request.num_computed_tokens, |
| 725 | num_new_tokens, |
| 726 | encoder_compute_budget, |
| 727 | shift_computed_tokens=self.num_prefill_lookahead, |
| 728 | ) |
| 729 | |
| 730 | # Multi-module MTP: avoid ending a prefill chunk within |
| 731 | # num_prefill_lookahead of the prefill end. |
| 732 | num_new_tokens = self._reserve_prefill_lookahead( |
| 733 | request, request.num_computed_tokens, num_new_tokens |
| 734 | ) |
| 735 | |
| 736 | if num_new_tokens == 0: |
| 737 | # The request cannot be scheduled because one of the following |
| 738 | # reasons: |
| 739 | # 1. No new tokens to schedule. This may happen when |
| 740 | # (1) PP>1 and we have already scheduled all prompt tokens |
| 741 | # but they are not finished yet. |
| 742 | # (2) Async scheduling and the request has reached to either |
| 743 | # its max_total_tokens or max_model_len. |
| 744 | # 2. The encoder budget is exhausted. |
| 745 | # 3. The encoder cache is exhausted. |
| 746 | # 4. Insufficient budget for a block-aligned chunk in hybrid |
| 747 | # models with mamba cache mode \"align\". |
| 748 | # 5. Insufficient budget to keep a multi-module MTP prefill |
| 749 | # chunk out of the prefill-lookahead window. |
| 750 | # NOTE(woosuk): Here, by doing `continue` instead of `break`, |
| 751 | # we do not strictly follow the FCFS scheduling policy and |
| 752 | # allow the lower-priority requests to be scheduled. |
| 753 | req_index += 1 |
| 754 | continue |
| 755 | |
| 756 | # Schedule newly needed KV blocks for the request. |
| 757 | with record_function_or_nullcontext("schedule: allocate_slots"): |
| 758 | while True: |
| 759 | new_blocks = self.kv_cache_manager.allocate_slots( |
| 760 | request, |
| 761 | num_new_tokens, |
| 762 | num_lookahead_tokens=self.num_lookahead_tokens, |
| 763 | ) |
| 764 | |
| 765 | if new_blocks is not None: |
| 766 | # The request can be scheduled. |
| 767 | break |
| 768 | |
| 769 | if ( |
| 770 | self.connector is not None |
| 771 | and self.connector.has_pending_block_frees() |
| 772 | ): |
| 773 | break |
| 774 | |
| 775 | # The request cannot be scheduled. |
| 776 | # Preempt the lowest-priority request. |
| 777 | if self.policy == SchedulingPolicy.PRIORITY: |
| 778 | preempted_req = max( |
| 779 | self.running, |
| 780 | key=lambda r: (r.priority, r.arrival_time), |
| 781 | ) |
| 782 | else: |
| 783 | preempted_req = self.running[-1] |
| 784 | |
| 785 | # A deferred free will not help with immediate allocation. |
| 786 | if not self._request_blocks_can_be_freed(preempted_req): |
| 787 | break |
| 788 | |
| 789 | if self.policy == SchedulingPolicy.PRIORITY: |
| 790 | victim_index = self.running.index(preempted_req) |
| 791 | del self.running[victim_index] |
| 792 | if victim_index < req_index: |
| 793 | req_index -= 1 |
| 794 | |
| 795 | if preempted_req in scheduled_running_reqs: |
| 796 | preempted_req_id = preempted_req.request_id |
| 797 | scheduled_running_reqs.remove(preempted_req) |
| 798 | restored = num_scheduled_tokens.pop(preempted_req_id) |
| 799 | token_budget += restored |
| 800 | input_budget += restored + draft_slots |
| 801 | req_to_new_blocks.pop(preempted_req_id) |
| 802 | scheduled_spec_decode_tokens.pop(preempted_req_id, None) |
| 803 | preempted_encoder_inputs = scheduled_encoder_inputs.pop( |
| 804 | preempted_req_id, None |
| 805 | ) |
| 806 | if preempted_encoder_inputs: |
| 807 | # Restore encoder compute budget if the preempted |
| 808 | # request had encoder inputs scheduled in this step. |
| 809 | num_embeds_to_restore = sum( |
| 810 | preempted_req.get_num_encoder_embeds(i) |
| 811 | for i in preempted_encoder_inputs |
| 812 | ) |
| 813 | encoder_compute_budget += num_embeds_to_restore |
| 814 | else: |
| 815 | preempted_req = self.running.pop() |
| 816 | |
| 817 | self._preempt_request( |
| 818 | preempted_req, |
| 819 | scheduled_timestamp, |
| 820 | drop_stale_output=self.requires_kv_delivery, |
| 821 | ) |
| 822 | preempted_reqs.append(preempted_req) |
| 823 | if preempted_req == request: |
| 824 | # No more request to preempt. Cannot schedule this request. |
| 825 | break |
| 826 | |
| 827 | if new_blocks is None: |
| 828 | # Cannot schedule this request. |
| 829 | break |
| 830 | |
| 831 | # Schedule the request. |
| 832 | scheduled_running_reqs.append(request) |
| 833 | prefill_scheduled |= request.is_prefill_chunk |
| 834 | request_id = request.request_id |
| 835 | req_to_new_blocks[request_id] = new_blocks |
| 836 | num_scheduled_tokens[request_id] = num_new_tokens |
| 837 | token_budget -= num_new_tokens |
| 838 | input_budget -= num_new_tokens + draft_slots |
| 839 | req_index += 1 |
| 840 | |
| 841 | # Speculative decode related. |
| 842 | if request.spec_token_ids: |
| 843 | num_scheduled_spec_tokens = ( |
| 844 | num_new_tokens |
| 845 | + request.num_computed_tokens |
| 846 | - request.num_tokens |
| 847 | - request.num_output_placeholders |
| 848 | ) |
| 849 | if num_scheduled_spec_tokens > 0: |
| 850 | spec_token_ids = request.spec_token_ids |
| 851 | if len(spec_token_ids) > num_scheduled_spec_tokens: |
| 852 | spec_token_ids = spec_token_ids[:num_scheduled_spec_tokens] |
| 853 | scheduled_spec_decode_tokens[request.request_id] = spec_token_ids |
| 854 | |
| 855 | # New spec tokens will be set in `update_draft_token_ids` before the |
| 856 | # next step when applicable. |
| 857 | request.spec_token_ids = [] |
| 858 | |
| 859 | # Encoder-related. |
| 860 | if encoder_inputs_to_schedule: |
| 861 | scheduled_encoder_inputs[request_id] = encoder_inputs_to_schedule |
| 862 | # Allocate the encoder cache. |
| 863 | for i in encoder_inputs_to_schedule: |
| 864 | self.encoder_cache_manager.allocate(request, i) |
| 865 | if self.ec_connector is not None: |
| 866 | self.ec_connector.update_state_after_alloc(request, i) |
| 867 | encoder_compute_budget = new_encoder_compute_budget |
| 868 | if external_load_encoder_input: |
| 869 | for i in external_load_encoder_input: |
| 870 | self.encoder_cache_manager.allocate(request, i) |
| 871 | if self.ec_connector is not None: |
| 872 | self.ec_connector.update_state_after_alloc(request, i) |
| 873 | |
| 874 | # Record the LoRAs in scheduled_running_reqs |
| 875 | scheduled_loras: set[int] = set() |
| 876 | if self.lora_config: |
| 877 | scheduled_loras = set( |
| 878 | req.lora_request.lora_int_id |
| 879 | for req in scheduled_running_reqs |
| 880 | if req.lora_request and req.lora_request.lora_int_id > 0 |
| 881 | ) |
| 882 | assert len(scheduled_loras) <= self.lora_config.max_loras |
| 883 | |
| 884 | # Next, schedule the WAITING requests. |
| 885 | if not preempted_reqs and self._pause_state == PauseState.UNPAUSED: |
| 886 | step_skipped_waiting: deque[Request] = deque() |
| 887 | step_skipped_kv_holding: deque[Request] = deque() |
| 888 | |
| 889 | def skip_request(from_queue: RequestQueue) -> None: |
| 890 | request_to_skip = from_queue.pop_request() |
| 891 | self.deferred_waiting.add(request_to_skip) |
| 892 | if self._holds_kv_blocks(request_to_skip): |
| 893 | step_skipped_kv_holding.appendleft(request_to_skip) |
| 894 | else: |
| 895 | step_skipped_waiting.appendleft(request_to_skip) |
| 896 | |
| 897 | while token_budget > 0: |
| 898 | # Requests holding KV blocks are always drained first. |
| 899 | request_queue = self.kv_holding_waiting or self.waiting |
| 900 | if not request_queue or input_budget <= draft_slots: |
| 901 | break |
| 902 | # Paused streaming sessions (WAITING_FOR_STREAMING_REQ) are not |
| 903 | # in `running` but still hold a model-runner request slot. |
| 904 | num_running = len(self.running) + self.num_waiting_for_streaming_input |
| 905 | if num_running >= self.max_num_active_reqs: |
| 906 | break |
| 907 | |
| 908 | request = request_queue.peek_request() |
| 909 | request_id = request.request_id |
| 910 | |
| 911 | ready_to_schedule = self._handle_blocked_waiting_request(request) |
| 912 | if not ready_to_schedule: |
| 913 | skip_request(request_queue) |
| 914 | continue |
| 915 | |
| 916 | if request.num_stale_output_tokens and not request.drop_stale_output: |
| 917 | # Deliverable stale output still in flight: resuming now |
| 918 | # could resample a position that output later delivers. |
| 919 | # It drains within the pipeline depth. |
| 920 | skip_request(request_queue) |
| 921 | continue |
| 922 | |
| 923 | # Check that adding the request still respects the max_loras constraint. |
| 924 | if ( |
| 925 | self.lora_config |
| 926 | and request.lora_request |
| 927 | and ( |
| 928 | len(scheduled_loras) == self.lora_config.max_loras |
| 929 | and request.lora_request.lora_int_id not in scheduled_loras |
| 930 | ) |
| 931 | ): |
| 932 | # Scheduling would exceed max_loras, skip. |
| 933 | skip_request(request_queue) |
| 934 | continue |
| 935 | |
| 936 | num_external_computed_tokens = 0 |
| 937 | load_kv_async = False |
| 938 | connector_prefix_cache_queries, connector_prefix_cache_hits = 0, 0 |
| 939 | did_prefix_cache_lookup = False |
| 940 | |
| 941 | # Get already-cached tokens. |
| 942 | if request.num_computed_tokens == 0: |
| 943 | did_prefix_cache_lookup = True |
| 944 | ( |
| 945 | new_computed_blocks, |
| 946 | num_new_local_computed_tokens, |
| 947 | request.shared_prefix_boundary, |
| 948 | hit_diverged, |
| 949 | ) = self._get_local_prefix_cache_hit(request) |
| 950 | |
| 951 | # Get externally-cached tokens if using a KVConnector. |
| 952 | if self.connector is not None: |
| 953 | # Present a block-aligned local hit to the connector so |
| 954 | # a strictly longer remote hit can supersede a local |
| 955 | # sub-block tail without racing its copy-on-write. |
| 956 | partial_tail = num_new_local_computed_tokens % self.block_size |
| 957 | block_aligned_local = ( |
| 958 | num_new_local_computed_tokens - partial_tail |
| 959 | ) |
| 960 | ext_tokens, load_kv_async = ( |
| 961 | self.connector.get_num_new_matched_tokens( |
| 962 | request, block_aligned_local |
| 963 | ) |
| 964 | ) |
| 965 | if request.skip_reading_prefix_cache: |
| 966 | ext_tokens, load_kv_async = 0, False |
| 967 | |
| 968 | if ext_tokens is None: |
| 969 | # The request cannot be scheduled because |
| 970 | # the KVConnector couldn't determine |
| 971 | # the number of matched tokens. |
| 972 | skip_request(request_queue) |
| 973 | continue |
| 974 | |
| 975 | if self.prefix_replay_tokens: |
| 976 | # SWA bounded replay recomputes the hit's last |
| 977 | # window, from the block holding its first |
| 978 | # token; the sliding-window groups retire |
| 979 | # whole blocks below the window of the hit's |
| 980 | # next token. Hits end on a block boundary, or |
| 981 | # a hit ending one token short of one would |
| 982 | # retire the block the replay starts in. |
| 983 | ext_tokens -= ext_tokens % self.block_size |
| 984 | load_kv_async = load_kv_async and ext_tokens > 0 |
| 985 | |
| 986 | if partial_tail and ext_tokens > partial_tail: |
| 987 | # Remote strictly exceeds the full local hit: drop the |
| 988 | # sub-block tail so no CoW is needed, and let the load |
| 989 | # cover it. Trim the partial block out of the local |
| 990 | # computed blocks so it is not adopted from the cache. |
| 991 | new_computed_blocks = ( |
| 992 | self.kv_cache_manager.truncate_computed_blocks( |
| 993 | new_computed_blocks, block_aligned_local |
| 994 | ) |
| 995 | ) |
| 996 | num_new_local_computed_tokens = block_aligned_local |
| 997 | num_external_computed_tokens = ext_tokens |
| 998 | elif partial_tail: |
| 999 | # Remote does not exceed the full local hit: keep the |
| 1000 | # local sub-block tail and load nothing external. |
| 1001 | num_external_computed_tokens = 0 |
| 1002 | # Nothing to load remotely -> not an async-load step; |
| 1003 | # clearing avoids the `load_kv_async` assert below. |
| 1004 | load_kv_async = False |
| 1005 | else: |
| 1006 | num_external_computed_tokens = ext_tokens |
| 1007 | |
| 1008 | if hit_diverged and num_external_computed_tokens == 0: |
| 1009 | # No external tokens back the deeper local hit, so its |
| 1010 | # resume boundary would have no valid Mamba state. |
| 1011 | # Reconcile to the boundary every group agrees on. |
| 1012 | ( |
| 1013 | new_computed_blocks, |
| 1014 | num_new_local_computed_tokens, |
| 1015 | request.shared_prefix_boundary, |
| 1016 | ) = self.kv_cache_manager.get_computed_blocks(request) |
| 1017 | |
| 1018 | connector_prefix_cache_queries = ( |
| 1019 | request.num_tokens - num_new_local_computed_tokens |
| 1020 | ) |
| 1021 | connector_prefix_cache_hits = num_external_computed_tokens |
| 1022 | |
| 1023 | # Total computed tokens (local + external). |
| 1024 | num_computed_tokens = ( |
| 1025 | num_new_local_computed_tokens + num_external_computed_tokens |
| 1026 | ) |
| 1027 | assert num_computed_tokens <= request.num_tokens |
| 1028 | if 0 < num_computed_tokens <= self.prefix_replay_tokens: |
| 1029 | # SWA bounded replay: a hit no longer than the replayed |
| 1030 | # window would be recomputed in full and save nothing. |
| 1031 | new_computed_blocks = ( |
| 1032 | self.kv_cache_manager.empty_kv_cache_blocks |
| 1033 | ) |
| 1034 | num_new_local_computed_tokens = 0 |
| 1035 | num_external_computed_tokens = 0 |
| 1036 | connector_prefix_cache_hits = 0 |
| 1037 | num_computed_tokens = 0 |
| 1038 | load_kv_async = False |
| 1039 | |
| 1040 | # Skip request with pending mm encoding prefetches |
| 1041 | if self._ec_transfer_pending(request, num_computed_tokens): |
| 1042 | skip_request(request_queue) |
| 1043 | continue |
| 1044 | |
| 1045 | # Track first scheduled prefill, not post-preemption repeat prefills |
| 1046 | if request.prefill_stats and request.num_preemptions <= 0: |
| 1047 | assert num_computed_tokens <= request.num_prompt_tokens |
| 1048 | request.prefill_stats.set( |
| 1049 | num_prompt_tokens=request.num_prompt_tokens, |
| 1050 | num_local_cached_tokens=num_new_local_computed_tokens, |
| 1051 | num_external_cached_tokens=num_external_computed_tokens, |
| 1052 | ) |
| 1053 | self._set_kv_fetch_stage( |
| 1054 | request, KV_FETCH_WAITING_TO_START if load_kv_async else None |
| 1055 | ) |
| 1056 | else: |
| 1057 | # KVTransfer: WAITING reqs have num_computed_tokens > 0 |
| 1058 | # after async KV recvs are completed. A streaming-input |
| 1059 | # session resumes here too, carrying whatever media its |
| 1060 | # latest chunk added, so this branch needs the same gate. |
| 1061 | new_computed_blocks = self.kv_cache_manager.empty_kv_cache_blocks |
| 1062 | num_new_local_computed_tokens = 0 |
| 1063 | num_computed_tokens = request.num_computed_tokens |
| 1064 | |
| 1065 | if self._ec_transfer_pending(request, num_computed_tokens): |
| 1066 | skip_request(request_queue) |
| 1067 | continue |
| 1068 | |
| 1069 | encoder_inputs_to_schedule = None |
| 1070 | external_load_encoder_input = [] |
| 1071 | new_encoder_compute_budget = encoder_compute_budget |
| 1072 | pad_spec_decode = False |
| 1073 | |
| 1074 | # SWA bounded replay: recompute the tail of the hit without |
| 1075 | # rewriting its cached KV. An async load replays once the KV |
| 1076 | # has arrived (_update_waiting_for_remote_kv). |
| 1077 | num_replay_tokens = 0 |
| 1078 | |
| 1079 | if load_kv_async: |
| 1080 | # KVTransfer: loading remote KV, do not allocate for new work. |
| 1081 | assert num_external_computed_tokens > 0 |
| 1082 | num_new_tokens = 0 |
| 1083 | elif defer_prefills and num_computed_tokens < request.num_tokens - 1: |
| 1084 | # DP prefill balancing: defer this step's local prefill |
| 1085 | # compute to a cadence-aligned step. |
| 1086 | break |
| 1087 | else: |
| 1088 | request_token_budget = min(token_budget, input_budget - draft_slots) |
| 1089 | if did_prefix_cache_lookup: |
| 1090 | # Fresh admission: local and/or sync external hit. |
| 1091 | num_replay_tokens = self._mark_prefix_replay( |
| 1092 | request, num_computed_tokens |
| 1093 | ) |
| 1094 | num_computed_tokens -= num_replay_tokens |
| 1095 | # Number of tokens to be scheduled. |
| 1096 | # We use `request.num_tokens` instead of |
| 1097 | # `request.num_prompt_tokens` to consider the resumed |
| 1098 | # requests, which have output tokens. |
| 1099 | num_new_tokens = request.num_tokens - num_computed_tokens |
| 1100 | |
| 1101 | # Pad new decode requests to uniform spec decoding size to |
| 1102 | # preserve full cudagraph for this step. |
| 1103 | # Not for diffusion where draft tokens can't be padded. |
| 1104 | if ( |
| 1105 | (self.num_spec_tokens > 0 and self.dynamic_sd_lookup is None) |
| 1106 | and self.num_sampled_tokens_per_step > 0 |
| 1107 | and num_new_tokens == 1 |
| 1108 | and not prefill_scheduled |
| 1109 | and (scheduled_running_reqs or num_computed_tokens > 0) |
| 1110 | ): |
| 1111 | padded_num_tokens = 1 + self.num_spec_tokens |
| 1112 | # Pad only when there is room for the sampled token(s). |
| 1113 | if ( |
| 1114 | num_computed_tokens |
| 1115 | + padded_num_tokens |
| 1116 | + self.num_sampled_tokens_per_step |
| 1117 | <= self.max_model_len |
| 1118 | ): |
| 1119 | if padded_num_tokens > request_token_budget: |
| 1120 | # Prefer to not schedule than schedule un-padded. |
| 1121 | break |
| 1122 | num_new_tokens = padded_num_tokens |
| 1123 | pad_spec_decode = True |
| 1124 | |
| 1125 | if 0 < long_prefill_token_threshold < num_new_tokens: |
| 1126 | num_new_tokens = long_prefill_token_threshold |
| 1127 | |
| 1128 | # chunked prefill has to be enabled explicitly to allow |
| 1129 | # pooling requests to be chunked |
| 1130 | if ( |
| 1131 | not self.scheduler_config.enable_chunked_prefill |
| 1132 | and num_new_tokens > request_token_budget |
| 1133 | and not self.is_mm_encoder_only |
| 1134 | ): |
| 1135 | # If chunked_prefill is disabled, |
| 1136 | # we can stop the scheduling here. |
| 1137 | break |
| 1138 | |
| 1139 | num_new_tokens = min(num_new_tokens, request_token_budget) |
| 1140 | assert num_new_tokens > 0 |
| 1141 | |
| 1142 | # Apply Mamba alignment before encoder caps. |
| 1143 | if self.need_mamba_block_aligned_split: |
| 1144 | num_new_tokens = self._mamba_block_aligned_split( |
| 1145 | request, |
| 1146 | num_new_tokens, |
| 1147 | num_new_local_computed_tokens, |
| 1148 | num_external_computed_tokens, |
| 1149 | ) |
| 1150 | if num_new_tokens == 0: |
| 1151 | break |
| 1152 | if ( |
| 1153 | pad_spec_decode |
| 1154 | and num_new_tokens != 1 + self.num_spec_tokens |
| 1155 | ): |
| 1156 | # Alignment clipped the placeholder rows. The split |
| 1157 | # aligns prefill chunks, but the padded tail rows are |
| 1158 | # speculative positions, not prefill tokens. A padded |
| 1159 | # request must keep all 1 + num_spec rows or the |
| 1160 | # sampler's row count stops matching its query rows, |
| 1161 | # so drop the padding instead of shortening it. |
| 1162 | num_new_tokens = 1 |
| 1163 | pad_spec_decode = False |
| 1164 | |
| 1165 | # Schedule encoder inputs. |
| 1166 | if request.has_encoder_inputs: |
| 1167 | ( |
| 1168 | encoder_inputs_to_schedule, |
| 1169 | num_new_tokens, |
| 1170 | new_encoder_compute_budget, |
| 1171 | external_load_encoder_input, |
| 1172 | ) = self._try_schedule_encoder_inputs( |
| 1173 | request, |
| 1174 | num_computed_tokens, |
| 1175 | num_new_tokens, |
| 1176 | encoder_compute_budget, |
| 1177 | shift_computed_tokens=self.num_prefill_lookahead, |
| 1178 | ) |
| 1179 | |
| 1180 | # Multi-module MTP: avoid ending a prefill chunk within |
| 1181 | # num_prefill_lookahead of the prefill end. |
| 1182 | num_new_tokens = self._reserve_prefill_lookahead( |
| 1183 | request, num_computed_tokens, num_new_tokens |
| 1184 | ) |
| 1185 | |
| 1186 | if num_new_tokens == 0: |
| 1187 | # The request cannot be scheduled. |
| 1188 | break |
| 1189 | |
| 1190 | # During async KV load, no forward pass is run yet. |
| 1191 | # Allocate speculative lookahead slots later to avoid |
| 1192 | # mismatching local and remote block counts. |
| 1193 | limit_lookahead_tokens = load_kv_async and self.num_lookahead_tokens > 0 |
| 1194 | effective_lookahead_tokens = ( |
| 1195 | 0 if limit_lookahead_tokens else self.num_lookahead_tokens |
| 1196 | ) |
| 1197 | |
| 1198 | # Determine if we need to allocate cross-attention blocks. |
| 1199 | num_encoder_tokens = 0 |
| 1200 | if ( |
| 1201 | self.is_encoder_decoder |
| 1202 | and request.has_encoder_inputs |
| 1203 | and encoder_inputs_to_schedule |
| 1204 | ): |
| 1205 | num_encoder_tokens = sum( |
| 1206 | request.get_num_encoder_embeds(i) |
| 1207 | for i in encoder_inputs_to_schedule |
| 1208 | ) |
| 1209 | |
| 1210 | reserved_blocks = 0 |
| 1211 | if load_kv_async: |
| 1212 | # An async load holds its blocks for the whole transfer with |
| 1213 | # no forward progress and isn't preemptible here. Admit it |
| 1214 | # only if it fits in (free - other in-flight reservations) |
| 1215 | # plus its own spec decode step blocks, to avoid deadlock and |
| 1216 | # predictable preemptions. |
| 1217 | reserved_blocks = ( |
| 1218 | self._inflight_prefill_reserved_blocks() |
| 1219 | + self._spec_decode_step_blocks() |
| 1220 | ) |
| 1221 | |
| 1222 | # Replayed tokens are already counted in the adopted hit; a |
| 1223 | # chunk that ends inside the replayed range needs no new slots. |
| 1224 | num_tokens_past_hit = max(num_new_tokens - num_replay_tokens, 0) |
| 1225 | new_blocks = self.kv_cache_manager.allocate_slots( |
| 1226 | request, |
| 1227 | num_tokens_past_hit, |
| 1228 | num_new_computed_tokens=num_new_local_computed_tokens, |
| 1229 | new_computed_blocks=new_computed_blocks, |
| 1230 | num_lookahead_tokens=effective_lookahead_tokens, |
| 1231 | num_external_computed_tokens=num_external_computed_tokens, |
| 1232 | delay_cache_blocks=load_kv_async, |
| 1233 | num_encoder_tokens=num_encoder_tokens, |
| 1234 | full_sequence_must_fit=self.scheduler_reserve_full_isl, |
| 1235 | reserved_blocks=reserved_blocks, |
| 1236 | has_scheduled_reqs=bool(self.running), |
| 1237 | ) |
| 1238 | |
| 1239 | if new_blocks is None: |
| 1240 | # The request cannot be scheduled. |
| 1241 | |
| 1242 | # NOTE: we need to untouch the request from the encode cache |
| 1243 | # manager |
| 1244 | if request.has_encoder_inputs: |
| 1245 | self.encoder_cache_manager.free(request) |
| 1246 | break |
| 1247 | |
| 1248 | # KVTransfer: the connector uses this info to determine |
| 1249 | # if a load is needed. Note that |
| 1250 | # This information is used to determine if a load is |
| 1251 | # needed for this request. |
| 1252 | if self.connector is not None: |
| 1253 | self.connector.update_state_after_alloc( |
| 1254 | request, |
| 1255 | self.kv_cache_manager.get_blocks(request_id), |
| 1256 | num_external_computed_tokens, |
| 1257 | ) |
| 1258 | if ( |
| 1259 | self.connector_prefix_cache_stats is not None |
| 1260 | and connector_prefix_cache_queries != 0 |
| 1261 | ): |
| 1262 | self.connector_prefix_cache_stats.record( |
| 1263 | num_tokens=connector_prefix_cache_queries, |
| 1264 | num_hits=connector_prefix_cache_hits, |
| 1265 | preempted=request.num_preemptions > 0, |
| 1266 | ) |
| 1267 | |
| 1268 | # Record at admission so unscheduled lookups are not counted. |
| 1269 | if did_prefix_cache_lookup: |
| 1270 | self.kv_cache_manager.record_prefix_cache_stats( |
| 1271 | request, num_new_local_computed_tokens |
| 1272 | ) |
| 1273 | |
| 1274 | if load_kv_async: |
| 1275 | # If loading async, allocate memory and put request |
| 1276 | # into the WAITING_FOR_REMOTE_KV state. |
| 1277 | request.status = RequestStatus.WAITING_FOR_REMOTE_KVS |
| 1278 | # Set num_computed_tokens even though KVs are not yet loaded. |
| 1279 | # request.num_computed_tokens will not be used anywhere until |
| 1280 | # the request finished the KV transfer. |
| 1281 | # |
| 1282 | # If a transfer error is reported by the connector, |
| 1283 | # request.num_computed_tokens will be re-set accordingly in |
| 1284 | # _update_requests_with_invalid_blocks. |
| 1285 | # |
| 1286 | # When the transfer is finished, either successfully or not, |
| 1287 | # request.num_computed_tokens will correctly reflect the number |
| 1288 | # of computed tokens. |
| 1289 | # _update_waiting_for_remote_kv will then cache |
| 1290 | # only the successfully loaded tokens. |
| 1291 | request.num_computed_tokens = num_computed_tokens |
| 1292 | self._inflight_prefills.add(request) |
| 1293 | self._set_kv_fetch_stage(request, KV_FETCH_IN_PROGRESS) |
| 1294 | if self.needs_kv_cache_zeroing: |
| 1295 | # Skip zeroing of the blocks the async load will |
| 1296 | # overwrite; the zeroing could race the write. |
| 1297 | self._skip_zero_block_ids.update( |
| 1298 | self.kv_cache_manager.get_zeroing_block_ids_in_range( |
| 1299 | request_id, |
| 1300 | num_new_local_computed_tokens, |
| 1301 | num_computed_tokens, |
| 1302 | ) |
| 1303 | ) |
| 1304 | skip_request(request_queue) |
| 1305 | continue |
| 1306 | |
| 1307 | request = request_queue.pop_request() |
| 1308 | self.deferred_waiting.discard(request) |
| 1309 | self.running.append(request) |
| 1310 | if num_external_computed_tokens > 0: |
| 1311 | # load_kv_async is False here |
| 1312 | has_sync_kv_loads = True |
| 1313 | if self.log_stats: |
| 1314 | request.record_event( |
| 1315 | EngineCoreEventType.SCHEDULED, scheduled_timestamp |
| 1316 | ) |
| 1317 | if request.status == RequestStatus.WAITING: |
| 1318 | scheduled_new_reqs.append(request) |
| 1319 | elif request.status == RequestStatus.PREEMPTED: |
| 1320 | scheduled_resumed_reqs.append(request) |
| 1321 | else: |
| 1322 | raise RuntimeError(f"Invalid request status: {request.status}") |
| 1323 | |
| 1324 | if self.lora_config and request.lora_request: |
| 1325 | scheduled_loras.add(request.lora_request.lora_int_id) |
| 1326 | req_to_new_blocks[request_id] = self.kv_cache_manager.get_blocks( |
| 1327 | request_id |
| 1328 | ) |
| 1329 | num_scheduled_tokens[request_id] = num_new_tokens |
| 1330 | token_budget -= num_new_tokens |
| 1331 | input_budget -= num_new_tokens + draft_slots |
| 1332 | request.status = RequestStatus.RUNNING |
| 1333 | request.num_computed_tokens = num_computed_tokens |
| 1334 | self._set_kv_fetch_stage(request, None) |
| 1335 | if pad_spec_decode: |
| 1336 | assert num_new_tokens == 1 + self.num_spec_tokens |
| 1337 | scheduled_spec_decode_tokens[request_id] = [ |
| 1338 | -1 |
| 1339 | ] * self.num_spec_tokens |
| 1340 | # Only track requests that will still be prefilling after this chunk. |
| 1341 | if num_computed_tokens + num_new_tokens < request.num_tokens: |
| 1342 | self._inflight_prefills.add(request) |
| 1343 | # Encoder-related. |
| 1344 | if encoder_inputs_to_schedule: |
| 1345 | scheduled_encoder_inputs[request_id] = encoder_inputs_to_schedule |
| 1346 | # Allocate the encoder cache. |
| 1347 | for i in encoder_inputs_to_schedule: |
| 1348 | self.encoder_cache_manager.allocate(request, i) |
| 1349 | if self.ec_connector is not None: |
| 1350 | self.ec_connector.update_state_after_alloc(request, i) |
| 1351 | encoder_compute_budget = new_encoder_compute_budget |
| 1352 | # Allocate for external load encoder cache |
| 1353 | if external_load_encoder_input: |
| 1354 | for i in external_load_encoder_input: |
| 1355 | self.encoder_cache_manager.allocate(request, i) |
| 1356 | if self.ec_connector is not None: |
| 1357 | self.ec_connector.update_state_after_alloc(request, i) |
| 1358 | |
| 1359 | # re-queue requests skipped in this pass ahead of older skipped items. |
| 1360 | if step_skipped_kv_holding: |
| 1361 | self.kv_holding_waiting.prepend_requests(step_skipped_kv_holding) |
| 1362 | if step_skipped_waiting: |
| 1363 | self.waiting.prepend_requests(step_skipped_waiting) |
| 1364 | |
| 1365 | # DP prefill balancing: on a step that admitted prefills (release), |
| 1366 | # record whether it was capacity-bound. |
| 1367 | if not defer_prefills: |
| 1368 | _, num_waiting = self.get_request_counts() |
| 1369 | self.prefill_capacity_bound = num_waiting > len(self.deferred_waiting) |
| 1370 | |
| 1371 | # Check if the scheduling constraints are satisfied. |
| 1372 | total_num_scheduled_tokens = sum(num_scheduled_tokens.values()) |
| 1373 | assert total_num_scheduled_tokens <= self.max_num_scheduled_tokens |
| 1374 | |
| 1375 | assert token_budget >= 0 |
| 1376 | assert input_budget >= 0 |
| 1377 | assert len(self.running) <= self.max_num_running_reqs |
| 1378 | # Since some requests in the RUNNING queue may not be scheduled in |
| 1379 | # this step, the total number of scheduled requests can be smaller than |
| 1380 | # len(self.running). |
| 1381 | assert len(scheduled_new_reqs) + len(scheduled_resumed_reqs) + len( |
| 1382 | scheduled_running_reqs |
| 1383 | ) <= len(self.running) |
| 1384 | |
| 1385 | # Get the longest common prefix among all requests in the running queue. |
| 1386 | # This can be potentially used for cascade attention. |
| 1387 | num_common_prefix_blocks = [0] * len(self.kv_cache_config.kv_cache_groups) |
| 1388 | with record_function_or_nullcontext("schedule: get_num_common_prefix_blocks"): |
| 1389 | if self.running: |
| 1390 | any_request_id = self.running[0].request_id |
| 1391 | num_common_prefix_blocks = ( |
| 1392 | self.kv_cache_manager.get_num_common_prefix_blocks(any_request_id) |
| 1393 | ) |
| 1394 | |
| 1395 | # Construct the scheduler output. |
| 1396 | if self.use_v2_model_runner: |
| 1397 | scheduled_new_reqs.extend(scheduled_resumed_reqs) |
| 1398 | scheduled_resumed_reqs.clear() |
| 1399 | new_reqs_data = [ |
| 1400 | NewRequestData.from_request( |
| 1401 | req, |
| 1402 | req_to_new_blocks[req.request_id].get_block_ids(), |
| 1403 | req._all_token_ids, |
| 1404 | uses_mrope=self.model_uses_mrope, |
| 1405 | ) |
| 1406 | for req in scheduled_new_reqs |
| 1407 | ] |
| 1408 | else: |
| 1409 | new_reqs_data = [ |
| 1410 | NewRequestData.from_request( |
| 1411 | req, |
| 1412 | req_to_new_blocks[req.request_id].get_block_ids(), |
| 1413 | uses_mrope=self.model_uses_mrope, |
| 1414 | ) |
| 1415 | for req in scheduled_new_reqs |
| 1416 | ] |
| 1417 | |
| 1418 | with record_function_or_nullcontext("schedule: make_cached_request_data"): |
| 1419 | cached_reqs_data = self._make_cached_request_data( |
| 1420 | scheduled_running_reqs, |
| 1421 | scheduled_resumed_reqs, |
| 1422 | num_scheduled_tokens, |
| 1423 | scheduled_spec_decode_tokens, |
| 1424 | req_to_new_blocks, |
| 1425 | ) |
| 1426 | |
| 1427 | # Record the request ids that were scheduled in this step (MRV1-only). |
| 1428 | if not self.use_v2_model_runner: |
| 1429 | self.prev_step_scheduled_req_ids.clear() |
| 1430 | self.prev_step_scheduled_req_ids.update(num_scheduled_tokens.keys()) |
| 1431 | |
| 1432 | # Mamba "align" boundary states must be handed off with exact block ids; |
| 1433 | # they cannot be reconstructed from a connector's append-only block |
| 1434 | # table. Drained every step so stale offers cannot accumulate. |
| 1435 | boundary_state_offloads = self.kv_cache_manager.take_boundary_state_offloads() |
| 1436 | |
| 1437 | kv_connector_block_state = None |
| 1438 | if self.connector is not None: |
| 1439 | # Any request scheduled this step can become a connector job now, |
| 1440 | # not only the ones that were allocated blocks: a store save lands |
| 1441 | # on the step that fills a block, which allocated none. |
| 1442 | block_state_req_ids = set(num_scheduled_tokens) |
| 1443 | block_state_req_ids.update( |
| 1444 | req_id for req_id in boundary_state_offloads if req_id in self.requests |
| 1445 | ) |
| 1446 | kv_connector_block_state = KVConnectorBlockState( |
| 1447 | req_ids=block_state_req_ids, |
| 1448 | resolve_block_ids=self.kv_cache_manager.get_block_ids, |
| 1449 | boundary_state_offloads=boundary_state_offloads, |
| 1450 | ) |
| 1451 | |
| 1452 | kv_cache_block_copies, cow_retained_blocks = ( |
| 1453 | self.kv_cache_manager.take_kv_cache_block_copies() |
| 1454 | ) |
| 1455 | if kv_cache_block_copies: |
| 1456 | # The copies run with this step's execution; the first non-empty |
| 1457 | # step at or after it gets seq `sched_step_seq + 1` (0-token steps |
| 1458 | # do not advance the seq), and its completion implies the copies |
| 1459 | # have run. |
| 1460 | self._free_cow_retained_blocks(cow_retained_blocks, self.sched_step_seq + 1) |
| 1461 | pending_kv_cache_block_copies = kv_cache_block_copies or None |
| 1462 | |
| 1463 | # Dynamic speculative decoding: compute optimal K |
| 1464 | num_spec_tokens_to_schedule = self.num_spec_tokens |
| 1465 | if self.dynamic_sd_lookup is not None and len(num_scheduled_tokens) > 0: |
| 1466 | num_spec_tokens_to_schedule = self.dynamic_sd_lookup[ |
| 1467 | len(num_scheduled_tokens) |
| 1468 | ] |
| 1469 | |
| 1470 | scheduled_encoder_input_stats = None |
| 1471 | if ( |
| 1472 | self.log_stats |
| 1473 | and self.observability_config.enable_logging_iteration_details |
| 1474 | ): |
| 1475 | scheduled_encoder_input_stats = self._make_scheduled_encoder_input_stats( |
| 1476 | scheduled_encoder_inputs |
| 1477 | ) |
| 1478 | |
| 1479 | scheduler_output = SchedulerOutput( |
| 1480 | scheduled_new_reqs=new_reqs_data, |
| 1481 | scheduled_cached_reqs=cached_reqs_data, |
| 1482 | num_scheduled_tokens=num_scheduled_tokens, |
| 1483 | total_num_scheduled_tokens=total_num_scheduled_tokens, |
| 1484 | scheduled_spec_decode_tokens=scheduled_spec_decode_tokens, |
| 1485 | scheduled_encoder_inputs=scheduled_encoder_inputs, |
| 1486 | scheduled_encoder_input_stats=scheduled_encoder_input_stats, |
| 1487 | num_common_prefix_blocks=num_common_prefix_blocks, |
| 1488 | preempted_req_ids=self.reset_preempted_req_ids, |
| 1489 | # finished_req_ids is an existing state in the scheduler, |
| 1490 | # instead of being newly scheduled in this step. |
| 1491 | # It contains the request IDs that are finished in between |
| 1492 | # the previous and the current steps. |
| 1493 | finished_req_ids=self.finished_req_ids, |
| 1494 | free_encoder_mm_hashes=self.encoder_cache_manager.get_freed_mm_hashes(), |
| 1495 | new_block_ids_to_zero=self._get_new_block_ids_to_zero(), |
| 1496 | has_sync_kv_loads=has_sync_kv_loads, |
| 1497 | kv_cache_block_copies=pending_kv_cache_block_copies, |
| 1498 | kv_connector_block_state=kv_connector_block_state, |
| 1499 | num_spec_tokens_to_schedule=num_spec_tokens_to_schedule, |
| 1500 | ec_manager_metadata=self.encoder_cache_manager.get_manager_metadata(), |
| 1501 | ) |
| 1502 | |
| 1503 | # NOTE(Kuntai): this function is designed for multiple purposes: |
| 1504 | # 1. Plan the KV cache store |
| 1505 | # 2. Wrap up all the KV cache load / save ops into an opaque object |
| 1506 | # 3. Clear the internal states of the connector |
| 1507 | if self.connector is not None: |
| 1508 | meta = self._build_kv_connector_meta(self.connector, scheduler_output) |
| 1509 | scheduler_output.kv_connector_metadata = meta |
| 1510 | |
| 1511 | if self.aux_output_connector is not None: |
| 1512 | scheduler_output.aux_output_connector_metadata = ( |
| 1513 | self.aux_output_connector.build_connector_meta( |
| 1514 | scheduler_output, self.requests |
| 1515 | ) |
| 1516 | ) |
| 1517 | |
| 1518 | # Build the connector meta for ECConnector |
| 1519 | if self.ec_connector is not None: |
| 1520 | ec_meta: ECConnectorMetadata = self.ec_connector.build_connector_meta( |
| 1521 | scheduler_output |
| 1522 | ) |
| 1523 | scheduler_output.ec_connector_metadata = ec_meta |
| 1524 | |
| 1525 | # Connector-only block state must not be dispatched to workers. |
| 1526 | scheduler_output.kv_connector_block_state = None |
| 1527 | |
| 1528 | # Advance the fence only for non-empty steps (those that actually |
| 1529 | # write KV and have their output processed later in update_from_output). |
| 1530 | if self.defer_block_free and total_num_scheduled_tokens > 0: |
| 1531 | self.sched_step_seq += 1 |
| 1532 | |
| 1533 | with record_function_or_nullcontext("schedule: update_after_schedule"): |
| 1534 | self._update_after_schedule(scheduler_output) |
| 1535 | return scheduler_output |
| 1536 | |
| 1537 | def _build_kv_connector_meta( |
| 1538 | self, connector: KVConnectorBase_V1, scheduler_output: SchedulerOutput |
| 1539 | ) -> KVConnectorMetadata: |
| 1540 | return connector.build_connector_meta(scheduler_output) |
| 1541 | |
| 1542 | def _get_new_block_ids_to_zero(self) -> list[int] | None: |
| 1543 | # Drain new attention block ids every step so the manager-side list |
| 1544 | # does not grow unbounded; only kv-cache zeroing consumes them. |
| 1545 | new_block_ids_to_zero = self.kv_cache_manager.take_new_block_ids() |
| 1546 | if not self.needs_kv_cache_zeroing: |
| 1547 | return None |
| 1548 | |
| 1549 | if self._skip_zero_block_ids: |
| 1550 | skip = self._skip_zero_block_ids |
| 1551 | new_block_ids_to_zero = [b for b in new_block_ids_to_zero if b not in skip] |
| 1552 | skip.clear() |
| 1553 | |
| 1554 | return new_block_ids_to_zero or None |
| 1555 | |
| 1556 | def _preempt_request( |
| 1557 | self, request: Request, timestamp: float, drop_stale_output: bool = False |
| 1558 | ) -> None: |
| 1559 | """Preempt a request and put it back to the waiting queue. |
| 1560 | |
| 1561 | NOTE: The request should be popped from the running queue outside of this |
| 1562 | method. |
| 1563 | |
| 1564 | drop_stale_output: drop (rather than deliver) any in-flight output; used |
| 1565 | by reset_prefix_cache, whose same-step resume would otherwise deliver |
| 1566 | tokens out of order, and for connectors with a pending KV hand-off, |
| 1567 | which the preemption's block free would leave without valid KV. |
| 1568 | """ |
| 1569 | assert request.status == RequestStatus.RUNNING, ( |
| 1570 | "Only running requests can be preempted" |
| 1571 | ) |
| 1572 | if self.aux_output_connector is not None: |
| 1573 | self.aux_output_connector.request_finished(request) |
| 1574 | self._free_request_blocks(request) |
| 1575 | self.encoder_cache_manager.free(request) |
| 1576 | self._inflight_prefills.discard(request) |
| 1577 | request.status = RequestStatus.PREEMPTED |
| 1578 | request.num_computed_tokens = 0 |
| 1579 | if request.spec_token_ids: |
| 1580 | request.spec_token_ids = [] |
| 1581 | # Async scheduling: mark all in-flight output as stale. Its tokens are |
| 1582 | # still delivered on return (dropping them would perturb spec-decode |
| 1583 | # acceptance) but must not mutate the reset counters; each step drains |
| 1584 | # its share in update_from_output. num_in_flight_tokens already |
| 1585 | # includes any undrained stale share, so assign rather than accumulate. |
| 1586 | # An undrained drop-mode share stays dropped: its positions have |
| 1587 | # already been resampled. |
| 1588 | request.drop_stale_output = drop_stale_output or ( |
| 1589 | request.drop_stale_output and request.num_stale_output_tokens > 0 |
| 1590 | ) |
| 1591 | request.num_stale_output_tokens = request.num_in_flight_tokens |
| 1592 | request.num_output_placeholders = 0 |
| 1593 | request.num_preemptions += 1 |
| 1594 | if self.log_stats: |
| 1595 | request.record_event(EngineCoreEventType.PREEMPTED, timestamp) |
| 1596 | |
| 1597 | # Put the request back to the waiting queue. |
| 1598 | self.waiting.prepend_request(request) |
| 1599 | self.reset_preempted_req_ids.add(request.request_id) |
| 1600 | |
| 1601 | def _update_after_schedule(self, scheduler_output: SchedulerOutput) -> None: |
| 1602 | # Advance the number of computed tokens for the request AFTER |
| 1603 | # the request is scheduled. |
| 1604 | # 1. The scheduler_output of the current step has to include the |
| 1605 | # original number of scheduled tokens to determine input IDs. |
| 1606 | # 2. Advance the number of computed tokens here allowing us to |
| 1607 | # schedule the prefill request again immediately in the next |
| 1608 | # scheduling step. |
| 1609 | # 3. If some tokens (e.g. spec tokens) are rejected later, the number of |
| 1610 | # computed tokens will be adjusted in update_from_output. |
| 1611 | num_scheduled_tokens = scheduler_output.num_scheduled_tokens |
| 1612 | for req_id, num_scheduled_token in num_scheduled_tokens.items(): |
| 1613 | request = self.requests[req_id] |
| 1614 | request.num_computed_tokens += num_scheduled_token |
| 1615 | request.num_in_flight_tokens += num_scheduled_token |
| 1616 | if self.defer_block_free: |
| 1617 | # Record the in-flight step, to fence deferred block freeing. |
| 1618 | request.last_sched_seq = self.sched_step_seq |
| 1619 | request.is_prefill_chunk = request.num_computed_tokens < ( |
| 1620 | request.num_tokens + request.num_output_placeholders |
| 1621 | ) |
| 1622 | scheduler_output.has_structured_output_requests |= ( |
| 1623 | request.use_structured_output and not request.is_prefill_chunk |
| 1624 | ) |
| 1625 | # Drop from the in-flight-prefill set once it's no longer prefilling. |
| 1626 | if not request.is_prefill_chunk: |
| 1627 | self._inflight_prefills.discard(request) |
| 1628 | |
| 1629 | # Clear the finished and preempted request IDs. |
| 1630 | # NOTE: We shouldn't just clear() here because it will also affect |
| 1631 | # the scheduler output. |
| 1632 | self.finished_req_ids = set() |
| 1633 | self.reset_preempted_req_ids = set() |
| 1634 | |
| 1635 | def _update_request_as_session( |
| 1636 | self, session: Request, update: StreamingUpdate |
| 1637 | ) -> None: |
| 1638 | """Updates the waiting session with the next streaming update. |
| 1639 | |
| 1640 | Discards the last sampled output token from the prior input chunk. |
| 1641 | """ |
| 1642 | # Current streaming input behaviour: Keep only computed output tokens |
| 1643 | # (discard final sampled output token). |
| 1644 | num_computed_tokens = session.num_computed_tokens |
| 1645 | kept_output_tokens = session._all_token_ids[ |
| 1646 | session.num_prompt_tokens : num_computed_tokens |
| 1647 | ] |
| 1648 | del session._all_token_ids[num_computed_tokens:] |
| 1649 | session._output_token_ids.clear() |
| 1650 | assert session.prompt_token_ids is not None |
| 1651 | # Extend prompt with kept output tokens. |
| 1652 | session.prompt_token_ids.extend(kept_output_tokens) |
| 1653 | |
| 1654 | if update.mm_features: |
| 1655 | base = session.num_tokens |
| 1656 | for mm_feature in update.mm_features: |
| 1657 | mm_feature.mm_position = replace( |
| 1658 | mm_feature.mm_position, offset=mm_feature.mm_position.offset + base |
| 1659 | ) |
| 1660 | session.mm_features.extend(update.mm_features) |
| 1661 | |
| 1662 | session._all_token_ids.extend(update.prompt_token_ids or ()) |
| 1663 | session.prompt_token_ids.extend(update.prompt_token_ids or ()) |
| 1664 | # Update block hashes for the new tokens. |
| 1665 | session.update_block_hashes() |
| 1666 | session.num_prompt_tokens = len(session.prompt_token_ids) |
| 1667 | session.arrival_time = update.arrival_time |
| 1668 | session.sampling_params = update.sampling_params |
| 1669 | session.max_tokens = update.max_tokens |
| 1670 | if session.status == RequestStatus.WAITING_FOR_STREAMING_REQ: |
| 1671 | self.num_waiting_for_streaming_input -= 1 |
| 1672 | session.status = RequestStatus.WAITING |
| 1673 | |
| 1674 | if self.log_stats: |
| 1675 | session.record_event(EngineCoreEventType.QUEUED) |
| 1676 | |
| 1677 | def _make_cached_request_data( |
| 1678 | self, |
| 1679 | running_reqs: list[Request], |
| 1680 | resumed_reqs: list[Request], |
| 1681 | num_scheduled_tokens: dict[str, int], |
| 1682 | spec_decode_tokens: dict[str, list[int]], |
| 1683 | req_to_new_blocks: dict[str, KVCacheBlocks], |
| 1684 | ) -> CachedRequestData: |
| 1685 | req_ids: list[str] = [] |
| 1686 | new_token_ids: list[list[int]] = [] |
| 1687 | new_block_ids: list[tuple[list[int], ...] | None] = [] |
| 1688 | all_token_ids: dict[str, list[int]] = {} |
| 1689 | num_computed_tokens: list[int] = [] |
| 1690 | num_output_tokens: list[int] = [] |
| 1691 | resumed_req_ids = set() |
| 1692 | |
| 1693 | num_running_reqs = len(running_reqs) |
| 1694 | for idx, req in enumerate(itertools.chain(running_reqs, resumed_reqs)): |
| 1695 | req_id = req.request_id |
| 1696 | req_ids.append(req_id) |
| 1697 | # NOTE: In PP+async scheduling, we consume token ids via a direct GPU |
| 1698 | # broadcast path (`input_batch.prev_sampled_token_ids`), so we can |
| 1699 | # omit this payload. |
| 1700 | if self.use_pp and not self.scheduler_config.async_scheduling: |
| 1701 | # When using PP, the scheduler sends the sampled tokens back, |
| 1702 | # because there's no direct communication between the first- |
| 1703 | # stage worker and the last-stage worker. Otherwise, we don't |
| 1704 | # need to send the sampled tokens back because the model runner |
| 1705 | # will cache them. |
| 1706 | num_tokens = num_scheduled_tokens[req_id] - len( |
| 1707 | spec_decode_tokens.get(req_id, ()) |
| 1708 | ) |
| 1709 | token_ids = req.all_token_ids[ |
| 1710 | req.num_computed_tokens : req.num_computed_tokens + num_tokens |
| 1711 | ] |
| 1712 | new_token_ids.append(token_ids) |
| 1713 | if idx >= num_running_reqs: |
| 1714 | resumed_req_ids.add(req_id) |
| 1715 | if not self.use_v2_model_runner: # noqa: SIM102 |
| 1716 | if req_id not in self.prev_step_scheduled_req_ids: |
| 1717 | all_token_ids[req_id] = req.all_token_ids.copy() |
| 1718 | new_block_ids.append( |
| 1719 | req_to_new_blocks[req_id].get_block_ids(allow_none=True) |
| 1720 | ) |
| 1721 | num_computed_tokens.append(req.num_computed_tokens) |
| 1722 | num_output_tokens.append( |
| 1723 | req.num_output_tokens + req.num_output_placeholders |
| 1724 | ) |
| 1725 | |
| 1726 | return CachedRequestData( |
| 1727 | req_ids=req_ids, |
| 1728 | resumed_req_ids=resumed_req_ids, |
| 1729 | new_token_ids=new_token_ids, |
| 1730 | all_token_ids=all_token_ids, |
| 1731 | new_block_ids=new_block_ids, |
| 1732 | num_computed_tokens=num_computed_tokens, |
| 1733 | num_output_tokens=num_output_tokens, |
| 1734 | ) |
| 1735 | |
| 1736 | def _reject_on_encoder_cache_embed_mismatch( |
| 1737 | self, |
| 1738 | request: Request, |
| 1739 | input_id: int, |
| 1740 | identifier: str, |
| 1741 | cached_num_encoder_embeds: int, |
| 1742 | requested_num_encoder_embeds: int, |
| 1743 | ) -> bool: |
| 1744 | if cached_num_encoder_embeds == requested_num_encoder_embeds: |
| 1745 | return False |
| 1746 | logger.warning( |
| 1747 | "Encoder cache entry for multimodal identifier %r holds %d " |
| 1748 | "embeddings, but request %s input %d expects %d. A client-provided " |
| 1749 | "multimodal UUID may have been reused for different content.", |
| 1750 | identifier, |
| 1751 | cached_num_encoder_embeds, |
| 1752 | request.request_id, |
| 1753 | input_id, |
| 1754 | requested_num_encoder_embeds, |
| 1755 | ) |
| 1756 | self.encoder_cache_mismatch_reqs.add(request.request_id) |
| 1757 | return True |
| 1758 | |
| 1759 | def _try_schedule_encoder_inputs( |
| 1760 | self, |
| 1761 | request: Request, |
| 1762 | num_computed_tokens: int, |
| 1763 | num_new_tokens: int, |
| 1764 | encoder_compute_budget: int, |
| 1765 | shift_computed_tokens: int = 0, |
| 1766 | ) -> tuple[list[int], int, int, list[int]]: |
| 1767 | """Determine which encoder inputs need to be scheduled in the current step, |
| 1768 | and update `num_new_tokens` and encoder token budget accordingly. |
| 1769 | |
| 1770 | An encoder input will be scheduled if: |
| 1771 | - Its output tokens overlap with the range of tokens being computed |
| 1772 | in this step, i.e., |
| 1773 | [num_computed_tokens, num_computed_tokens + num_new_tokens). |
| 1774 | - It is not already computed and stored in the encoder cache. |
| 1775 | - It is not exist on remote encoder cache (via ECConnector) |
| 1776 | - There is sufficient encoder token budget to process it. |
| 1777 | - The encoder cache has space to store it. |
| 1778 | |
| 1779 | If an encoder input cannot be scheduled due to cache or budget |
| 1780 | limitations, the method adjusts `num_new_tokens` to schedule only the |
| 1781 | decoder tokens up to just before the unschedulable encoder input. |
| 1782 | |
| 1783 | Note that num_computed_tokens includes both locally cached |
| 1784 | blocks and externally cached blocks (via KVConnector). |
| 1785 | """ |
| 1786 | if num_new_tokens == 0 or not request.has_encoder_inputs: |
| 1787 | return [], num_new_tokens, encoder_compute_budget, [] |
| 1788 | encoder_inputs_to_schedule: list[int] = [] |
| 1789 | mm_features = request.mm_features |
| 1790 | assert mm_features is not None |
| 1791 | assert len(mm_features) > 0 |
| 1792 | external_load_encoder_input = [] |
| 1793 | |
| 1794 | # NOTE: since scheduler operates on the request level (possibly with |
| 1795 | # multiple encoder inputs per request), we need to create temporary |
| 1796 | # trackers for accounting at the encoder input level. |
| 1797 | mm_hashes_to_schedule: dict[str, int] = {} |
| 1798 | num_embeds_to_schedule = 0 |
| 1799 | |
| 1800 | encoder_window_end = ( |
| 1801 | num_computed_tokens + num_new_tokens + shift_computed_tokens |
| 1802 | ) |
| 1803 | lo, hi = get_mm_features_in_window( |
| 1804 | mm_features, |
| 1805 | start=num_computed_tokens, |
| 1806 | end=encoder_window_end, |
| 1807 | ) |
| 1808 | # For encoder-decoder, all inputs sit at start_pos=0, so lo=0 always. |
| 1809 | if self.is_encoder_decoder: |
| 1810 | lo = 0 |
| 1811 | |
| 1812 | for i in range(lo, hi): |
| 1813 | mm_feature = mm_features[i] |
| 1814 | start_pos = mm_feature.mm_position.offset |
| 1815 | num_encoder_tokens = mm_feature.mm_position.length |
| 1816 | num_encoder_embeds = mm_feature.mm_position.get_num_embeds() |
| 1817 | item_identifier = mm_feature.identifier |
| 1818 | |
| 1819 | if self.is_encoder_decoder and num_computed_tokens > 0: |
| 1820 | assert start_pos == 0, ( |
| 1821 | "Encoder input should be processed at the beginning of " |
| 1822 | "the sequence when encoder-decoder models are used." |
| 1823 | ) |
| 1824 | # Encoder input has already been computed |
| 1825 | # The calculation here is a bit different. We don't turn encoder |
| 1826 | # output into tokens that get processed by the decoder and |
| 1827 | # reflected in num_computed_tokens. Instead, start_pos reflects |
| 1828 | # the position where we need to ensure we calculate encoder |
| 1829 | # inputs. This should always be 0 to ensure we calculate encoder |
| 1830 | # inputs before running the decoder. Once we've calculated some |
| 1831 | # decoder tokens (num_computed_tokens > 0), then we know we |
| 1832 | # already calculated encoder inputs and can skip here. |
| 1833 | continue |
| 1834 | |
| 1835 | if not self.is_encoder_decoder: |
| 1836 | # We are not using the encoder cache for encoder-decoder models, |
| 1837 | # yet. |
| 1838 | if item_identifier in mm_hashes_to_schedule: |
| 1839 | scheduled_num_encoder_embeds = mm_hashes_to_schedule[ |
| 1840 | item_identifier |
| 1841 | ] |
| 1842 | if self._reject_on_encoder_cache_embed_mismatch( |
| 1843 | request, |
| 1844 | i, |
| 1845 | item_identifier, |
| 1846 | scheduled_num_encoder_embeds, |
| 1847 | num_encoder_embeds, |
| 1848 | ): |
| 1849 | return [], 0, encoder_compute_budget, [] |
| 1850 | # The same encoder input has already been scheduled in the |
| 1851 | # current step. |
| 1852 | continue |
| 1853 | |
| 1854 | cached_num_encoder_embeds = ( |
| 1855 | self.encoder_cache_manager.get_cached_num_encoder_embeds(request, i) |
| 1856 | ) |
| 1857 | if cached_num_encoder_embeds is not None and ( |
| 1858 | self._reject_on_encoder_cache_embed_mismatch( |
| 1859 | request, |
| 1860 | i, |
| 1861 | item_identifier, |
| 1862 | cached_num_encoder_embeds, |
| 1863 | num_encoder_embeds, |
| 1864 | ) |
| 1865 | ): |
| 1866 | return [], 0, encoder_compute_budget, [] |
| 1867 | |
| 1868 | if self.encoder_cache_manager.check_and_update_cache(request, i): |
| 1869 | # The encoder input is already computed and cached from a |
| 1870 | # previous step. |
| 1871 | continue |
| 1872 | |
| 1873 | # If no encoder input chunking is allowed, we do not want to |
| 1874 | # partially schedule a multimodal item. If the scheduled range would |
| 1875 | # only cover part of the mm input, roll back to before the mm item. |
| 1876 | if ( |
| 1877 | self.scheduler_config.disable_chunked_mm_input |
| 1878 | and num_computed_tokens < start_pos |
| 1879 | and (num_computed_tokens + num_new_tokens) |
| 1880 | < (start_pos + num_encoder_tokens) |
| 1881 | ): |
| 1882 | # Account for EAGLE shift when rolling back to avoid |
| 1883 | # encoder cache miss. This ensures the scheduled range |
| 1884 | # stops before start_pos even with the shift. |
| 1885 | num_new_tokens = max( |
| 1886 | 0, start_pos - (num_computed_tokens + shift_computed_tokens) |
| 1887 | ) |
| 1888 | break |
| 1889 | if not self.encoder_cache_manager.can_allocate( |
| 1890 | request, i, encoder_compute_budget, num_embeds_to_schedule |
| 1891 | ): |
| 1892 | # The encoder cache is full or the encoder budget is exhausted. |
| 1893 | # NOTE(woosuk): We assume that the encoder input tokens should |
| 1894 | # be processed altogether, as the encoder usually uses |
| 1895 | # bidirectional attention. |
| 1896 | if num_computed_tokens + shift_computed_tokens < start_pos: |
| 1897 | # We only schedule the decoder tokens just before the |
| 1898 | # encoder input. |
| 1899 | num_new_tokens = start_pos - ( |
| 1900 | num_computed_tokens + shift_computed_tokens |
| 1901 | ) |
| 1902 | else: |
| 1903 | # Because of prefix caching, num_computed_tokens is greater |
| 1904 | # than start_pos even though its encoder input is not |
| 1905 | # available. In this case, we can't schedule any token for |
| 1906 | # the request in this step. |
| 1907 | num_new_tokens = 0 |
| 1908 | break |
| 1909 | |
| 1910 | # Calculate the number of embeddings to schedule in the current range |
| 1911 | # of scheduled encoder placeholder tokens. |
| 1912 | start_idx_rel = max(0, num_computed_tokens - start_pos) |
| 1913 | end_idx_rel = min(num_encoder_tokens, encoder_window_end - start_pos) |
| 1914 | curr_embeds_start, curr_embeds_end = ( |
| 1915 | mm_feature.mm_position.get_embeds_indices_in_range( |
| 1916 | start_idx_rel, end_idx_rel |
| 1917 | ) |
| 1918 | ) |
| 1919 | # There's no embeddings in the current range of encoder placeholder tokens |
| 1920 | # so we can skip the encoder input. |
| 1921 | if curr_embeds_end - curr_embeds_start == 0: |
| 1922 | continue |
| 1923 | |
| 1924 | if self.ec_connector is not None and self.ec_connector.has_cache_item( |
| 1925 | item_identifier |
| 1926 | ): |
| 1927 | mm_hashes_to_schedule[item_identifier] = num_encoder_embeds |
| 1928 | external_load_encoder_input.append(i) |
| 1929 | num_embeds_to_schedule += num_encoder_embeds |
| 1930 | continue |
| 1931 | |
| 1932 | num_embeds_to_schedule += num_encoder_embeds |
| 1933 | encoder_compute_budget -= num_encoder_embeds |
| 1934 | mm_hashes_to_schedule[item_identifier] = num_encoder_embeds |
| 1935 | encoder_inputs_to_schedule.append(i) |
| 1936 | |
| 1937 | return ( |
| 1938 | encoder_inputs_to_schedule, |
| 1939 | num_new_tokens, |
| 1940 | encoder_compute_budget, |
| 1941 | external_load_encoder_input, |
| 1942 | ) |
| 1943 | |
| 1944 | def _make_scheduled_encoder_input_stats( |
| 1945 | self, scheduled_encoder_inputs: dict[str, list[int]] |
| 1946 | ) -> ScheduledEncoderInputStats | None: |
| 1947 | stats = ScheduledEncoderInputStats() |
| 1948 | |
| 1949 | for req_id, input_ids in scheduled_encoder_inputs.items(): |
| 1950 | request = self.requests.get(req_id) |
| 1951 | if request is None: |
| 1952 | continue |
| 1953 | |
| 1954 | for input_id in input_ids: |
| 1955 | mm_feature = request.mm_features[input_id] |
| 1956 | stats.num_inputs += 1 |
| 1957 | stats.output_tokens += mm_feature.mm_position.get_num_embeds() |
| 1958 | |
| 1959 | return stats if stats.num_inputs else None |
| 1960 | |
| 1961 | def get_grammar_bitmask( |
| 1962 | self, scheduler_output: SchedulerOutput |
| 1963 | ) -> GrammarOutput | None: |
| 1964 | # Collect list of scheduled request ids that use structured output. |
| 1965 | # The corresponding rows of the bitmask will be in this order. |
| 1966 | if not scheduler_output.has_structured_output_requests: |
| 1967 | return None |
| 1968 | |
| 1969 | structured_output_request_ids = [ |
| 1970 | req_id |
| 1971 | for req_id in scheduler_output.num_scheduled_tokens |
| 1972 | if (req := self.requests.get(req_id)) |
| 1973 | and (req.use_structured_output and not req.is_prefill_chunk) |
| 1974 | ] |
| 1975 | if not structured_output_request_ids: |
| 1976 | return None |
| 1977 | |
| 1978 | bitmask = self.structured_output_manager.grammar_bitmask( |
| 1979 | self.requests, |
| 1980 | structured_output_request_ids, |
| 1981 | scheduler_output.scheduled_spec_decode_tokens, |
| 1982 | ) |
| 1983 | spec_tokens = scheduler_output.scheduled_spec_decode_tokens |
| 1984 | num_acceptable_drafts = [ |
| 1985 | len(strip_speculative_padding(spec_tokens.get(req_id, []))) |
| 1986 | for req_id in structured_output_request_ids |
| 1987 | ] |
| 1988 | return GrammarOutput( |
| 1989 | structured_output_request_ids, bitmask, num_acceptable_drafts |
| 1990 | ) |
| 1991 | |
| 1992 | def update_from_output( |
| 1993 | self, |
| 1994 | scheduler_output: SchedulerOutput, |
| 1995 | model_runner_output: ModelRunnerOutput, |
| 1996 | ) -> dict[int, EngineCoreOutputs]: |
| 1997 | sampled_token_ids = model_runner_output.sampled_token_ids |
| 1998 | logprobs = model_runner_output.logprobs |
| 1999 | prompt_logprobs_dict = model_runner_output.prompt_logprobs_dict |
| 2000 | prompt_token_id_logprobs_dict = ( |
| 2001 | model_runner_output.prompt_token_id_logprobs_dict |
| 2002 | ) |
| 2003 | num_scheduled_tokens = scheduler_output.num_scheduled_tokens |
| 2004 | pooler_outputs = model_runner_output.pooler_output |
| 2005 | num_nans_in_logits = model_runner_output.num_nans_in_logits |
| 2006 | kv_connector_output = model_runner_output.kv_connector_output |
| 2007 | ec_connector_output = model_runner_output.ec_connector_output |
| 2008 | cudagraph_stats = model_runner_output.cudagraph_stats |
| 2009 | |
| 2010 | # Every GPU write enqueued by this and earlier steps has completed, so it is |
| 2011 | # safe to return deferred-free blocks to the pool. |
| 2012 | if self.defer_block_free and scheduler_output.total_num_scheduled_tokens > 0: |
| 2013 | self.processed_step_seq += 1 |
| 2014 | self._drain_deferred_frees() |
| 2015 | |
| 2016 | perf_stats: PerfStats | None = None |
| 2017 | if self.perf_metrics and self.perf_metrics.is_enabled(): |
| 2018 | perf_stats = self.perf_metrics.get_step_perf_stats_per_gpu(scheduler_output) |
| 2019 | |
| 2020 | outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list) |
| 2021 | spec_decoding_stats: SpecDecodingStats | None = None |
| 2022 | |
| 2023 | failed_kv_load_req_ids: set[str] = set() |
| 2024 | if kv_connector_output and kv_connector_output.failed_recving: |
| 2025 | failed_kv_load_req_ids.update( |
| 2026 | self._handle_failed_recving(kv_connector_output.failed_recving) |
| 2027 | ) |
| 2028 | if kv_connector_output and kv_connector_output.invalid_block_ids: |
| 2029 | # These blocks contain externally computed tokens that failed to |
| 2030 | # load. Identify affected requests and adjust their computed token |
| 2031 | # count to trigger recomputation of the invalid blocks. |
| 2032 | failed_kv_load_req_ids.update( |
| 2033 | self._handle_invalid_blocks( |
| 2034 | kv_connector_output.invalid_block_ids, |
| 2035 | num_scheduled_tokens, |
| 2036 | ) |
| 2037 | ) |
| 2038 | # NOTE(woosuk): As len(num_scheduled_tokens) can be up to 1K or more, |
| 2039 | # the below loop can be a performance bottleneck. We should do our best |
| 2040 | # to avoid expensive operations inside the loop. |
| 2041 | stopped_running_reqs: set[Request] = set() |
| 2042 | stopped_preempted_reqs: set[Request] = set() |
| 2043 | for req_id, num_tokens_scheduled in num_scheduled_tokens.items(): |
| 2044 | assert num_tokens_scheduled > 0 |
| 2045 | request = self.requests.get(req_id) |
| 2046 | output_is_stale = False |
| 2047 | if request is not None: |
| 2048 | request.num_in_flight_tokens -= num_tokens_scheduled |
| 2049 | # Drain any stale share (see _preempt_request) in lockstep. |
| 2050 | if request.num_stale_output_tokens > 0: |
| 2051 | output_is_stale = True |
| 2052 | request.num_stale_output_tokens -= num_tokens_scheduled |
| 2053 | assert request.num_stale_output_tokens >= 0 |
| 2054 | if failed_kv_load_req_ids and req_id in failed_kv_load_req_ids: |
| 2055 | # skip failed or rescheduled requests from KV load failure |
| 2056 | continue |
| 2057 | if request is None or request.is_finished(): |
| 2058 | # The request is already finished. This can happen if the |
| 2059 | # request is aborted while the model is executing it (e.g., |
| 2060 | # in pipeline parallelism or in async scheduling). |
| 2061 | # NOTE(Kuntai): When delay_free_blocks=True (for async KV |
| 2062 | # cache transfer in KV connector), the aborted request will not |
| 2063 | # be set to None (in order to finish async KV transfer). |
| 2064 | # In this case, we use is_finished() to check. |
| 2065 | continue |
| 2066 | |
| 2067 | # Drop-mode stale output (same-step resume) is discarded entirely. |
| 2068 | if output_is_stale and request.drop_stale_output: |
| 2069 | continue |
| 2070 | |
| 2071 | req_index = model_runner_output.req_id_to_index[req_id] |
| 2072 | generated_token_ids = ( |
| 2073 | sampled_token_ids[req_index] if sampled_token_ids else [] |
| 2074 | ) |
| 2075 | |
| 2076 | scheduled_spec_token_ids = ( |
| 2077 | scheduler_output.scheduled_spec_decode_tokens.get(req_id) |
| 2078 | ) |
| 2079 | if scheduled_spec_token_ids and ( |
| 2080 | generated_token_ids or self.num_sampled_tokens_per_step == 0 |
| 2081 | ): |
| 2082 | num_draft_tokens = len(scheduled_spec_token_ids) |
| 2083 | num_sampled = self.num_sampled_tokens_per_step |
| 2084 | num_accepted = max(len(generated_token_ids) - num_sampled, 0) |
| 2085 | num_rejected = num_draft_tokens - num_accepted |
| 2086 | # Rejections roll back num_computed_tokens (and, under async |
| 2087 | # scheduling, num_output_placeholders, which covers the spec |
| 2088 | # tokens). A stale rejection count predates the preemption |
| 2089 | # rollback and must not apply. |
| 2090 | if not output_is_stale: |
| 2091 | if request.num_computed_tokens > 0: |
| 2092 | request.num_computed_tokens -= num_rejected |
| 2093 | if request.num_output_placeholders > 0: |
| 2094 | request.num_output_placeholders -= num_rejected |
| 2095 | spec_decoding_stats = self.make_spec_decoding_stats( |
| 2096 | spec_decoding_stats, |
| 2097 | num_draft_tokens=num_draft_tokens, |
| 2098 | num_accepted_tokens=num_accepted, |
| 2099 | num_invalid_spec_tokens=scheduler_output.num_invalid_spec_tokens, |
| 2100 | request_id=req_id, |
| 2101 | ) |
| 2102 | if request.spec_decode_metrics is not None: |
| 2103 | # Exclude grammar-invalidated drafts from the proposed |
| 2104 | # count, mirroring make_spec_decoding_stats; the accepted |
| 2105 | # bucket (j) is unaffected. |
| 2106 | adj_draft_tokens = num_draft_tokens |
| 2107 | if scheduler_output.num_invalid_spec_tokens: |
| 2108 | adj_draft_tokens -= ( |
| 2109 | scheduler_output.num_invalid_spec_tokens.get(req_id, 0) |
| 2110 | ) |
| 2111 | request.spec_decode_metrics.observe( |
| 2112 | num_draft_tokens=adj_draft_tokens, |
| 2113 | num_accepted=num_accepted, |
| 2114 | detailed=self.spec_decode_metrics_level == "detailed", |
| 2115 | ) |
| 2116 | |
| 2117 | # Free encoder inputs only after the step has actually executed. |
| 2118 | if request.has_encoder_inputs: |
| 2119 | self._free_encoder_inputs(request) |
| 2120 | |
| 2121 | stopped = False |
| 2122 | new_logprobs = None |
| 2123 | new_sampling_mask = None |
| 2124 | new_token_ids = generated_token_ids |
| 2125 | pooler_output = pooler_outputs[req_index] if pooler_outputs else None |
| 2126 | kv_transfer_params = None |
| 2127 | ec_transfer_params = None |
| 2128 | prefill_stats = None |
| 2129 | status_before_stop = request.status |
| 2130 | |
| 2131 | # Check for stop and update request status. |
| 2132 | if new_token_ids: |
| 2133 | new_token_ids, stopped = self._update_request_with_output( |
| 2134 | request, new_token_ids, is_stale=output_is_stale |
| 2135 | ) |
| 2136 | elif request.pooling_params and pooler_output is not None: |
| 2137 | # Pooling stops as soon as there is output. |
| 2138 | request.status = RequestStatus.FINISHED_STOPPED |
| 2139 | stopped = True |
| 2140 | elif ( |
| 2141 | self.is_mm_encoder_only |
| 2142 | and request.num_computed_tokens >= request.num_prompt_tokens |
| 2143 | ): |
| 2144 | # An encoder instance runs the encoder and publishes the |
| 2145 | # embeddings instead of sampling, so it stops as soon as the |
| 2146 | # whole prompt is consumed. Encoder inputs are never scheduled |
| 2147 | # past a multi-modal item the encoder cache could not admit, so |
| 2148 | # a consumed prompt also means every item in it was encoded. |
| 2149 | request.status = RequestStatus.FINISHED_STOPPED |
| 2150 | stopped = True |
| 2151 | |
| 2152 | if new_token_ids and not self.structured_output_manager.accept_tokens( |
| 2153 | request, new_token_ids |
| 2154 | ): |
| 2155 | logger.error( |
| 2156 | "Unexpected: grammar rejected tokens %s for request %s. " |
| 2157 | "Terminating request.", |
| 2158 | new_token_ids, |
| 2159 | req_id, |
| 2160 | ) |
| 2161 | request.status = RequestStatus.FINISHED_ERROR |
| 2162 | request.resumable = False |
| 2163 | stopped = True |
| 2164 | |
| 2165 | routed_experts = None |
| 2166 | should_emit_output = bool( |
| 2167 | new_token_ids or pooler_output is not None or stopped |
| 2168 | ) |
| 2169 | if self.aux_output_connector is not None and should_emit_output: |
| 2170 | routed_experts = self.aux_output_connector.take_output( |
| 2171 | request, model_runner_output.aux_output_connector_output |
| 2172 | ) |
| 2173 | if should_emit_output: |
| 2174 | prefill_stats = request.take_prefill_stats() |
| 2175 | if prefill_stats is not None: |
| 2176 | prefill_stats.finalize( |
| 2177 | self.kv_cache_manager.estimate_cached_tokens(request) |
| 2178 | ) |
| 2179 | |
| 2180 | # Extract sample logprobs before stop handling can replace the |
| 2181 | # sampling parameters for a streaming continuation. |
| 2182 | if ( |
| 2183 | request.sampling_params is not None |
| 2184 | and request.sampling_params.num_logprobs is not None |
| 2185 | and logprobs |
| 2186 | ): |
| 2187 | new_logprobs = logprobs.slice_request(req_index, len(new_token_ids)) |
| 2188 | |
| 2189 | finish_reason = None |
| 2190 | if stopped: |
| 2191 | # Capture finish_reason BEFORE _handle_stopped_request, which may |
| 2192 | # reset the status to WAITING for streaming requests that continue. |
| 2193 | finish_reason = request.get_finished_reason() |
| 2194 | finished = self._handle_stopped_request(request) |
| 2195 | if finished: |
| 2196 | kv_transfer_params, ec_transfer_params = self._free_request(request) |
| 2197 | if ( |
| 2198 | prefill_stats is not None |
| 2199 | and kv_transfer_params is not None |
| 2200 | and kv_transfer_params.get("do_remote_prefill") |
| 2201 | ): |
| 2202 | # P-side cache hits, so D can report them in |
| 2203 | # prompt_tokens_details instead of its own (~100%) |
| 2204 | # hit rate from the KV transfer. |
| 2205 | kv_transfer_params["remote_prefill_cached_tokens"] = ( |
| 2206 | prefill_stats.num_cached_tokens |
| 2207 | ) |
| 2208 | |
| 2209 | if status_before_stop == RequestStatus.RUNNING: |
| 2210 | stopped_running_reqs.add(request) |
| 2211 | else: |
| 2212 | stopped_preempted_reqs.add(request) |
| 2213 | |
| 2214 | if self.return_sampling_mask: |
| 2215 | sampling_masks = model_runner_output.sampling_masks |
| 2216 | if new_token_ids and sampling_masks is not None: |
| 2217 | new_sampling_mask = sampling_masks.slice_request( |
| 2218 | req_index, len(new_token_ids) |
| 2219 | ) |
| 2220 | |
| 2221 | if num_nans_in_logits is not None and req_id in num_nans_in_logits: |
| 2222 | request.num_nans_in_logits = num_nans_in_logits[req_id] |
| 2223 | |
| 2224 | # Get prompt logprobs for this request. |
| 2225 | prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id) |
| 2226 | prompt_token_id_logprobs = prompt_token_id_logprobs_dict.get(req_id) |
| 2227 | if should_emit_output: |
| 2228 | # Add EngineCoreOutput for this Request. |
| 2229 | outputs[request.client_index].append( |
| 2230 | EngineCoreOutput( |
| 2231 | request_id=req_id, |
| 2232 | new_token_ids=new_token_ids, |
| 2233 | finish_reason=finish_reason, |
| 2234 | new_logprobs=new_logprobs, |
| 2235 | new_sampling_mask=new_sampling_mask, |
| 2236 | new_prompt_logprobs_tensors=prompt_logprobs_tensors, |
| 2237 | prompt_token_id_logprobs=prompt_token_id_logprobs, |
| 2238 | pooling_output=pooler_output, |
| 2239 | stop_reason=request.stop_reason, |
| 2240 | events=request.take_events(), |
| 2241 | prefill_stats=prefill_stats, |
| 2242 | spec_decode_metrics=( |
| 2243 | request.spec_decode_metrics |
| 2244 | if finish_reason is not None |
| 2245 | else None |
| 2246 | ), |
| 2247 | kv_transfer_params=kv_transfer_params, |
| 2248 | ec_transfer_params=ec_transfer_params, |
| 2249 | trace_headers=request.trace_headers, |
| 2250 | routed_experts=routed_experts, |
| 2251 | num_nans_in_logits=request.num_nans_in_logits, |
| 2252 | ) |
| 2253 | ) |
| 2254 | else: |
| 2255 | # Invariant: EngineCore returns no partial prefill outputs. |
| 2256 | assert not prompt_logprobs_tensors |
| 2257 | assert prompt_token_id_logprobs is None |
| 2258 | |
| 2259 | # Remove the stopped requests from the running and waiting queues. |
| 2260 | if stopped_running_reqs: |
| 2261 | self.running = remove_all(self.running, stopped_running_reqs) |
| 2262 | if stopped_preempted_reqs: |
| 2263 | # This is a rare case and unlikely to impact performance. |
| 2264 | self.waiting.remove_requests(stopped_preempted_reqs) |
| 2265 | self.kv_holding_waiting.remove_requests(stopped_preempted_reqs) |
| 2266 | self.deferred_waiting.difference_update(stopped_preempted_reqs) |
| 2267 | |
| 2268 | error_req_ids = set(self.grammar_compile_error_reqs) |
| 2269 | self.grammar_compile_error_reqs.clear() |
| 2270 | error_req_ids.update(self.encoder_cache_mismatch_reqs) |
| 2271 | self.encoder_cache_mismatch_reqs.clear() |
| 2272 | if failed_kv_load_req_ids and not self.recompute_kv_load_failures: |
| 2273 | error_req_ids.update(failed_kv_load_req_ids) |
| 2274 | if self.ec_connector is not None: |
| 2275 | # An encoder input the connector can no longer obtain. Failing is |
| 2276 | # retryable: re-issuing the request re-runs the encode. |
| 2277 | error_req_ids.update(self.ec_connector.take_unavailable_requests()) |
| 2278 | |
| 2279 | if error_req_ids: |
| 2280 | error_reqs = self.finish_requests( |
| 2281 | error_req_ids, RequestStatus.FINISHED_ERROR |
| 2282 | ) |
| 2283 | for request in error_reqs: |
| 2284 | outputs[request.client_index].append( |
| 2285 | EngineCoreOutput( |
| 2286 | request_id=request.request_id, |
| 2287 | new_token_ids=[], |
| 2288 | finish_reason=request.get_finished_reason(), |
| 2289 | events=request.take_events(), |
| 2290 | trace_headers=request.trace_headers, |
| 2291 | ) |
| 2292 | ) |
| 2293 | |
| 2294 | # KV Connector: update state for finished KV Transfers. |
| 2295 | if kv_connector_output: |
| 2296 | self._update_from_kv_xfer_finished(kv_connector_output) |
| 2297 | |
| 2298 | # EC Connector: update state from worker-side EC connector output. |
| 2299 | if self.ec_connector is not None and ec_connector_output: |
| 2300 | self.ec_connector.update_connector_output(ec_connector_output) |
| 2301 | |
| 2302 | # Worker-side KV connector stats from the model runner output. |
| 2303 | kv_connector_stats: KVConnectorStats | None = ( |
| 2304 | kv_connector_output.kv_connector_stats if kv_connector_output else None |
| 2305 | ) |
| 2306 | if self.connector: |
| 2307 | # Scheduler-side KV connector stats collected after connector update. |
| 2308 | scheduler_kv_connector_stats = self.connector.get_kv_connector_stats() |
| 2309 | if ( |
| 2310 | scheduler_kv_connector_stats is not None |
| 2311 | and not scheduler_kv_connector_stats.is_empty() |
| 2312 | ): |
| 2313 | kv_connector_stats = ( |
| 2314 | kv_connector_stats.aggregate(scheduler_kv_connector_stats) |
| 2315 | if kv_connector_stats is not None |
| 2316 | else scheduler_kv_connector_stats |
| 2317 | ) |
| 2318 | |
| 2319 | # Worker-side EC connector stats from the model runner output. |
| 2320 | ec_connector_stats: ECConnectorStats | None = ( |
| 2321 | ec_connector_output.ec_connector_stats if ec_connector_output else None |
| 2322 | ) |
| 2323 | if self.ec_connector: |
| 2324 | # Scheduler-side EC connector stats collected after connector update. |
| 2325 | scheduler_ec_connector_stats = self.ec_connector.get_ec_connector_stats() |
| 2326 | if ( |
| 2327 | scheduler_ec_connector_stats is not None |
| 2328 | and not scheduler_ec_connector_stats.is_empty() |
| 2329 | ): |
| 2330 | ec_connector_stats = ( |
| 2331 | ec_connector_stats.aggregate(scheduler_ec_connector_stats) |
| 2332 | if ec_connector_stats is not None |
| 2333 | else scheduler_ec_connector_stats |
| 2334 | ) |
| 2335 | |
| 2336 | # collect KV cache events from KV cache manager |
| 2337 | events = self.kv_cache_manager.take_events() |
| 2338 | |
| 2339 | # collect KV cache events from connector |
| 2340 | if self.connector is not None: |
| 2341 | connector_events = self.connector.take_events() |
| 2342 | if connector_events: |
| 2343 | if events is None: |
| 2344 | events = list(connector_events) |
| 2345 | else: |
| 2346 | events.extend(connector_events) |
| 2347 | |
| 2348 | # publish collected KV cache events |
| 2349 | if events: |
| 2350 | batch = KVEventBatch(ts=time.time(), events=events) |
| 2351 | self.kv_event_publisher.publish(batch) |
| 2352 | |
| 2353 | # Create EngineCoreOutputs for all clients that have requests with |
| 2354 | # outputs in this step. |
| 2355 | engine_core_outputs = { |
| 2356 | client_index: EngineCoreOutputs(outputs=outs) |
| 2357 | for client_index, outs in outputs.items() |
| 2358 | } |
| 2359 | |
| 2360 | finished_req_ids = self.finished_req_ids_dict |
| 2361 | if finished_req_ids: |
| 2362 | # Include ids of requests that finished since last outputs |
| 2363 | # were sent. |
| 2364 | for client_index, finished_set in finished_req_ids.items(): |
| 2365 | # Set finished request set in EngineCoreOutputs for this client. |
| 2366 | if (eco := engine_core_outputs.get(client_index)) is not None: |
| 2367 | eco.finished_requests = finished_set |
| 2368 | else: |
| 2369 | engine_core_outputs[client_index] = EngineCoreOutputs( |
| 2370 | finished_requests=finished_set |
| 2371 | ) |
| 2372 | finished_req_ids.clear() |
| 2373 | |
| 2374 | if ( |
| 2375 | stats := self.make_stats( |
| 2376 | spec_decoding_stats, |
| 2377 | kv_connector_stats, |
| 2378 | cudagraph_stats, |
| 2379 | perf_stats, |
| 2380 | ec_connector_stats=ec_connector_stats, |
| 2381 | ) |
| 2382 | ) is not None: |
| 2383 | # Return stats to only one of the front-ends. |
| 2384 | if (eco := next(iter(engine_core_outputs.values()), None)) is None: |
| 2385 | # We must return the stats even if there are no request |
| 2386 | # outputs this step. |
| 2387 | engine_core_outputs[0] = eco = EngineCoreOutputs() |
| 2388 | eco.scheduler_stats = stats |
| 2389 | |
| 2390 | return engine_core_outputs |
| 2391 | |
| 2392 | def _ec_transfer_pending(self, request: Request, num_computed_tokens: int) -> bool: |
| 2393 | """Whether an encoder input this request needs is still in transit.""" |
| 2394 | return ( |
| 2395 | self.ec_connector is not None |
| 2396 | and bool(request.mm_features) |
| 2397 | and not self.ec_connector.ensure_cache_available( |
| 2398 | request, num_computed_tokens |
| 2399 | ) |
| 2400 | ) |
| 2401 | |
| 2402 | @staticmethod |
| 2403 | def _is_blocked_waiting_status(status: RequestStatus) -> bool: |
| 2404 | return status in ( |
| 2405 | RequestStatus.WAITING_FOR_STRUCTURED_OUTPUT_GRAMMAR, |
| 2406 | RequestStatus.WAITING_FOR_REMOTE_KVS, |
| 2407 | RequestStatus.WAITING_FOR_STREAMING_REQ, |
| 2408 | ) |
| 2409 | |
| 2410 | @staticmethod |
| 2411 | def _holds_kv_blocks(request: Request) -> bool: |
| 2412 | # Whether a request currently has kv blocks allocated. |
| 2413 | # While always the case when status is WAITING_FOR_REMOTE_KVS |
| 2414 | # or WAITING_FOR_STREAMING_REQ, there are such cases where it's |
| 2415 | # WAITING or PREEMPTED. |
| 2416 | return ( |
| 2417 | request.num_computed_tokens > 0 |
| 2418 | # num_computed_tokens is reset to 0 if kv transfer fails for requests in |
| 2419 | # WAITING_FOR_REMOTE_KVS state, before their allocated blocks are freed. |
| 2420 | or request.status == RequestStatus.WAITING_FOR_REMOTE_KVS |
| 2421 | ) |
| 2422 | |
| 2423 | def _enqueue_waiting_request(self, request: Request) -> None: |
| 2424 | if self._is_blocked_waiting_status(request.status): |
| 2425 | self.deferred_waiting.add(request) |
| 2426 | if self._holds_kv_blocks(request): |
| 2427 | self.kv_holding_waiting.add_request(request) |
| 2428 | else: |
| 2429 | self.waiting.add_request(request) |
| 2430 | |
| 2431 | def _handle_stopped_request(self, request: Request) -> bool: |
| 2432 | """Return True if finished (can be False for resumable requests).""" |
| 2433 | if not request.resumable: |
| 2434 | return True |
| 2435 | |
| 2436 | # Drop the finished turn's in-flight work and resume from the |
| 2437 | # materialized frontier. |
| 2438 | safe_frontier = request.num_computed_tokens - request.num_output_placeholders |
| 2439 | assert safe_frontier >= 0 |
| 2440 | request.num_computed_tokens = safe_frontier |
| 2441 | request.spec_token_ids = [] |
| 2442 | request.drop_stale_output = True |
| 2443 | request.num_stale_output_tokens = request.num_in_flight_tokens |
| 2444 | request.num_output_placeholders = 0 |
| 2445 | |
| 2446 | if request.streaming_queue: |
| 2447 | update = request.streaming_queue.popleft() |
| 2448 | if update is None: |
| 2449 | # Streaming request finished. |
| 2450 | return True |
| 2451 | self._update_request_as_session(request, update) |
| 2452 | else: |
| 2453 | request.status = RequestStatus.WAITING_FOR_STREAMING_REQ |
| 2454 | self.num_waiting_for_streaming_input += 1 |
| 2455 | |
| 2456 | self._enqueue_waiting_request(request) |
| 2457 | return False |
| 2458 | |
| 2459 | def _update_request_with_output( |
| 2460 | self, request: Request, new_token_ids: list[int], is_stale: bool = False |
| 2461 | ) -> tuple[list[int], bool]: |
| 2462 | # is_stale is only used by the AsyncScheduler override. |
| 2463 | # Append generated tokens and check for stop. Note that if |
| 2464 | # a request is still being prefilled, we expect the model runner |
| 2465 | # to return empty token ids for the request. |
| 2466 | stopped = False |
| 2467 | for num_new, output_token_id in enumerate(new_token_ids, 1): |
| 2468 | request.append_output_token_ids(output_token_id) |
| 2469 | |
| 2470 | # Check for stop and update request state. |
| 2471 | # This must be called before we make the EngineCoreOutput. |
| 2472 | stopped = check_stop(request, self.max_model_len) |
| 2473 | if stopped: |
| 2474 | del new_token_ids[num_new:] # Trim new tokens if needed. |
| 2475 | break |
| 2476 | return new_token_ids, stopped |
| 2477 | |
| 2478 | def _free_encoder_inputs(self, request: Request) -> None: |
| 2479 | cached_encoder_input_ids = self.encoder_cache_manager.get_cached_input_ids( |
| 2480 | request |
| 2481 | ) |
| 2482 | # OPTIMIZATION: Avoid list(set) if the set is empty. |
| 2483 | if not cached_encoder_input_ids: |
| 2484 | return |
| 2485 | |
| 2486 | # Defer the free by the drafter's look-ahead so an entry stays |
| 2487 | # referenced until the drafter's read-ahead has also passed it, |
| 2488 | # mirroring the shift the encoder scheduling path applies. |
| 2489 | spec_lookahead = self.num_prefill_lookahead |
| 2490 | |
| 2491 | # Here, we use list(set) to avoid modifying the set while iterating |
| 2492 | # over it. |
| 2493 | for input_id in list(cached_encoder_input_ids): |
| 2494 | mm_feature = request.mm_features[input_id] |
| 2495 | start_pos = mm_feature.mm_position.offset |
| 2496 | num_tokens = mm_feature.mm_position.length |
| 2497 | if self.is_encoder_decoder and request.num_computed_tokens > 0: |
| 2498 | # With Whisper, as soon as we've generated a single token, |
| 2499 | # we know we're done with the encoder input. Cross Attention |
| 2500 | # KVs have been calculated and cached already. |
| 2501 | self._free_encoder_input(request, input_id) |
| 2502 | elif ( |
| 2503 | start_pos + num_tokens + spec_lookahead |
| 2504 | <= request.num_computed_tokens - request.num_output_placeholders |
| 2505 | ): |
| 2506 | # Processed, stored in the decoder KV cache, and far enough past |
| 2507 | # the placeholder range (plus the drafter's look-ahead) that no |
| 2508 | # rejection or drafter gather can reference it. |
| 2509 | self._free_encoder_input(request, input_id) |
| 2510 | |
| 2511 | def _free_encoder_input(self, request: Request, input_id: int) -> None: |
| 2512 | self.encoder_cache_manager.free_encoder_input(request, input_id) |
| 2513 | if self.ec_connector is not None: |
| 2514 | self.ec_connector.update_state_after_free(request, input_id) |
| 2515 | |
| 2516 | def update_draft_token_ids(self, draft_token_ids: DraftTokenIds) -> None: |
| 2517 | for req_id, spec_token_ids in zip( |
| 2518 | draft_token_ids.req_ids, |
| 2519 | draft_token_ids.draft_token_ids, |
| 2520 | ): |
| 2521 | request = self.requests.get(req_id) |
| 2522 | if request is None or request.is_finished(): |
| 2523 | # The request may have been finished. Skip. |
| 2524 | continue |
| 2525 | |
| 2526 | if request.is_prefill_chunk: |
| 2527 | # Ignore draft tokens for prefill chunks. |
| 2528 | if request.spec_token_ids: |
| 2529 | request.spec_token_ids = [] |
| 2530 | continue |
| 2531 | |
| 2532 | # Add newly generated spec token ids to the request. |
| 2533 | request.spec_token_ids = self.structured_output_manager.validate_tokens( |
| 2534 | request, spec_token_ids |
| 2535 | ) |
| 2536 | |
| 2537 | def update_draft_token_ids_in_output( |
| 2538 | self, draft_token_ids: DraftTokenIds, scheduler_output: SchedulerOutput |
| 2539 | ) -> None: |
| 2540 | num_invalid_spec_tokens: dict[str, int] = {} |
| 2541 | |
| 2542 | sched_spec_tokens = scheduler_output.scheduled_spec_decode_tokens |
| 2543 | for req_id, spec_token_ids in zip( |
| 2544 | draft_token_ids.req_ids, |
| 2545 | draft_token_ids.draft_token_ids, |
| 2546 | ): |
| 2547 | request = self.requests.get(req_id) |
| 2548 | if request is None or request.is_finished(): |
| 2549 | # The request may have been finished. Skip. |
| 2550 | continue |
| 2551 | |
| 2552 | placeholder_spec_tokens = sched_spec_tokens.get(req_id) |
| 2553 | if not placeholder_spec_tokens: |
| 2554 | continue |
| 2555 | |
| 2556 | orig_num_spec_tokens = len(placeholder_spec_tokens) |
| 2557 | # Trim drafts to scheduled number of spec tokens |
| 2558 | # (needed for chunked prefill case for example). |
| 2559 | del spec_token_ids[orig_num_spec_tokens:] |
| 2560 | # Filter out spec tokens which do not adhere to the grammar. |
| 2561 | spec_token_ids = self.structured_output_manager.validate_tokens( |
| 2562 | request, spec_token_ids |
| 2563 | ) |
| 2564 | # Pad to original number of spec tokens. |
| 2565 | num_invalid_tokens = orig_num_spec_tokens - len(spec_token_ids) |
| 2566 | if num_invalid_tokens: |
| 2567 | spec_token_ids.extend([-1] * num_invalid_tokens) |
| 2568 | num_invalid_spec_tokens[req_id] = num_invalid_tokens |
| 2569 | |
| 2570 | sched_spec_tokens[req_id] = spec_token_ids |
| 2571 | |
| 2572 | scheduler_output.num_invalid_spec_tokens = num_invalid_spec_tokens |
| 2573 | |
| 2574 | def get_request_counts(self) -> tuple[int, int]: |
| 2575 | """Returns (num_running_reqs, num_waiting_reqs).""" |
| 2576 | return len(self.running), len(self.waiting) + len(self.kv_holding_waiting) |
| 2577 | |
| 2578 | def get_kv_cache_usage(self) -> float: |
| 2579 | """Returns the fraction of the KV cache currently in use (0.0-1.0).""" |
| 2580 | return self.kv_cache_manager.usage |
| 2581 | |
| 2582 | def add_request(self, request: Request) -> None: |
| 2583 | existing = self.requests.get(request.request_id) |
| 2584 | if existing is not None: |
| 2585 | update = StreamingUpdate.from_request(request) |
| 2586 | if existing.status != RequestStatus.WAITING_FOR_STREAMING_REQ: |
| 2587 | assert existing.streaming_queue is not None, "duplicate request id" |
| 2588 | # Queue next input chunk (or finished sentinel). |
| 2589 | existing.streaming_queue.append(update) |
| 2590 | elif update is not None: |
| 2591 | # Commence next input chunk. |
| 2592 | self._update_request_as_session(existing, update) |
| 2593 | else: |
| 2594 | # Streaming-input session finished. |
| 2595 | self.finish_requests(request.request_id, RequestStatus.FINISHED_ABORTED) |
| 2596 | else: |
| 2597 | if request.resumable: |
| 2598 | request.streaming_queue = deque() |
| 2599 | self._enqueue_waiting_request(request) |
| 2600 | self.requests[request.request_id] = request |
| 2601 | if self.spec_decode_metrics_level != "none": |
| 2602 | request.spec_decode_metrics = RequestSpecDecodeMetrics.new( |
| 2603 | self.num_spec_tokens |
| 2604 | ) |
| 2605 | if self.connector is not None: |
| 2606 | self.connector.on_new_request(request) |
| 2607 | kv_transfer_params = request.kv_transfer_params |
| 2608 | if kv_transfer_params and kv_transfer_params.get("do_remote_prefill"): |
| 2609 | # Count P/D requests on arrival: connector matching is |
| 2610 | # skipped while all run slots are taken. |
| 2611 | self._set_kv_fetch_stage(request, KV_FETCH_WAITING_TO_START) |
| 2612 | if self.log_stats: |
| 2613 | request.record_event(EngineCoreEventType.QUEUED) |
| 2614 | |
| 2615 | def finish_requests( |
| 2616 | self, request_ids: str | Iterable[str] | None, finished_status: RequestStatus |
| 2617 | ) -> list[Request]: |
| 2618 | """Handles the finish signal from outside the scheduler. |
| 2619 | |
| 2620 | For example, the API server can abort a request when the client |
| 2621 | disconnects. |
| 2622 | |
| 2623 | If request_ids is None, all requests will be finished. |
| 2624 | |
| 2625 | Returns: |
| 2626 | List of requests that were aborted. Will not include any that were |
| 2627 | already finished. |
| 2628 | |
| 2629 | """ |
| 2630 | assert RequestStatus.is_finished(finished_status) |
| 2631 | if isinstance(request_ids, str): |
| 2632 | request_ids = (request_ids,) |
| 2633 | elif request_ids is not None: |
| 2634 | request_ids = set(request_ids) |
| 2635 | else: |
| 2636 | request_ids = self.requests.keys() |
| 2637 | |
| 2638 | running_requests_to_remove = set() |
| 2639 | waiting_requests_to_remove = [] |
| 2640 | valid_requests = [] |
| 2641 | |
| 2642 | # First pass: collect requests to remove from queues |
| 2643 | for req_id in request_ids: |
| 2644 | request = self.requests.get(req_id) |
| 2645 | if request is None or request.is_finished(): |
| 2646 | # Invalid request ID. |
| 2647 | continue |
| 2648 | |
| 2649 | valid_requests.append(request) |
| 2650 | if request.status == RequestStatus.RUNNING: |
| 2651 | running_requests_to_remove.add(request) |
| 2652 | else: |
| 2653 | if request.status == RequestStatus.WAITING_FOR_STREAMING_REQ: |
| 2654 | self.num_waiting_for_streaming_input -= 1 |
| 2655 | waiting_requests_to_remove.append(request) |
| 2656 | |
| 2657 | # Remove all requests from queues at once for better efficiency |
| 2658 | if running_requests_to_remove: |
| 2659 | self.running = remove_all(self.running, running_requests_to_remove) |
| 2660 | if waiting_requests_to_remove: |
| 2661 | self.waiting.remove_requests(waiting_requests_to_remove) |
| 2662 | self.kv_holding_waiting.remove_requests(waiting_requests_to_remove) |
| 2663 | self.deferred_waiting.difference_update(waiting_requests_to_remove) |
| 2664 | |
| 2665 | # Second pass: set status and free requests |
| 2666 | for request in valid_requests: |
| 2667 | delay_free_blocks = False |
| 2668 | if request.status == RequestStatus.WAITING_FOR_REMOTE_KVS: |
| 2669 | delay_free_blocks = ( |
| 2670 | request.request_id not in self.finished_recving_kv_req_ids |
| 2671 | ) |
| 2672 | self.finished_recving_kv_req_ids.discard(request.request_id) |
| 2673 | self.failed_recving_kv_req_ids.discard(request.request_id) |
| 2674 | |
| 2675 | request.status = finished_status |
| 2676 | self._free_request(request, delay_free_blocks=delay_free_blocks) |
| 2677 | |
| 2678 | return valid_requests |
| 2679 | |
| 2680 | def _free_request( |
| 2681 | self, request: Request, delay_free_blocks: bool = False |
| 2682 | ) -> tuple[dict[str, Any] | None, dict[str, Any] | None]: |
| 2683 | assert request.is_finished() |
| 2684 | |
| 2685 | if self.aux_output_connector is not None: |
| 2686 | self.aux_output_connector.request_finished(request) |
| 2687 | self._inflight_prefills.discard(request) |
| 2688 | self._set_kv_fetch_stage(request, None) |
| 2689 | connector_delay_free_blocks, kv_xfer_params = self._connector_finished(request) |
| 2690 | |
| 2691 | # EC Connector: mirror the KV hook. The contract requires firing |
| 2692 | # before the encoder cache is freed so the connector can inspect |
| 2693 | # per-request state (e.g. which mm_hashes it recorded during |
| 2694 | # save_caches()) and emit ec_transfer_params for the response body. |
| 2695 | ec_xfer_params: dict[str, Any] | None = None |
| 2696 | if self.ec_connector is not None: |
| 2697 | ec_delay_free, ec_xfer_params = self.ec_connector.request_finished(request) |
| 2698 | connector_delay_free_blocks |= ec_delay_free |
| 2699 | |
| 2700 | self.encoder_cache_manager.free(request) |
| 2701 | request_id = request.request_id |
| 2702 | self.finished_req_ids.add(request_id) |
| 2703 | if self.finished_req_ids_dict is not None: |
| 2704 | self.finished_req_ids_dict[request.client_index].add(request_id) |
| 2705 | |
| 2706 | delay_free_blocks |= connector_delay_free_blocks |
| 2707 | if not delay_free_blocks: |
| 2708 | self._free_blocks(request) |
| 2709 | |
| 2710 | return kv_xfer_params, ec_xfer_params |
| 2711 | |
| 2712 | def _free_blocks(self, request: Request): |
| 2713 | assert request.is_finished() |
| 2714 | self._free_request_blocks(request) |
| 2715 | del self.requests[request.request_id] |
| 2716 | |
| 2717 | @property |
| 2718 | def pause_state(self) -> PauseState: |
| 2719 | return self._pause_state |
| 2720 | |
| 2721 | def set_pause_state(self, pause_state: PauseState) -> None: |
| 2722 | logger.info("setting pause state to %s", pause_state.name) |
| 2723 | self._pause_state = pause_state |
| 2724 | |
| 2725 | def _request_blocks_can_be_freed(self, request: Request) -> bool: |
| 2726 | # We must defer freeing blocks if an async kv connector may |
| 2727 | # write to them immediately (not ordered with GPU stream). |
| 2728 | return not self.defer_block_free or ( |
| 2729 | request.last_sched_seq <= self.processed_step_seq |
| 2730 | ) |
| 2731 | |
| 2732 | def _free_request_blocks(self, request: Request): |
| 2733 | """Free the request's KV blocks, deferring the return to the block |
| 2734 | pool when an in-flight GPU step may still write them. |
| 2735 | """ |
| 2736 | if self._request_blocks_can_be_freed(request): |
| 2737 | self.kv_cache_manager.free(request) |
| 2738 | return |
| 2739 | blocks = self.kv_cache_manager.pop_blocks_for_free(request) |
| 2740 | if blocks: |
| 2741 | self.deferred_frees.append((self.sched_step_seq, blocks)) |
| 2742 | |
| 2743 | def _free_cow_retained_blocks( |
| 2744 | self, blocks: list[KVCacheBlock], fence_seq: int |
| 2745 | ) -> None: |
| 2746 | """Release CoW copy retentions, deferring their return to the block |
| 2747 | pool while the step that runs the copy may still be in flight. |
| 2748 | """ |
| 2749 | if not self.defer_block_free or fence_seq <= self.processed_step_seq: |
| 2750 | self.kv_cache_manager.block_pool.free_blocks(blocks) |
| 2751 | return |
| 2752 | self.deferred_frees.append((fence_seq, blocks[::-1])) |
| 2753 | |
| 2754 | def _drain_deferred_frees(self): |
| 2755 | """Return deferred blocks whose fence step has completed. |
| 2756 | |
| 2757 | Fences are appended in near-monotonic order (a CoW retention fence |
| 2758 | can lead request-free fences by one step), so stop at the first |
| 2759 | pending one; any satisfied entry behind it is merely freed later. |
| 2760 | """ |
| 2761 | while self.deferred_frees: |
| 2762 | fence, _ = self.deferred_frees[0] |
| 2763 | if fence > self.processed_step_seq: |
| 2764 | break |
| 2765 | _, blocks = self.deferred_frees.popleft() |
| 2766 | # Free in reverse order so that the tail blocks are evicted first. |
| 2767 | self.kv_cache_manager.block_pool.free_blocks(reversed(blocks)) |
| 2768 | |
| 2769 | def get_num_unfinished_requests(self) -> int: |
| 2770 | if self._pause_state == PauseState.PAUSED_ALL: |
| 2771 | return 0 |
| 2772 | num_running, num_waiting = self.get_request_counts() |
| 2773 | if self._pause_state == PauseState.PAUSED_NEW: |
| 2774 | return num_running |
| 2775 | num_waiting -= self.num_waiting_for_streaming_input |
| 2776 | return num_waiting + num_running |
| 2777 | |
| 2778 | def has_finished_requests(self) -> bool: |
| 2779 | if self.finished_req_ids: |
| 2780 | return True |
| 2781 | if self.connector is None: |
| 2782 | return False |
| 2783 | # Finished requests waiting on delayed connector cleanup remain in |
| 2784 | # self.requests after they have been removed from scheduling queues. |
| 2785 | num_running, num_waiting = self.get_request_counts() |
| 2786 | return len(self.requests) > (num_running + num_waiting) |
| 2787 | |
| 2788 | def has_requests(self) -> bool: |
| 2789 | # Override the interface default to also keep the engine alive while a |
| 2790 | # connector still has pending push work (e.g. push-mode WRITE transfers |
| 2791 | # in flight after all "live" requests have finished). Without this hook |
| 2792 | # the engine would quiesce before the connector can drain completions. |
| 2793 | # TODO: replace with a more general mechanism for connectors to keep |
| 2794 | # the scheduler alive. |
| 2795 | return ( |
| 2796 | self.has_unfinished_requests() |
| 2797 | or self.has_finished_requests() |
| 2798 | or (self.connector is not None and self.connector.has_pending_push_work()) |
| 2799 | or ( |
| 2800 | self.ec_connector is not None |
| 2801 | and self.ec_connector.has_pending_push_work() |
| 2802 | ) |
| 2803 | ) |
| 2804 | |
| 2805 | def reset_prefix_cache( |
| 2806 | self, reset_running_requests: bool = False, reset_connector: bool = False |
| 2807 | ) -> bool: |
| 2808 | """Reset the KV prefix cache. |
| 2809 | |
| 2810 | If reset_running_requests is True, all the running requests will be |
| 2811 | preempted and moved to the waiting queue. |
| 2812 | Otherwise, this method will only reset the KV prefix cache when there |
| 2813 | is no running requests taking KV cache. |
| 2814 | """ |
| 2815 | if reset_running_requests: |
| 2816 | # For logging. |
| 2817 | timestamp = time.monotonic() |
| 2818 | # Invalidate all the current running requests KV's by pushing them to |
| 2819 | # the waiting queue. In this case, we can reduce the ref count of all |
| 2820 | # the kv blocks to 0 and thus we can make sure the reset is successful. |
| 2821 | # Preempt in reverse order so the requests will be added back to the |
| 2822 | # running queue in FIFO order. |
| 2823 | while self.running: |
| 2824 | request = self.running.pop() |
| 2825 | self._preempt_request(request, timestamp, drop_stale_output=True) |
| 2826 | |
| 2827 | # Clear scheduled request ids cache. Since we are forcing preemption |
| 2828 | # + resumption in the same step, we must act as if these requests were |
| 2829 | # not scheduled in the prior step. They will be flushed from the |
| 2830 | # persistent batch in the model runner. |
| 2831 | self.prev_step_scheduled_req_ids.clear() |
| 2832 | |
| 2833 | reset_successful = self.kv_cache_manager.reset_prefix_cache() |
| 2834 | if reset_running_requests and not reset_successful: |
| 2835 | raise RuntimeError( |
| 2836 | "Failed to reset KV cache even when all the running requests are " |
| 2837 | "preempted and moved to the waiting queue. This is likely due to " |
| 2838 | "the presence of running requests waiting for remote KV transfer, " |
| 2839 | "which is not supported yet." |
| 2840 | ) |
| 2841 | |
| 2842 | if reset_connector: |
| 2843 | reset_successful = self.reset_connector_cache() and reset_successful |
| 2844 | |
| 2845 | if reset_successful and self.aux_output_connector is not None: |
| 2846 | self.aux_output_connector.reset() |
| 2847 | |
| 2848 | return reset_successful |
| 2849 | |
| 2850 | def reset_connector_cache(self) -> bool: |
| 2851 | if self.connector is None: |
| 2852 | # No connector attached -> nothing to reset, treat as success so |
| 2853 | # callers that unconditionally request a connector reset (e.g. as |
| 2854 | # part of a cache-clearing cascade after a weight update) don't |
| 2855 | # see reset_prefix_cache() flip to False purely because they |
| 2856 | # didn't configure a connector. |
| 2857 | logger.debug( |
| 2858 | "reset_connector requested but no KV connector is configured; " |
| 2859 | "treating as no-op success." |
| 2860 | ) |
| 2861 | return True |
| 2862 | |
| 2863 | if self.connector.reset_cache() is False: |
| 2864 | return False |
| 2865 | |
| 2866 | if self.log_stats: |
| 2867 | assert self.connector_prefix_cache_stats is not None |
| 2868 | self.connector_prefix_cache_stats.reset = True |
| 2869 | |
| 2870 | return True |
| 2871 | |
| 2872 | def reset_encoder_cache(self) -> None: |
| 2873 | """Reset the encoder cache to invalidate all cached encoder outputs. |
| 2874 | |
| 2875 | This should be called when model weights are updated to ensure |
| 2876 | stale vision embeddings are not reused. |
| 2877 | """ |
| 2878 | self.encoder_cache_manager.reset() |
| 2879 | |
| 2880 | def make_stats( |
| 2881 | self, |
| 2882 | spec_decoding_stats: SpecDecodingStats | None = None, |
| 2883 | kv_connector_stats: KVConnectorStats | None = None, |
| 2884 | cudagraph_stats: CUDAGraphStat | None = None, |
| 2885 | perf_stats: PerfStats | None = None, |
| 2886 | ec_connector_stats: ECConnectorStats | None = None, |
| 2887 | ) -> SchedulerStats | None: |
| 2888 | if not self.log_stats: |
| 2889 | return None |
| 2890 | prefix_cache_stats = self.kv_cache_manager.make_prefix_cache_stats() |
| 2891 | assert prefix_cache_stats is not None |
| 2892 | connector_prefix_cache_stats: PrefixCacheStats | None = None |
| 2893 | if self.connector_prefix_cache_stats is not None: |
| 2894 | connector_prefix_cache_stats = self.connector_prefix_cache_stats |
| 2895 | self.connector_prefix_cache_stats = PrefixCacheStats() |
| 2896 | eviction_events = ( |
| 2897 | self.kv_metrics_collector.drain_events() |
| 2898 | if self.kv_metrics_collector is not None |
| 2899 | else [] |
| 2900 | ) |
| 2901 | spec_stats = spec_decoding_stats |
| 2902 | connector_stats_payload = ( |
| 2903 | kv_connector_stats.to_dict() if kv_connector_stats else None |
| 2904 | ) |
| 2905 | ec_connector_stats_payload = ( |
| 2906 | ec_connector_stats.data if ec_connector_stats else None |
| 2907 | ) |
| 2908 | num_running, num_waiting = self.get_request_counts() |
| 2909 | num_deferred = len(self.deferred_waiting) |
| 2910 | return SchedulerStats( |
| 2911 | num_running_reqs=num_running, |
| 2912 | num_waiting_reqs=num_waiting - num_deferred, |
| 2913 | num_skipped_waiting_reqs=num_deferred, |
| 2914 | num_kv_fetch_reqs_by_stage=self._kv_fetch_counts.copy(), |
| 2915 | kv_cache_usage=self.kv_cache_manager.usage, |
| 2916 | prefix_cache_stats=prefix_cache_stats, |
| 2917 | connector_prefix_cache_stats=connector_prefix_cache_stats, |
| 2918 | kv_cache_eviction_events=eviction_events, |
| 2919 | spec_decoding_stats=spec_stats, |
| 2920 | kv_connector_stats=connector_stats_payload, |
| 2921 | cudagraph_stats=cudagraph_stats, |
| 2922 | perf_stats=perf_stats, |
| 2923 | ec_connector_stats=ec_connector_stats_payload, |
| 2924 | ) |
| 2925 | |
| 2926 | def make_spec_decoding_stats( |
| 2927 | self, |
| 2928 | spec_decoding_stats: SpecDecodingStats | None, |
| 2929 | num_draft_tokens: int, |
| 2930 | num_accepted_tokens: int, |
| 2931 | num_invalid_spec_tokens: dict[str, int] | None, |
| 2932 | request_id: str, |
| 2933 | ) -> SpecDecodingStats | None: |
| 2934 | if not self.log_stats or not num_draft_tokens: |
| 2935 | return None |
| 2936 | if spec_decoding_stats is None: |
| 2937 | spec_decoding_stats = SpecDecodingStats.new(self.num_spec_tokens) |
| 2938 | if num_invalid_spec_tokens: |
| 2939 | num_draft_tokens -= num_invalid_spec_tokens.get(request_id, 0) |
| 2940 | spec_decoding_stats.observe_draft( |
| 2941 | num_draft_tokens=num_draft_tokens, num_accepted_tokens=num_accepted_tokens |
| 2942 | ) |
| 2943 | return spec_decoding_stats |
| 2944 | |
| 2945 | def shutdown(self) -> None: |
| 2946 | logger.debug_once("[shutdown] Scheduler: start") |
| 2947 | if self.kv_event_publisher: |
| 2948 | self.kv_event_publisher.shutdown() |
| 2949 | if self.connector is not None: |
| 2950 | self.connector.shutdown() |
| 2951 | |
| 2952 | if self.ec_connector is not None: |
| 2953 | self.ec_connector.shutdown() |
| 2954 | |
| 2955 | logger.debug_once("[shutdown] Scheduler: complete") |
| 2956 | |
| 2957 | ######################################################################## |
| 2958 | # KV Connector Related Methods |
| 2959 | ######################################################################## |
| 2960 | |
| 2961 | def get_kv_connector(self) -> KVConnectorBase_V1 | None: |
| 2962 | return self.connector |
| 2963 | |
| 2964 | def get_ec_connector(self) -> ECConnectorBase | None: |
| 2965 | return self.ec_connector |
| 2966 | |
| 2967 | def get_kv_event_publisher_config(self) -> KVEventsConfig | None: |
| 2968 | return self.kv_event_publisher.get_publisher_config() |
| 2969 | |
| 2970 | def _connector_finished( |
| 2971 | self, request: Request |
| 2972 | ) -> tuple[bool, dict[str, Any] | None]: |
| 2973 | """Invoke the KV connector request_finished() method if applicable. |
| 2974 | |
| 2975 | Returns optional kv transfer parameters to be included with the |
| 2976 | request outputs. |
| 2977 | """ |
| 2978 | if self.connector is None: |
| 2979 | return False, None |
| 2980 | |
| 2981 | finished_partial_tails: list[tuple[int, int, int]] = [] |
| 2982 | kv_transfer_config = self.vllm_config.kv_transfer_config |
| 2983 | if kv_transfer_config is not None and kv_transfer_config.is_kv_producer: |
| 2984 | finished_partial_tails = ( |
| 2985 | self.kv_cache_manager.finalize_partial_tail_offloads(request) |
| 2986 | ) |
| 2987 | |
| 2988 | # Free any out-of-window prefix blocks before we hand the block table to |
| 2989 | # the connector, on the processed-token basis (see `allocate_slots`). |
| 2990 | self.kv_cache_manager.remove_skipped_blocks( |
| 2991 | request_id=request.request_id, |
| 2992 | processed_computed_tokens=max( |
| 2993 | 0, request.num_computed_tokens - request.num_in_flight_tokens |
| 2994 | ), |
| 2995 | num_prompt_tokens=request.num_prompt_tokens, |
| 2996 | ) |
| 2997 | |
| 2998 | block_ids = self.kv_cache_manager.get_block_ids_for_computed_tokens( |
| 2999 | request_id=request.request_id, |
| 3000 | num_computed_tokens=request.num_computed_tokens, |
| 3001 | ) |
| 3002 | partial_tail_delay = False |
| 3003 | if finished_partial_tails: |
| 3004 | partial_tail_delay = self.connector.register_finished_partial_tail( |
| 3005 | request, |
| 3006 | block_ids, |
| 3007 | finished_partial_tails, |
| 3008 | ) |
| 3009 | |
| 3010 | if not isinstance(self.connector, SupportsHMA): |
| 3011 | # NOTE(Kuntai): We should deprecate this code path after we enforce |
| 3012 | # all connectors to support HMA. |
| 3013 | # Hybrid memory allocator should be already turned off for this |
| 3014 | # code path, but let's double-check here. |
| 3015 | assert len(self.kv_cache_config.kv_cache_groups) == 1 |
| 3016 | delay_free, kv_xfer_params = self.connector.request_finished( |
| 3017 | request, block_ids[0] |
| 3018 | ) |
| 3019 | else: |
| 3020 | delay_free, kv_xfer_params = self.connector.request_finished_all_groups( |
| 3021 | request, block_ids |
| 3022 | ) |
| 3023 | return delay_free or partial_tail_delay, kv_xfer_params |
| 3024 | |
| 3025 | def _request_remaining_blocks(self, request: Request) -> int: |
| 3026 | """Blocks `request` still needs to hold its full sequence and be promoted.""" |
| 3027 | full_num_tokens = min(request.num_tokens, self.max_model_len) |
| 3028 | num_blocks = self.kv_cache_manager.coordinator.get_num_blocks_to_allocate( |
| 3029 | request_id=request.request_id, |
| 3030 | num_tokens=full_num_tokens, |
| 3031 | new_computed_blocks=self.kv_cache_manager.empty_kv_cache_blocks.blocks, |
| 3032 | num_encoder_tokens=0, |
| 3033 | total_computed_tokens=request.num_computed_tokens, |
| 3034 | num_local_computed_tokens=request.num_computed_tokens, |
| 3035 | num_tokens_main_model=full_num_tokens, |
| 3036 | apply_admission_cap=True, |
| 3037 | prefill_end=max(request.num_prompt_tokens, request.num_tokens - 1), |
| 3038 | ) |
| 3039 | return num_blocks + self._spec_decode_step_blocks() |
| 3040 | |
| 3041 | def _spec_decode_step_blocks(self) -> int: |
| 3042 | """Number of blocks for the extra KV slots a spec decode step needs. |
| 3043 | |
| 3044 | When using async kv load, scheduler must reserve enough blocks for |
| 3045 | full sequence + the spec decode step, otherwise request cannot be |
| 3046 | able to run after the async_load if we are out of kv blocks. |
| 3047 | """ |
| 3048 | if not self.num_spec_tokens: |
| 3049 | return 0 |
| 3050 | return cdiv( |
| 3051 | 1 + self.num_spec_tokens + self.num_lookahead_tokens, self.block_size |
| 3052 | ) |
| 3053 | |
| 3054 | def _set_kv_fetch_stage(self, request: Request, stage: str | None) -> None: |
| 3055 | kv_fetch_stages = self._kv_fetch_stages |
| 3056 | if kv_fetch_stages is not None: |
| 3057 | prev_stage = kv_fetch_stages.pop(request, None) |
| 3058 | if prev_stage is not None: |
| 3059 | self._kv_fetch_counts[prev_stage] -= 1 |
| 3060 | if stage is not None: |
| 3061 | kv_fetch_stages[request] = stage |
| 3062 | self._kv_fetch_counts[stage] += 1 |
| 3063 | |
| 3064 | def _inflight_prefill_reserved_blocks(self) -> int: |
| 3065 | """Num blocks in-flight prefills still need to finish (their reservation).""" |
| 3066 | return sum( |
| 3067 | self._request_remaining_blocks(req) for req in self._inflight_prefills |
| 3068 | ) |
| 3069 | |
| 3070 | def _mark_prefix_replay(self, request: Request, num_hit_tokens: int) -> int: |
| 3071 | """Record where a prefix hit's replay starts and return the number of |
| 3072 | hit tokens to recompute; the caller rewinds the computed count by it. |
| 3073 | |
| 3074 | The replayed tokens are the hit's last window (see |
| 3075 | ``Request.replay_start``): the worker rebuilds their sliding-window KV |
| 3076 | and leaves their cached KV alone. |
| 3077 | """ |
| 3078 | if not self.prefix_replay_tokens: |
| 3079 | return 0 |
| 3080 | num_replay_tokens = min(self.prefix_replay_tokens, num_hit_tokens) |
| 3081 | request.replay_start = num_hit_tokens - num_replay_tokens |
| 3082 | return num_replay_tokens |
| 3083 | |
| 3084 | def _update_waiting_for_remote_kv(self, request: Request) -> None: |
| 3085 | """KV Connector: update request state after async recv is finished. |
| 3086 | |
| 3087 | When the kv transfer is ready, we cache the blocks |
| 3088 | and the request state will be moved back to WAITING from |
| 3089 | WAITING_FOR_REMOTE_KV. |
| 3090 | """ |
| 3091 | assert self.connector is not None |
| 3092 | |
| 3093 | if request.request_id in self.failed_recving_kv_req_ids: |
| 3094 | # Request had KV load failures; num_computed_tokens was already |
| 3095 | # updated in _update_requests_with_invalid_blocks |
| 3096 | if request.num_computed_tokens: |
| 3097 | # Cache any valid computed tokens. |
| 3098 | self.kv_cache_manager.cache_blocks(request, request.num_computed_tokens) |
| 3099 | if self.needs_kv_cache_zeroing: |
| 3100 | # The failed load left the blocks beyond the valid |
| 3101 | # prefix unwritten and their zeroing was skipped; zero |
| 3102 | # them before they are recomputed locally. |
| 3103 | self.kv_cache_manager.record_blocks_for_zeroing( |
| 3104 | request.request_id, request.num_computed_tokens |
| 3105 | ) |
| 3106 | else: |
| 3107 | # No valid computed tokens, release allocated blocks. |
| 3108 | # There may be a local cache hit on retry. |
| 3109 | # (Freed blocks are re-recorded for zeroing when |
| 3110 | # reallocated, so the skipped blocks need no handling.) |
| 3111 | self.kv_cache_manager.free(request) |
| 3112 | |
| 3113 | self.failed_recving_kv_req_ids.remove(request.request_id) |
| 3114 | else: |
| 3115 | # Now that the blocks are ready, actually cache them. |
| 3116 | # This will cache the blocks iff caching is enabled. |
| 3117 | self.kv_cache_manager.cache_blocks(request, request.num_computed_tokens) |
| 3118 | |
| 3119 | # SWA bounded replay recomputes the tail of the hit, which covers the |
| 3120 | # last token; otherwise a full prompt hit re-computes that token so the |
| 3121 | # next one can be sampled. |
| 3122 | num_replay_tokens = self._mark_prefix_replay( |
| 3123 | request, request.num_computed_tokens |
| 3124 | ) |
| 3125 | if num_replay_tokens > 0: |
| 3126 | request.num_computed_tokens -= num_replay_tokens |
| 3127 | elif request.num_computed_tokens == request.num_tokens: |
| 3128 | request.num_computed_tokens = request.num_tokens - 1 |
| 3129 | self.finished_recving_kv_req_ids.remove(request.request_id) |
| 3130 | |
| 3131 | def _handle_blocked_waiting_request(self, request: Request) -> bool: |
| 3132 | """Returns True if scheduling of this request should proceed, False if |
| 3133 | it should be skipped for this step.""" |
| 3134 | if request.status == RequestStatus.WAITING_FOR_REMOTE_KVS: |
| 3135 | # Check whether async loading requests are finished. |
| 3136 | request_id = request.request_id |
| 3137 | if request_id not in self.finished_recving_kv_req_ids: |
| 3138 | logger.debug("%s is still in WAITING_FOR_REMOTE_KVS state.", request_id) |
| 3139 | return False |
| 3140 | # Promote request to waiting or preempted. |
| 3141 | self._update_waiting_for_remote_kv(request) |
| 3142 | if request.num_preemptions: |
| 3143 | request.status = RequestStatus.PREEMPTED |
| 3144 | else: |
| 3145 | request.status = RequestStatus.WAITING |
| 3146 | |
| 3147 | elif request.status == RequestStatus.WAITING_FOR_STRUCTURED_OUTPUT_GRAMMAR: |
| 3148 | # Check whether grammar compilation has finished. |
| 3149 | struct_output_req = request.structured_output_request |
| 3150 | if not struct_output_req or struct_output_req.grammar is None: |
| 3151 | return False |
| 3152 | if isinstance(struct_output_req.grammar, Exception): |
| 3153 | self.grammar_compile_error_reqs.add(request.request_id) |
| 3154 | return False |
| 3155 | # Grammar is ready, promote to waiting state. |
| 3156 | request.status = RequestStatus.WAITING |
| 3157 | |
| 3158 | elif request.status == RequestStatus.WAITING_FOR_STREAMING_REQ: |
| 3159 | # Streaming input request waiting for more input. |
| 3160 | assert not request.streaming_queue |
| 3161 | return False |
| 3162 | |
| 3163 | return True |
| 3164 | |
| 3165 | def _update_from_kv_xfer_finished(self, kv_connector_output: KVConnectorOutput): |
| 3166 | """KV Connector: update the scheduler state based on the output. |
| 3167 | |
| 3168 | The Worker side connectors add finished_recving and |
| 3169 | finished_sending reqs to the output. |
| 3170 | * if finished_sending: free the blocks |
| 3171 | # if finished_recving: add to state so we can |
| 3172 | schedule the request during the next step. |
| 3173 | """ |
| 3174 | if self.connector is not None: |
| 3175 | self.connector.update_connector_output(kv_connector_output) |
| 3176 | |
| 3177 | # KV Connector:: update recv and send status from last step. |
| 3178 | for req_id in kv_connector_output.finished_recving or (): |
| 3179 | logger.debug("Finished recving KV transfer for request %s", req_id) |
| 3180 | assert req_id in self.requests |
| 3181 | req = self.requests[req_id] |
| 3182 | if req.status == RequestStatus.WAITING_FOR_REMOTE_KVS: |
| 3183 | self.finished_recving_kv_req_ids.add(req_id) |
| 3184 | self._set_kv_fetch_stage(req, KV_FETCH_COMPLETED_WAITING) |
| 3185 | else: |
| 3186 | assert RequestStatus.is_finished(req.status) |
| 3187 | self._free_blocks(self.requests[req_id]) |
| 3188 | for req_id in kv_connector_output.finished_sending or (): |
| 3189 | logger.debug("Finished sending KV transfer for request %s", req_id) |
| 3190 | assert req_id in self.requests |
| 3191 | self._free_blocks(self.requests[req_id]) |
| 3192 | |
| 3193 | def _update_requests_with_invalid_blocks( |
| 3194 | self, |
| 3195 | requests: Iterable[Request], |
| 3196 | invalid_block_ids: set[int], |
| 3197 | num_scheduled_tokens: dict[str, int], |
| 3198 | evict_blocks: bool = True, |
| 3199 | ) -> tuple[set[str], int, set[int]]: |
| 3200 | """Identify and update requests affected by invalid KV cache blocks. |
| 3201 | |
| 3202 | This method scans the given requests, detects those with invalid blocks |
| 3203 | and adjusts their `num_computed_tokens` to the longest valid prefix. |
| 3204 | For observability, it also accumulates the total number of tokens that |
| 3205 | will need to be recomputed across all affected requests. |
| 3206 | |
| 3207 | Args: |
| 3208 | requests: The set of requests to scan for invalid blocks. |
| 3209 | invalid_block_ids: IDs of invalid blocks. |
| 3210 | num_scheduled_tokens: req_id -> number of scheduled tokens. |
| 3211 | evict_blocks: Whether to collect blocks for eviction (False for |
| 3212 | async requests which aren't cached yet). |
| 3213 | |
| 3214 | Returns: |
| 3215 | tuple: |
| 3216 | - affected_req_ids (set[str]): IDs of requests impacted by |
| 3217 | invalid blocks. |
| 3218 | - total_affected_tokens (int): Total number of tokens that must |
| 3219 | be recomputed across all affected requests. |
| 3220 | - blocks_to_evict (set[int]): Block IDs to evict from cache, |
| 3221 | including invalid blocks and downstream dependent blocks. |
| 3222 | |
| 3223 | """ |
| 3224 | affected_req_ids: set[str] = set() |
| 3225 | total_affected_tokens = 0 |
| 3226 | blocks_to_evict: set[int] = set() |
| 3227 | # If a block is invalid and shared by multiple requests in the batch, |
| 3228 | # these requests must be rescheduled, but only the first will recompute |
| 3229 | # it. This set tracks blocks already marked for recomputation. |
| 3230 | marked_invalid_block_ids: set[int] = set() |
| 3231 | for request in requests: |
| 3232 | is_affected = False |
| 3233 | marked_invalid_block = False |
| 3234 | req_id = request.request_id |
| 3235 | # One block-id list per KV cache group; single-group layouts are |
| 3236 | # enforced by _handle_invalid_blocks. |
| 3237 | (req_block_ids,) = self.kv_cache_manager.get_block_ids(req_id) |
| 3238 | # We iterate only over blocks that may contain externally computed |
| 3239 | # tokens |
| 3240 | req_num_computed_tokens = ( |
| 3241 | request.num_computed_tokens - num_scheduled_tokens.get(req_id, 0) |
| 3242 | ) |
| 3243 | |
| 3244 | req_num_computed_blocks = ( |
| 3245 | req_num_computed_tokens + self.block_size - 1 |
| 3246 | ) // self.block_size |
| 3247 | for idx, block_id in zip(range(req_num_computed_blocks), req_block_ids): |
| 3248 | if block_id not in invalid_block_ids: |
| 3249 | continue |
| 3250 | |
| 3251 | is_affected = True |
| 3252 | |
| 3253 | if block_id in marked_invalid_block_ids: |
| 3254 | # This invalid block is shared with a previous request |
| 3255 | # and was already marked for recomputation. |
| 3256 | # This means this request can still consider this block |
| 3257 | # as computed when rescheduled. |
| 3258 | # Currently this only applies to sync loading; Async |
| 3259 | # loading does not yet support block sharing |
| 3260 | continue |
| 3261 | |
| 3262 | marked_invalid_block_ids.add(block_id) |
| 3263 | |
| 3264 | if marked_invalid_block: |
| 3265 | # This request has already marked an invalid block for |
| 3266 | # recomputation and updated its num_computed_tokens. |
| 3267 | continue |
| 3268 | |
| 3269 | marked_invalid_block = True |
| 3270 | # Truncate the computed tokens at the first failed block |
| 3271 | request.num_computed_tokens = idx * self.block_size |
| 3272 | num_affected_tokens = ( |
| 3273 | req_num_computed_tokens - request.num_computed_tokens |
| 3274 | ) |
| 3275 | total_affected_tokens += num_affected_tokens |
| 3276 | |
| 3277 | # collect invalid block and all downstream dependent blocks |
| 3278 | if evict_blocks: |
| 3279 | blocks_to_evict.update(req_block_ids[idx:]) |
| 3280 | |
| 3281 | if is_affected: |
| 3282 | if not marked_invalid_block: |
| 3283 | # All invalid blocks of this request are shared with |
| 3284 | # previous requests and will be recomputed by them. |
| 3285 | # Revert to considering only cached tokens as computed. |
| 3286 | # Currently this only applies to sync loading; Async |
| 3287 | # loading does not yet support block sharing |
| 3288 | total_affected_tokens += ( |
| 3289 | request.num_computed_tokens - req_num_computed_tokens |
| 3290 | ) |
| 3291 | request.num_computed_tokens = req_num_computed_tokens |
| 3292 | |
| 3293 | affected_req_ids.add(request.request_id) |
| 3294 | |
| 3295 | return affected_req_ids, total_affected_tokens, blocks_to_evict |
| 3296 | |
| 3297 | def _handle_failed_recving(self, failed_req_ids: set[str]) -> set[str]: |
| 3298 | """Fail closed for layouts whose block IDs are not globally unique.""" |
| 3299 | affected_req_ids: set[str] = set() |
| 3300 | for req_id in failed_req_ids: |
| 3301 | request = self.requests.get(req_id) |
| 3302 | if ( |
| 3303 | request is not None |
| 3304 | and request.status == RequestStatus.WAITING_FOR_REMOTE_KVS |
| 3305 | ): |
| 3306 | affected_req_ids.add(req_id) |
| 3307 | if self.recompute_kv_load_failures: |
| 3308 | request.num_computed_tokens = 0 |
| 3309 | self.failed_recving_kv_req_ids.add(req_id) |
| 3310 | return set() if self.recompute_kv_load_failures else affected_req_ids |
| 3311 | |
| 3312 | def _handle_invalid_blocks( |
| 3313 | self, invalid_block_ids: set[int], num_scheduled_tokens: dict[str, int] |
| 3314 | ) -> set[str]: |
| 3315 | """Handle requests affected by invalid KV cache blocks. |
| 3316 | |
| 3317 | Returns: |
| 3318 | Set of affected request IDs to skip in update_from_output main loop. |
| 3319 | |
| 3320 | """ |
| 3321 | if len(self.kv_cache_config.kv_cache_groups) > 1: |
| 3322 | # Block IDs are only unique within a group, so a flat set of |
| 3323 | # invalid block IDs cannot be mapped back to requests. |
| 3324 | raise RuntimeError( |
| 3325 | "A KV connector reported block-level load failures " |
| 3326 | "(invalid_block_ids) which is not supported for models " |
| 3327 | "with multiple KV cache groups. " |
| 3328 | "Connectors must report failed requests via " |
| 3329 | "KVConnectorTransferResults.failed_recving instead." |
| 3330 | ) |
| 3331 | |
| 3332 | should_fail = not self.recompute_kv_load_failures |
| 3333 | |
| 3334 | # handle async KV loads (not cached yet, evict_blocks=False) |
| 3335 | async_load_reqs = ( |
| 3336 | req |
| 3337 | for req in self.kv_holding_waiting |
| 3338 | if req.status == RequestStatus.WAITING_FOR_REMOTE_KVS |
| 3339 | ) |
| 3340 | async_failed_req_ids, num_failed_tokens, _ = ( |
| 3341 | self._update_requests_with_invalid_blocks( |
| 3342 | async_load_reqs, |
| 3343 | invalid_block_ids, |
| 3344 | num_scheduled_tokens, |
| 3345 | evict_blocks=False, |
| 3346 | ) |
| 3347 | ) |
| 3348 | |
| 3349 | total_failed_requests = len(async_failed_req_ids) |
| 3350 | total_failed_tokens = num_failed_tokens |
| 3351 | |
| 3352 | # handle sync loads (may be cached, collect blocks for eviction) |
| 3353 | sync_failed_req_ids, num_failed_tokens, sync_blocks_to_evict = ( |
| 3354 | self._update_requests_with_invalid_blocks( |
| 3355 | self.running, invalid_block_ids, num_scheduled_tokens, evict_blocks=True |
| 3356 | ) |
| 3357 | ) |
| 3358 | |
| 3359 | total_failed_requests += len(sync_failed_req_ids) |
| 3360 | total_failed_tokens += num_failed_tokens |
| 3361 | |
| 3362 | if not total_failed_requests: |
| 3363 | return set() |
| 3364 | |
| 3365 | # evict invalid blocks and downstream dependent blocks from cache |
| 3366 | # only when not using recompute policy (where blocks will be recomputed |
| 3367 | # and reused by other requests sharing them) |
| 3368 | if sync_blocks_to_evict and not self.recompute_kv_load_failures: |
| 3369 | self.kv_cache_manager.evict_blocks(sync_blocks_to_evict) |
| 3370 | |
| 3371 | if should_fail: |
| 3372 | all_failed_req_ids = async_failed_req_ids | sync_failed_req_ids |
| 3373 | logger.error( |
| 3374 | "Failing %d request(s) due to KV load failure " |
| 3375 | "(failure_policy=fail, %d tokens affected). Request IDs: %s", |
| 3376 | total_failed_requests, |
| 3377 | total_failed_tokens, |
| 3378 | all_failed_req_ids, |
| 3379 | ) |
| 3380 | return all_failed_req_ids |
| 3381 | |
| 3382 | logger.warning( |
| 3383 | "Recovered from KV load failure: " |
| 3384 | "%d request(s) rescheduled (%d tokens affected).", |
| 3385 | total_failed_requests, |
| 3386 | total_failed_tokens, |
| 3387 | ) |
| 3388 | |
| 3389 | # Mark async requests with KV load failures for retry once loading completes |
| 3390 | self.failed_recving_kv_req_ids |= async_failed_req_ids |
| 3391 | # Return sync affected IDs to skip in update_from_output |
| 3392 | return sync_failed_req_ids |