diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index 294a2b505a..ff68486917 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -237,6 +237,37 @@ class KVCacheManager: # whenever a running request would otherwise fail to allocate # (``yield_pinned_hits``), so a pin never causes a preemption. self._pinned_hits: dict[str, list[KVCacheBlock]] = {} + # A finished request's hashed blocks (its attention prefix and + # the GDN states it holds) stay referenced as a "grace" pin until the + # conversation's next turn is admitted on them. With waiter pins alone, + # every remaining total miss had lost its top attention block BEFORE + # the next turn was first looked up -- in the gap between the + # previous turn's free and the next turn's arrival, or while its + # admission kept failing -- when nothing referenced the prefix and + # the running set's allocations popped it; once the hybrid hit is 0 + # a waiter pin has nothing to pin. Grace pins yield first (oldest + # first) whenever anyone needs blocks, so they never cost a running + # request or an admission anything a free queue would have given. + # A grace is one extra ref per block, taken on whichever free path + # the request takes (immediate or deferred behind an in-flight step). + self._grace_pins: dict[str, list[KVCacheBlock]] = {} + self._grace_block_ids: dict[str, frozenset[int]] = {} + self._grace_prompt_tokens: dict[str, int] = {} + # Only for fine-grained hybrid hits (align Mamba hashed finer than its + # block): there the next turn's key sits in the top attention block + # and one lost block is a total miss; at a block-sized unit the top + # block is a spare and nothing changes. The scheduler also turns + # graces off when a KV connector is configured (as for waiter pins). + self.enable_grace_pins = ( + self.enable_caching and self.coordinator.enable_partial_hash_hits + ) + # Shortfall (blocks) of the last allocate_slots that returned None. + self.last_alloc_shortfall = 0 + # A producer that moves its live GDN tail's hash to a CoW copy keeps + # writing the source; pins and graces follow the hash to the copy. + for manager in self.coordinator.single_type_managers: + if isinstance(manager, MambaManager): + manager.moved_hash_listeners.append(self._pins_follow_moved_hash) @property def usage(self) -> float: @@ -551,6 +582,9 @@ class KVCacheManager: ) required_blocks = num_blocks_to_allocate + watermark_blocks if required_blocks > self.block_pool.get_num_free_blocks(): + self.last_alloc_shortfall = ( + required_blocks - self.block_pool.get_num_free_blocks() + ) return None num_tokens_main_model = total_computed_tokens + num_new_tokens @@ -591,7 +625,9 @@ class KVCacheManager: required_blocks = num_blocks_to_allocate + watermark_blocks if required_blocks > available_blocks: # Cannot allocate new blocks + self.last_alloc_shortfall = required_blocks - available_blocks return None + self.last_alloc_shortfall = 0 if ( new_computed_block_list is not self.empty_kv_cache_blocks.blocks @@ -618,6 +654,10 @@ class KVCacheManager: # remote. Update state for locally cached blocks. if delay_cache_blocks: self.release_pinned_hit(request.request_id) + if self._grace_pins and num_new_computed_tokens > 0: + self._release_grace_for_hit( + new_computed_block_list, num_local_computed_tokens + ) return self.create_kv_cache_blocks(new_blocks) # NOTE(woosuk): We want to commit (cache) up to num_local_computed_tokens @@ -631,8 +671,13 @@ class KVCacheManager: ) self.coordinator.cache_blocks(request, num_tokens_to_cache) - # Admitted: the table now references the hit, the pin is redundant. + # Admitted: the table now references the hit, the pins are redundant. self.release_pinned_hit(request.request_id) + if self._grace_pins and num_new_computed_tokens > 0: + # ... and so does the grace pin of the turn it resumes. + self._release_grace_for_hit( + new_computed_block_list, num_local_computed_tokens + ) return self.create_kv_cache_blocks(new_blocks) def free(self, request: Request) -> None: @@ -645,7 +690,220 @@ class KVCacheManager: """ self.release_pinned_hit(request.request_id) - self.coordinator.free(request.request_id) + if not self._grace_eligible(request): + self.coordinator.free(request.request_id) + return + # Same blocks and order as coordinator.free (the deferred path frees + # this list as-is too); the grace takes its ref before the table's + # goes, so the hashed blocks never pass through the free queue. + blocks = self.coordinator.pop_blocks_for_free(request.request_id) + self._take_grace_pin(request, blocks) + self.block_pool.free_blocks(blocks) + + # Finished requests whose hashed blocks stay pinned for the next turn + # (oldest released first beyond this). One per conversation is the + # need; twice the deployed max_num_seqs. + MAX_GRACE_PINS = 16 + + @property + def num_grace_pins(self) -> int: + return len(self._grace_pins) + + @property + def num_grace_pinned_blocks(self) -> int: + return sum(len(blocks) for blocks in self._grace_pins.values()) + + def grace_pinned_blocks(self, request_id: str) -> list[KVCacheBlock]: + return list(self._grace_pins.get(request_id, ())) + + def _grace_eligible(self, request: Request) -> bool: + """A turn that finished with output (stop, length, repetition, or + an abort after decoding) may have a next turn that resumes it. A + preemption (still RUNNING), an aborted waiter (no output), an error + or an ignored request leaves nothing worth holding.""" + return ( + self.enable_grace_pins + and self.enable_caching + and self.prefix_cache_lookup_enabled(request) + and request.is_finished() + and request.status + not in (RequestStatus.FINISHED_ERROR, RequestStatus.FINISHED_IGNORED) + and request.num_output_tokens > 0 + ) + + def _take_grace_pin(self, request: Request, blocks: list[KVCacheBlock]) -> int: + """Take one extra ref on every hashed block of a finished request's + just-popped blocks (``blocks`` is in eviction order: each group's + table tail first, the held GDN states last). The caller then drops + the table refs, now or after the in-flight fence; the hashed blocks + stay out of the free queue until the grace is released. Unhashed + blocks (running state, MTP scratch, an unkeyed last block) go back + with the table refs as before. Returns the number of blocks held.""" + keep: list[KVCacheBlock] = [] + seen: set[int] = set() + for block in blocks: + if block.is_null or block.block_hash is None or id(block) in seen: + continue + # Still referenced by the request's table or hold: the grace ref + # never moves a block in or out of a free queue (any pool). + assert block.ref_cnt > 0 + seen.add(id(block)) + keep.append(block) + if not keep: + return 0 + request_id = request.request_id + self.release_grace_pin(request_id) # a repeated free: replace + for block in keep: + block.ref_cnt += 1 + self._grace_pins[request_id] = keep + self._grace_block_ids[request_id] = frozenset(seen) + self._grace_prompt_tokens[request_id] = request.num_prompt_tokens + while len(self._grace_pins) > self.MAX_GRACE_PINS: + self.release_grace_pin(next(iter(self._grace_pins))) + return len(keep) + + def release_grace_pin(self, request_id: str) -> int: + """Return a finished request's grace-pinned blocks to the pool (in + the eviction order a free would have used). Returns the count.""" + blocks = self._grace_pins.pop(request_id, None) + self._grace_block_ids.pop(request_id, None) + self._grace_prompt_tokens.pop(request_id, None) + if not blocks: + return 0 + self.block_pool.free_blocks(blocks) + return len(blocks) + + def _release_grace_for_hit( + self, + computed_blocks: tuple[Sequence[KVCacheBlock], ...], + num_hit_tokens: int, + ) -> None: + """An admitted hit now references its blocks through the table, so + the grace of the turn it resumes has done its job. That turn is a + grace that holds the deepest hit block of every group whose deepest + hit block any grace holds (the GDN state it resumes from, the + attention block holding the key), and whose prompt the hit reaches + to within two prefix units (the EAGLE-dropped tail key). A new turn + that only shares a system prompt with finished conversations + releases none of theirs.""" + candidates: set[str] | None = None + for group in computed_blocks: + deepest = next((b for b in reversed(group) if not b.is_null), None) + if deepest is None: + continue + owners = { + request_id + for request_id, block_ids in self._grace_block_ids.items() + if id(deepest) in block_ids + } + if owners: + candidates = owners if candidates is None else candidates & owners + if not candidates: + return + reach = num_hit_tokens + 2 * self.block_pool.hash_block_size + for request_id in candidates: + if reach >= self._grace_prompt_tokens.get(request_id, 0): + self.release_grace_pin(request_id) + + def _pins_follow_moved_hash( + self, source_block: KVCacheBlock, cow_block: KVCacheBlock + ) -> None: + """A Mamba producer moved ``source_block``'s hashes to ``cow_block`` + and keeps writing ``source_block``. A grace or waiter pin on the + source follows to the copy (as a hold does), so it protects the + hashed state and never leaves an extra ref on a block its owner must + hold exclusively (speculative relocation asserts on that).""" + for request_id, blocks in self._grace_pins.items(): + if id(source_block) in self._grace_block_ids[request_id]: + self._swap_pinned_block(blocks, source_block, cow_block) + self._grace_block_ids[request_id] = frozenset( + id(block) for block in blocks + ) + for blocks in self._pinned_hits.values(): + self._swap_pinned_block(blocks, source_block, cow_block) + + def _swap_pinned_block( + self, + blocks: list[KVCacheBlock], + source_block: KVCacheBlock, + cow_block: KVCacheBlock, + ) -> None: + for i, block in enumerate(blocks): + if block is source_block: + if any(other is cow_block for other in blocks): + del blocks[i] + else: + assert cow_block.ref_cnt > 0 + cow_block.ref_cnt += 1 + blocks[i] = cow_block + self.block_pool.free_blocks([source_block]) + return + + def num_yieldable_blocks( + self, + keep: str | None = None, + keep_hit: tuple[Sequence[KVCacheBlock], ...] | None = None, + ) -> int: + """Blocks that would return to the free queue if every grace and + every waiter pin except ``keep``'s were released now. With + ``keep_hit`` (the hit ``keep`` is being admitted on), ``keep``'s + pinned blocks outside that hit (the proof block past it) count too: + ``release_pinned_extras`` can give them back. Blocks of that hit + never count: freed, they are taken straight back by the admission + (a hit that grew past the pin snapshot would otherwise be counted + through the grace that still holds it).""" + pins: dict[int, list] = {} + holders = [*self._grace_pins.values()] + [ + blocks + for request_id, blocks in self._pinned_hits.items() + if request_id != keep + ] + hit_ids: set[int] = set() + if keep is not None and keep_hit is not None: + holders.append(self._pinned_extras(keep, keep_hit)) + hit_ids = {id(block) for group in keep_hit for block in group} + for blocks in holders: + for block in blocks: + entry = pins.setdefault(id(block), [block, 0]) + entry[1] += 1 + return sum( + 1 + for block_id, (block, n) in pins.items() + if block.ref_cnt == n and block_id not in hit_ids + ) + + def _pinned_extras( + self, request_id: str, hit: tuple[Sequence[KVCacheBlock], ...] + ) -> list[KVCacheBlock]: + hit_ids = {id(block) for group in hit for block in group} + return [ + block + for block in self._pinned_hits.get(request_id, ()) + if id(block) not in hit_ids + ] + + def release_pinned_extras( + self, request_id: str, hit: tuple[Sequence[KVCacheBlock], ...] + ) -> int: + """Drop from the request's pin the blocks outside ``hit``: the proof + block past it, and any hit block a later lookup added. The request + is being admitted on ``hit`` now, so they protect nothing it still + needs; they cost it blocks. Returns the number dropped.""" + extras = self._pinned_extras(request_id, hit) + if not extras: + return 0 + extra_ids = {id(block) for block in extras} + rest = [ + block + for block in self._pinned_hits[request_id] + if id(block) not in extra_ids + ] + if rest: + self._pinned_hits[request_id] = rest + else: + del self._pinned_hits[request_id] + self.block_pool.free_blocks(extras) + return len(extras) # Waiting requests whose prefix-cache hit is pinned at once (arrival # order). Pins never block anyone (they yield before a running request @@ -671,11 +929,10 @@ class KVCacheManager: Known gap: a pin is a snapshot of the hit at pin time. It is not refreshed if the hit grows while the request waits (a sibling - computing more of the same prefix) or if a producer moves the hash of - a pinned live block into its CoW copy (a resume at a running - producer's tail, unreachable under the EAGLE drop); the admission - lookup still finds whatever is cached then, so the only cost is a - smaller hit, never a wrong one.""" + computing more of the same prefix); the admission lookup still finds + whatever is cached then, so the only cost is a smaller hit, never a + wrong one. A producer moving the hash of a pinned live GDN block into + its CoW copy is followed (``_pins_follow_moved_hash``).""" request_id = request.request_id if ( request_id in self._pinned_hits @@ -734,13 +991,21 @@ class KVCacheManager: self, keep: str | None = None, limit: int | None = None ) -> int: """Drop pins (except ``keep``'s) so a running request can allocate, - or a waiter ahead in line can be admitted: the most recently pinned - waiters first, at most ``limit`` of them per call (all by default; - the scheduler passes 1 and retries, so a decode step that needs a - block costs one waiter its pin, not every waiter). Returns the number - of blocks released; 0 means there was nothing to yield.""" + or a waiter ahead in line can be admitted: grace pins of finished + turns first (oldest first), then the most recently pinned waiters, at + most ``limit`` pins per call (all by default; the scheduler passes 1 + and retries, so a decode step that needs a block costs one pin, not + every pin). Returns the number of blocks released; 0 means there was + nothing to yield.""" released = 0 yielded = 0 + # Grace pins first (oldest finished conversation first), then the + # waiters' pins (most recently pinned first). + for request_id in list(self._grace_pins): + if limit is not None and yielded >= limit: + return released + released += self.release_grace_pin(request_id) + yielded += 1 for request_id in reversed(list(self._pinned_hits)): if request_id == keep: continue @@ -787,7 +1052,12 @@ class KVCacheManager: # A pin is never written by an in-flight step (the waiter has no # table), so it is released now, not with the deferred blocks. self.release_pinned_hit(request.request_id) - return self.coordinator.pop_blocks_for_free(request.request_id) + blocks = self.coordinator.pop_blocks_for_free(request.request_id) + # A finished turn whose last step is still in flight: the grace takes + # its ref now; the table refs go when the fence drains. + if self._grace_eligible(request): + self._take_grace_pin(request, blocks) + return blocks def evict_blocks(self, block_ids: set[int]) -> None: """Evict blocks from the prefix cache by their block IDs. @@ -850,6 +1120,12 @@ class KVCacheManager: group. """ + if self._pinned_hits or self._grace_pins: + # A pinned or graced block carries a ref that is not a running + # request's, so the ref count can equal the running count for a + # block only some of them share. Report no common prefix (cascade + # attention is an optimisation; a wrong prefix would be read). + return [0] * len(self.kv_cache_config.kv_cache_groups) return self.coordinator.get_num_common_prefix_blocks(running_request_id) def take_events(self) -> list[KVCacheEvent]: diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 671417a105..c7cf3b3fe5 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -366,6 +366,9 @@ class Scheduler(SchedulerInterface): # Bind after construction so connectors can access the cache manager. if self.connector is not None: self.connector.bind_kv_cache_manager(self.kv_cache_manager) + # Finished turns' grace pins are off with a connector, as + # waiter pins are (_pin_waiting_hits). + self.kv_cache_manager.enable_grace_pins = False # Ring blocks an async KV load is about to fill. Their tag clear # is skipped for the same reason as zeroing (it could land after the # out-of-band write and wipe the producer's committed tags). Tracked @@ -583,6 +586,15 @@ class Scheduler(SchedulerInterface): self.hash_block_size, self.block_size, ) + if manager.enable_grace_pins: + logger.info( + "Prefix cache: a finished turn's cached prefix stays referenced " + "until its next turn is admitted (up to %d turns; hash unit %d, " + "block %d)", + manager.MAX_GRACE_PINS, + self.hash_block_size, + self.block_size, + ) def _pin_waiting_hits(self) -> None: """Pin the prefix-cache hits of the requests at the front of the @@ -1326,12 +1338,89 @@ class Scheduler(SchedulerInterface): has_scheduled_reqs=bool(self.running), ) - if new_blocks is None and self.kv_cache_manager.yield_pinned_hits( - keep=request_id + # Finished turns' grace pins (oldest first), then the + # pins of waiters behind this one, yield to it one at a time + # (arrival order). Its own hit is pinned first (a turn that + # arrived this step is not pinned yet and its hit may sit in + # a grace that yields), so only blocks it does not resume + # from go back. And only when releasing them all would cover + # the shortfall: otherwise yielding costs other conversations + # their prefix and still does not admit this one. + kv_cache_manager = self.kv_cache_manager + yield_for_head = new_blocks is None and ( + kv_cache_manager.last_alloc_shortfall > 0 + and ( + kv_cache_manager.num_grace_pins > 0 + or kv_cache_manager.num_pinned_waiters > 0 + ) + ) + if yield_for_head: + if ( + self.connector is None + and request.num_computed_tokens == 0 + and kv_cache_manager.pin_waiting_hit(request) > 0 + ): + # The pin just took blocks out of the free queue that + # the failed attempt counted as free: the proof block + # past the hit, and any hit block the lookup above did + # not return. Recompute the shortfall on the pinned + # state before measuring it against the others' pins. + new_blocks = self.kv_cache_manager.allocate_slots( + request, + num_new_tokens, + num_new_computed_tokens=num_new_local_computed_tokens, + new_computed_blocks=new_computed_blocks, + num_lookahead_tokens=effective_lookahead_tokens, + num_external_computed_tokens=num_external_computed_tokens, + delay_cache_blocks=load_kv_async, + num_encoder_tokens=num_encoder_tokens, + full_sequence_must_fit=self.scheduler_reserve_full_isl, + reserved_blocks=reserved_blocks, + has_scheduled_reqs=bool(self.running), + ) + yield_for_head = new_blocks is None and ( + kv_cache_manager.last_alloc_shortfall + <= kv_cache_manager.num_yieldable_blocks( + keep=request_id, keep_hit=new_computed_blocks.blocks + ) + ) + if ( + yield_for_head + and kv_cache_manager.release_pinned_extras( + request_id, new_computed_blocks.blocks + ) + > 0 + ): + # Its own pinned blocks outside the hit it is admitted + # on (the proof block past it) go first: they cost + # nobody else anything, and at an exact fit they are + # the blocks it needs. Retry on them before the others + # yield. + new_blocks = self.kv_cache_manager.allocate_slots( + request, + num_new_tokens, + num_new_computed_tokens=num_new_local_computed_tokens, + new_computed_blocks=new_computed_blocks, + num_lookahead_tokens=effective_lookahead_tokens, + num_external_computed_tokens=num_external_computed_tokens, + delay_cache_blocks=load_kv_async, + num_encoder_tokens=num_encoder_tokens, + full_sequence_must_fit=self.scheduler_reserve_full_isl, + reserved_blocks=reserved_blocks, + has_scheduled_reqs=bool(self.running), + ) + # Coverage is re-read before every yield: a retry can fail + # by more than the attempt before it (the full-sequence gate + # passes, the lookahead gate does not). + while ( + yield_for_head + and new_blocks is None + and kv_cache_manager.last_alloc_shortfall + <= kv_cache_manager.num_yieldable_blocks( + keep=request_id, keep_hit=new_computed_blocks.blocks + ) + and kv_cache_manager.yield_pinned_hits(keep=request_id, limit=1) ): - # Pins of waiters behind this one yield to it - # (arrival order); its own pin stays, those blocks are - # its hit. One retry with them back in the free queue. new_blocks = self.kv_cache_manager.allocate_slots( request, num_new_tokens, diff --git a/vllm/v1/core/sched/scheduler.py.orig b/vllm/v1/core/sched/scheduler.py.orig new file mode 100644 index 0000000000..671417a105 --- /dev/null +++ b/vllm/v1/core/sched/scheduler.py.orig @@ -0,0 +1,3551 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import itertools +import os +import time +from collections import defaultdict, deque +from collections.abc import Iterable +from dataclasses import replace +from typing import Any + +from vllm.compilation.cuda_graph import CUDAGraphStat +from vllm.config import KVEventsConfig, VllmConfig +from vllm.distributed.aux_output_connector.connector import AuxOutputSchedulerConnector +from vllm.distributed.ec_transfer.ec_connector.base import ( + ECConnectorBase, + ECConnectorMetadata, + ECConnectorRole, +) +from vllm.distributed.ec_transfer.ec_connector.factory import ECConnectorFactory +from vllm.distributed.ec_transfer.ec_connector.metrics import ECConnectorStats +from vllm.distributed.kv_events import EventPublisherFactory, KVEventBatch +from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory +from vllm.distributed.kv_transfer.kv_connector.v1 import ( + KVConnectorBase_V1, + KVConnectorRole, + SupportsHMA, +) +from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata +from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats +from vllm.logger import init_logger +from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry +from vllm.multimodal.encoder_budget import MultiModalBudget +from vllm.multimodal.utils import get_mm_features_in_window +from vllm.utils.math_utils import cdiv +from vllm.v1.core.encoder_cache_manager import ( + EncoderCacheManager, + EncoderDecoderCacheManager, +) +from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager +from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector +from vllm.v1.core.kv_cache_utils import KVCacheBlock +from vllm.v1.core.sched import e2_alloc_log as _e2_alloc_log +from vllm.v1.core.sched.interface import PauseState, SchedulerInterface +from vllm.v1.core.sched.output import ( + CachedRequestData, + GrammarOutput, + KVConnectorBlockState, + NewRequestData, + ScheduledEncoderInputStats, + SchedulerOutput, +) +from vllm.v1.core.sched.request_queue import ( + RequestQueue, + SchedulingPolicy, + create_request_queue, +) +from vllm.v1.core.sched.utils import check_stop, remove_all +from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs +from vllm.v1.kv_cache_interface import ( + CircularBufferSpec, + KVCacheConfig, + MambaSpec, + get_mamba_prefill_checkpoint_position, + is_mamba_prefill_checkpoint_valid, + iter_layer_specs, +) +from vllm.v1.metrics.perf import ModelMetrics, PerfStats +from vllm.v1.metrics.stats import ( + KV_FETCH_COMPLETED_WAITING, + KV_FETCH_IN_PROGRESS, + KV_FETCH_STAGES, + KV_FETCH_WAITING_TO_START, + PrefixCacheStats, + RequestSpecDecodeMetrics, + SchedulerStats, +) +from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput +from vllm.v1.request import Request, RequestStatus, StreamingUpdate +from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup +from vllm.v1.spec_decode.metrics import SpecDecodingStats +from vllm.v1.structured_output import StructuredOutputManager +from vllm.v1.structured_output.utils import strip_speculative_padding +from vllm.v1.structured_output.utils import strip_speculative_padding +from vllm.v1.utils import record_function_or_nullcontext + +logger = init_logger(__name__) + +# Bounds-guard instrumentation: verify KVBlockZeroer's exclusive-ownership precondition before a +# block is zeroed. Off by default, so it costs nothing. +_ZERO_OWNERSHIP_CHECK = os.getenv("VLLM_E2_ZERO_OWNERSHIP_CHECK") == "1" + + +class Scheduler(SchedulerInterface): + def __init__( + self, + vllm_config: VllmConfig, + kv_cache_config: KVCacheConfig, + structured_output_manager: StructuredOutputManager, + block_size: int, + hash_block_size: int | None = None, + mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY, + include_finished_set: bool = False, + log_stats: bool = False, + ) -> None: + self.vllm_config = vllm_config + self.scheduler_config = vllm_config.scheduler_config + self.cache_config = vllm_config.cache_config + self.lora_config = vllm_config.lora_config + self.model_uses_mrope = vllm_config.model_config.uses_mrope + self.kv_cache_config = kv_cache_config + self.kv_events_config = vllm_config.kv_events_config + self.parallel_config = vllm_config.parallel_config + self.log_stats = log_stats + self.observability_config = vllm_config.observability_config + self.spec_decode_metrics_level = ( + self.observability_config.per_request_spec_decode_metrics + ) + self.kv_metrics_collector: KVCacheMetricsCollector | None = None + if self.observability_config.kv_cache_metrics: + self.kv_metrics_collector = KVCacheMetricsCollector( + self.observability_config.kv_cache_metrics_sample, + ) + self.structured_output_manager = structured_output_manager + self.is_encoder_decoder = vllm_config.model_config.is_encoder_decoder + self.is_mm_encoder_only = vllm_config.is_mm_encoder_only + + # include_finished_set controls whether a separate set of finished + # request ids should be included in the EngineCoreOutputs returned + # by update_from_outputs(). This is currently used in the multi-engine + # case to track request lifetimes efficiently. + self.finished_req_ids_dict: dict[int, set[str]] | None = ( + defaultdict(set) if include_finished_set else None + ) + # Track requests scheduled in prior step (MRV1-only). + self.prev_step_scheduled_req_ids: set[str] = set() + + # Scheduling constraints. + self.max_num_running_reqs = self.scheduler_config.max_num_seqs + # Admission into RUNNING. Separate from max_num_running_reqs, which + # is the model-runner slot count. + self.max_num_active_reqs = ( + self.scheduler_config.max_num_active_seqs + if self.scheduler_config.max_num_active_seqs is not None + else self.max_num_running_reqs + ) + self.max_num_scheduled_tokens = ( + self.scheduler_config.max_num_scheduled_tokens + if self.scheduler_config.max_num_scheduled_tokens is not None + else self.scheduler_config.max_num_batched_tokens + ) + self.adaptive_long_prefill_threshold = ( + self.scheduler_config.long_prefill_token_threshold_adaptive + ) + self.max_model_len = vllm_config.model_config.max_model_len + self.enable_kv_cache_events = ( + self.kv_events_config is not None + and self.kv_events_config.enable_kv_cache_events + ) + # Diffusion models may not sample any tokens for a denoising step. + self.num_sampled_tokens_per_step = ( + 1 if not vllm_config.model_config.is_diffusion else 0 + ) + + # Create KVConnector for the Scheduler. Note that each Worker + # will have a corresponding KVConnector with Role=WORKER. + # KV Connector pushes/pull of remote KVs for P/D and offloading. + self.connector = None + self.connector_prefix_cache_stats: PrefixCacheStats | None = None + self.recompute_kv_load_failures = True + self.defer_block_free = False + # Whether a preempted request's in-flight output must be dropped; see + # KVConnectorBase_V1.requires_kv_delivery. + self.requires_kv_delivery = False + kv_transfer_config = self.vllm_config.kv_transfer_config + if kv_transfer_config is not None: + assert not self.is_encoder_decoder, ( + "Encoder-decoder models are not currently supported with KV connectors" + ) + self.connector = KVConnectorFactory.create_connector( + config=self.vllm_config, + role=KVConnectorRole.SCHEDULER, + kv_cache_config=self.kv_cache_config, + ) + if self.log_stats: + self.connector_prefix_cache_stats = PrefixCacheStats() + kv_load_failure_policy = kv_transfer_config.kv_load_failure_policy + self.recompute_kv_load_failures = kv_load_failure_policy == "recompute" + + # With overlapping batches (async scheduling or PP), a step may + # still be writing a freed request's KV blocks. A consumer KV + # Connector can reallocate and fill those blocks via a load that + # isn't ordered against that write, so defer freeing them. + multiple_inflight_batches = self.vllm_config.max_concurrent_batches > 1 + if multiple_inflight_batches and kv_transfer_config.is_kv_consumer: + self.defer_block_free = True + + self.requires_kv_delivery = self.connector.requires_kv_delivery + + # Ring-skip telemetry: the QSA ring block (CircularBufferSpec) is rewritten + # by every step that touches the request, so with overlapping batches + # (async scheduling or PP) an in-flight step may still be committing a + # freed/preempted request's suffix rows when the scheduler hands that + # block to another request. A KV connector is not the only writer that + # can race the stale in-flight write; the next request's own first + # step is. Defer the return of freed blocks to the pool until the + # newest scheduled step's output has been processed whenever a ring + # group exists. Same mechanism and fence as the connector gate above. + if ( + not self.defer_block_free + and self.vllm_config.max_concurrent_batches > 1 + and any( + isinstance(spec, CircularBufferSpec) + for group in kv_cache_config.kv_cache_groups + for spec in iter_layer_specs(group.kv_cache_spec) + ) + ): + self.defer_block_free = True + + self.kv_event_publisher = EventPublisherFactory.create( + self.kv_events_config, + self.parallel_config.data_parallel_index, + ) + self.ec_connector = None + if self.vllm_config.ec_transfer_config is not None: + self.ec_connector = ECConnectorFactory.create_connector( + config=self.vllm_config, role=ECConnectorRole.SCHEDULER + ) + + num_gpu_blocks = self.cache_config.num_gpu_blocks + assert num_gpu_blocks is not None and num_gpu_blocks > 0 + + self.block_size = block_size + self.dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size + self.pcp_world_size = vllm_config.parallel_config.prefill_context_parallel_size + + # req_id -> Request + self.requests: dict[str, Request] = {} + # Scheduling policy + try: + self.policy = SchedulingPolicy(self.scheduler_config.policy) + except ValueError as e: + raise ValueError( + f"Unknown scheduling policy: {self.scheduler_config.policy}" + ) from e + # Priority queues for requests. + self.waiting = create_request_queue(self.policy) + # Waiting requests that hold KV blocks are always drained before + # self.waiting: a block-holder must never sit behind a request whose + # failed allocation would stop the scheduling scan. + self.kv_holding_waiting = create_request_queue(self.policy) + # Waiting requests that are deferred, i.e. were skipped by the scheduling + # scan or enqueued in a blocked status. + self.deferred_waiting: set[Request] = set() + self.running: list[Request] = [] + + # The request IDs that are finished in between the previous and the + # current steps. This is used to notify the workers about the finished + # requests so that they can free the cached states for those requests. + # This is flushed at the end of each scheduling step. + self.finished_req_ids: set[str] = set() + + # IDs of requests preempted since the last call to schedule(). + self.reset_preempted_req_ids: set[str] = set() + + # Counter for requests waiting for streaming input. Used to calculate + # number of unfinished requests + self.num_waiting_for_streaming_input: int = 0 + + # KV Connector: requests in process of async KV loading or recving + self.finished_recving_kv_req_ids: set[str] = set() + self.failed_recving_kv_req_ids: set[str] = set() + + # Grammar compilation failures to finish as per-request errors in + # update_from_output. + self.grammar_compile_error_reqs: set[str] = set() + self.encoder_cache_mismatch_reqs: set[str] = set() + + # Encoder-related. + # Calculate encoder cache size if applicable + supports_mm_inputs = vllm_config.model_config.supports_multimodal_inputs + mm_budget = ( + MultiModalBudget(vllm_config, mm_registry) if supports_mm_inputs else None + ) + + # NOTE: Text-only encoder-decoder models are implemented as + # multi-modal models for convenience + # Example: https://github.com/vllm-project/bart-plugin + if self.is_encoder_decoder: + assert mm_budget and len(mm_budget.mm_max_toks_per_item) <= 1, ( + "Encoder-decoder models are expected to implement the " + "multimodal interface with at most one modality." + ) + + self.max_num_encoder_input_tokens = ( + mm_budget.encoder_compute_budget if mm_budget else 0 + ) + encoder_cache_size = mm_budget.encoder_cache_size if mm_budget else 0 + manager_cls_obj = vllm_config.ec_manager_config.get_encoder_cache_manager_obj() + if manager_cls_obj is None: + manager_cls_obj = ( + EncoderDecoderCacheManager + if self.is_encoder_decoder + else EncoderCacheManager + ) + self.encoder_cache_manager = manager_cls_obj.create_manager( + cache_size=encoder_cache_size, vllm_config=vllm_config + ) + speculative_config = vllm_config.speculative_config + self.use_eagle = False + self.use_eagle_block_drop = False + self.num_spec_tokens = vllm_config.num_speculative_tokens + self.num_lookahead_tokens = vllm_config.num_lookahead_tokens + # DSV41 SWA bounded replay: groups that declare a replay window are rebuilt + # after a prefix hit by recomputing its trailing tokens. One window + # for all such groups, so the rewind matches every group's allocation. + replay_windows = { + group.kv_cache_spec.prefix_replay_tokens + for group in kv_cache_config.kv_cache_groups + if group.kv_cache_spec.prefix_replay_tokens > 0 + } + assert len(replay_windows) <= 1, ( + f"Prefix replay windows should agree: {sorted(replay_windows)}" + ) + self.prefix_replay_tokens = replay_windows.pop() if replay_windows else 0 + self.num_prefill_lookahead = vllm_config.num_prefill_lookahead_tokens + self.dynamic_sd_lookup: list[int] | None = None + if speculative_config is not None: + if speculative_config.num_speculative_tokens_per_batch_size: + self.dynamic_sd_lookup = build_dynamic_sd_schedule_lookup( + speculative_config.num_speculative_tokens_per_batch_size, + vllm_max_batch_size=self.scheduler_config.max_num_seqs, + vllm_num_speculative_tokens=self.num_spec_tokens, + ) + self.use_eagle = speculative_config.use_eagle() + self.use_eagle_block_drop = speculative_config.use_eagle_block_drop() + if self.use_eagle and not self.use_eagle_block_drop: + logger.warning( + "EAGLE trailing prefix-cache block dropping is disabled. " + "This is experimental and may affect speculative-token " + "acceptance rates." + ) + + # Create the KV cache manager. + if hash_block_size is None: + hash_block_size = block_size + self.hash_block_size = hash_block_size + self.kv_cache_manager = KVCacheManager( + kv_cache_config=kv_cache_config, + max_model_len=self.max_model_len, + max_in_flight_tokens=vllm_config.max_in_flight_tokens, + enable_caching=self.cache_config.enable_prefix_caching, + use_eagle=self.use_eagle_block_drop, + num_prefill_lookahead=self.num_prefill_lookahead, + log_stats=self.log_stats, + enable_kv_cache_events=self.enable_kv_cache_events, + dcp_world_size=self.dcp_world_size, + pcp_world_size=1, + scheduler_block_size=self.block_size, + hash_block_size=hash_block_size, + metrics_collector=self.kv_metrics_collector, + watermark=self.scheduler_config.watermark, + enable_mamba_shared_prefix_checkpoint=( + self.cache_config.enable_mamba_shared_prefix_checkpoint + ), + ) + # Bind after construction so connectors can access the cache manager. + if self.connector is not None: + self.connector.bind_kv_cache_manager(self.kv_cache_manager) + # Ring blocks an async KV load is about to fill. Their tag clear + # is skipped for the same reason as zeroing (it could land after the + # out-of-band write and wipe the producer's committed tags). Tracked + # whatever needs_kv_cache_zeroing says, like the clear itself. + self._skip_ring_clear_block_ids: set[int] = set() + + self._log_prefix_cache_features() + + self.use_pp = self.parallel_config.pipeline_parallel_size > 1 + self.use_v2_model_runner = vllm_config.use_v2_model_runner + # Scheduler iteration counter. Drives the V2+PP+async decode-throttle + # cadence (`next_decode_eligible_step`). + self.current_step = 0 + # DP prefill balancing: Flag to track whether the last cadence-aligned + # prefill batch fully drained the waiting queue. Prefill throttling + # is disabled in this case. + self.prefill_capacity_bound = False + self.scheduler_reserve_full_isl = ( + self.scheduler_config.scheduler_reserve_full_isl + ) + + self.has_mamba_layers = kv_cache_config.has_mamba_layers + self.needs_kv_cache_zeroing = kv_cache_config.needs_kv_cache_zeroing + self.need_mamba_block_aligned_split = ( + self.has_mamba_layers and self.cache_config.mamba_cache_mode == "align" + ) + # TODO: Support models with multiple Mamba specs that require different + # prefill checkpoint alignments instead of selecting the first one. + self.mamba_prefill_checkpoint_alignment = next( + ( + group.kv_cache_spec.prefill_checkpoint_alignment + for group in kv_cache_config.kv_cache_groups + if isinstance(group.kv_cache_spec, MambaSpec) + ), + None, + ) + self.mamba_has_prefill_checkpoint_blocks = self.has_mamba_layers and all( + not isinstance(group.kv_cache_spec, MambaSpec) + or group.kv_cache_spec.num_prefill_checkpoint_blocks > 0 + for group in kv_cache_config.kv_cache_groups + ) + # A finer prefix_match_unit is configured: a mamba partial tail entry + # can only be registered by a step ending exactly at the prompt's last + # hash boundary, so the split adds that stop. + self.mamba_partial_cache_hit = ( + self.need_mamba_block_aligned_split + and self.hash_block_size < self.block_size + and self.kv_cache_manager.coordinator.enable_partial_hash_hits + ) + # Opt-in: also stop at the junction, where an eagle sibling resumes. The + # manager decides whether it can check-point there (per-group eagle bit, + # no MTP re-prefill tail); splitting for a stop it would refuse costs a + # forward pass and displaces the block-boundary stop. + self.mamba_shared_prefix_checkpoint = ( + self.mamba_partial_cache_hit + and self.kv_cache_manager.mamba_shared_prefix_checkpoint + ) + + # Counts of non-empty steps scheduled / processed. update_from_output + # is called once per scheduled step in FIFO order, so these stay in sync. + self.sched_step_seq = 0 + self.processed_step_seq = 0 + # FIFO of (fence_seq, blocks): blocks become safe to free once + # processed_step_seq >= fence_seq. + self.deferred_frees: deque[tuple[int, list[KVCacheBlock]]] = deque() + + self.perf_metrics: ModelMetrics | None = None + if self.log_stats and vllm_config.observability_config.enable_mfu_metrics: + self.perf_metrics = ModelMetrics(vllm_config) + + self.aux_output_connector = ( + AuxOutputSchedulerConnector() + if vllm_config.aux_output_config.enabled + else None + ) + self.return_sampling_mask = vllm_config.model_config.return_sampling_mask + + self._pause_state: PauseState = PauseState.UNPAUSED + + # In-flight requests still prefilling (prefill chunks + in-progress + # async KV loads). Their remaining-block reservation gates async loads. + self._inflight_prefills: set[Request] = set() + + # Requests with an async KV load, by stage (KV_FETCH_STAGES). + # None when not tracked (stats disabled or no KV connector). + self._kv_fetch_stages: dict[Request, str] | None = ( + {} if self.log_stats and self.connector is not None else None + ) + self._kv_fetch_counts: dict[str, int] = dict.fromkeys(KV_FETCH_STAGES, 0) + + def _mamba_block_aligned_split( + self, + request: Request, + num_new_tokens: int, + num_new_local_computed_tokens: int = 0, + num_external_computed_tokens: int = 0, + ) -> int: + """Clip a prefill chunk so it ends where Mamba state must be cached. + + In "align" cache mode reusable SSM states are materialized at block + boundaries, plus mandatory early stops (the prompt's partial-tail hash + boundary, a detected shared-prefix junction). If a block is larger + than the configured prefill chunk limit, intermediate chunks keep + private running state until they reach the next cacheable position. + """ + start = ( + request.num_computed_tokens + + num_new_local_computed_tokens + + num_external_computed_tokens + ) + # Split only during prefill: `request.num_tokens - 1` extends this to + # resumed requests replaying their output tokens. + prefill_end = max(request.num_prompt_tokens, request.num_tokens - 1) + if start >= prefill_end: + return num_new_tokens + + block_size = self.cache_config.block_size + # The last block-aligned position whose state can be cached. With + # Eagle, FullAttn prunes the last matching block, so back off one + # block to avoid a Mamba cache miss. + last_cache_position = request.num_tokens - request.num_tokens % block_size + if self.use_eagle_block_drop: + last_cache_position = max(last_cache_position - block_size, 0) + + end = start + num_new_tokens + checkpoint_position = get_mamba_prefill_checkpoint_position( + prefill_end, + self.hash_block_size, + drop_eagle_block=self.use_eagle_block_drop, + ) + use_internal_checkpoint = ( + self.mamba_has_prefill_checkpoint_blocks + and end >= prefill_end + and is_mamba_prefill_checkpoint_valid( + query_start=start, + query_end=end, + checkpoint_position=checkpoint_position, + hash_block_size=self.hash_block_size, + mamba_block_size=block_size, + checkpoint_alignment=self.mamba_prefill_checkpoint_alignment, + ) + ) + if use_internal_checkpoint: + last_cache_position = 0 + # Invariant: slot p holds the state after exactly (p + 1) * block_size + # tokens. State is written at chunk ends, so chunk ends must be block + # aligned. Exempt: the prompt's last chunk, whose slot decode advances + # to the boundary. A block too wide for one chunk advances sub-block + # and re-aligns at the next boundary. + if end < prefill_end: + max_prefill_tokens = self.max_num_scheduled_tokens + long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold + if long_prefill_threshold > 0: + max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) + aligned_end = end // block_size * block_size + if aligned_end > start or block_size <= max_prefill_tokens: + end = aligned_end + + next_block_boundary = (start // block_size + 1) * block_size + tail_boundary = ( + request.num_prompt_tokens // self.hash_block_size * self.hash_block_size + if self.mamba_partial_cache_hit and not use_internal_checkpoint + else 0 + ) + if tail_boundary and self.use_eagle_block_drop: + # Eagle matches one hash unit past the candidate and drops it, so + # nothing proves the prompt's own last hash boundary. Materialize + # the state one unit lower, where the hit can actually land. Keyed on + # the block-drop bit, not plain use_eagle: this shift exists only to + # compensate for the drop, and the Mamba manager's matching gate + # reads the same bit (the coordinator is handed use_eagle_block_drop). + tail_boundary = max(tail_boundary - self.hash_block_size, 0) + junction = request.shared_prefix_boundary + # Block-floored: a sub-block junction's state is not separately cacheable. + block_floored = start + (junction - start) // block_size * block_size + # Past the prompt the manager writes nothing, so fall back to the + # block-floored stop rather than dropping it: a resumed request replaying + # output tokens can still observe a junction there. + junction_stop = ( + junction + if self.mamba_shared_prefix_checkpoint + and junction <= request.num_prompt_tokens + else block_floored + ) + stops = ( + # Same invariant: a chunk starting mid-block stops at the boundary + # rather than running past it. + next_block_boundary + if start % block_size != 0 and not use_internal_checkpoint + else 0, + # Never run past the last cacheable block boundary mid-chunk. + last_cache_position, + # Fine-grained hits: the prompt's partial-tail entry can only be + # registered by a chunk ending exactly at its last hash boundary. + tail_boundary + if last_cache_position < tail_boundary < request.num_prompt_tokens + else 0, + # Marconi shared-prefix junction: cache its state so sibling + # requests sharing the prefix can reuse it. + junction_stop if start < junction < end else 0, + ) + # Stop at the earliest mandatory position strictly inside the chunk. + end = min((s for s in stops if start < s < end), default=end) + return max(end - start, 0) + + def _log_prefix_cache_features(self) -> None: + """One startup line per prefix-cache retention feature that is active + (nothing when it is off).""" + manager = self.kv_cache_manager + if manager.enable_caching and self.connector is None: + logger.info( + "Prefix cache: waiting requests' cached prefixes are pinned " + "until admission (up to %d waiters; hash unit %d, block %d)", + manager.MAX_PINNED_WAITERS, + self.hash_block_size, + self.block_size, + ) + + def _pin_waiting_hits(self) -> None: + """Pin the prefix-cache hits of the requests at the front of the + waiting queue (arrival order, bounded by the manager) so the + admissions and decode growth of other requests cannot evict them + while they wait -- also while the running set is full, when they are + not even looked up. Released on admission, free, and whenever a + running request would otherwise fail to allocate. + """ + manager = self.kv_cache_manager + if ( + not manager.enable_caching + or self.connector is not None + or manager.num_pinned_waiters >= manager.MAX_PINNED_WAITERS + ): + return + for request in self.waiting: + if manager.num_pinned_waiters >= manager.MAX_PINNED_WAITERS: + break + if ( + request.num_computed_tokens != 0 + or request.status + not in (RequestStatus.WAITING, RequestStatus.PREEMPTED) + ): + continue + manager.pin_waiting_hit(request) + + def _get_local_prefix_cache_hit( + self, request: Request + ) -> tuple[KVCacheBlocks, int, int, bool]: + connector = self.connector + if connector is not None and connector.supports_divergent_local_hybrid_hits: + return self.kv_cache_manager.get_computed_blocks_for_connector(request) + + blocks, num_local, shared_prefix_boundary = ( + self.kv_cache_manager.get_computed_blocks(request) + ) + return blocks, num_local, shared_prefix_boundary, False + + def _reserve_prefill_lookahead( + self, + request: Request, + num_computed_tokens: int, + num_new_tokens: int, + ) -> int: + """Never end a prefill chunk within num_prefill_lookahead of the + prefill end. + + At a chunked-prefill boundary, the multi-module MTP drafter consumes + the next num_prefill_lookahead known prefill tokens as draft inputs. A + boundary closer to the end than that would make it fall back to + sampled drafts, permanently polluting the trailing modules' KV caches. + Either finish the prefill or leave at least num_prefill_lookahead for + the next chunk. No-op for eagle-family drafters (lookahead 1). + """ + remaining = request.num_tokens - num_computed_tokens - num_new_tokens + if 0 < remaining < self.num_prefill_lookahead: + num_new_tokens -= self.num_prefill_lookahead - remaining + return max(num_new_tokens, 0) + + def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: + self.current_step += 1 + # NOTE(woosuk) on the scheduling algorithm: + # There's no "decoding phase" nor "prefill phase" in the scheduler. + # Each request just has the num_computed_tokens and + # num_tokens_with_spec. num_tokens_with_spec = + # len(prompt_token_ids) + len(output_token_ids) + len(spec_token_ids). + # At each step, the scheduler tries to assign tokens to the requests + # so that each request's num_computed_tokens can catch up its + # num_tokens_with_spec. This is general enough to cover + # chunked prefills, prefix caching, speculative decoding, + # and the "jump decoding" optimization in the future. + + scheduled_new_reqs: list[Request] = [] + scheduled_resumed_reqs: list[Request] = [] + scheduled_running_reqs: list[Request] = [] + preempted_reqs: list[Request] = [] + + req_to_new_blocks: dict[str, KVCacheBlocks] = {} + num_scheduled_tokens: dict[str, int] = {} + token_budget = self.max_num_scheduled_tokens + spec = self.vllm_config.speculative_config + draft_slots = spec.max_num_new_slots_for_drafting if spec is not None else 0 + input_budget = self.scheduler_config.max_num_batched_tokens + if self._pause_state == PauseState.PAUSED_ALL: + # Do not schedule any requests when paused. + token_budget = 0 + + # Encoder-related. + scheduled_encoder_inputs: dict[str, list[int]] = {} + encoder_compute_budget = self.max_num_encoder_input_tokens + # Spec decode-related. + scheduled_spec_decode_tokens: dict[str, list[int]] = {} + # Whether the running batch contains any prefill requests. + prefill_scheduled = False + # Whether any scheduled request has a synchronous connector KV load. + has_sync_kv_loads = False + + # For logging. + scheduled_timestamp = time.monotonic() + + self.kv_cache_manager.new_step_starts() + + # DP prefill balancing: on a throttled (non-cadence-aligned) step, defer + # all prefill compute unless saturated. + defer_prefills = ( + throttle_prefills and not self.prefill_capacity_bound + ) and any(not r.is_prefill_chunk for r in self.running) + + # `long_prefill_token_threshold` exists to stop a long prefill from + # starving other requests of the token budget. When it is the only + # request there is nobody to starve, so let it use the whole budget. + num_running, num_waiting = self.get_request_counts() + num_eligible_reqs = num_running + num_waiting + long_prefill_token_threshold = ( + self.scheduler_config.long_prefill_token_threshold + if num_eligible_reqs > 1 + else 0 + ) + if long_prefill_token_threshold > 0 and self.adaptive_long_prefill_threshold: + # Floor the cap at a fair share of the input budget so it never + # cuts a request below max_num_batched_tokens / num requests. + long_prefill_token_threshold = max( + long_prefill_token_threshold, input_budget // num_eligible_reqs + ) + + # First, schedule the RUNNING requests. + req_index = 0 + while req_index < len(self.running) and token_budget > 0: + request = self.running[req_index] + if input_budget <= draft_slots: + break + + if ( + request.num_output_placeholders > 0 + # This is (num_computed_tokens + 1) - (num_output_placeholders - 1). + # Since output placeholders are also included in the computed tokens + # count, we subtract (num_output_placeholders - 1) to remove any draft + # tokens, so that we can be sure no further steps are needed even if + # they are all rejected. + and request.num_computed_tokens + 2 - request.num_output_placeholders + >= request.num_prompt_tokens + request.max_tokens + ): + # Async scheduling: Avoid scheduling an extra step when we are sure that + # the previous step has reached request.max_tokens. We don't schedule + # partial draft tokens since this prevents uniform decode optimizations. + req_index += 1 + continue + + if self.current_step < request.next_decode_eligible_step: + # V2+PP+async: enforce `pp_size` steps between same-req decodes + # to match worker-side sampled-tokens broadcast slot ring cadence. + req_index += 1 + continue + + if defer_prefills and request.is_prefill_chunk: + # DP prefill balancing: defer this in-progress prefill chunk to a + # cadence-aligned step; decodes still run to fill this step. + req_index += 1 + continue + + if ( + self.ec_connector is not None + and request.mm_features + and not self.ec_connector.ensure_cache_available( + request, + request.num_computed_tokens - request.num_output_placeholders, + ) + ): + req_index += 1 + continue + + num_new_tokens = ( + request.num_tokens_with_spec + + request.num_output_placeholders + - request.num_computed_tokens + ) + if 0 < long_prefill_token_threshold < num_new_tokens: + num_new_tokens = long_prefill_token_threshold + num_new_tokens = min( + num_new_tokens, token_budget, input_budget - draft_slots + ) + + # Make sure the input position does not exceed the max model len. + # This is necessary when using spec decoding. + num_new_tokens = min( + num_new_tokens, + self.max_model_len + - request.num_computed_tokens + - self.num_sampled_tokens_per_step, + ) + + # Apply Mamba alignment before encoder caps. + if self.need_mamba_block_aligned_split: + num_new_tokens = self._mamba_block_aligned_split( + request, num_new_tokens + ) + + # Schedule encoder inputs. + encoder_inputs_to_schedule = None + external_load_encoder_input: list[int] = [] + new_encoder_compute_budget = encoder_compute_budget + if request.has_encoder_inputs: + ( + encoder_inputs_to_schedule, + num_new_tokens, + new_encoder_compute_budget, + external_load_encoder_input, + ) = self._try_schedule_encoder_inputs( + request, + request.num_computed_tokens, + num_new_tokens, + encoder_compute_budget, + shift_computed_tokens=self.num_prefill_lookahead, + ) + + # Multi-module MTP: avoid ending a prefill chunk within + # num_prefill_lookahead of the prefill end. + num_new_tokens = self._reserve_prefill_lookahead( + request, request.num_computed_tokens, num_new_tokens + ) + + if num_new_tokens == 0: + # The request cannot be scheduled because one of the following + # reasons: + # 1. No new tokens to schedule. This may happen when + # (1) PP>1 and we have already scheduled all prompt tokens + # but they are not finished yet. + # (2) Async scheduling and the request has reached to either + # its max_total_tokens or max_model_len. + # 2. The encoder budget is exhausted. + # 3. The encoder cache is exhausted. + # 4. Insufficient budget for a block-aligned chunk in hybrid + # models with mamba cache mode \"align\". + # 5. Insufficient budget to keep a multi-module MTP prefill + # chunk out of the prefill-lookahead window. + # NOTE(woosuk): Here, by doing `continue` instead of `break`, + # we do not strictly follow the FCFS scheduling policy and + # allow the lower-priority requests to be scheduled. + req_index += 1 + continue + + # Schedule newly needed KV blocks for the request. + with record_function_or_nullcontext("schedule: allocate_slots"): + while True: + new_blocks = self.kv_cache_manager.allocate_slots( + request, + num_new_tokens, + num_lookahead_tokens=self.num_lookahead_tokens, + ) + + if new_blocks is not None: + # The request can be scheduled. + break + + if ( + self.connector is not None + and self.connector.has_pending_block_frees() + ): + break + + # Waiting requests' pinned prefix hits yield + # before anyone is preempted (a pin never costs a running + # request its blocks); retry with them in the free queue. + if self.kv_cache_manager.yield_pinned_hits(limit=1): + continue + + # The request cannot be scheduled. + # Preempt the lowest-priority request. + if self.policy == SchedulingPolicy.PRIORITY: + preempted_req = max( + self.running, + key=lambda r: (r.priority, r.arrival_time), + ) + else: + preempted_req = self.running[-1] + + # A deferred free will not help with immediate allocation. + # (Upstream #49675. It subsumes our earlier post-hoc guard, + # which stopped after ONE deferred victim: a victim that + # passes this check is freed immediately by + # _preempt_request, so deferred_frees cannot grow below. + # The requester stays running, unscheduled, and retries + # after the fence drains; with every running request in + # flight (PP / async) that is zero victims, not the batch.) + if not self._request_blocks_can_be_freed(preempted_req): + break + + if self.policy == SchedulingPolicy.PRIORITY: + victim_index = self.running.index(preempted_req) + del self.running[victim_index] + if victim_index < req_index: + req_index -= 1 + + if preempted_req in scheduled_running_reqs: + preempted_req_id = preempted_req.request_id + scheduled_running_reqs.remove(preempted_req) + restored = num_scheduled_tokens.pop(preempted_req_id) + token_budget += restored + input_budget += restored + draft_slots + req_to_new_blocks.pop(preempted_req_id) + scheduled_spec_decode_tokens.pop(preempted_req_id, None) + preempted_encoder_inputs = scheduled_encoder_inputs.pop( + preempted_req_id, None + ) + if preempted_encoder_inputs: + # Restore encoder compute budget if the preempted + # request had encoder inputs scheduled in this step. + num_embeds_to_restore = sum( + preempted_req.get_num_encoder_embeds(i) + for i in preempted_encoder_inputs + ) + encoder_compute_budget += num_embeds_to_restore + else: + preempted_req = self.running.pop() + + self._preempt_request( + preempted_req, + scheduled_timestamp, + drop_stale_output=self.requires_kv_delivery, + ) + preempted_reqs.append(preempted_req) + if preempted_req == request: + # No more request to preempt. Cannot schedule this request. + break + + if new_blocks is None: + # Cannot schedule this request. + break + + # Schedule the request. + scheduled_running_reqs.append(request) + prefill_scheduled |= request.is_prefill_chunk + request_id = request.request_id + req_to_new_blocks[request_id] = new_blocks + num_scheduled_tokens[request_id] = num_new_tokens + token_budget -= num_new_tokens + input_budget -= num_new_tokens + draft_slots + req_index += 1 + + # Speculative decode related. + if request.spec_token_ids: + num_scheduled_spec_tokens = ( + num_new_tokens + + request.num_computed_tokens + - request.num_tokens + - request.num_output_placeholders + ) + if num_scheduled_spec_tokens > 0: + spec_token_ids = request.spec_token_ids + if len(spec_token_ids) > num_scheduled_spec_tokens: + spec_token_ids = spec_token_ids[:num_scheduled_spec_tokens] + scheduled_spec_decode_tokens[request.request_id] = spec_token_ids + + # New spec tokens will be set in `update_draft_token_ids` before the + # next step when applicable. + request.spec_token_ids = [] + + # Encoder-related. + if encoder_inputs_to_schedule: + scheduled_encoder_inputs[request_id] = encoder_inputs_to_schedule + # Allocate the encoder cache. + for i in encoder_inputs_to_schedule: + self.encoder_cache_manager.allocate(request, i) + if self.ec_connector is not None: + self.ec_connector.update_state_after_alloc(request, i) + encoder_compute_budget = new_encoder_compute_budget + if external_load_encoder_input: + for i in external_load_encoder_input: + self.encoder_cache_manager.allocate(request, i) + if self.ec_connector is not None: + self.ec_connector.update_state_after_alloc(request, i) + + # Record the LoRAs in scheduled_running_reqs + scheduled_loras: set[int] = set() + if self.lora_config: + scheduled_loras = set( + req.lora_request.lora_int_id + for req in scheduled_running_reqs + if req.lora_request and req.lora_request.lora_int_id > 0 + ) + assert len(scheduled_loras) <= self.lora_config.max_loras + + # Next, schedule the WAITING requests. + if not preempted_reqs and self._pause_state == PauseState.UNPAUSED: + step_skipped_waiting: deque[Request] = deque() + step_skipped_kv_holding: deque[Request] = deque() + + def skip_request(from_queue: RequestQueue) -> None: + request_to_skip = from_queue.pop_request() + self.deferred_waiting.add(request_to_skip) + if self._holds_kv_blocks(request_to_skip): + step_skipped_kv_holding.appendleft(request_to_skip) + else: + step_skipped_waiting.appendleft(request_to_skip) + + while token_budget > 0: + # Requests holding KV blocks are always drained first. + request_queue = self.kv_holding_waiting or self.waiting + if not request_queue or input_budget <= draft_slots: + break + # Paused streaming sessions (WAITING_FOR_STREAMING_REQ) are not + # in `running` but still hold a model-runner request slot. + num_running = len(self.running) + self.num_waiting_for_streaming_input + if num_running >= self.max_num_active_reqs: + break + + request = request_queue.peek_request() + request_id = request.request_id + + ready_to_schedule = self._handle_blocked_waiting_request(request) + if not ready_to_schedule: + skip_request(request_queue) + continue + + if request.num_stale_output_tokens and not request.drop_stale_output: + # Deliverable stale output still in flight: resuming now + # could resample a position that output later delivers. + # It drains within the pipeline depth. + skip_request(request_queue) + continue + + # Check that adding the request still respects the max_loras constraint. + if ( + self.lora_config + and request.lora_request + and ( + len(scheduled_loras) == self.lora_config.max_loras + and request.lora_request.lora_int_id not in scheduled_loras + ) + ): + # Scheduling would exceed max_loras, skip. + skip_request(request_queue) + continue + + num_external_computed_tokens = 0 + load_kv_async = False + connector_prefix_cache_queries, connector_prefix_cache_hits = 0, 0 + did_prefix_cache_lookup = False + + # Get already-cached tokens. + if request.num_computed_tokens == 0: + did_prefix_cache_lookup = True + ( + new_computed_blocks, + num_new_local_computed_tokens, + request.shared_prefix_boundary, + hit_diverged, + ) = self._get_local_prefix_cache_hit(request) + + # Get externally-cached tokens if using a KVConnector. + if self.connector is not None: + # Present a block-aligned local hit to the connector so + # a strictly longer remote hit can supersede a local + # sub-block tail without racing its copy-on-write. + partial_tail = num_new_local_computed_tokens % self.block_size + block_aligned_local = ( + num_new_local_computed_tokens - partial_tail + ) + ext_tokens, load_kv_async = ( + self.connector.get_num_new_matched_tokens( + request, block_aligned_local + ) + ) + if request.skip_reading_prefix_cache: + ext_tokens, load_kv_async = 0, False + + if ext_tokens is None: + # The request cannot be scheduled because + # the KVConnector couldn't determine + # the number of matched tokens. + skip_request(request_queue) + continue + + if self.prefix_replay_tokens: + # SWA bounded replay recomputes the hit's last + # window, from the block holding its first + # token; the sliding-window groups retire + # whole blocks below the window of the hit's + # next token. Hits end on a block boundary, or + # a hit ending one token short of one would + # retire the block the replay starts in. + ext_tokens -= ext_tokens % self.block_size + load_kv_async = load_kv_async and ext_tokens > 0 + + if partial_tail and ext_tokens > partial_tail: + # Remote strictly exceeds the full local hit: drop the + # sub-block tail so no CoW is needed, and let the load + # cover it. Trim the partial block out of the local + # computed blocks so it is not adopted from the cache. + new_computed_blocks = ( + self.kv_cache_manager.truncate_computed_blocks( + new_computed_blocks, block_aligned_local + ) + ) + num_new_local_computed_tokens = block_aligned_local + num_external_computed_tokens = ext_tokens + elif partial_tail: + # Remote does not exceed the full local hit: keep the + # local sub-block tail and load nothing external. + num_external_computed_tokens = 0 + # Nothing to load remotely -> not an async-load step; + # clearing avoids the `load_kv_async` assert below. + load_kv_async = False + else: + num_external_computed_tokens = ext_tokens + + if hit_diverged and num_external_computed_tokens == 0: + # No external tokens back the deeper local hit, so its + # resume boundary would have no valid Mamba state. + # Reconcile to the boundary every group agrees on. + ( + new_computed_blocks, + num_new_local_computed_tokens, + request.shared_prefix_boundary, + ) = self.kv_cache_manager.get_computed_blocks(request) + + connector_prefix_cache_queries = ( + request.num_tokens - num_new_local_computed_tokens + ) + connector_prefix_cache_hits = num_external_computed_tokens + + # Total computed tokens (local + external). + num_computed_tokens = ( + num_new_local_computed_tokens + num_external_computed_tokens + ) + assert num_computed_tokens <= request.num_tokens + if 0 < num_computed_tokens <= self.prefix_replay_tokens: + # SWA bounded replay: a hit no longer than the replayed + # window would be recomputed in full and save nothing. + new_computed_blocks = ( + self.kv_cache_manager.empty_kv_cache_blocks + ) + num_new_local_computed_tokens = 0 + num_external_computed_tokens = 0 + connector_prefix_cache_hits = 0 + num_computed_tokens = 0 + load_kv_async = False + + # Skip request with pending mm encoding prefetches + if self._ec_transfer_pending(request, num_computed_tokens): + skip_request(request_queue) + continue + + # Track first scheduled prefill, not post-preemption repeat prefills + if request.prefill_stats and request.num_preemptions <= 0: + assert num_computed_tokens <= request.num_prompt_tokens + request.prefill_stats.set( + num_prompt_tokens=request.num_prompt_tokens, + num_local_cached_tokens=num_new_local_computed_tokens, + num_external_cached_tokens=num_external_computed_tokens, + ) + self._set_kv_fetch_stage( + request, KV_FETCH_WAITING_TO_START if load_kv_async else None + ) + else: + # KVTransfer: WAITING reqs have num_computed_tokens > 0 + # after async KV recvs are completed. A streaming-input + # session resumes here too, carrying whatever media its + # latest chunk added, so this branch needs the same gate. + new_computed_blocks = self.kv_cache_manager.empty_kv_cache_blocks + num_new_local_computed_tokens = 0 + num_computed_tokens = request.num_computed_tokens + + if self._ec_transfer_pending(request, num_computed_tokens): + skip_request(request_queue) + continue + + encoder_inputs_to_schedule = None + external_load_encoder_input = [] + new_encoder_compute_budget = encoder_compute_budget + pad_spec_decode = False + + # SWA bounded replay: recompute the tail of the hit without + # rewriting its cached KV. An async load replays once the KV + # has arrived (_update_waiting_for_remote_kv). + num_replay_tokens = 0 + + if load_kv_async: + # KVTransfer: loading remote KV, do not allocate for new work. + assert num_external_computed_tokens > 0 + num_new_tokens = 0 + elif defer_prefills and num_computed_tokens < request.num_tokens - 1: + # DP prefill balancing: defer this step's local prefill + # compute to a cadence-aligned step. + break + else: + request_token_budget = min(token_budget, input_budget - draft_slots) + if did_prefix_cache_lookup: + # Fresh admission: local and/or sync external hit. + num_replay_tokens = self._mark_prefix_replay( + request, num_computed_tokens + ) + num_computed_tokens -= num_replay_tokens + # Number of tokens to be scheduled. + # We use `request.num_tokens` instead of + # `request.num_prompt_tokens` to consider the resumed + # requests, which have output tokens. + num_new_tokens = request.num_tokens - num_computed_tokens + + # Pad new decode requests to uniform spec decoding size to + # preserve full cudagraph for this step. + # Not for diffusion where draft tokens can't be padded. + if ( + (self.num_spec_tokens > 0 and self.dynamic_sd_lookup is None) + and self.num_sampled_tokens_per_step > 0 + and num_new_tokens == 1 + and not prefill_scheduled + and (scheduled_running_reqs or num_computed_tokens > 0) + ): + padded_num_tokens = 1 + self.num_spec_tokens + # Pad only when there is room for the sampled token(s). + if ( + num_computed_tokens + + padded_num_tokens + + self.num_sampled_tokens_per_step + <= self.max_model_len + ): + if padded_num_tokens > request_token_budget: + # Prefer to not schedule than schedule un-padded. + break + num_new_tokens = padded_num_tokens + pad_spec_decode = True + + if 0 < long_prefill_token_threshold < num_new_tokens: + num_new_tokens = long_prefill_token_threshold + + # chunked prefill has to be enabled explicitly to allow + # pooling requests to be chunked + if ( + not self.scheduler_config.enable_chunked_prefill + and num_new_tokens > request_token_budget + and not self.is_mm_encoder_only + ): + # If chunked_prefill is disabled, + # we can stop the scheduling here. + break + + num_new_tokens = min(num_new_tokens, request_token_budget) + assert num_new_tokens > 0 + + # Apply Mamba alignment before encoder caps. + if self.need_mamba_block_aligned_split: + num_new_tokens = self._mamba_block_aligned_split( + request, + num_new_tokens, + num_new_local_computed_tokens, + num_external_computed_tokens, + ) + if num_new_tokens == 0: + break + if ( + pad_spec_decode + and num_new_tokens != 1 + self.num_spec_tokens + ): + # Alignment clipped the placeholder rows. The split + # aligns prefill chunks, but the padded tail rows are + # speculative positions, not prefill tokens. A padded + # request must keep all 1 + num_spec rows or the + # sampler's row count stops matching its query rows, + # so drop the padding instead of shortening it. + num_new_tokens = 1 + pad_spec_decode = False + + # Schedule encoder inputs. + if request.has_encoder_inputs: + ( + encoder_inputs_to_schedule, + num_new_tokens, + new_encoder_compute_budget, + external_load_encoder_input, + ) = self._try_schedule_encoder_inputs( + request, + num_computed_tokens, + num_new_tokens, + encoder_compute_budget, + shift_computed_tokens=self.num_prefill_lookahead, + ) + + # Multi-module MTP: avoid ending a prefill chunk within + # num_prefill_lookahead of the prefill end. + num_new_tokens = self._reserve_prefill_lookahead( + request, num_computed_tokens, num_new_tokens + ) + + if num_new_tokens == 0: + # The request cannot be scheduled. + break + + # During async KV load, no forward pass is run yet. + # Allocate speculative lookahead slots later to avoid + # mismatching local and remote block counts. + limit_lookahead_tokens = load_kv_async and self.num_lookahead_tokens > 0 + effective_lookahead_tokens = ( + 0 if limit_lookahead_tokens else self.num_lookahead_tokens + ) + + # Determine if we need to allocate cross-attention blocks. + num_encoder_tokens = 0 + if ( + self.is_encoder_decoder + and request.has_encoder_inputs + and encoder_inputs_to_schedule + ): + num_encoder_tokens = sum( + request.get_num_encoder_embeds(i) + for i in encoder_inputs_to_schedule + ) + + reserved_blocks = 0 + if load_kv_async: + # An async load holds its blocks for the whole transfer with + # no forward progress and isn't preemptible here. Admit it + # only if it fits in (free - other in-flight reservations) + # plus its own spec decode step blocks, to avoid deadlock and + # predictable preemptions. + reserved_blocks = ( + self._inflight_prefill_reserved_blocks() + + self._spec_decode_step_blocks() + ) + + # Replayed tokens are already counted in the adopted hit; a + # chunk that ends inside the replayed range needs no new slots. + num_tokens_past_hit = max(num_new_tokens - num_replay_tokens, 0) + new_blocks = self.kv_cache_manager.allocate_slots( + request, + num_tokens_past_hit, + num_new_computed_tokens=num_new_local_computed_tokens, + new_computed_blocks=new_computed_blocks, + num_lookahead_tokens=effective_lookahead_tokens, + num_external_computed_tokens=num_external_computed_tokens, + delay_cache_blocks=load_kv_async, + skip_zeroing_group_ids=( + self.connector.get_loaded_kv_cache_group_ids(request) + if load_kv_async and self.connector is not None + else () + ), + num_encoder_tokens=num_encoder_tokens, + full_sequence_must_fit=self.scheduler_reserve_full_isl, + reserved_blocks=reserved_blocks, + has_scheduled_reqs=bool(self.running), + ) + + if new_blocks is None and self.kv_cache_manager.yield_pinned_hits( + keep=request_id + ): + # Pins of waiters behind this one yield to it + # (arrival order); its own pin stays, those blocks are + # its hit. One retry with them back in the free queue. + new_blocks = self.kv_cache_manager.allocate_slots( + request, + num_new_tokens, + num_new_computed_tokens=num_new_local_computed_tokens, + new_computed_blocks=new_computed_blocks, + num_lookahead_tokens=effective_lookahead_tokens, + num_external_computed_tokens=num_external_computed_tokens, + delay_cache_blocks=load_kv_async, + num_encoder_tokens=num_encoder_tokens, + full_sequence_must_fit=self.scheduler_reserve_full_isl, + reserved_blocks=reserved_blocks, + has_scheduled_reqs=bool(self.running), + ) + + if new_blocks is None: + # The request cannot be scheduled. + + # NOTE: we need to untouch the request from the encode cache + # manager + if request.has_encoder_inputs: + self.encoder_cache_manager.free(request) + break + + # KVTransfer: the connector uses this info to determine + # if a load is needed. Note that + # This information is used to determine if a load is + # needed for this request. + if self.connector is not None: + self.connector.update_state_after_alloc( + request, + self.kv_cache_manager.get_blocks(request_id), + num_external_computed_tokens, + ) + if ( + self.connector_prefix_cache_stats is not None + and connector_prefix_cache_queries != 0 + ): + self.connector_prefix_cache_stats.record( + num_tokens=connector_prefix_cache_queries, + num_hits=connector_prefix_cache_hits, + preempted=request.num_preemptions > 0, + ) + + # Record at admission so unscheduled lookups are not counted. + if did_prefix_cache_lookup: + self.kv_cache_manager.record_prefix_cache_stats( + request, num_new_local_computed_tokens + ) + + if load_kv_async: + # If loading async, allocate memory and put request + # into the WAITING_FOR_REMOTE_KV state. + request.status = RequestStatus.WAITING_FOR_REMOTE_KVS + # Set num_computed_tokens even though KVs are not yet loaded. + # request.num_computed_tokens will not be used anywhere until + # the request finished the KV transfer. + # + # If a transfer error is reported by the connector, + # request.num_computed_tokens will be re-set accordingly in + # _update_requests_with_invalid_blocks. + # + # When the transfer is finished, either successfully or not, + # request.num_computed_tokens will correctly reflect the number + # of computed tokens. + # _update_waiting_for_remote_kv will then cache + # only the successfully loaded tokens. + request.num_computed_tokens = num_computed_tokens + self._inflight_prefills.add(request) + self._set_kv_fetch_stage(request, KV_FETCH_IN_PROGRESS) + # Same race for the ring's tag clear; the load + # brings the producer's rows and tags (or none, and the + # reader then fails closed on the missing history). + self._skip_ring_clear_block_ids.update( + self.kv_cache_manager.get_ring_block_ids(request.request_id) + ) + skip_request(request_queue) + continue + + request = request_queue.pop_request() + self.deferred_waiting.discard(request) + self.running.append(request) + if num_external_computed_tokens > 0: + # load_kv_async is False here + has_sync_kv_loads = True + if self.log_stats: + request.record_event( + EngineCoreEventType.SCHEDULED, scheduled_timestamp + ) + if request.status == RequestStatus.WAITING: + scheduled_new_reqs.append(request) + elif request.status == RequestStatus.PREEMPTED: + scheduled_resumed_reqs.append(request) + else: + raise RuntimeError(f"Invalid request status: {request.status}") + + if self.lora_config and request.lora_request: + scheduled_loras.add(request.lora_request.lora_int_id) + req_to_new_blocks[request_id] = self.kv_cache_manager.get_blocks( + request_id + ) + num_scheduled_tokens[request_id] = num_new_tokens + token_budget -= num_new_tokens + input_budget -= num_new_tokens + draft_slots + request.status = RequestStatus.RUNNING + request.num_computed_tokens = num_computed_tokens + self._set_kv_fetch_stage(request, None) + if pad_spec_decode: + assert num_new_tokens == 1 + self.num_spec_tokens + scheduled_spec_decode_tokens[request_id] = [ + -1 + ] * self.num_spec_tokens + # Only track requests that will still be prefilling after this chunk. + if num_computed_tokens + num_new_tokens < request.num_tokens: + self._inflight_prefills.add(request) + # Encoder-related. + if encoder_inputs_to_schedule: + scheduled_encoder_inputs[request_id] = encoder_inputs_to_schedule + # Allocate the encoder cache. + for i in encoder_inputs_to_schedule: + self.encoder_cache_manager.allocate(request, i) + if self.ec_connector is not None: + self.ec_connector.update_state_after_alloc(request, i) + encoder_compute_budget = new_encoder_compute_budget + # Allocate for external load encoder cache + if external_load_encoder_input: + for i in external_load_encoder_input: + self.encoder_cache_manager.allocate(request, i) + if self.ec_connector is not None: + self.ec_connector.update_state_after_alloc(request, i) + + # re-queue requests skipped in this pass ahead of older skipped items. + if step_skipped_kv_holding: + self.kv_holding_waiting.prepend_requests(step_skipped_kv_holding) + if step_skipped_waiting: + self.waiting.prepend_requests(step_skipped_waiting) + + # Whoever is still waiting keeps its prefix-cache hit. + self._pin_waiting_hits() + + # DP prefill balancing: on a step that admitted prefills (release), + # record whether it was capacity-bound. + if not defer_prefills: + _, num_waiting = self.get_request_counts() + self.prefill_capacity_bound = num_waiting > len(self.deferred_waiting) + + # Check if the scheduling constraints are satisfied. + total_num_scheduled_tokens = sum(num_scheduled_tokens.values()) + assert total_num_scheduled_tokens <= self.max_num_scheduled_tokens + + assert token_budget >= 0 + assert input_budget >= 0 + assert len(self.running) <= self.max_num_running_reqs + # Since some requests in the RUNNING queue may not be scheduled in + # this step, the total number of scheduled requests can be smaller than + # len(self.running). + assert len(scheduled_new_reqs) + len(scheduled_resumed_reqs) + len( + scheduled_running_reqs + ) <= len(self.running) + + # Get the longest common prefix among all requests in the running queue. + # This can be potentially used for cascade attention. + num_common_prefix_blocks = [0] * len(self.kv_cache_config.kv_cache_groups) + with record_function_or_nullcontext("schedule: get_num_common_prefix_blocks"): + if self.running: + any_request_id = self.running[0].request_id + num_common_prefix_blocks = ( + self.kv_cache_manager.get_num_common_prefix_blocks(any_request_id) + ) + + # Construct the scheduler output. + if self.use_v2_model_runner: + scheduled_new_reqs.extend(scheduled_resumed_reqs) + scheduled_resumed_reqs.clear() + new_reqs_data = [ + NewRequestData.from_request( + req, + req_to_new_blocks[req.request_id].get_block_ids(), + req._all_token_ids, + uses_mrope=self.model_uses_mrope, + ) + for req in scheduled_new_reqs + ] + else: + new_reqs_data = [ + NewRequestData.from_request( + req, + req_to_new_blocks[req.request_id].get_block_ids(), + uses_mrope=self.model_uses_mrope, + ) + for req in scheduled_new_reqs + ] + + with record_function_or_nullcontext("schedule: make_cached_request_data"): + cached_reqs_data = self._make_cached_request_data( + scheduled_running_reqs, + scheduled_resumed_reqs, + num_scheduled_tokens, + scheduled_spec_decode_tokens, + req_to_new_blocks, + ) + + # Record the request ids that were scheduled in this step (MRV1-only). + if not self.use_v2_model_runner: + self.prev_step_scheduled_req_ids.clear() + self.prev_step_scheduled_req_ids.update(num_scheduled_tokens.keys()) + + # Mamba "align" boundary states must be handed off with exact block ids; + # they cannot be reconstructed from a connector's append-only block + # table. Drained every step so stale offers cannot accumulate. + boundary_state_offloads = self.kv_cache_manager.take_boundary_state_offloads() + + kv_connector_block_state = None + if self.connector is not None: + # Any request scheduled this step can become a connector job now, + # not only the ones that were allocated blocks: a store save lands + # on the step that fills a block, which allocated none. + block_state_req_ids = set(num_scheduled_tokens) + block_state_req_ids.update( + req_id for req_id in boundary_state_offloads if req_id in self.requests + ) + kv_connector_block_state = KVConnectorBlockState( + req_ids=block_state_req_ids, + resolve_block_ids=self.kv_cache_manager.get_block_ids, + boundary_state_offloads=boundary_state_offloads, + ) + + kv_cache_block_copies, cow_retained_blocks = ( + self.kv_cache_manager.take_kv_cache_block_copies() + ) + if kv_cache_block_copies: + # The copies run with this step's execution; the first non-empty + # step at or after it gets seq `sched_step_seq + 1` (0-token steps + # do not advance the seq), and its completion implies the copies + # have run. + self._free_cow_retained_blocks(cow_retained_blocks, self.sched_step_seq + 1) + pending_kv_cache_block_copies = kv_cache_block_copies or None + + # Dynamic speculative decoding: compute optimal K + num_spec_tokens_to_schedule = self.num_spec_tokens + if self.dynamic_sd_lookup is not None and len(num_scheduled_tokens) > 0: + num_spec_tokens_to_schedule = self.dynamic_sd_lookup[ + len(num_scheduled_tokens) + ] + + scheduled_encoder_input_stats = None + if ( + self.log_stats + and self.observability_config.enable_logging_iteration_details + ): + scheduled_encoder_input_stats = self._make_scheduled_encoder_input_stats( + scheduled_encoder_inputs + ) + + scheduler_output = SchedulerOutput( + scheduled_new_reqs=new_reqs_data, + scheduled_cached_reqs=cached_reqs_data, + num_scheduled_tokens=num_scheduled_tokens, + total_num_scheduled_tokens=total_num_scheduled_tokens, + scheduled_spec_decode_tokens=scheduled_spec_decode_tokens, + scheduled_encoder_inputs=scheduled_encoder_inputs, + scheduled_encoder_input_stats=scheduled_encoder_input_stats, + num_common_prefix_blocks=num_common_prefix_blocks, + preempted_req_ids=self.reset_preempted_req_ids, + # finished_req_ids is an existing state in the scheduler, + # instead of being newly scheduled in this step. + # It contains the request IDs that are finished in between + # the previous and the current steps. + finished_req_ids=self.finished_req_ids, + free_encoder_mm_hashes=self.encoder_cache_manager.get_freed_mm_hashes(), + new_block_ids_to_zero=self._get_new_block_ids_to_zero(), + # Drained every step, never gated on KV zeroing. + new_ring_block_ids_to_clear=self._get_new_ring_block_ids_to_clear(), + has_sync_kv_loads=has_sync_kv_loads, + kv_cache_block_copies=pending_kv_cache_block_copies, + kv_connector_block_state=kv_connector_block_state, + num_spec_tokens_to_schedule=num_spec_tokens_to_schedule, + ec_manager_metadata=self.encoder_cache_manager.get_manager_metadata(), + scheduler_step=self.current_step, + ) + + # NOTE(Kuntai): this function is designed for multiple purposes: + # 1. Plan the KV cache store + # 2. Wrap up all the KV cache load / save ops into an opaque object + # 3. Clear the internal states of the connector + if self.connector is not None: + meta = self._build_kv_connector_meta(self.connector, scheduler_output) + scheduler_output.kv_connector_metadata = meta + + if self.aux_output_connector is not None: + scheduler_output.aux_output_connector_metadata = ( + self.aux_output_connector.build_connector_meta( + scheduler_output, self.requests + ) + ) + + # Build the connector meta for ECConnector + if self.ec_connector is not None: + ec_meta: ECConnectorMetadata = self.ec_connector.build_connector_meta( + scheduler_output + ) + scheduler_output.ec_connector_metadata = ec_meta + + # Connector-only block state must not be dispatched to workers. + scheduler_output.kv_connector_block_state = None + + # Advance the fence only for non-empty steps (those that actually + # write KV and have their output processed later in update_from_output). + if self.defer_block_free and total_num_scheduled_tokens > 0: + self.sched_step_seq += 1 + + with record_function_or_nullcontext("schedule: update_after_schedule"): + self._update_after_schedule(scheduler_output) + return scheduler_output + + def _build_kv_connector_meta( + self, connector: KVConnectorBase_V1, scheduler_output: SchedulerOutput + ) -> KVConnectorMetadata: + return connector.build_connector_meta(scheduler_output) + + def _get_new_ring_block_ids_to_clear(self) -> list[int] | None: + """Ring claims of this step whose tag words the worker clears. + Unconditional (not under needs_kv_cache_zeroing); only blocks an async + KV load is about to write are left out.""" + ids = self.kv_cache_manager.take_new_ring_block_ids() + if self._skip_ring_clear_block_ids: + skip = self._skip_ring_clear_block_ids + ids = [b for b in ids if b not in skip] + skip.clear() + return ids or None + + def _get_new_block_ids_to_zero(self) -> list[int] | None: + # Drain new attention block ids every step so the manager-side list + # does not grow unbounded; only kv-cache zeroing consumes them. + new_block_ids_to_zero = self.kv_cache_manager.take_new_block_ids() + if not self.needs_kv_cache_zeroing: + return None + if new_block_ids_to_zero and _ZERO_OWNERSHIP_CHECK: + self._check_zero_block_ownership(new_block_ids_to_zero) + + return new_block_ids_to_zero or None + + def _check_zero_block_ownership(self, block_ids: list[int]) -> None: + """Test KVBlockZeroer's stated precondition, host side, off by default. + + worker/utils.py justifies zeroing a whole tile with "block IDs are + global pool indices and a newly allocated block owns its whole tile". + Ownership is visible only here, in the scheduler process, so the check + lives here rather than in the zeroer. A block about to be zeroed should + be referenced by exactly one live request. More than one means a stale + reference survived into the zeroing, and the kernel is about to destroy + state another request still points at -- the stale-reference shape. + + VLLM_E2_ZERO_OWNERSHIP_CHECK=1. O(live blocks) per step. + """ + pending = set(block_ids) + owners: dict[int, set[str]] = {} + for mgr in self.kv_cache_manager.coordinator.single_type_managers: + if getattr(mgr, "slot_pool", None) is not None: + # v2.5 slot-packed ring: its block ids are slot ids, not pool + # block ids, and its pages are never zeroed. + continue + for req_id, blocks in mgr.req_to_blocks.items(): + for blk in blocks: + if blk.block_id in pending: + owners.setdefault(blk.block_id, set()).add(req_id) + for block_id, req_ids in sorted(owners.items()): + if len(req_ids) > 1: + logger.warning( + "KV zeroing ownership violation: block %d is about to be " + "zeroed but is referenced by %d live requests (%s).", + block_id, + len(req_ids), + ", ".join(sorted(req_ids)[:4]), + ) + + def _preempt_request( + self, request: Request, timestamp: float, drop_stale_output: bool = False + ) -> None: + """Preempt a request and put it back to the waiting queue. + + NOTE: The request should be popped from the running queue outside of this + method. + + drop_stale_output: drop (rather than deliver) any in-flight output; used + by reset_prefix_cache, whose same-step resume would otherwise deliver + tokens out of order, and for connectors with a pending KV hand-off, + which the preemption's block free would leave without valid KV. + """ + assert request.status == RequestStatus.RUNNING, ( + "Only running requests can be preempted" + ) + if self.aux_output_connector is not None: + self.aux_output_connector.request_finished(request) + self._free_request_blocks(request) + self.encoder_cache_manager.free(request) + self._inflight_prefills.discard(request) + request.status = RequestStatus.PREEMPTED + request.num_computed_tokens = 0 + if request.spec_token_ids: + request.spec_token_ids = [] + # Async scheduling: mark all in-flight output as stale. Its tokens are + # still delivered on return (dropping them would perturb spec-decode + # acceptance) but must not mutate the reset counters; each step drains + # its share in update_from_output. num_in_flight_tokens already + # includes any undrained stale share, so assign rather than accumulate. + # An undrained drop-mode share stays dropped: its positions have + # already been resampled. + request.drop_stale_output = drop_stale_output or ( + request.drop_stale_output and request.num_stale_output_tokens > 0 + ) + request.num_stale_output_tokens = request.num_in_flight_tokens + request.num_output_placeholders = 0 + request.num_preemptions += 1 + if self.log_stats: + request.record_event(EngineCoreEventType.PREEMPTED, timestamp) + + # Put the request back to the waiting queue. + self.waiting.prepend_request(request) + self.reset_preempted_req_ids.add(request.request_id) + + def _update_after_schedule(self, scheduler_output: SchedulerOutput) -> None: + # Advance the number of computed tokens for the request AFTER + # the request is scheduled. + # 1. The scheduler_output of the current step has to include the + # original number of scheduled tokens to determine input IDs. + # 2. Advance the number of computed tokens here allowing us to + # schedule the prefill request again immediately in the next + # scheduling step. + # 3. If some tokens (e.g. spec tokens) are rejected later, the number of + # computed tokens will be adjusted in update_from_output. + num_scheduled_tokens = scheduler_output.num_scheduled_tokens + for req_id, num_scheduled_token in num_scheduled_tokens.items(): + request = self.requests[req_id] + request.num_computed_tokens += num_scheduled_token + request.num_in_flight_tokens += num_scheduled_token + if self.defer_block_free: + # Record the in-flight step, to fence deferred block freeing. + request.last_sched_seq = self.sched_step_seq + request.is_prefill_chunk = request.num_computed_tokens < ( + request.num_tokens + request.num_output_placeholders + ) + scheduler_output.has_structured_output_requests |= ( + request.use_structured_output and not request.is_prefill_chunk + ) + # Drop from the in-flight-prefill set once it's no longer prefilling. + if not request.is_prefill_chunk: + self._inflight_prefills.discard(request) + + # Clear the finished and preempted request IDs. + # NOTE: We shouldn't just clear() here because it will also affect + # the scheduler output. + self.finished_req_ids = set() + self.reset_preempted_req_ids = set() + + def _update_request_as_session( + self, session: Request, update: StreamingUpdate + ) -> None: + """Updates the waiting session with the next streaming update. + + Discards the last sampled output token from the prior input chunk. + """ + # Current streaming input behaviour: Keep only computed output tokens + # (discard final sampled output token). + num_computed_tokens = session.num_computed_tokens + kept_output_tokens = session._all_token_ids[ + session.num_prompt_tokens : num_computed_tokens + ] + del session._all_token_ids[num_computed_tokens:] + session._output_token_ids.clear() + assert session.prompt_token_ids is not None + # Extend prompt with kept output tokens. + session.prompt_token_ids.extend(kept_output_tokens) + + if update.mm_features: + base = session.num_tokens + for mm_feature in update.mm_features: + mm_feature.mm_position = replace( + mm_feature.mm_position, offset=mm_feature.mm_position.offset + base + ) + session.mm_features.extend(update.mm_features) + + session._all_token_ids.extend(update.prompt_token_ids or ()) + session.prompt_token_ids.extend(update.prompt_token_ids or ()) + # Update block hashes for the new tokens. + session.update_block_hashes() + session.num_prompt_tokens = len(session.prompt_token_ids) + session.arrival_time = update.arrival_time + session.sampling_params = update.sampling_params + session.max_tokens = update.max_tokens + if session.status == RequestStatus.WAITING_FOR_STREAMING_REQ: + self.num_waiting_for_streaming_input -= 1 + session.status = RequestStatus.WAITING + + if self.log_stats: + session.record_event(EngineCoreEventType.QUEUED) + + def _make_cached_request_data( + self, + running_reqs: list[Request], + resumed_reqs: list[Request], + num_scheduled_tokens: dict[str, int], + spec_decode_tokens: dict[str, list[int]], + req_to_new_blocks: dict[str, KVCacheBlocks], + ) -> CachedRequestData: + req_ids: list[str] = [] + new_token_ids: list[list[int]] = [] + new_block_ids: list[tuple[list[int], ...] | None] = [] + all_token_ids: dict[str, list[int]] = {} + num_computed_tokens: list[int] = [] + num_output_tokens: list[int] = [] + resumed_req_ids = set() + + num_running_reqs = len(running_reqs) + for idx, req in enumerate(itertools.chain(running_reqs, resumed_reqs)): + req_id = req.request_id + req_ids.append(req_id) + # NOTE: In PP+async scheduling, we consume token ids via a direct GPU + # broadcast path (`input_batch.prev_sampled_token_ids`), so we can + # omit this payload. + if self.use_pp and not self.scheduler_config.async_scheduling: + # When using PP, the scheduler sends the sampled tokens back, + # because there's no direct communication between the first- + # stage worker and the last-stage worker. Otherwise, we don't + # need to send the sampled tokens back because the model runner + # will cache them. + num_tokens = num_scheduled_tokens[req_id] - len( + spec_decode_tokens.get(req_id, ()) + ) + token_ids = req.all_token_ids[ + req.num_computed_tokens : req.num_computed_tokens + num_tokens + ] + new_token_ids.append(token_ids) + if idx >= num_running_reqs: + resumed_req_ids.add(req_id) + if not self.use_v2_model_runner: # noqa: SIM102 + if req_id not in self.prev_step_scheduled_req_ids: + all_token_ids[req_id] = req.all_token_ids.copy() + new_block_ids.append( + req_to_new_blocks[req_id].get_block_ids(allow_none=True) + ) + num_computed_tokens.append(req.num_computed_tokens) + num_output_tokens.append( + req.num_output_tokens + req.num_output_placeholders + ) + + return CachedRequestData( + req_ids=req_ids, + resumed_req_ids=resumed_req_ids, + new_token_ids=new_token_ids, + all_token_ids=all_token_ids, + new_block_ids=new_block_ids, + num_computed_tokens=num_computed_tokens, + num_output_tokens=num_output_tokens, + ) + + def _reject_on_encoder_cache_embed_mismatch( + self, + request: Request, + input_id: int, + identifier: str, + cached_num_encoder_embeds: int, + requested_num_encoder_embeds: int, + ) -> bool: + if cached_num_encoder_embeds == requested_num_encoder_embeds: + return False + logger.warning( + "Encoder cache entry for multimodal identifier %r holds %d " + "embeddings, but request %s input %d expects %d. A client-provided " + "multimodal UUID may have been reused for different content.", + identifier, + cached_num_encoder_embeds, + request.request_id, + input_id, + requested_num_encoder_embeds, + ) + self.encoder_cache_mismatch_reqs.add(request.request_id) + return True + + def _try_schedule_encoder_inputs( + self, + request: Request, + num_computed_tokens: int, + num_new_tokens: int, + encoder_compute_budget: int, + shift_computed_tokens: int = 0, + ) -> tuple[list[int], int, int, list[int]]: + """Determine which encoder inputs need to be scheduled in the current step, + and update `num_new_tokens` and encoder token budget accordingly. + + An encoder input will be scheduled if: + - Its output tokens overlap with the range of tokens being computed + in this step, i.e., + [num_computed_tokens, num_computed_tokens + num_new_tokens). + - It is not already computed and stored in the encoder cache. + - It is not exist on remote encoder cache (via ECConnector) + - There is sufficient encoder token budget to process it. + - The encoder cache has space to store it. + + If an encoder input cannot be scheduled due to cache or budget + limitations, the method adjusts `num_new_tokens` to schedule only the + decoder tokens up to just before the unschedulable encoder input. + + Note that num_computed_tokens includes both locally cached + blocks and externally cached blocks (via KVConnector). + """ + if num_new_tokens == 0 or not request.has_encoder_inputs: + return [], num_new_tokens, encoder_compute_budget, [] + encoder_inputs_to_schedule: list[int] = [] + mm_features = request.mm_features + assert mm_features is not None + assert len(mm_features) > 0 + external_load_encoder_input = [] + + # NOTE: since scheduler operates on the request level (possibly with + # multiple encoder inputs per request), we need to create temporary + # trackers for accounting at the encoder input level. + mm_hashes_to_schedule: dict[str, int] = {} + num_embeds_to_schedule = 0 + + encoder_window_end = ( + num_computed_tokens + num_new_tokens + shift_computed_tokens + ) + lo, hi = get_mm_features_in_window( + mm_features, + start=num_computed_tokens, + end=encoder_window_end, + ) + # For encoder-decoder, all inputs sit at start_pos=0, so lo=0 always. + if self.is_encoder_decoder: + lo = 0 + + for i in range(lo, hi): + mm_feature = mm_features[i] + start_pos = mm_feature.mm_position.offset + num_encoder_tokens = mm_feature.mm_position.length + num_encoder_embeds = mm_feature.mm_position.get_num_embeds() + item_identifier = mm_feature.identifier + + if self.is_encoder_decoder and num_computed_tokens > 0: + assert start_pos == 0, ( + "Encoder input should be processed at the beginning of " + "the sequence when encoder-decoder models are used." + ) + # Encoder input has already been computed + # The calculation here is a bit different. We don't turn encoder + # output into tokens that get processed by the decoder and + # reflected in num_computed_tokens. Instead, start_pos reflects + # the position where we need to ensure we calculate encoder + # inputs. This should always be 0 to ensure we calculate encoder + # inputs before running the decoder. Once we've calculated some + # decoder tokens (num_computed_tokens > 0), then we know we + # already calculated encoder inputs and can skip here. + continue + + if not self.is_encoder_decoder: + # We are not using the encoder cache for encoder-decoder models, + # yet. + if item_identifier in mm_hashes_to_schedule: + scheduled_num_encoder_embeds = mm_hashes_to_schedule[ + item_identifier + ] + if self._reject_on_encoder_cache_embed_mismatch( + request, + i, + item_identifier, + scheduled_num_encoder_embeds, + num_encoder_embeds, + ): + return [], 0, encoder_compute_budget, [] + # The same encoder input has already been scheduled in the + # current step. + continue + + cached_num_encoder_embeds = ( + self.encoder_cache_manager.get_cached_num_encoder_embeds(request, i) + ) + if cached_num_encoder_embeds is not None and ( + self._reject_on_encoder_cache_embed_mismatch( + request, + i, + item_identifier, + cached_num_encoder_embeds, + num_encoder_embeds, + ) + ): + return [], 0, encoder_compute_budget, [] + + if self.encoder_cache_manager.check_and_update_cache(request, i): + # The encoder input is already computed and cached from a + # previous step. + continue + + # If no encoder input chunking is allowed, we do not want to + # partially schedule a multimodal item. If the scheduled range would + # only cover part of the mm input, roll back to before the mm item. + if ( + self.scheduler_config.disable_chunked_mm_input + and num_computed_tokens < start_pos + and (num_computed_tokens + num_new_tokens) + < (start_pos + num_encoder_tokens) + ): + # Account for EAGLE shift when rolling back to avoid + # encoder cache miss. This ensures the scheduled range + # stops before start_pos even with the shift. + num_new_tokens = max( + 0, start_pos - (num_computed_tokens + shift_computed_tokens) + ) + break + if not self.encoder_cache_manager.can_allocate( + request, i, encoder_compute_budget, num_embeds_to_schedule + ): + # The encoder cache is full or the encoder budget is exhausted. + # NOTE(woosuk): We assume that the encoder input tokens should + # be processed altogether, as the encoder usually uses + # bidirectional attention. + if num_computed_tokens + shift_computed_tokens < start_pos: + # We only schedule the decoder tokens just before the + # encoder input. + num_new_tokens = start_pos - ( + num_computed_tokens + shift_computed_tokens + ) + else: + # Because of prefix caching, num_computed_tokens is greater + # than start_pos even though its encoder input is not + # available. In this case, we can't schedule any token for + # the request in this step. + num_new_tokens = 0 + break + + # Calculate the number of embeddings to schedule in the current range + # of scheduled encoder placeholder tokens. + start_idx_rel = max(0, num_computed_tokens - start_pos) + end_idx_rel = min(num_encoder_tokens, encoder_window_end - start_pos) + curr_embeds_start, curr_embeds_end = ( + mm_feature.mm_position.get_embeds_indices_in_range( + start_idx_rel, end_idx_rel + ) + ) + # There's no embeddings in the current range of encoder placeholder tokens + # so we can skip the encoder input. + if curr_embeds_end - curr_embeds_start == 0: + continue + + if self.ec_connector is not None and self.ec_connector.has_cache_item( + item_identifier + ): + mm_hashes_to_schedule[item_identifier] = num_encoder_embeds + external_load_encoder_input.append(i) + num_embeds_to_schedule += num_encoder_embeds + continue + + num_embeds_to_schedule += num_encoder_embeds + encoder_compute_budget -= num_encoder_embeds + mm_hashes_to_schedule[item_identifier] = num_encoder_embeds + encoder_inputs_to_schedule.append(i) + + return ( + encoder_inputs_to_schedule, + num_new_tokens, + encoder_compute_budget, + external_load_encoder_input, + ) + + def _make_scheduled_encoder_input_stats( + self, scheduled_encoder_inputs: dict[str, list[int]] + ) -> ScheduledEncoderInputStats | None: + stats = ScheduledEncoderInputStats() + + for req_id, input_ids in scheduled_encoder_inputs.items(): + request = self.requests.get(req_id) + if request is None: + continue + + for input_id in input_ids: + mm_feature = request.mm_features[input_id] + stats.num_inputs += 1 + stats.output_tokens += mm_feature.mm_position.get_num_embeds() + + return stats if stats.num_inputs else None + + def get_grammar_bitmask( + self, scheduler_output: SchedulerOutput + ) -> GrammarOutput | None: + # Collect list of scheduled request ids that use structured output. + # The corresponding rows of the bitmask will be in this order. + if not scheduler_output.has_structured_output_requests: + return None + + structured_output_request_ids = [ + req_id + for req_id in scheduler_output.num_scheduled_tokens + if (req := self.requests.get(req_id)) + and (req.use_structured_output and not req.is_prefill_chunk) + ] + if not structured_output_request_ids: + return None + + bitmask = self.structured_output_manager.grammar_bitmask( + self.requests, + structured_output_request_ids, + scheduler_output.scheduled_spec_decode_tokens, + ) + spec_tokens = scheduler_output.scheduled_spec_decode_tokens + num_acceptable_drafts = [ + len(strip_speculative_padding(spec_tokens.get(req_id, []))) + for req_id in structured_output_request_ids + ] + return GrammarOutput( + structured_output_request_ids, bitmask, num_acceptable_drafts + ) + + def update_from_output( + self, + scheduler_output: SchedulerOutput, + model_runner_output: ModelRunnerOutput, + ) -> dict[int, EngineCoreOutputs]: + sampled_token_ids = model_runner_output.sampled_token_ids + logprobs = model_runner_output.logprobs + prompt_logprobs_dict = model_runner_output.prompt_logprobs_dict + prompt_token_id_logprobs_dict = ( + model_runner_output.prompt_token_id_logprobs_dict + ) + num_scheduled_tokens = scheduler_output.num_scheduled_tokens + pooler_outputs = model_runner_output.pooler_output + num_nans_in_logits = model_runner_output.num_nans_in_logits + kv_connector_output = model_runner_output.kv_connector_output + ec_connector_output = model_runner_output.ec_connector_output + cudagraph_stats = model_runner_output.cudagraph_stats + + # Every GPU write enqueued by this and earlier steps has completed, so it is + # safe to return deferred-free blocks to the pool. + if self.defer_block_free and scheduler_output.total_num_scheduled_tokens > 0: + self.processed_step_seq += 1 + self._drain_deferred_frees() + + perf_stats: PerfStats | None = None + if self.perf_metrics and self.perf_metrics.is_enabled(): + perf_stats = self.perf_metrics.get_step_perf_stats_per_gpu(scheduler_output) + + outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list) + spec_decoding_stats: SpecDecodingStats | None = None + + failed_kv_load_req_ids: set[str] = set() + if kv_connector_output and kv_connector_output.failed_recving: + failed_kv_load_req_ids.update( + self._handle_failed_recving(kv_connector_output.failed_recving) + ) + if kv_connector_output and kv_connector_output.invalid_block_ids: + # These blocks contain externally computed tokens that failed to + # load. Identify affected requests and adjust their computed token + # count to trigger recomputation of the invalid blocks. + failed_kv_load_req_ids.update( + self._handle_invalid_blocks( + kv_connector_output.invalid_block_ids, + num_scheduled_tokens, + ) + ) + # NOTE(woosuk): As len(num_scheduled_tokens) can be up to 1K or more, + # the below loop can be a performance bottleneck. We should do our best + # to avoid expensive operations inside the loop. + stopped_running_reqs: set[Request] = set() + stopped_preempted_reqs: set[Request] = set() + for req_id, num_tokens_scheduled in num_scheduled_tokens.items(): + assert num_tokens_scheduled > 0 + request = self.requests.get(req_id) + output_is_stale = False + if request is not None: + request.num_in_flight_tokens -= num_tokens_scheduled + # Drain any stale share (see _preempt_request) in lockstep. + if request.num_stale_output_tokens > 0: + output_is_stale = True + request.num_stale_output_tokens -= num_tokens_scheduled + assert request.num_stale_output_tokens >= 0 + if failed_kv_load_req_ids and req_id in failed_kv_load_req_ids: + # skip failed or rescheduled requests from KV load failure + continue + if request is None or request.is_finished(): + # The request is already finished. This can happen if the + # request is aborted while the model is executing it (e.g., + # in pipeline parallelism or in async scheduling). + # NOTE(Kuntai): When delay_free_blocks=True (for async KV + # cache transfer in KV connector), the aborted request will not + # be set to None (in order to finish async KV transfer). + # In this case, we use is_finished() to check. + continue + + # Drop-mode stale output (same-step resume) is discarded entirely. + if output_is_stale and request.drop_stale_output: + continue + + req_index = model_runner_output.req_id_to_index[req_id] + generated_token_ids = ( + sampled_token_ids[req_index] if sampled_token_ids else [] + ) + + scheduled_spec_token_ids = ( + scheduler_output.scheduled_spec_decode_tokens.get(req_id) + ) + if scheduled_spec_token_ids and ( + generated_token_ids or self.num_sampled_tokens_per_step == 0 + ): + num_draft_tokens = len(scheduled_spec_token_ids) + num_sampled = self.num_sampled_tokens_per_step + num_accepted = max(len(generated_token_ids) - num_sampled, 0) + num_rejected = num_draft_tokens - num_accepted + # Rejections roll back num_computed_tokens (and, under async + # scheduling, num_output_placeholders, which covers the spec + # tokens). A stale rejection count predates the preemption + # rollback and must not apply. + if not output_is_stale: + if request.num_computed_tokens > 0: + request.num_computed_tokens -= num_rejected + if request.num_output_placeholders > 0: + request.num_output_placeholders -= num_rejected + spec_decoding_stats = self.make_spec_decoding_stats( + spec_decoding_stats, + num_draft_tokens=num_draft_tokens, + num_accepted_tokens=num_accepted, + num_invalid_spec_tokens=scheduler_output.num_invalid_spec_tokens, + request_id=req_id, + ) + if request.spec_decode_metrics is not None: + # Exclude grammar-invalidated drafts from the proposed + # count, mirroring make_spec_decoding_stats; the accepted + # bucket (j) is unaffected. + adj_draft_tokens = num_draft_tokens + if scheduler_output.num_invalid_spec_tokens: + adj_draft_tokens -= ( + scheduler_output.num_invalid_spec_tokens.get(req_id, 0) + ) + request.spec_decode_metrics.observe( + num_draft_tokens=adj_draft_tokens, + num_accepted=num_accepted, + detailed=self.spec_decode_metrics_level == "detailed", + ) + + # Free encoder inputs only after the step has actually executed. + if request.has_encoder_inputs: + self._free_encoder_inputs(request) + + stopped = False + new_logprobs = None + new_sampling_mask = None + new_token_ids = generated_token_ids + pooler_output = pooler_outputs[req_index] if pooler_outputs else None + kv_transfer_params = None + ec_transfer_params = None + prefill_stats = None + status_before_stop = request.status + + # Check for stop and update request status. + if new_token_ids: + new_token_ids, stopped = self._update_request_with_output( + request, new_token_ids, is_stale=output_is_stale + ) + elif request.pooling_params and pooler_output is not None: + # Pooling stops as soon as there is output. + request.status = RequestStatus.FINISHED_STOPPED + stopped = True + elif ( + self.is_mm_encoder_only + and request.num_computed_tokens >= request.num_prompt_tokens + ): + # An encoder instance runs the encoder and publishes the + # embeddings instead of sampling, so it stops as soon as the + # whole prompt is consumed. Encoder inputs are never scheduled + # past a multi-modal item the encoder cache could not admit, so + # a consumed prompt also means every item in it was encoded. + request.status = RequestStatus.FINISHED_STOPPED + stopped = True + + if new_token_ids and not self.structured_output_manager.accept_tokens( + request, new_token_ids + ): + logger.error( + "Unexpected: grammar rejected tokens %s for request %s. " + "Terminating request.", + new_token_ids, + req_id, + ) + request.status = RequestStatus.FINISHED_ERROR + request.resumable = False + stopped = True + + routed_experts = None + should_emit_output = bool( + new_token_ids or pooler_output is not None or stopped + ) + if self.aux_output_connector is not None and should_emit_output: + routed_experts = self.aux_output_connector.take_output( + request, model_runner_output.aux_output_connector_output + ) + if should_emit_output: + prefill_stats = request.take_prefill_stats() + if prefill_stats is not None: + prefill_stats.finalize( + self.kv_cache_manager.estimate_cached_tokens(request) + ) + + # Extract sample logprobs before stop handling can replace the + # sampling parameters for a streaming continuation. + if ( + request.sampling_params is not None + and request.sampling_params.num_logprobs is not None + and logprobs + ): + new_logprobs = logprobs.slice_request(req_index, len(new_token_ids)) + + finish_reason = None + if stopped: + # Capture finish_reason BEFORE _handle_stopped_request, which may + # reset the status to WAITING for streaming requests that continue. + finish_reason = request.get_finished_reason() + finished = self._handle_stopped_request(request) + if finished: + kv_transfer_params, ec_transfer_params = self._free_request(request) + if ( + prefill_stats is not None + and kv_transfer_params is not None + and kv_transfer_params.get("do_remote_prefill") + ): + # P-side cache hits, so D can report them in + # prompt_tokens_details instead of its own (~100%) + # hit rate from the KV transfer. + kv_transfer_params["remote_prefill_cached_tokens"] = ( + prefill_stats.num_cached_tokens + ) + + if status_before_stop == RequestStatus.RUNNING: + stopped_running_reqs.add(request) + else: + stopped_preempted_reqs.add(request) + + if self.return_sampling_mask: + sampling_masks = model_runner_output.sampling_masks + if new_token_ids and sampling_masks is not None: + new_sampling_mask = sampling_masks.slice_request( + req_index, len(new_token_ids) + ) + + if num_nans_in_logits is not None and req_id in num_nans_in_logits: + request.num_nans_in_logits = num_nans_in_logits[req_id] + + # Get prompt logprobs for this request. + prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id) + prompt_token_id_logprobs = prompt_token_id_logprobs_dict.get(req_id) + if should_emit_output: + # Add EngineCoreOutput for this Request. + outputs[request.client_index].append( + EngineCoreOutput( + request_id=req_id, + new_token_ids=new_token_ids, + finish_reason=finish_reason, + new_logprobs=new_logprobs, + new_sampling_mask=new_sampling_mask, + new_prompt_logprobs_tensors=prompt_logprobs_tensors, + prompt_token_id_logprobs=prompt_token_id_logprobs, + pooling_output=pooler_output, + stop_reason=request.stop_reason, + events=request.take_events(), + prefill_stats=prefill_stats, + spec_decode_metrics=( + request.spec_decode_metrics + if finish_reason is not None + else None + ), + kv_transfer_params=kv_transfer_params, + ec_transfer_params=ec_transfer_params, + trace_headers=request.trace_headers, + routed_experts=routed_experts, + num_nans_in_logits=request.num_nans_in_logits, + ) + ) + else: + # Invariant: EngineCore returns no partial prefill outputs. + assert not prompt_logprobs_tensors + assert prompt_token_id_logprobs is None + + # Remove the stopped requests from the running and waiting queues. + if stopped_running_reqs: + self.running = remove_all(self.running, stopped_running_reqs) + if stopped_preempted_reqs: + # This is a rare case and unlikely to impact performance. + self.waiting.remove_requests(stopped_preempted_reqs) + self.kv_holding_waiting.remove_requests(stopped_preempted_reqs) + self.deferred_waiting.difference_update(stopped_preempted_reqs) + + error_req_ids = set(self.grammar_compile_error_reqs) + self.grammar_compile_error_reqs.clear() + error_req_ids.update(self.encoder_cache_mismatch_reqs) + self.encoder_cache_mismatch_reqs.clear() + if failed_kv_load_req_ids and not self.recompute_kv_load_failures: + error_req_ids.update(failed_kv_load_req_ids) + if self.ec_connector is not None: + # An encoder input the connector can no longer obtain. Failing is + # retryable: re-issuing the request re-runs the encode. + error_req_ids.update(self.ec_connector.take_unavailable_requests()) + + if error_req_ids: + error_reqs = self.finish_requests( + error_req_ids, RequestStatus.FINISHED_ERROR + ) + for request in error_reqs: + outputs[request.client_index].append( + EngineCoreOutput( + request_id=request.request_id, + new_token_ids=[], + finish_reason=request.get_finished_reason(), + events=request.take_events(), + trace_headers=request.trace_headers, + ) + ) + + # KV Connector: update state for finished KV Transfers. + if kv_connector_output: + self._update_from_kv_xfer_finished(kv_connector_output) + + # EC Connector: update state from worker-side EC connector output. + if self.ec_connector is not None and ec_connector_output: + self.ec_connector.update_connector_output(ec_connector_output) + + # Worker-side KV connector stats from the model runner output. + kv_connector_stats: KVConnectorStats | None = ( + kv_connector_output.kv_connector_stats if kv_connector_output else None + ) + if self.connector: + # Scheduler-side KV connector stats collected after connector update. + scheduler_kv_connector_stats = self.connector.get_kv_connector_stats() + if ( + scheduler_kv_connector_stats is not None + and not scheduler_kv_connector_stats.is_empty() + ): + kv_connector_stats = ( + kv_connector_stats.aggregate(scheduler_kv_connector_stats) + if kv_connector_stats is not None + else scheduler_kv_connector_stats + ) + + # Worker-side EC connector stats from the model runner output. + ec_connector_stats: ECConnectorStats | None = ( + ec_connector_output.ec_connector_stats if ec_connector_output else None + ) + if self.ec_connector: + # Scheduler-side EC connector stats collected after connector update. + scheduler_ec_connector_stats = self.ec_connector.get_ec_connector_stats() + if ( + scheduler_ec_connector_stats is not None + and not scheduler_ec_connector_stats.is_empty() + ): + ec_connector_stats = ( + ec_connector_stats.aggregate(scheduler_ec_connector_stats) + if ec_connector_stats is not None + else scheduler_ec_connector_stats + ) + + # collect KV cache events from KV cache manager + events = self.kv_cache_manager.take_events() + + # collect KV cache events from connector + if self.connector is not None: + connector_events = self.connector.take_events() + if connector_events: + if events is None: + events = list(connector_events) + else: + events.extend(connector_events) + + # publish collected KV cache events + if events: + batch = KVEventBatch(ts=time.time(), events=events) + self.kv_event_publisher.publish(batch) + + # Create EngineCoreOutputs for all clients that have requests with + # outputs in this step. + engine_core_outputs = { + client_index: EngineCoreOutputs(outputs=outs) + for client_index, outs in outputs.items() + } + + finished_req_ids = self.finished_req_ids_dict + if finished_req_ids: + # Include ids of requests that finished since last outputs + # were sent. + for client_index, finished_set in finished_req_ids.items(): + # Set finished request set in EngineCoreOutputs for this client. + if (eco := engine_core_outputs.get(client_index)) is not None: + eco.finished_requests = finished_set + else: + engine_core_outputs[client_index] = EngineCoreOutputs( + finished_requests=finished_set + ) + finished_req_ids.clear() + + if ( + stats := self.make_stats( + spec_decoding_stats, + kv_connector_stats, + cudagraph_stats, + perf_stats, + ec_connector_stats=ec_connector_stats, + ) + ) is not None: + # Return stats to only one of the front-ends. + if (eco := next(iter(engine_core_outputs.values()), None)) is None: + # We must return the stats even if there are no request + # outputs this step. + engine_core_outputs[0] = eco = EngineCoreOutputs() + eco.scheduler_stats = stats + + if _e2_alloc_log.enabled(): + _e2_alloc_log.log_step(self, scheduler_output) + return engine_core_outputs + + def _ec_transfer_pending(self, request: Request, num_computed_tokens: int) -> bool: + """Whether an encoder input this request needs is still in transit.""" + return ( + self.ec_connector is not None + and bool(request.mm_features) + and not self.ec_connector.ensure_cache_available( + request, num_computed_tokens + ) + ) + + @staticmethod + def _is_blocked_waiting_status(status: RequestStatus) -> bool: + return status in ( + RequestStatus.WAITING_FOR_STRUCTURED_OUTPUT_GRAMMAR, + RequestStatus.WAITING_FOR_REMOTE_KVS, + RequestStatus.WAITING_FOR_STREAMING_REQ, + ) + + @staticmethod + def _holds_kv_blocks(request: Request) -> bool: + # Whether a request currently has kv blocks allocated. + # While always the case when status is WAITING_FOR_REMOTE_KVS + # or WAITING_FOR_STREAMING_REQ, there are such cases where it's + # WAITING or PREEMPTED. + return ( + request.num_computed_tokens > 0 + # num_computed_tokens is reset to 0 if kv transfer fails for requests in + # WAITING_FOR_REMOTE_KVS state, before their allocated blocks are freed. + or request.status == RequestStatus.WAITING_FOR_REMOTE_KVS + ) + + def _enqueue_waiting_request(self, request: Request) -> None: + if self._is_blocked_waiting_status(request.status): + self.deferred_waiting.add(request) + if self._holds_kv_blocks(request): + self.kv_holding_waiting.add_request(request) + else: + self.waiting.add_request(request) + + def _handle_stopped_request(self, request: Request) -> bool: + """Return True if finished (can be False for resumable requests).""" + if not request.resumable: + return True + + # Drop the finished turn's in-flight work and resume from the + # materialized frontier. + safe_frontier = request.num_computed_tokens - request.num_output_placeholders + assert safe_frontier >= 0 + request.num_computed_tokens = safe_frontier + request.spec_token_ids = [] + request.drop_stale_output = True + request.num_stale_output_tokens = request.num_in_flight_tokens + request.num_output_placeholders = 0 + + if request.streaming_queue: + update = request.streaming_queue.popleft() + if update is None: + # Streaming request finished. + return True + self._update_request_as_session(request, update) + else: + request.status = RequestStatus.WAITING_FOR_STREAMING_REQ + self.num_waiting_for_streaming_input += 1 + + self._enqueue_waiting_request(request) + return False + + def _update_request_with_output( + self, request: Request, new_token_ids: list[int], is_stale: bool = False + ) -> tuple[list[int], bool]: + # is_stale is only used by the AsyncScheduler override. + # Append generated tokens and check for stop. Note that if + # a request is still being prefilled, we expect the model runner + # to return empty token ids for the request. + stopped = False + for num_new, output_token_id in enumerate(new_token_ids, 1): + request.append_output_token_ids(output_token_id) + + # Check for stop and update request state. + # This must be called before we make the EngineCoreOutput. + stopped = check_stop(request, self.max_model_len) + if stopped: + del new_token_ids[num_new:] # Trim new tokens if needed. + break + return new_token_ids, stopped + + def _free_encoder_inputs(self, request: Request) -> None: + cached_encoder_input_ids = self.encoder_cache_manager.get_cached_input_ids( + request + ) + # OPTIMIZATION: Avoid list(set) if the set is empty. + if not cached_encoder_input_ids: + return + + # Defer the free by the drafter's look-ahead so an entry stays + # referenced until the drafter's read-ahead has also passed it, + # mirroring the shift the encoder scheduling path applies. + spec_lookahead = self.num_prefill_lookahead + + # Here, we use list(set) to avoid modifying the set while iterating + # over it. + for input_id in list(cached_encoder_input_ids): + mm_feature = request.mm_features[input_id] + start_pos = mm_feature.mm_position.offset + num_tokens = mm_feature.mm_position.length + if self.is_encoder_decoder and request.num_computed_tokens > 0: + # With Whisper, as soon as we've generated a single token, + # we know we're done with the encoder input. Cross Attention + # KVs have been calculated and cached already. + self._free_encoder_input(request, input_id) + elif ( + start_pos + num_tokens + spec_lookahead + <= request.num_computed_tokens - request.num_output_placeholders + ): + # Processed, stored in the decoder KV cache, and far enough past + # the placeholder range (plus the drafter's look-ahead) that no + # rejection or drafter gather can reference it. + self._free_encoder_input(request, input_id) + + def _free_encoder_input(self, request: Request, input_id: int) -> None: + self.encoder_cache_manager.free_encoder_input(request, input_id) + if self.ec_connector is not None: + self.ec_connector.update_state_after_free(request, input_id) + + def update_draft_token_ids(self, draft_token_ids: DraftTokenIds) -> None: + for req_id, spec_token_ids in zip( + draft_token_ids.req_ids, + draft_token_ids.draft_token_ids, + ): + request = self.requests.get(req_id) + if request is None or request.is_finished(): + # The request may have been finished. Skip. + continue + + if request.is_prefill_chunk: + # Ignore draft tokens for prefill chunks. + if request.spec_token_ids: + request.spec_token_ids = [] + continue + + # Add newly generated spec token ids to the request. + request.spec_token_ids = self.structured_output_manager.validate_tokens( + request, spec_token_ids + ) + + def update_draft_token_ids_in_output( + self, draft_token_ids: DraftTokenIds, scheduler_output: SchedulerOutput + ) -> None: + num_invalid_spec_tokens: dict[str, int] = {} + + sched_spec_tokens = scheduler_output.scheduled_spec_decode_tokens + for req_id, spec_token_ids in zip( + draft_token_ids.req_ids, + draft_token_ids.draft_token_ids, + ): + request = self.requests.get(req_id) + if request is None or request.is_finished(): + # The request may have been finished. Skip. + continue + + placeholder_spec_tokens = sched_spec_tokens.get(req_id) + if not placeholder_spec_tokens: + continue + + orig_num_spec_tokens = len(placeholder_spec_tokens) + # Trim drafts to scheduled number of spec tokens + # (needed for chunked prefill case for example). + del spec_token_ids[orig_num_spec_tokens:] + # Filter out spec tokens which do not adhere to the grammar. + spec_token_ids = self.structured_output_manager.validate_tokens( + request, spec_token_ids + ) + # Pad to original number of spec tokens. + num_invalid_tokens = orig_num_spec_tokens - len(spec_token_ids) + if num_invalid_tokens: + spec_token_ids.extend([-1] * num_invalid_tokens) + num_invalid_spec_tokens[req_id] = num_invalid_tokens + + sched_spec_tokens[req_id] = spec_token_ids + + scheduler_output.num_invalid_spec_tokens = num_invalid_spec_tokens + + def get_request_counts(self) -> tuple[int, int]: + """Returns (num_running_reqs, num_waiting_reqs).""" + return len(self.running), len(self.waiting) + len(self.kv_holding_waiting) + + def get_kv_cache_usage(self) -> float: + """Returns the fraction of the KV cache currently in use (0.0-1.0).""" + return self.kv_cache_manager.usage + + def add_request(self, request: Request) -> None: + existing = self.requests.get(request.request_id) + if existing is not None: + update = StreamingUpdate.from_request(request) + if existing.status != RequestStatus.WAITING_FOR_STREAMING_REQ: + assert existing.streaming_queue is not None, "duplicate request id" + # Queue next input chunk (or finished sentinel). + existing.streaming_queue.append(update) + elif update is not None: + # Commence next input chunk. + self._update_request_as_session(existing, update) + else: + # Streaming-input session finished. + self.finish_requests(request.request_id, RequestStatus.FINISHED_ABORTED) + else: + if request.resumable: + request.streaming_queue = deque() + self._enqueue_waiting_request(request) + self.requests[request.request_id] = request + if self.spec_decode_metrics_level != "none": + request.spec_decode_metrics = RequestSpecDecodeMetrics.new( + self.num_spec_tokens + ) + if self.connector is not None: + self.connector.on_new_request(request) + kv_transfer_params = request.kv_transfer_params + if kv_transfer_params and kv_transfer_params.get("do_remote_prefill"): + # Count P/D requests on arrival: connector matching is + # skipped while all run slots are taken. + self._set_kv_fetch_stage(request, KV_FETCH_WAITING_TO_START) + if self.log_stats: + request.record_event(EngineCoreEventType.QUEUED) + + def finish_requests( + self, request_ids: str | Iterable[str] | None, finished_status: RequestStatus + ) -> list[Request]: + """Handles the finish signal from outside the scheduler. + + For example, the API server can abort a request when the client + disconnects. + + If request_ids is None, all requests will be finished. + + Returns: + List of requests that were aborted. Will not include any that were + already finished. + + """ + assert RequestStatus.is_finished(finished_status) + if isinstance(request_ids, str): + request_ids = (request_ids,) + elif request_ids is not None: + request_ids = set(request_ids) + else: + request_ids = self.requests.keys() + + running_requests_to_remove = set() + waiting_requests_to_remove = [] + valid_requests = [] + + # First pass: collect requests to remove from queues + for req_id in request_ids: + request = self.requests.get(req_id) + if request is None or request.is_finished(): + # Invalid request ID. + continue + + valid_requests.append(request) + if request.status == RequestStatus.RUNNING: + running_requests_to_remove.add(request) + else: + if request.status == RequestStatus.WAITING_FOR_STREAMING_REQ: + self.num_waiting_for_streaming_input -= 1 + waiting_requests_to_remove.append(request) + + # Remove all requests from queues at once for better efficiency + if running_requests_to_remove: + self.running = remove_all(self.running, running_requests_to_remove) + if waiting_requests_to_remove: + self.waiting.remove_requests(waiting_requests_to_remove) + self.kv_holding_waiting.remove_requests(waiting_requests_to_remove) + self.deferred_waiting.difference_update(waiting_requests_to_remove) + + # Second pass: set status and free requests + for request in valid_requests: + delay_free_blocks = False + if request.status == RequestStatus.WAITING_FOR_REMOTE_KVS: + delay_free_blocks = ( + request.request_id not in self.finished_recving_kv_req_ids + ) + self.finished_recving_kv_req_ids.discard(request.request_id) + self.failed_recving_kv_req_ids.discard(request.request_id) + + request.status = finished_status + self._free_request(request, delay_free_blocks=delay_free_blocks) + + return valid_requests + + def _free_request( + self, request: Request, delay_free_blocks: bool = False + ) -> tuple[dict[str, Any] | None, dict[str, Any] | None]: + assert request.is_finished() + + if self.aux_output_connector is not None: + self.aux_output_connector.request_finished(request) + self._inflight_prefills.discard(request) + self._set_kv_fetch_stage(request, None) + connector_delay_free_blocks, kv_xfer_params = self._connector_finished(request) + + # EC Connector: mirror the KV hook. The contract requires firing + # before the encoder cache is freed so the connector can inspect + # per-request state (e.g. which mm_hashes it recorded during + # save_caches()) and emit ec_transfer_params for the response body. + ec_xfer_params: dict[str, Any] | None = None + if self.ec_connector is not None: + ec_delay_free, ec_xfer_params = self.ec_connector.request_finished(request) + connector_delay_free_blocks |= ec_delay_free + + self.encoder_cache_manager.free(request) + request_id = request.request_id + self.finished_req_ids.add(request_id) + if self.finished_req_ids_dict is not None: + self.finished_req_ids_dict[request.client_index].add(request_id) + + delay_free_blocks |= connector_delay_free_blocks + if not delay_free_blocks: + self._free_blocks(request) + + return kv_xfer_params, ec_xfer_params + + def _free_blocks(self, request: Request): + assert request.is_finished() + self._free_request_blocks(request) + del self.requests[request.request_id] + + @property + def pause_state(self) -> PauseState: + return self._pause_state + + def set_pause_state(self, pause_state: PauseState) -> None: + logger.info("setting pause state to %s", pause_state.name) + self._pause_state = pause_state + + def _request_blocks_can_be_freed(self, request: Request) -> bool: + # We must defer freeing blocks if an async kv connector may + # write to them immediately (not ordered with GPU stream). + return not self.defer_block_free or ( + request.last_sched_seq <= self.processed_step_seq + ) + + def _free_request_blocks(self, request: Request): + """Free the request's KV blocks, deferring the return to the block + pool when an in-flight GPU step may still write them. + """ + if self._request_blocks_can_be_freed(request): + self.kv_cache_manager.free(request) + return + blocks = self.kv_cache_manager.pop_blocks_for_free(request) + if blocks: + self.deferred_frees.append((self.sched_step_seq, blocks)) + + def _free_cow_retained_blocks( + self, blocks: list[KVCacheBlock], fence_seq: int + ) -> None: + """Release CoW copy retentions, deferring their return to the block + pool while the step that runs the copy may still be in flight. + """ + if not self.defer_block_free or fence_seq <= self.processed_step_seq: + self.kv_cache_manager.block_pool.free_blocks(blocks) + return + # The immediate path above frees these as-is, so enqueue them as-is: + # the drain no longer reverses. + self.deferred_frees.append((fence_seq, blocks)) + + def _drain_deferred_frees(self): + """Return deferred blocks whose fence step has completed. + + Fences are appended in near-monotonic order (a CoW retention fence + can lead request-free fences by one step), so stop at the first + pending one; any satisfied entry behind it is merely freed later. + """ + while self.deferred_frees: + fence, _ = self.deferred_frees[0] + if fence > self.processed_step_seq: + break + _, blocks = self.deferred_frees.popleft() + # Already in eviction order (reversed per KV cache group by + # `pop_blocks_for_free`), so free as-is. + self.kv_cache_manager.block_pool.free_blocks(blocks) + + def get_num_unfinished_requests(self) -> int: + if self._pause_state == PauseState.PAUSED_ALL: + return 0 + num_running, num_waiting = self.get_request_counts() + if self._pause_state == PauseState.PAUSED_NEW: + return num_running + num_waiting -= self.num_waiting_for_streaming_input + return num_waiting + num_running + + def has_finished_requests(self) -> bool: + if self.finished_req_ids: + return True + if self.connector is None: + return False + # Finished requests waiting on delayed connector cleanup remain in + # self.requests after they have been removed from scheduling queues. + num_running, num_waiting = self.get_request_counts() + return len(self.requests) > (num_running + num_waiting) + + def has_requests(self) -> bool: + # Override the interface default to also keep the engine alive while a + # connector still has pending push work (e.g. push-mode WRITE transfers + # in flight after all "live" requests have finished). Without this hook + # the engine would quiesce before the connector can drain completions. + # TODO: replace with a more general mechanism for connectors to keep + # the scheduler alive. + return ( + self.has_unfinished_requests() + or self.has_finished_requests() + or (self.connector is not None and self.connector.has_pending_push_work()) + or ( + self.ec_connector is not None + and self.ec_connector.has_pending_push_work() + ) + ) + + def reset_prefix_cache( + self, reset_running_requests: bool = False, reset_connector: bool = False + ) -> bool: + """Reset the KV prefix cache. + + If reset_running_requests is True, all the running requests will be + preempted and moved to the waiting queue. + Otherwise, this method will only reset the KV prefix cache when there + is no running requests taking KV cache. + """ + if reset_running_requests: + # For logging. + timestamp = time.monotonic() + # Invalidate all the current running requests KV's by pushing them to + # the waiting queue. In this case, we can reduce the ref count of all + # the kv blocks to 0 and thus we can make sure the reset is successful. + # Preempt in reverse order so the requests will be added back to the + # running queue in FIFO order. + while self.running: + request = self.running.pop() + self._preempt_request(request, timestamp, drop_stale_output=True) + + # Clear scheduled request ids cache. Since we are forcing preemption + # + resumption in the same step, we must act as if these requests were + # not scheduled in the prior step. They will be flushed from the + # persistent batch in the model runner. + self.prev_step_scheduled_req_ids.clear() + + reset_successful = self.kv_cache_manager.reset_prefix_cache() + if reset_running_requests and not reset_successful: + raise RuntimeError( + "Failed to reset KV cache even when all the running requests are " + "preempted and moved to the waiting queue. This is likely due to " + "the presence of running requests waiting for remote KV transfer, " + "which is not supported yet." + ) + + if reset_connector: + reset_successful = self.reset_connector_cache() and reset_successful + + if reset_successful and self.aux_output_connector is not None: + self.aux_output_connector.reset() + + return reset_successful + + def reset_connector_cache(self) -> bool: + if self.connector is None: + # No connector attached -> nothing to reset, treat as success so + # callers that unconditionally request a connector reset (e.g. as + # part of a cache-clearing cascade after a weight update) don't + # see reset_prefix_cache() flip to False purely because they + # didn't configure a connector. + logger.debug( + "reset_connector requested but no KV connector is configured; " + "treating as no-op success." + ) + return True + + if self.connector.reset_cache() is False: + return False + + if self.log_stats: + assert self.connector_prefix_cache_stats is not None + self.connector_prefix_cache_stats.reset = True + + return True + + def reset_encoder_cache(self) -> None: + """Reset the encoder cache to invalidate all cached encoder outputs. + + This should be called when model weights are updated to ensure + stale vision embeddings are not reused. + """ + self.encoder_cache_manager.reset() + + def make_stats( + self, + spec_decoding_stats: SpecDecodingStats | None = None, + kv_connector_stats: KVConnectorStats | None = None, + cudagraph_stats: CUDAGraphStat | None = None, + perf_stats: PerfStats | None = None, + ec_connector_stats: ECConnectorStats | None = None, + ) -> SchedulerStats | None: + if not self.log_stats: + return None + prefix_cache_stats = self.kv_cache_manager.make_prefix_cache_stats() + assert prefix_cache_stats is not None + connector_prefix_cache_stats: PrefixCacheStats | None = None + if self.connector_prefix_cache_stats is not None: + connector_prefix_cache_stats = self.connector_prefix_cache_stats + self.connector_prefix_cache_stats = PrefixCacheStats() + eviction_events = ( + self.kv_metrics_collector.drain_events() + if self.kv_metrics_collector is not None + else [] + ) + spec_stats = spec_decoding_stats + connector_stats_payload = ( + kv_connector_stats.to_dict() if kv_connector_stats else None + ) + ec_connector_stats_payload = ( + ec_connector_stats.data if ec_connector_stats else None + ) + num_running, num_waiting = self.get_request_counts() + num_deferred = len(self.deferred_waiting) + return SchedulerStats( + num_running_reqs=num_running, + num_waiting_reqs=num_waiting - num_deferred, + num_skipped_waiting_reqs=num_deferred, + num_kv_fetch_reqs_by_stage=self._kv_fetch_counts.copy(), + kv_cache_usage=self.kv_cache_manager.usage, + prefix_cache_stats=prefix_cache_stats, + connector_prefix_cache_stats=connector_prefix_cache_stats, + kv_cache_eviction_events=eviction_events, + spec_decoding_stats=spec_stats, + kv_connector_stats=connector_stats_payload, + cudagraph_stats=cudagraph_stats, + perf_stats=perf_stats, + ec_connector_stats=ec_connector_stats_payload, + ) + + def make_spec_decoding_stats( + self, + spec_decoding_stats: SpecDecodingStats | None, + num_draft_tokens: int, + num_accepted_tokens: int, + num_invalid_spec_tokens: dict[str, int] | None, + request_id: str, + ) -> SpecDecodingStats | None: + if not self.log_stats or not num_draft_tokens: + return None + if spec_decoding_stats is None: + spec_decoding_stats = SpecDecodingStats.new(self.num_spec_tokens) + if num_invalid_spec_tokens: + num_draft_tokens -= num_invalid_spec_tokens.get(request_id, 0) + spec_decoding_stats.observe_draft( + num_draft_tokens=num_draft_tokens, num_accepted_tokens=num_accepted_tokens + ) + return spec_decoding_stats + + def shutdown(self) -> None: + logger.debug_once("[shutdown] Scheduler: start") + if self.kv_event_publisher: + self.kv_event_publisher.shutdown() + if self.connector is not None: + self.connector.shutdown() + + if self.ec_connector is not None: + self.ec_connector.shutdown() + + logger.debug_once("[shutdown] Scheduler: complete") + + ######################################################################## + # KV Connector Related Methods + ######################################################################## + + def get_kv_connector(self) -> KVConnectorBase_V1 | None: + return self.connector + + def get_ec_connector(self) -> ECConnectorBase | None: + return self.ec_connector + + def get_kv_event_publisher_config(self) -> KVEventsConfig | None: + return self.kv_event_publisher.get_publisher_config() + + def _connector_finished( + self, request: Request + ) -> tuple[bool, dict[str, Any] | None]: + """Invoke the KV connector request_finished() method if applicable. + + Returns optional kv transfer parameters to be included with the + request outputs. + """ + if self.connector is None: + return False, None + + finished_partial_tails: list[tuple[int, int, int]] = [] + kv_transfer_config = self.vllm_config.kv_transfer_config + if kv_transfer_config is not None and kv_transfer_config.is_kv_producer: + finished_partial_tails = ( + self.kv_cache_manager.finalize_partial_tail_offloads(request) + ) + + # Free any out-of-window prefix blocks before we hand the block table to + # the connector, on the processed-token basis (see `allocate_slots`). + self.kv_cache_manager.remove_skipped_blocks( + request_id=request.request_id, + processed_computed_tokens=max( + 0, request.num_computed_tokens - request.num_in_flight_tokens + ), + num_prompt_tokens=request.num_prompt_tokens, + ) + + block_ids = self.kv_cache_manager.get_block_ids_for_computed_tokens( + request_id=request.request_id, + num_computed_tokens=request.num_computed_tokens, + ) + partial_tail_delay = False + if finished_partial_tails: + partial_tail_delay = self.connector.register_finished_partial_tail( + request, + block_ids, + finished_partial_tails, + ) + + if not isinstance(self.connector, SupportsHMA): + # NOTE(Kuntai): We should deprecate this code path after we enforce + # all connectors to support HMA. + # Hybrid memory allocator should be already turned off for this + # code path, but let's double-check here. + assert len(self.kv_cache_config.kv_cache_groups) == 1 + delay_free, kv_xfer_params = self.connector.request_finished( + request, block_ids[0] + ) + else: + delay_free, kv_xfer_params = self.connector.request_finished_all_groups( + request, block_ids + ) + return delay_free or partial_tail_delay, kv_xfer_params + + def _request_remaining_blocks(self, request: Request) -> int: + """Blocks `request` still needs to hold its full sequence and be promoted.""" + full_num_tokens = min(request.num_tokens, self.max_model_len) + num_blocks = self.kv_cache_manager.coordinator.get_num_blocks_to_allocate( + request_id=request.request_id, + num_tokens=full_num_tokens, + new_computed_blocks=self.kv_cache_manager.empty_kv_cache_blocks.blocks, + num_encoder_tokens=0, + total_computed_tokens=request.num_computed_tokens, + num_local_computed_tokens=request.num_computed_tokens, + num_tokens_main_model=full_num_tokens, + apply_admission_cap=True, + prefill_end=max(request.num_prompt_tokens, request.num_tokens - 1), + ) + return num_blocks + self._spec_decode_step_blocks() + + def _spec_decode_step_blocks(self) -> int: + """Number of blocks for the extra KV slots a spec decode step needs. + + When using async kv load, scheduler must reserve enough blocks for + full sequence + the spec decode step, otherwise request cannot be + able to run after the async_load if we are out of kv blocks. + """ + if not self.num_spec_tokens: + return 0 + return cdiv( + 1 + self.num_spec_tokens + self.num_lookahead_tokens, self.block_size + ) + + def _set_kv_fetch_stage(self, request: Request, stage: str | None) -> None: + kv_fetch_stages = self._kv_fetch_stages + if kv_fetch_stages is not None: + prev_stage = kv_fetch_stages.pop(request, None) + if prev_stage is not None: + self._kv_fetch_counts[prev_stage] -= 1 + if stage is not None: + kv_fetch_stages[request] = stage + self._kv_fetch_counts[stage] += 1 + + def _inflight_prefill_reserved_blocks(self) -> int: + """Num blocks in-flight prefills still need to finish (their reservation).""" + return sum( + self._request_remaining_blocks(req) for req in self._inflight_prefills + ) + + def _mark_prefix_replay(self, request: Request, num_hit_tokens: int) -> int: + """Record where a prefix hit's replay starts and return the number of + hit tokens to recompute; the caller rewinds the computed count by it. + + The replayed tokens are the hit's last window (see + ``Request.replay_start``): the worker rebuilds their sliding-window KV + and leaves their cached KV alone. + """ + if not self.prefix_replay_tokens: + return 0 + num_replay_tokens = min(self.prefix_replay_tokens, num_hit_tokens) + request.replay_start = num_hit_tokens - num_replay_tokens + return num_replay_tokens + + def _update_waiting_for_remote_kv(self, request: Request) -> None: + """KV Connector: update request state after async recv is finished. + + When the kv transfer is ready, we cache the blocks + and the request state will be moved back to WAITING from + WAITING_FOR_REMOTE_KV. + """ + assert self.connector is not None + + if request.request_id in self.failed_recving_kv_req_ids: + # Request had KV load failures; num_computed_tokens was already + # updated in _update_requests_with_invalid_blocks + if request.num_computed_tokens: + # Cache any valid computed tokens. + self.kv_cache_manager.cache_blocks(request, request.num_computed_tokens) + if self.needs_kv_cache_zeroing: + # The failed load left the blocks beyond the valid + # prefix unwritten and their zeroing was skipped; zero + # them before they are recomputed locally. + self.kv_cache_manager.record_blocks_for_zeroing( + request.request_id, request.num_computed_tokens + ) + else: + # No valid computed tokens, release allocated blocks. + # There may be a local cache hit on retry. + # (Freed blocks are re-recorded for zeroing when + # reallocated, so the skipped blocks need no handling.) + self.kv_cache_manager.free(request) + + self.failed_recving_kv_req_ids.remove(request.request_id) + else: + # Now that the blocks are ready, actually cache them. + # This will cache the blocks iff caching is enabled. + self.kv_cache_manager.cache_blocks(request, request.num_computed_tokens) + + # SWA bounded replay recomputes the tail of the hit, which covers the + # last token; otherwise a full prompt hit re-computes that token so the + # next one can be sampled. + num_replay_tokens = self._mark_prefix_replay( + request, request.num_computed_tokens + ) + if num_replay_tokens > 0: + request.num_computed_tokens -= num_replay_tokens + elif request.num_computed_tokens == request.num_tokens: + request.num_computed_tokens = request.num_tokens - 1 + self.finished_recving_kv_req_ids.remove(request.request_id) + + def _handle_blocked_waiting_request(self, request: Request) -> bool: + """Returns True if scheduling of this request should proceed, False if + it should be skipped for this step.""" + if request.status == RequestStatus.WAITING_FOR_REMOTE_KVS: + # Check whether async loading requests are finished. + request_id = request.request_id + if request_id not in self.finished_recving_kv_req_ids: + logger.debug("%s is still in WAITING_FOR_REMOTE_KVS state.", request_id) + return False + # Promote request to waiting or preempted. + self._update_waiting_for_remote_kv(request) + if request.num_preemptions: + request.status = RequestStatus.PREEMPTED + else: + request.status = RequestStatus.WAITING + + elif request.status == RequestStatus.WAITING_FOR_STRUCTURED_OUTPUT_GRAMMAR: + # Check whether grammar compilation has finished. + struct_output_req = request.structured_output_request + if not struct_output_req or struct_output_req.grammar is None: + return False + if isinstance(struct_output_req.grammar, Exception): + self.grammar_compile_error_reqs.add(request.request_id) + return False + # Grammar is ready, promote to waiting state. + request.status = RequestStatus.WAITING + + elif request.status == RequestStatus.WAITING_FOR_STREAMING_REQ: + # Streaming input request waiting for more input. + assert not request.streaming_queue + return False + + return True + + def _update_from_kv_xfer_finished(self, kv_connector_output: KVConnectorOutput): + """KV Connector: update the scheduler state based on the output. + + The Worker side connectors add finished_recving and + finished_sending reqs to the output. + * if finished_sending: free the blocks + # if finished_recving: add to state so we can + schedule the request during the next step. + """ + if self.connector is not None: + self.connector.update_connector_output(kv_connector_output) + + # KV Connector:: update recv and send status from last step. + for req_id in kv_connector_output.finished_recving or (): + logger.debug("Finished recving KV transfer for request %s", req_id) + assert req_id in self.requests + req = self.requests[req_id] + if req.status == RequestStatus.WAITING_FOR_REMOTE_KVS: + self.finished_recving_kv_req_ids.add(req_id) + self._set_kv_fetch_stage(req, KV_FETCH_COMPLETED_WAITING) + else: + assert RequestStatus.is_finished(req.status) + self._free_blocks(self.requests[req_id]) + for req_id in kv_connector_output.finished_sending or (): + logger.debug("Finished sending KV transfer for request %s", req_id) + assert req_id in self.requests + self._free_blocks(self.requests[req_id]) + + def _update_requests_with_invalid_blocks( + self, + requests: Iterable[Request], + invalid_block_ids: set[int], + num_scheduled_tokens: dict[str, int], + evict_blocks: bool = True, + ) -> tuple[set[str], int, set[int]]: + """Identify and update requests affected by invalid KV cache blocks. + + This method scans the given requests, detects those with invalid blocks + and adjusts their `num_computed_tokens` to the longest valid prefix. + For observability, it also accumulates the total number of tokens that + will need to be recomputed across all affected requests. + + Args: + requests: The set of requests to scan for invalid blocks. + invalid_block_ids: IDs of invalid blocks. + num_scheduled_tokens: req_id -> number of scheduled tokens. + evict_blocks: Whether to collect blocks for eviction (False for + async requests which aren't cached yet). + + Returns: + tuple: + - affected_req_ids (set[str]): IDs of requests impacted by + invalid blocks. + - total_affected_tokens (int): Total number of tokens that must + be recomputed across all affected requests. + - blocks_to_evict (set[int]): Block IDs to evict from cache, + including invalid blocks and downstream dependent blocks. + + """ + affected_req_ids: set[str] = set() + total_affected_tokens = 0 + blocks_to_evict: set[int] = set() + # If a block is invalid and shared by multiple requests in the batch, + # these requests must be rescheduled, but only the first will recompute + # it. This set tracks blocks already marked for recomputation. + marked_invalid_block_ids: set[int] = set() + for request in requests: + is_affected = False + marked_invalid_block = False + req_id = request.request_id + # One block-id list per KV cache group; single-group layouts are + # enforced by _handle_invalid_blocks. + (req_block_ids,) = self.kv_cache_manager.get_block_ids(req_id) + # We iterate only over blocks that may contain externally computed + # tokens + req_num_computed_tokens = ( + request.num_computed_tokens - num_scheduled_tokens.get(req_id, 0) + ) + + req_num_computed_blocks = ( + req_num_computed_tokens + self.block_size - 1 + ) // self.block_size + for idx, block_id in zip(range(req_num_computed_blocks), req_block_ids): + if block_id not in invalid_block_ids: + continue + + is_affected = True + + if block_id in marked_invalid_block_ids: + # This invalid block is shared with a previous request + # and was already marked for recomputation. + # This means this request can still consider this block + # as computed when rescheduled. + # Currently this only applies to sync loading; Async + # loading does not yet support block sharing + continue + + marked_invalid_block_ids.add(block_id) + + if marked_invalid_block: + # This request has already marked an invalid block for + # recomputation and updated its num_computed_tokens. + continue + + marked_invalid_block = True + # Truncate the computed tokens at the first failed block + request.num_computed_tokens = idx * self.block_size + num_affected_tokens = ( + req_num_computed_tokens - request.num_computed_tokens + ) + total_affected_tokens += num_affected_tokens + + # collect invalid block and all downstream dependent blocks + if evict_blocks: + blocks_to_evict.update(req_block_ids[idx:]) + + if is_affected: + if not marked_invalid_block: + # All invalid blocks of this request are shared with + # previous requests and will be recomputed by them. + # Revert to considering only cached tokens as computed. + # Currently this only applies to sync loading; Async + # loading does not yet support block sharing + total_affected_tokens += ( + request.num_computed_tokens - req_num_computed_tokens + ) + request.num_computed_tokens = req_num_computed_tokens + + affected_req_ids.add(request.request_id) + + return affected_req_ids, total_affected_tokens, blocks_to_evict + + def _handle_failed_recving(self, failed_req_ids: set[str]) -> set[str]: + """Fail closed for layouts whose block IDs are not globally unique.""" + affected_req_ids: set[str] = set() + for req_id in failed_req_ids: + request = self.requests.get(req_id) + if ( + request is not None + and request.status == RequestStatus.WAITING_FOR_REMOTE_KVS + ): + affected_req_ids.add(req_id) + if self.recompute_kv_load_failures: + request.num_computed_tokens = 0 + self.failed_recving_kv_req_ids.add(req_id) + return set() if self.recompute_kv_load_failures else affected_req_ids + + def _handle_invalid_blocks( + self, invalid_block_ids: set[int], num_scheduled_tokens: dict[str, int] + ) -> set[str]: + """Handle requests affected by invalid KV cache blocks. + + Returns: + Set of affected request IDs to skip in update_from_output main loop. + + """ + if len(self.kv_cache_config.kv_cache_groups) > 1: + # Block IDs are only unique within a group, so a flat set of + # invalid block IDs cannot be mapped back to requests. + raise RuntimeError( + "A KV connector reported block-level load failures " + "(invalid_block_ids) which is not supported for models " + "with multiple KV cache groups. " + "Connectors must report failed requests via " + "KVConnectorTransferResults.failed_recving instead." + ) + + should_fail = not self.recompute_kv_load_failures + + # handle async KV loads (not cached yet, evict_blocks=False) + async_load_reqs = ( + req + for req in self.kv_holding_waiting + if req.status == RequestStatus.WAITING_FOR_REMOTE_KVS + ) + async_failed_req_ids, num_failed_tokens, _ = ( + self._update_requests_with_invalid_blocks( + async_load_reqs, + invalid_block_ids, + num_scheduled_tokens, + evict_blocks=False, + ) + ) + + total_failed_requests = len(async_failed_req_ids) + total_failed_tokens = num_failed_tokens + + # handle sync loads (may be cached, collect blocks for eviction) + sync_failed_req_ids, num_failed_tokens, sync_blocks_to_evict = ( + self._update_requests_with_invalid_blocks( + self.running, invalid_block_ids, num_scheduled_tokens, evict_blocks=True + ) + ) + + total_failed_requests += len(sync_failed_req_ids) + total_failed_tokens += num_failed_tokens + + if not total_failed_requests: + return set() + + # evict invalid blocks and downstream dependent blocks from cache + # only when not using recompute policy (where blocks will be recomputed + # and reused by other requests sharing them) + if sync_blocks_to_evict and not self.recompute_kv_load_failures: + self.kv_cache_manager.evict_blocks(sync_blocks_to_evict) + + if should_fail: + all_failed_req_ids = async_failed_req_ids | sync_failed_req_ids + logger.error( + "Failing %d request(s) due to KV load failure " + "(failure_policy=fail, %d tokens affected). Request IDs: %s", + total_failed_requests, + total_failed_tokens, + all_failed_req_ids, + ) + return all_failed_req_ids + + logger.warning( + "Recovered from KV load failure: " + "%d request(s) rescheduled (%d tokens affected).", + total_failed_requests, + total_failed_tokens, + ) + + # Mark async requests with KV load failures for retry once loading completes + self.failed_recving_kv_req_ids |= async_failed_req_ids + # Return sync affected IDs to skip in update_from_output + return sync_failed_req_ids diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index cf3c9c1b61..096093036a 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -3,7 +3,7 @@ import itertools from abc import ABC, abstractmethod from collections import defaultdict -from collections.abc import Sequence +from collections.abc import Callable, Sequence from typing import TYPE_CHECKING, ClassVar from vllm import envs @@ -1705,6 +1705,12 @@ class MambaManager(SingleTypeKVCacheManager): # therefore hold their hashed endpoint here too: the moved tail for # a running request, and the hit source a new request copies from. self._retained_checkpoints: dict[str, tuple[int, KVCacheBlock]] = {} + # Callbacks (source_block, cow_block) run after a running + # request moved its live tail's hash to the CoW copy, for refs held + # outside this manager (the KVCacheManager's grace and waiter pins). + self.moved_hash_listeners: list[ + Callable[[KVCacheBlock, KVCacheBlock], None] + ] = [] if self.mamba_cache_mode == "align" and self.enable_caching: # The A/B evidence line: boot logs must show which arm ran. logger.info_once( @@ -1801,6 +1807,8 @@ class MambaManager(SingleTypeKVCacheManager): self.block_pool.touch([cow_block]) self._retained_checkpoints[request_id] = (position, cow_block) self.block_pool.free_blocks([source_block]) + for listener in self.moved_hash_listeners: + listener(source_block, cow_block) def pop_retained_checkpoints(self, request_id: str) -> list[KVCacheBlock]: if self.mamba_cache_mode != "align":