iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity
Files
e1a5086256a00c47a968aac6a583c25d040100a7
vLLM/vllm/v1/core/sched/scheduler.py
T
Woosuk KwonandClaude Opus 5.5 e1a5086256 [Attention] Remove cascade attention (#61126) 
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-12 10:41:07 -07:00

3436 lines
158 KiB
Python
Raw Blame History

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import itertools
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.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 (
KVCacheConfig,
MambaSpec,
get_mamba_prefill_checkpoint_position,
is_mamba_prefill_checkpoint_valid,
)
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.utils import record_function_or_nullcontext
logger = init_logger(__name__)
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
)
# Pooling and diffusion models do not append a sampled token after
# prefill, so they should not reserve one context slot while scheduling.
self.num_sampled_tokens_per_step = (
0
if (
vllm_config.model_config.runner_type == "pooling"
or vllm_config.model_config.is_diffusion
)
else 1
)
# 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
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.
self.prefix_replay_group_ids = tuple(
group_id
for group_id, group in enumerate(kv_cache_config.kv_cache_groups)
if group.kv_cache_spec.prefix_replay_tokens > 0
)
replay_windows = {
kv_cache_config.kv_cache_groups[group_id].kv_cache_spec.prefix_replay_tokens
for group_id in self.prefix_replay_group_ids
}
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
),
)
# Mamba managers already account for their speculative state blocks.
num_spec_decode_tokens = (
1 + self.num_spec_tokens + self.num_lookahead_tokens
if self.num_spec_tokens
else 0
)
self._spec_decode_step_blocks = sum(
cdiv(num_spec_decode_tokens, manager.block_size)
for manager in self.kv_cache_manager.coordinator.single_type_managers
if not isinstance(manager.kv_cache_spec, MambaSpec)
)
# 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)
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
)
# Fine-grained hits reuse the prompt's hash-aligned Mamba checkpoint.
# Split there only if the backend cannot save it within the final chunk.
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.
ckpt_block_size = (
self.hash_block_size if self.mamba_partial_cache_hit else self.block_size
)
last_cache_position = get_mamba_prefill_checkpoint_position(
request.num_tokens,
block_size,
self.use_eagle_block_drop,
)
end = start + num_new_tokens
checkpoint_position = get_mamba_prefill_checkpoint_position(
prefill_end,
ckpt_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 = (
get_mamba_prefill_checkpoint_position(
request.num_prompt_tokens,
self.hash_block_size,
self.use_eagle_block_drop,
)
if self.mamba_partial_cache_hit and not use_internal_checkpoint
else 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,
# Stop at the prompt checkpoint unless the backend saves it internally.
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 _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] = []
duplicate_encoder_inputs: 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,
duplicate_encoder_inputs,
) = 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
# 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.
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
encoder_compute_budget = new_encoder_compute_budget
if encoder_inputs_to_schedule or external_load_encoder_input:
self._allocate_encoder_inputs(
request,
(encoder_inputs_to_schedule or []) + external_load_encoder_input,
duplicate_encoder_inputs,
)
# 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.PAUSED_ALL:
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, and are
# the only ones admitted under PAUSED_NEW (e.g. async KV loads).
request_queue = self.kv_holding_waiting
if not request_queue and self._pause_state == PauseState.UNPAUSED:
request_queue = 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 and not (
ext_tokens and self._load_restores_replay_window(request)
):
# DSV41 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. A load
# that carries the window replays nothing.
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
# DSV41 SWA bounded replay: a KV load that carries the
# window is taken as is; any other hit no longer than the
# replayed window would be recomputed in full and save
# nothing.
window_loaded = bool(
num_external_computed_tokens
) and self._load_restores_replay_window(request)
if (
not window_loaded
and 0 < num_computed_tokens <= self.prefix_replay_tokens
):
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 = []
duplicate_encoder_inputs = []
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, window_loaded
)
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,
duplicate_encoder_inputs,
) = 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:
# 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
):
preempted = request.num_preemptions > 0
self.connector_prefix_cache_stats.record(
num_tokens=connector_prefix_cache_queries,
num_hits=connector_prefix_cache_hits,
preempted=preempted,
# The connector knows the tier only after the load
# plan above exists.
hits_by_source=(
self.connector.get_external_cache_hit_sources(
request, connector_prefix_cache_hits
)
if connector_prefix_cache_hits and not preempted
else None
),
)
# 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)
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
encoder_compute_budget = new_encoder_compute_budget
if encoder_inputs_to_schedule or external_load_encoder_input:
self._allocate_encoder_inputs(
request,
(encoder_inputs_to_schedule or [])
+ external_load_encoder_input,
duplicate_encoder_inputs,
)
# 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)
# 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)
# 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,
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(),
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(),
)
# 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_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
return new_block_ids_to_zero or None
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:]
del session.block_hashes[num_computed_tokens // self.hash_block_size :]
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 _allocate_encoder_inputs(
self,
request: Request,
input_ids: list[int],
duplicate_input_ids: list[int],
) -> None:
for input_id in input_ids:
self.encoder_cache_manager.allocate(request, input_id)
if self.ec_connector is not None:
self.ec_connector.update_state_after_alloc(request, input_id)
for input_id in duplicate_input_ids:
self.encoder_cache_manager.check_and_update_cache(request, input_id)
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], 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 = []
duplicate_encoder_inputs = []
# 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.
duplicate_encoder_inputs.append(i)
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,
duplicate_encoder_inputs,
)
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
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
self.deferred_frees.append((fence_seq, blocks[::-1]))
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()
# Free in reverse order so that the tail blocks are evicted first.
self.kv_cache_manager.block_pool.free_blocks(reversed(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:
# Requests holding KV blocks still drain; queued ones stay queued.
num_waiting = len(self.kv_holding_waiting)
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_tokens = full_num_tokens
if self.num_spec_tokens:
num_tokens += 1 + self.num_spec_tokens + self.num_lookahead_tokens
return self.kv_cache_manager.coordinator.get_num_blocks_to_allocate(
request_id=request.request_id,
num_tokens=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),
)
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 _load_restores_replay_window(self, request: Request) -> bool:
"""Whether the KV connector's load for ``request`` restores every DSV41
SWA bounded-replay group, as a P/D transfer of the request's own blocks
does; a prefix-cache store holds no window."""
assert self.connector is not None
return bool(self.prefix_replay_group_ids) and set(
self.prefix_replay_group_ids
).issubset(self.connector.get_loaded_kv_cache_group_ids(request))
def _mark_prefix_replay(
self, request: Request, num_hit_tokens: int, window_loaded: bool
) -> 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. A hit whose KV load carries the
window replays nothing.
"""
if not self.prefix_replay_tokens:
return 0
if window_loaded:
request.replay_start = 0
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
load_failed = request.request_id in self.failed_recving_kv_req_ids
if load_failed:
# 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)
# DSV41 SWA bounded replay recomputes the tail of a hit loaded without
# its window, or of a failed load, which may have lost it; the replay
# 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,
window_loaded=not load_failed
and self._load_restores_replay_window(request),
)
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
Reference in New Issue View Git Blame Copy Permalink