# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import itertools from collections.abc import Sequence from dataclasses import dataclass from typing import Literal, overload from vllm.distributed.kv_events import BlockStored, KVCacheEvent from vllm.logger import init_logger from vllm.utils.math_utils import cdiv from vllm.v1.core.kv_cache_coordinator import get_kv_cache_coordinator from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector from vllm.v1.core.kv_cache_utils import KVCacheBlock, KVCacheBlockCopy from vllm.v1.kv_cache_interface import ( AttentionSpec, CrossAttentionSpec, EncoderOnlyAttentionSpec, KVCacheConfig, get_kv_cache_spec_kind, get_kv_cache_spec_sliding_window, ) from vllm.v1.metrics.stats import PrefixCacheStats from vllm.v1.request import Request, RequestStatus logger = init_logger(__name__) @dataclass class KVCacheBlocks: """ The allocation result of KVCacheManager, work as the interface between Scheduler and KVCacheManager, to hide KVCacheManager's internal data structure from the Scheduler. """ blocks: tuple[Sequence[KVCacheBlock], ...] """ `blocks[i][j]` refers to the i-th kv_cache_group and the j-th block of tokens.We don't use block of tokens as the outer dimension because it assumes all kv_cache_groups have the same number of blocks, which is true for now but will be broken if we want to give different block_size to different kv_cache_groups in the future. Each single type KVCacheBlocks could be represented as: - list[KVCacheBlock] for more than one KVCacheBlock - an empty tuple for requests without KVCacheBlock (a precomputed KVCacheBlocks is in KVCacheManager to avoid GC overhead) """ def __add__(self, other: "KVCacheBlocks") -> "KVCacheBlocks": """Adds two KVCacheBlocks instances.""" return KVCacheBlocks( tuple( list(itertools.chain(blk1, blk2)) for blk1, blk2 in zip(self.blocks, other.blocks) ) ) @overload def get_block_ids( self, allow_none: Literal[False] = False, ) -> tuple[list[int], ...]: ... @overload def get_block_ids( self, allow_none: Literal[True] = True, ) -> tuple[list[int], ...] | None: ... def get_block_ids( self, allow_none: bool = False, ) -> tuple[list[int], ...] | None: """ Converts the KVCacheBlocks instance to block_ids. Returns: tuple[list[int], ...]: A tuple of lists where: - the outer tuple corresponds to KV cache groups - each inner list contains the block_ids of the blocks in that group """ if allow_none and all(len(group) == 0 for group in self.blocks): return None return tuple([blk.block_id for blk in group] for group in self.blocks) def get_unhashed_block_ids(self) -> list[int]: """Get block_ids of unhashed blocks from KVCacheBlocks instance.""" assert len(self.blocks) == 1, "Only one group is supported" return [block.block_id for block in self.blocks[0] if block.block_hash is None] def get_unhashed_block_ids_all_groups(self) -> list[list[int]]: """Get block_ids of unhashed blocks from KVCacheBlocks instance.""" # Skip padding blocks. return [ [ block.block_id for block in group if block.block_hash is None and not block.is_null ] for group in self.blocks ] def new_empty(self) -> "KVCacheBlocks": """ Creates a new KVCacheBlocks instance with no blocks. """ return KVCacheBlocks(tuple(() for _ in range(len(self.blocks)))) class KVCacheManager: def __init__( self, kv_cache_config: KVCacheConfig, max_model_len: int, scheduler_block_size: int, hash_block_size: int, max_in_flight_tokens: int | None = None, enable_caching: bool = True, use_eagle: bool = False, log_stats: bool = False, enable_kv_cache_events: bool = False, dcp_world_size: int = 1, pcp_world_size: int = 1, metrics_collector: KVCacheMetricsCollector | None = None, watermark: float = 0.0, ) -> None: self.max_model_len = max_model_len # When unset, fall back to `max_model_len` so the recycling-aware cap # collapses to the prior (uncapped) admission behavior. The scheduler # always supplies the real value at runtime. if max_in_flight_tokens is None: max_in_flight_tokens = max_model_len self.enable_caching = enable_caching self.enable_kv_cache_events = enable_kv_cache_events self.use_eagle = use_eagle self.log_stats = log_stats self.metrics_collector = metrics_collector # FIXME: make prefix cache stats conditional on log_stats. We still need # this comment because when the log stats is enabled there are still # potential configs we could expose in the future. self.prefix_cache_stats = PrefixCacheStats() if log_stats else None self.coordinator = get_kv_cache_coordinator( kv_cache_config=kv_cache_config, max_model_len=self.max_model_len, max_in_flight_tokens=max_in_flight_tokens, use_eagle=self.use_eagle, enable_caching=self.enable_caching, enable_kv_cache_events=enable_kv_cache_events, dcp_world_size=dcp_world_size, pcp_world_size=pcp_world_size, scheduler_block_size=scheduler_block_size, hash_block_size=hash_block_size, metrics_collector=self.metrics_collector, ) self.num_kv_cache_groups = len(kv_cache_config.kv_cache_groups) self.block_pool = self.coordinator.block_pool self.kv_cache_config = kv_cache_config # Watermark: minimum number of KV cache blocks to keep free when # admitting waiting/preempted requests, to avoid frequent preemptions. assert watermark >= 0.0, "watermark must be non-negative" self.watermark_blocks = int(watermark * kv_cache_config.num_blocks) self.kv_cache_event_metadata = tuple( ( get_kv_cache_spec_kind(group.kv_cache_spec).value, get_kv_cache_spec_sliding_window(group.kv_cache_spec), ) for group in kv_cache_config.kv_cache_groups ) # Pre-constructed KVCacheBlocks with no blocks, callers should use this # via create_kv_cache_blocks instead of creating new ones to avoid GC # overhead. # # We use nested tuples to ensure the empty KVCacheBlocks is immutable. self.empty_kv_cache_blocks = KVCacheBlocks( tuple(() for _ in range(self.num_kv_cache_groups)) ) @property def usage(self) -> float: """Get the KV cache usage. Returns: The KV cache usage (between 0.0 and 1.0). """ return self.block_pool.get_usage() def make_prefix_cache_stats(self) -> PrefixCacheStats | None: """Get (and reset) the prefix cache stats. Returns: The current prefix caching stats, or None if logging is disabled. """ if not self.log_stats: return None stats = self.prefix_cache_stats self.prefix_cache_stats = PrefixCacheStats() return stats def get_computed_blocks(self, request: Request) -> tuple[KVCacheBlocks, int]: """Get the computed (cached) blocks for the request. Note that the computed blocks must be full. Args: request: The request to get the computed blocks. Returns: A tuple containing: - A list of blocks that are computed for the request. - The number of computed tokens. """ # We skip finding the prefix cache hit when prefix caching is # disabled or the request is marked as skipping kv cache read # (which happens when the request requires prompt logprobs # or calls a pooling model with all pooling). if not self.enable_caching or request.skip_reading_prefix_cache: return self.empty_kv_cache_blocks, 0 # NOTE: When all tokens hit the cache, we must recompute the last token # to obtain logits. Thus, set max_cache_hit_length to prompt_length - 1. # This can trigger recomputation of an entire block, rather than just # the single last token, because allocate_slots() requires # num_computed_tokens to be block-size aligned. Removing this limitation # could slightly improve performance in the future. max_cache_hit_length = request.num_tokens - 1 computed_blocks, num_new_computed_tokens = ( self.coordinator.find_longest_cache_hit( request.block_hashes, max_cache_hit_length ) ) # When kv_cache_report_mode is "full", emit BlockStored events # for the reused prefix cache blocks so that external consumers # (e.g. gateway) can learn about them. if ( num_new_computed_tokens > 0 and self.enable_kv_cache_events and getattr(request, "kv_cache_report_mode", "incremental") == "full" ): for group_idx, group_blocks in enumerate(computed_blocks): num_blocks = len(group_blocks) if num_blocks > 0: group = self.kv_cache_config.kv_cache_groups[group_idx] block_size = group.kv_cache_spec.block_size self.block_pool.emit_cached_block_events( request, num_blocks, block_size, group_idx, ) if self.log_stats: assert self.prefix_cache_stats is not None self.prefix_cache_stats.record( num_tokens=request.num_tokens, num_hits=num_new_computed_tokens, preempted=request.num_preemptions > 0, ) return self.create_kv_cache_blocks(computed_blocks), num_new_computed_tokens def allocate_slots( self, request: Request, num_new_tokens: int, num_new_computed_tokens: int = 0, new_computed_blocks: KVCacheBlocks | None = None, num_lookahead_tokens: int = 0, num_external_computed_tokens: int = 0, delay_cache_blocks: bool = False, num_encoder_tokens: int = 0, full_sequence_must_fit: bool = False, reserved_blocks: int = 0, has_scheduled_reqs: bool = True, ) -> KVCacheBlocks | None: """Add slots for a request with new tokens to append. Args: request: The request to allocate slots. num_new_tokens: The number of new tokens to be allocated and computed. num_new_computed_tokens: The number of new computed tokens just hitting the prefix caching, excluding external tokens. new_computed_blocks: The cached blocks for the above new computed tokens, grouped as a tuple by kv cache groups. num_lookahead_tokens: The number of speculative tokens to allocate. This is used by spec decode proposers with kv-cache such as eagle. num_external_computed_tokens: The number of tokens that their KV caches are not cached by vLLM but cached by the connector. delay_cache_blocks: Whether to skip caching the blocks. This is used by P/D when allocating blocks used in a KV transfer which will complete in a future step. num_encoder_tokens: The number of encoder tokens to allocate for cross-attention in encoder-decoder models(e.g., Whisper). For decoder-only models, this should be 0. full_sequence_must_fit: Only allocate blocks if the KV cache has enough free blocks to hold the full sequence, accounting for prefix cache hits and sliding window. Used as an admission gate to prevent over-admitting requests when chunked prefill would otherwise only check the first chunk reserved_blocks: Number of free blocks that must be left available for other in-flight sequences to complete. The actual allocation is only made if it fits within (free blocks - reserved_blocks). Used to gate async KV-connector loads so their initial allocation cannot consume blocks an already in-flight (prefilling) sequence is relying on. has_scheduled_reqs: Whether any requests are already scheduled to run this step, controls whether watermark is applied. Blocks layout: ``` ---------------------------------------------------------------------- | < comp > | < new_comp > | < ext_comp > | < new > | < lookahead > | ---------------------------------------------------------------------- | < to be computed > | ---------------------------------------------------------------------- | < to be allocated > | ---------------------------------------------------------------------- | < to be cached (roughly, | | details below)> | ---------------------------------------------------------------------- | Prefix-cached tokens from either vLLM | | or connector. Can be safely removed if | | they are outside sliding window. | ---------------------------------------------------------------------- | < cached by vLLM > | not cached by | | vLLM, but | | ref_cnt | ref_cnt not | cached by | | increased| increased yet| connector | ---------------------------------------------------------------------- ``` Abbrivations: ``` comp = request.num_computed_tokens new_comp = num_new_computed_tokens = len(new_computed_blocks) * block_size ext_comp = num_external_computed_tokens, cached by the connector new = num_new_tokens, including unverified draft tokens lookahead = num_lookahead_tokens ``` NOTE: for new tokens which include both verified and unverified draft tokens, we only cache the verified tokens (by capping the number at `request.num_tokens`). The allocation has three stages: - Free unnecessary blocks in `comp` and check if we have sufficient free blocks (return None if not). - Handle prefix tokens (`comp + new_comp + ext_comp`): - Free unnecessary blocks (e.g. outside sliding window) - Allocate new blocks for `ext_comp` tokens inside sliding window - Allocate new blocks for tokens to be computed (`new + lookahead`) Returns: A list of new allocated blocks. """ # When loading KV data asynchronously, we may have zero new tokens to # compute while still allocating slots for externally computed tokens. if num_new_tokens == 0 and num_external_computed_tokens == 0: raise ValueError( "num_new_tokens must be greater than 0 when there are no " "external computed tokens" ) if new_computed_blocks is not None: new_computed_block_list = new_computed_blocks.blocks else: new_computed_block_list = self.empty_kv_cache_blocks.blocks # The number of computed tokens is the number of computed tokens plus # the new prefix caching hits num_local_computed_tokens = ( request.num_computed_tokens + num_new_computed_tokens ) total_computed_tokens = min( num_local_computed_tokens + num_external_computed_tokens, self.max_model_len, ) watermark_blocks = 0 # The watermark is applied to waiting/preempted requests only, and only # when there's at least one request already scheduled. if has_scheduled_reqs and request.status in ( RequestStatus.WAITING, RequestStatus.PREEMPTED, ): watermark_blocks = self.watermark_blocks if full_sequence_must_fit: # First check and fail if the full request sequence won't fit. full_num_tokens = min(request.num_tokens, self.max_model_len) num_blocks_to_allocate = self.coordinator.get_num_blocks_to_allocate( request_id=request.request_id, num_tokens=full_num_tokens, new_computed_blocks=new_computed_block_list, num_encoder_tokens=num_encoder_tokens, total_computed_tokens=total_computed_tokens, num_local_computed_tokens=num_local_computed_tokens, num_tokens_main_model=full_num_tokens, apply_admission_cap=True, ) required_blocks = num_blocks_to_allocate + watermark_blocks if required_blocks > self.block_pool.get_num_free_blocks(): return None num_tokens_main_model = total_computed_tokens + num_new_tokens num_tokens_need_slot = min( num_tokens_main_model + num_lookahead_tokens, self.max_model_len ) # Free the blocks that are skipped during the attention computation # (e.g., tokens outside the sliding window). # We can do this even if we cannot schedule this request due to # insufficient free blocks. # Should call this function before allocating new blocks to reduce # the number of evicted blocks. # Free on the processed-token basis: in-flight steps' attention windows # still read blocks below the optimistic boundary, and rejected spec # tokens can roll it back. self.coordinator.remove_skipped_blocks( request.request_id, max(0, total_computed_tokens - request.num_in_flight_tokens), num_prompt_tokens=request.num_prompt_tokens, ) num_blocks_to_allocate = self.coordinator.get_num_blocks_to_allocate( request_id=request.request_id, num_tokens=num_tokens_need_slot, new_computed_blocks=new_computed_block_list, num_encoder_tokens=num_encoder_tokens, total_computed_tokens=num_local_computed_tokens + num_external_computed_tokens, num_local_computed_tokens=num_local_computed_tokens, num_tokens_main_model=num_tokens_main_model, ) # Keep `reserved_blocks` free for other in-flight sequences, and an # additional watermark of headroom for waiting/preempted admissions. available_blocks = self.block_pool.get_num_free_blocks() - reserved_blocks required_blocks = num_blocks_to_allocate + watermark_blocks if required_blocks > available_blocks: # Cannot allocate new blocks return None if ( new_computed_block_list is not self.empty_kv_cache_blocks.blocks or num_external_computed_tokens > 0 ): # Append the new computed blocks to the request blocks until now to # avoid the case where the new blocks cannot be allocated. self.coordinator.allocate_new_computed_blocks( request_id=request.request_id, new_computed_blocks=new_computed_block_list, num_local_computed_tokens=num_local_computed_tokens, num_external_computed_tokens=num_external_computed_tokens, ) new_blocks = self.coordinator.allocate_new_blocks( request.request_id, num_tokens_need_slot, num_tokens_main_model, num_encoder_tokens, ) # P/D: delay caching blocks if we have to recv from # remote. Update state for locally cached blocks. if not self.enable_caching or delay_cache_blocks: return self.create_kv_cache_blocks(new_blocks) # NOTE(woosuk): We want to commit (cache) up to num_local_computed_tokens # + num_external_computed_tokens + num_new_tokens, but must exclude # "non-committable" tokens (e.g., draft tokens that could be rejected). # Therefore, we cap the number at `request.num_tokens`, ensuring only # "finalized" tokens are cached. num_tokens_to_cache = min( total_computed_tokens + num_new_tokens, request.num_tokens, ) self.coordinator.cache_blocks(request, num_tokens_to_cache) return self.create_kv_cache_blocks(new_blocks) def free(self, request: Request) -> None: """Free the blocks allocated for the request. We free the blocks in reverse order so that the tail blocks are evicted first when caching is enabled. Args: request: The request to free the blocks. """ self.coordinator.free(request.request_id) def remove_skipped_blocks( self, request_id: str, processed_computed_tokens: int, num_prompt_tokens: int | None = None, ) -> None: """Remove the blocks that are no longer needed from `blocks` and replace the removed blocks with null_block. Args: request_id: The request ID. processed_computed_tokens: Computed-token prefix length covering fully processed and committed tokens only (safe to free). num_prompt_tokens: Optional prompt length for R-SWA gap eviction. """ self.coordinator.remove_skipped_blocks( request_id, processed_computed_tokens, num_prompt_tokens ) def pop_blocks_for_free(self, request: Request) -> list[KVCacheBlock]: """Pop the request's bookkeeping and return its blocks without returning them to the block pool. The caller must eventually free them in reverse order (so that tail blocks are evicted first). Args: request: The request to pop the blocks for. Returns: The request's blocks in allocation order. """ return self.coordinator.pop_blocks_for_free(request.request_id) def evict_blocks(self, block_ids: set[int]) -> None: """evict blocks from the prefix cache by their block IDs. Args: block_ids: Set of block IDs to evict from cache. """ self.block_pool.evict_blocks(block_ids) def reset_prefix_cache(self) -> bool: """Reset prefix cache. This function may be used in RLHF flows to invalidate prefix caching after the weights are updated, or used for resetting prefix caching status for benchmarking. Returns: bool: True if the prefix cache is successfully reset, False otherwise. """ if not self.block_pool.reset_prefix_cache(): return False if self.log_stats: assert self.prefix_cache_stats is not None self.prefix_cache_stats.reset = True return True def get_num_common_prefix_blocks(self, running_request_id: str) -> list[int]: """Calculate the number of common prefix blocks for each kv cache group. The function selects a running request and iterates through its blocks. A block is considered a common prefix block if ALL requests with allocated KV cache share it (i.e., ref_cnt equals the number of entries in req_to_blocks). NOTE(woosuk): The number of requests with allocated KV cache is **greater than or equal to** the number of requests scheduled in the current step. This is because having allocated KV cache only indicates that: 1. The request has not yet finished, and 2. The request holds its blocks unfreed. While all scheduled requests must have allocated KV cache, the inverse is not necessarily true. There may be requests with allocated KV cache that are not scheduled in the current step. This can result in an edge case where the number of common prefix blocks is 0, even though all scheduled requests share a common prefix. This occurs because there may be unscheduled requests that do not share the common prefix. Currently, this case cannot be easily detected, so the function returns 0 in such cases. Args: running_request_id: The request ID of any running request, used to identify the common prefix blocks. Returns: list[int]: The number of common prefix blocks for each kv cache group. """ return self.coordinator.get_num_common_prefix_blocks(running_request_id) def take_events(self) -> list[KVCacheEvent]: """Take the KV cache events from the block pool. Returns: A list of KV cache events. """ events = self.block_pool.take_events() for event in events: if not isinstance(event, BlockStored): continue if event.group_idx is None: continue if event.group_idx < 0 or event.group_idx >= len( self.kv_cache_event_metadata ): logger.warning( "Group index `%s` not in KV cache metadata", event.group_idx ) continue # Annotate here so BlockPool can keep emitting structural cache # events without owning semantic KV cache spec metadata. kind, sliding_window = self.kv_cache_event_metadata[event.group_idx] event.kv_cache_spec_kind = kind event.kv_cache_spec_sliding_window = sliding_window return events def get_blocks(self, request_id: str) -> KVCacheBlocks: """Get the blocks of a request.""" return self.create_kv_cache_blocks(self.coordinator.get_blocks(request_id)) def get_block_ids(self, request_id: str) -> tuple[list[int], ...]: """Get the block ids of a request.""" return self.get_blocks(request_id).get_block_ids() def get_block_ids_for_computed_tokens( self, request_id: str, num_computed_tokens: int, ) -> tuple[list[int], ...]: """Get block ids covering the request's computed tokens.""" block_ids = self.get_block_ids(request_id) clipped_block_ids: list[list[int]] = [] for group, ids in zip(self.kv_cache_config.kv_cache_groups, block_ids): spec = group.kv_cache_spec if not isinstance(spec, AttentionSpec) or isinstance( spec, (CrossAttentionSpec, EncoderOnlyAttentionSpec) ): clipped_block_ids.append(ids) continue num_valid_blocks = cdiv(num_computed_tokens, spec.block_size) clipped_block_ids.append(ids[:num_valid_blocks]) return tuple(clipped_block_ids) def cache_blocks(self, request: Request, num_computed_tokens: int) -> None: """Cache the blocks for the request, if enabled. Args: request: The request to cache the blocks. num_computed_tokens: The number of computed tokens, including tokens that are already cached and tokens to be cached. """ if self.enable_caching: self.coordinator.cache_blocks(request, num_computed_tokens) def create_kv_cache_blocks( self, blocks: tuple[list[KVCacheBlock], ...] ) -> KVCacheBlocks: # Only create new KVCacheBlocks for non-empty blocks return KVCacheBlocks(blocks) if any(blocks) else self.empty_kv_cache_blocks def take_new_block_ids(self) -> list[int]: """Drain and return new attention block IDs for zeroing.""" ids: list[int] = [] for mgr in self.coordinator.single_type_managers: ids.extend(mgr.take_new_block_ids()) return ids def take_kv_cache_block_copies( self, ) -> tuple[list[KVCacheBlockCopy], list[KVCacheBlock]]: """Drain pending copies and return their retained endpoints.""" pending_copies: list[tuple[KVCacheBlock, KVCacheBlock]] = [] for mgr in self.coordinator.single_type_managers: pending_copies.extend(mgr.take_pending_cow_copies()) copies = [ KVCacheBlockCopy( src_block_id=source_block.block_id, dst_block_id=cow_block.block_id, ) for source_block, cow_block in pending_copies ] retained_blocks = [block for pair in pending_copies for block in pair] return copies, retained_blocks def new_step_starts(self) -> None: """Notify the coordinator that a new step is starting.""" self.coordinator.new_step_starts()