From bf3e4af6b1ecd86f1bd434069f247f04af6798b3 Mon Sep 17 00:00:00 2001 From: MOVZX Date: Tue, 6 Oct 2026 02:49:52 -0400 Subject: [PATCH] [KV Cache] Hold a finished turn's prefix until its next turn is admitted With a prefix unit finer than the block, a hybrid follow-up resumes from a single GDN state at the previous prompt's tail, and that key sits in the top attention block. Losing any one block of that prefix is a total miss. The waiting-request pins only protect a hit that is still whole when the follow-up is first looked up. Tracing the remaining misses under a long-context multi-turn agent load shows they lost their top attention block before that, in the gap between the previous turn finishing and the next turn arriving. In that gap nothing referenced the prefix. LRU then consumed it ahead of dead blocks freed later (other conversations' handed-over CoW sources and superseded states). A finished turn (stop, length, repetition, or an abort after decoding) now keeps one extra ref on each of its hashed blocks, including its GDN holds, as a grace pin. The ref is taken on both free paths, the immediate free and the deferred pop_blocks_for_free behind an in-flight step, so the prefix never passes through the free queue. A grace is released when: - the next turn is admitted on it. That is the grace holding the deepest hit block of every group, with a hit that reaches its prompt to within two prefix units, so a new turn sharing only a system prompt releases nothing. The table takes the blocks over, with no double ref. - anyone needs blocks. Graces yield before waiter pins, oldest first and one per failed allocation, so they never cause a preemption. - the bound (16) is exceeded. The oldest grace goes. - reset_prefix_cache runs. A head waiter whose admission fails pins its own hit first. It then yields graces and other pins one at a time, and only when everything yieldable covers its shortfall. Otherwise yielding would drop other conversations' prefixes and still not admit it. The shortfall is measured after the pin: the pin takes the proof block past the hit (and any hit block the lookup did not return) out of the free queue, so the shortfall of the failed attempt can be short by those blocks and the guard would yield another conversation's grace for nothing. Its own pinned blocks outside the hit it is admitted on, the proof block past it, go back first: once the hit is taken over they protect nothing, and they cost nobody else anything. At an exact fit they are the blocks it needs, and nothing else ever yields a head waiter's own pin, so without this a request that needs the whole pool waited forever. While any pin or grace is held the scheduler reports no common prefix to cascade attention. Those refs can make a block's count equal the number of running requests when only some of them share it, and cascade attention would then read it for all of them. Graces and waiter pins now follow a producer's GDN hash into its CoW copy, as holds already do. Otherwise they would keep a second ref on a running state block that its owner must own exclusively. Graces apply only to fine-grained hybrid hits (coordinator.enable_partial_hash_hits) and are off with a KV connector; a block-sized unit (the 4096 lane) is unchanged. When graces are active the scheduler logs one startup line saying so, with the bound, the hash unit and the block size. Tests: the arrival-gap reproduction fails on the previous commit (hit 0 instead of 168, both free paths). It passes here. Further tests cover the handover refcounts, yield order and bound, finish statuses, a running request at an exactly full pool, yieldable accounting, moved-hash following, the pin's extras contract, and six runs on the real scheduler: among them a head waiter whose pin raises its shortfall (it gives its proof block back and the one grace that covers the rest yields), a request that needs the whole pool (admitted on its own proof block; it waited forever before), and the common-prefix count under a grace (none reported; the count returns once the pins are gone). Co-Authored-By: Hinata Signed-off-by: MOVZX --- vllm/v1/core/kv_cache_manager.py | 302 +- vllm/v1/core/sched/scheduler.py | 99 +- vllm/v1/core/sched/scheduler.py.orig | 3551 ++++++++++++++++++ vllm/v1/core/single_type_kv_cache_manager.py | 10 +- 4 files changed, 3943 insertions(+), 19 deletions(-) create mode 100644 vllm/v1/core/sched/scheduler.py.orig 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":