| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from collections.abc import Callable, Iterable, Sequence |
| 4 | from typing import Any |
| 5 | |
| 6 | from vllm.distributed.kv_events import ( |
| 7 | MEDIUM_GPU, |
| 8 | AllBlocksCleared, |
| 9 | BlockRemoved, |
| 10 | BlockStored, |
| 11 | KVCacheEvent, |
| 12 | ) |
| 13 | from vllm.logger import init_logger |
| 14 | from vllm.multimodal.utils import get_mm_features_in_window |
| 15 | from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector |
| 16 | from vllm.v1.core.kv_cache_utils import ( |
| 17 | BlockHash, |
| 18 | BlockHashWithGroupId, |
| 19 | ExternalBlockHash, |
| 20 | FreeKVCacheBlockQueue, |
| 21 | KVCacheBlock, |
| 22 | generate_block_hash_extra_keys, |
| 23 | get_block_hash, |
| 24 | get_group_id, |
| 25 | make_block_hash_with_group_id, |
| 26 | maybe_convert_block_hash, |
| 27 | resolve_block_hashes, |
| 28 | to_event_extra_keys, |
| 29 | ) |
| 30 | from vllm.v1.request import Request |
| 31 | |
| 32 | logger = init_logger(__name__) |
| 33 | |
| 34 | |
| 35 | class BlockHashToBlockMap: |
| 36 | """Cache of blocks that are used for prefix caching. It caches blocks |
| 37 | from hash directly to a block or multiple blocks |
| 38 | (i.e. {block_hash: KVCacheBlocks}) |
| 39 | - Mostly block_hash maps to a single KVCacheBlock, and KVCacheBlocks |
| 40 | would simply be a KVCacheBlock. |
| 41 | - Otherwise, KVCacheBlocks is a dict from {block_id: KVCacheBlock} |
| 42 | |
| 43 | A cached block is a full block with a block hash that can be used |
| 44 | for prefix caching. |
| 45 | The cached block may be used by running requests or in the |
| 46 | free_block_queue that could potentially be evicted. |
| 47 | |
| 48 | NOTE #1: We currently don't de-duplicate the blocks in the cache, |
| 49 | meaning that if a block becomes full and is cached, we don't check |
| 50 | if there is already an identical block in the cache. This is because |
| 51 | we want to make sure the allocated block IDs won't change so that |
| 52 | block tables are append-only. |
| 53 | NOTE #2: The union type is introduced in order to reduce GC costs |
| 54 | from the inner dict. |
| 55 | """ |
| 56 | |
| 57 | def __init__(self): |
| 58 | self._cache: dict[ |
| 59 | BlockHashWithGroupId, KVCacheBlock | dict[int, KVCacheBlock] |
| 60 | ] = {} |
| 61 | |
| 62 | def get_one_block(self, key: BlockHashWithGroupId) -> KVCacheBlock | None: |
| 63 | """Gets any block with the given block hash key.""" |
| 64 | blocks = self._cache.get(key) |
| 65 | if blocks is not None: |
| 66 | if isinstance(blocks, KVCacheBlock): |
| 67 | return blocks |
| 68 | if isinstance(blocks, dict): |
| 69 | return next(iter(blocks.values())) |
| 70 | self._unexpected_blocks_type(blocks) |
| 71 | return None |
| 72 | |
| 73 | def contain(self, key: BlockHashWithGroupId, block_id: int) -> bool: |
| 74 | """Checks whether the key maps to the given block ID.""" |
| 75 | blocks = self._cache.get(key) |
| 76 | if blocks is None: |
| 77 | return False |
| 78 | if isinstance(blocks, KVCacheBlock): |
| 79 | return blocks.block_id == block_id |
| 80 | if isinstance(blocks, dict): |
| 81 | return block_id in blocks |
| 82 | self._unexpected_blocks_type(blocks) |
| 83 | return False |
| 84 | |
| 85 | def insert(self, key: BlockHashWithGroupId, block: KVCacheBlock) -> None: |
| 86 | """Inserts the KVCacheBlock to the cache""" |
| 87 | blocks = self._cache.get(key) |
| 88 | if blocks is None: |
| 89 | # When key is not found, attach a single block to the key |
| 90 | self._cache[key] = block |
| 91 | elif isinstance(blocks, KVCacheBlock): |
| 92 | # If there's a block with the same key, merge the original block |
| 93 | # and the new block into a dict |
| 94 | self._cache[key] = {blocks.block_id: blocks, block.block_id: block} |
| 95 | elif isinstance(blocks, dict): |
| 96 | # If it's already a dict, simply insert the block |
| 97 | blocks[block.block_id] = block |
| 98 | else: |
| 99 | self._unexpected_blocks_type(blocks) |
| 100 | |
| 101 | def pop(self, key: BlockHashWithGroupId, block_id: int) -> KVCacheBlock | None: |
| 102 | """Checks if block_hash exists and pop block_id from the cache""" |
| 103 | blocks = self._cache.pop(key, None) |
| 104 | if blocks is None: |
| 105 | # block_hash not found in the cache |
| 106 | return None |
| 107 | # TODO(Jialin): If key is found, block_id should always present |
| 108 | # in blocks. We currently keep the original behaviour for safety. |
| 109 | # |
| 110 | # Will add block_id == blocks.block_id assertion and |
| 111 | # use del blocks[block_id] instead as followup. |
| 112 | if isinstance(blocks, KVCacheBlock): |
| 113 | if blocks.block_id == block_id: |
| 114 | return blocks |
| 115 | # If the single block ID doesn't match, we should put the |
| 116 | # block back (it should happen rarely) |
| 117 | self._cache[key] = blocks |
| 118 | return None |
| 119 | if isinstance(blocks, dict): |
| 120 | # Try to pop block_id from the block dict, and if dict still |
| 121 | # contain blocks, put back to the cache. |
| 122 | block = blocks.pop(block_id, None) |
| 123 | if len(blocks) > 0: |
| 124 | self._cache[key] = blocks |
| 125 | return block |
| 126 | self._unexpected_blocks_type(blocks) |
| 127 | return None |
| 128 | |
| 129 | def __len__(self) -> int: |
| 130 | return len(self._cache) |
| 131 | |
| 132 | def _unexpected_blocks_type(self, blocks: Any) -> None: |
| 133 | raise AssertionError(f"Invalid KV cache block type {type(blocks)}") |
| 134 | |
| 135 | |
| 136 | class BlockPool: |
| 137 | """BlockPool that manages KVCacheBlocks. |
| 138 | It provides methods to allocate, free and cache the kv cache blocks. The |
| 139 | free_block_queue stores the free blocks in eviction order to enable |
| 140 | allocation, free, and cache eviction. The cached_block_hash_to_block |
| 141 | maps between block hash and cached block to support finding cached blocks |
| 142 | by their block hash. |
| 143 | |
| 144 | Args: |
| 145 | num_gpu_blocks: The number of blocks in the pool. |
| 146 | enable_caching: Whether to enable prefix caching. |
| 147 | hash_block_size: The block size of which the block hashes are computed. |
| 148 | The actual block size usually equals hash_block_size, but in cases |
| 149 | where different KV cache groups have different block sizes, the |
| 150 | actual block size can be a multiple of hash_block_size. |
| 151 | enable_kv_cache_events: Whether to enable kv cache events. |
| 152 | metrics_collector: Optional metrics collector for tracking block residency. |
| 153 | medium: Storage medium reported in KV cache events. |
| 154 | |
| 155 | """ |
| 156 | |
| 157 | def __init__( |
| 158 | self, |
| 159 | num_gpu_blocks: int, |
| 160 | enable_caching: bool, |
| 161 | hash_block_size: int, |
| 162 | enable_kv_cache_events: bool = False, |
| 163 | metrics_collector: KVCacheMetricsCollector | None = None, |
| 164 | medium: str = MEDIUM_GPU, |
| 165 | ): |
| 166 | assert isinstance(num_gpu_blocks, int) and num_gpu_blocks > 0 |
| 167 | self.num_gpu_blocks = num_gpu_blocks |
| 168 | self.medium = medium |
| 169 | self.enable_caching = enable_caching |
| 170 | self.hash_block_size = hash_block_size |
| 171 | # All kv-cache blocks. |
| 172 | self.blocks: list[KVCacheBlock] = [ |
| 173 | KVCacheBlock(idx, pool=self) for idx in range(num_gpu_blocks) |
| 174 | ] |
| 175 | # Free block queue that constructs and manipulates a doubly linked |
| 176 | # list of free blocks (including eviction candidates when caching is |
| 177 | # enabled). |
| 178 | self.free_block_queue = FreeKVCacheBlockQueue(self.blocks) |
| 179 | |
| 180 | # Cache for block lookup |
| 181 | self.cached_block_hash_to_block: BlockHashToBlockMap = BlockHashToBlockMap() |
| 182 | self.cached_block_hashes_by_block: dict[int, set[BlockHashWithGroupId]] = {} |
| 183 | |
| 184 | # To represent a placeholder block with block_id=0. |
| 185 | # The ref_cnt of null_block is not maintained, needs special care to |
| 186 | # avoid freeing it. |
| 187 | self.null_block = self.free_block_queue.popleft() |
| 188 | self.null_block.is_null = True |
| 189 | |
| 190 | self.enable_kv_cache_events = enable_kv_cache_events |
| 191 | self.kv_event_queue: list[KVCacheEvent] = [] |
| 192 | |
| 193 | self.metrics_collector = metrics_collector |
| 194 | # Callbacks for blocks released with ``unpin_blocks`` whose contents |
| 195 | # are still being read until the pool reuses them. |
| 196 | self._reuse_watchers: dict[int, Callable[[KVCacheBlock], None]] = {} |
| 197 | |
| 198 | def get_cached_block( |
| 199 | self, block_hash: BlockHash, kv_cache_group_ids: list[int] |
| 200 | ) -> list[KVCacheBlock] | None: |
| 201 | """Get the cached block by the block hash for each group in |
| 202 | `kv_cache_group_ids`, or None if cache miss for any group. |
| 203 | If there are duplicated blocks, we return the first block in the cache. |
| 204 | |
| 205 | Args: |
| 206 | block_hash: The hash value of the block. |
| 207 | kv_cache_group_ids: The ids of the KV cache groups. |
| 208 | |
| 209 | Returns: |
| 210 | The cached blocks if exists, or None. |
| 211 | |
| 212 | """ |
| 213 | cached_blocks = [] |
| 214 | for group_id in kv_cache_group_ids: |
| 215 | block_hash_with_group_id = make_block_hash_with_group_id( |
| 216 | block_hash, group_id |
| 217 | ) |
| 218 | block = self.cached_block_hash_to_block.get_one_block( |
| 219 | block_hash_with_group_id |
| 220 | ) |
| 221 | if not block: |
| 222 | return None |
| 223 | cached_blocks.append(block) |
| 224 | return cached_blocks |
| 225 | |
| 226 | def cache_full_blocks( |
| 227 | self, |
| 228 | request: Request, |
| 229 | blocks: list[KVCacheBlock], |
| 230 | num_cached_blocks: int, |
| 231 | num_full_blocks: int, |
| 232 | block_size: int, |
| 233 | kv_cache_group_id: int, |
| 234 | block_mask: list[bool] | None = None, |
| 235 | ) -> None: |
| 236 | """Cache a list of full blocks for prefix caching. |
| 237 | This function takes a list of blocks that will have their block hash |
| 238 | metadata to be updated and cached. Given a request, it updates the |
| 239 | metadata for each block and caching it in the |
| 240 | `cached_block_hash_to_block`. |
| 241 | The block hashes values are computed by the Request object immediately |
| 242 | when it is created and when new tokens are appended. |
| 243 | |
| 244 | Args: |
| 245 | request: The request to cache the blocks. |
| 246 | blocks: All blocks in the request. |
| 247 | num_cached_blocks: The number of blocks that are already cached. |
| 248 | num_full_blocks: The number of blocks that are full and should |
| 249 | be cached after this function. |
| 250 | block_size: Number of tokens in each block. |
| 251 | kv_cache_group_id: The id of the KV cache group. |
| 252 | block_mask: Optional mask aligned with |
| 253 | ``blocks[num_cached_blocks:num_full_blocks]``. When provided, |
| 254 | blocks where the mask is False are skipped (treated like null |
| 255 | blocks). Used by groups whose ``find_longest_cache_hit`` only |
| 256 | consults a subset of blocks (e.g. SWA tail-window), so blocks |
| 257 | that can never serve a hit stay out of the prefix-cache hash |
| 258 | map. |
| 259 | |
| 260 | """ |
| 261 | if num_cached_blocks >= num_full_blocks: |
| 262 | return |
| 263 | new_full_blocks = blocks[num_cached_blocks:num_full_blocks] |
| 264 | assert block_mask is None or len(block_mask) == len(new_full_blocks) |
| 265 | block_hashes = resolve_block_hashes( |
| 266 | request.block_hashes, self.hash_block_size, block_size |
| 267 | ) |
| 268 | |
| 269 | new_block_hashes = block_hashes[num_cached_blocks:] |
| 270 | new_hashes: list[ExternalBlockHash] | None = ( |
| 271 | [] if self.enable_kv_cache_events else None |
| 272 | ) |
| 273 | for i, blk in enumerate(new_full_blocks): |
| 274 | # Some blocks may be null or masked out when enabling sparse attention |
| 275 | # like sliding window attention, or Mamba models with prefix-caching |
| 276 | # in align mode. We skip null blocks here. |
| 277 | if blk.is_null or (block_mask is not None and not block_mask[i]): |
| 278 | continue |
| 279 | block_hash = new_block_hashes[i] |
| 280 | num_hash_tokens = (num_cached_blocks + i + 1) * block_size |
| 281 | |
| 282 | # Update and added the full block to the cache. |
| 283 | block_hash_with_group_id = make_block_hash_with_group_id( |
| 284 | block_hash, kv_cache_group_id |
| 285 | ) |
| 286 | if blk.block_hash is not None: |
| 287 | # The only valid case where a "new full block" already has a |
| 288 | # hash is partial->full promotion of the same cache block. |
| 289 | assert ( |
| 290 | blk.block_hash_num_tokens is not None |
| 291 | and blk.block_hash_num_tokens < num_hash_tokens |
| 292 | ) |
| 293 | removed_hashes = self._remove_cached_block_hashes(blk) |
| 294 | self._emit_block_removed_events(removed_hashes) |
| 295 | self._insert_block_hash( |
| 296 | block_hash_with_group_id, |
| 297 | blk, |
| 298 | num_tokens=num_hash_tokens, |
| 299 | ) |
| 300 | if new_hashes is not None: |
| 301 | new_hashes.append(maybe_convert_block_hash(block_hash)) |
| 302 | |
| 303 | if self.enable_kv_cache_events: |
| 304 | if num_cached_blocks == 0: |
| 305 | parent_block_hash: ExternalBlockHash | None = None |
| 306 | else: |
| 307 | parent_block_hash = maybe_convert_block_hash( |
| 308 | block_hashes[num_cached_blocks - 1] |
| 309 | ) |
| 310 | |
| 311 | # Calculate token range for the blocks being cached |
| 312 | start_token_idx = num_cached_blocks * block_size |
| 313 | end_token_idx = num_full_blocks * block_size |
| 314 | |
| 315 | # Generate extra keys for each block individually. |
| 316 | # Each block may have different extra_keys (e.g., different MM |
| 317 | # features, or cache_salt only for the first block). |
| 318 | # Skip null/masked-out blocks to match the length of new_hashes. |
| 319 | extra_keys_list: list[tuple[Any, ...] | None] = [] |
| 320 | curr_mm_idx = 0 |
| 321 | for i in range(num_cached_blocks, num_full_blocks): |
| 322 | if blocks[i].is_null: |
| 323 | continue |
| 324 | if block_mask is not None and not block_mask[i - num_cached_blocks]: |
| 325 | continue |
| 326 | block_start = i * block_size |
| 327 | block_end = block_start + block_size |
| 328 | extra_keys, curr_mm_idx = generate_block_hash_extra_keys( |
| 329 | request, block_start, block_end, curr_mm_idx |
| 330 | ) |
| 331 | extra_keys_list.append(extra_keys) |
| 332 | |
| 333 | self.kv_event_queue.append( |
| 334 | self._build_block_stored_event( |
| 335 | request, |
| 336 | block_hashes=new_hashes, |
| 337 | parent_block_hash=parent_block_hash, |
| 338 | start_token_idx=start_token_idx, |
| 339 | end_token_idx=end_token_idx, |
| 340 | block_size=block_size, |
| 341 | kv_cache_group_id=kv_cache_group_id, |
| 342 | extra_keys_list=extra_keys_list, |
| 343 | ) |
| 344 | ) |
| 345 | |
| 346 | def _build_block_stored_event( |
| 347 | self, |
| 348 | request: Request, |
| 349 | block_hashes: list[ExternalBlockHash] | None, |
| 350 | parent_block_hash: ExternalBlockHash | None, |
| 351 | start_token_idx: int, |
| 352 | end_token_idx: int, |
| 353 | block_size: int, |
| 354 | kv_cache_group_id: int, |
| 355 | extra_keys_list: list[tuple[Any, ...] | None], |
| 356 | ) -> BlockStored: |
| 357 | """Build a ``BlockStored`` KV event for ``request``. |
| 358 | |
| 359 | Shared by ``cache_full_blocks`` (newly cached blocks) and |
| 360 | ``emit_cached_block_events`` (prefix-cache-reused blocks) so both emit |
| 361 | identical event shapes for downstream consumers. |
| 362 | """ |
| 363 | return BlockStored( |
| 364 | block_hashes=block_hashes, |
| 365 | parent_block_hash=parent_block_hash, |
| 366 | token_ids=request.all_token_ids[start_token_idx:end_token_idx], |
| 367 | block_size=block_size, |
| 368 | lora_id=request.lora_request.adapter_id if request.lora_request else None, |
| 369 | medium=self.medium, |
| 370 | lora_name=request.lora_request.name if request.lora_request else None, |
| 371 | extra_keys=to_event_extra_keys(extra_keys_list), |
| 372 | group_idx=kv_cache_group_id, |
| 373 | session_id=request.session_id, |
| 374 | ) |
| 375 | |
| 376 | def emit_cached_block_events( |
| 377 | self, |
| 378 | request: Request, |
| 379 | num_cached_blocks: int, |
| 380 | block_size: int, |
| 381 | kv_cache_group_id: int, |
| 382 | ) -> None: |
| 383 | """Generate BlockStored events for blocks reused from prefix cache. |
| 384 | |
| 385 | Unlike cache_full_blocks(), this does NOT modify block state — |
| 386 | the blocks are already cached. It only generates events so that |
| 387 | external consumers (e.g. gateway) can learn about reused blocks. |
| 388 | |
| 389 | Args: |
| 390 | request: The request whose prefix cache blocks were reused. |
| 391 | num_cached_blocks: Number of blocks that were cache hits. |
| 392 | block_size: Number of tokens per block. |
| 393 | kv_cache_group_id: The KV cache group ID. |
| 394 | |
| 395 | """ |
| 396 | if not self.enable_kv_cache_events or num_cached_blocks == 0: |
| 397 | return |
| 398 | |
| 399 | block_hashes = resolve_block_hashes( |
| 400 | request.block_hashes, self.hash_block_size, block_size |
| 401 | ) |
| 402 | |
| 403 | # Collect external hashes and extra_keys for cached blocks. |
| 404 | cached_hashes: list[ExternalBlockHash] = [] |
| 405 | extra_keys_list: list[tuple[Any, ...] | None] = [] |
| 406 | curr_mm_idx = 0 |
| 407 | for i in range(num_cached_blocks): |
| 408 | block_start = i * block_size |
| 409 | block_end = block_start + block_size |
| 410 | cached_hashes.append(maybe_convert_block_hash(block_hashes[i])) |
| 411 | extra_keys, curr_mm_idx = generate_block_hash_extra_keys( |
| 412 | request, block_start, block_end, curr_mm_idx |
| 413 | ) |
| 414 | extra_keys_list.append(extra_keys) |
| 415 | |
| 416 | if not cached_hashes: |
| 417 | return |
| 418 | |
| 419 | # Prefix-cache hits always form a contiguous prefix starting at block 0, |
| 420 | # so the first (and thus the whole group's) parent block hash is None. |
| 421 | parent_block_hash: ExternalBlockHash | None = None |
| 422 | start_token_idx = 0 |
| 423 | end_token_idx = num_cached_blocks * block_size |
| 424 | |
| 425 | logger.debug( |
| 426 | "EmitCachedBlock event: block_size=%d, " |
| 427 | "num_cached_blocks=%d, parent_block_hash=%s, " |
| 428 | "token_ids_len=%d, group_idx=%s", |
| 429 | block_size, |
| 430 | num_cached_blocks, |
| 431 | parent_block_hash, |
| 432 | len(request.all_token_ids[start_token_idx:end_token_idx]), |
| 433 | kv_cache_group_id, |
| 434 | ) |
| 435 | |
| 436 | self.kv_event_queue.append( |
| 437 | self._build_block_stored_event( |
| 438 | request, |
| 439 | block_hashes=cached_hashes, |
| 440 | parent_block_hash=parent_block_hash, |
| 441 | start_token_idx=start_token_idx, |
| 442 | end_token_idx=end_token_idx, |
| 443 | block_size=block_size, |
| 444 | kv_cache_group_id=kv_cache_group_id, |
| 445 | extra_keys_list=extra_keys_list, |
| 446 | ) |
| 447 | ) |
| 448 | |
| 449 | def cache_partial_block( |
| 450 | self, |
| 451 | request: Request, |
| 452 | block: KVCacheBlock, |
| 453 | num_tokens: int, |
| 454 | kv_cache_group_id: int, |
| 455 | block_size: int, |
| 456 | replace_existing_hashes: bool = False, |
| 457 | ) -> BlockHashWithGroupId | None: |
| 458 | """Register a partial prefix-cache entry for an existing block. |
| 459 | |
| 460 | Prefix-cache keys normally identify full cache blocks. A partial entry |
| 461 | makes an existing cache block reachable from a fine-grained prefix |
| 462 | boundary inside that block without allocating or copying a new |
| 463 | ``KVCacheBlock``. |
| 464 | |
| 465 | The partial entry is lookup metadata owned by ``block``. If ``block`` |
| 466 | has no primary hash, the key becomes its primary hash. If the block |
| 467 | already has a primary hash, the partial entry is tracked in |
| 468 | ``cached_block_hashes_by_block`` so eviction, reset, and promotion can |
| 469 | remove every hash key that points to the block. |
| 470 | |
| 471 | Args: |
| 472 | request: Request whose token IDs and block hashes define the |
| 473 | partial entry. |
| 474 | block: Existing cache block to make reachable from the partial |
| 475 | prefix boundary. |
| 476 | num_tokens: Prefix length represented by the partial entry. It |
| 477 | must be a positive multiple of ``self.hash_block_size`` and |
| 478 | cannot exceed the request's computed block hashes. |
| 479 | kv_cache_group_id: KV cache group that owns the partial entry. |
| 480 | block_size: Cache block size for the owning group. The partial |
| 481 | entry hash itself is always the prefix-chain hash at |
| 482 | ``num_tokens``; ``block_size`` is used to assert that the |
| 483 | entry is partial within the owning cache block. |
| 484 | replace_existing_hashes: Whether the block contents were replaced |
| 485 | and all existing cache entries must be removed before the new |
| 486 | entry is registered. |
| 487 | |
| 488 | Returns: |
| 489 | The hash key with group ID if a partial entry can be registered; |
| 490 | otherwise ``None`` for null blocks. |
| 491 | |
| 492 | """ |
| 493 | if block.is_null: |
| 494 | return None |
| 495 | |
| 496 | assert block_size % self.hash_block_size == 0 |
| 497 | assert replace_existing_hashes or ( |
| 498 | block_size > self.hash_block_size and num_tokens % block_size != 0 |
| 499 | ) |
| 500 | block_hash = self._get_partial_block_hash(request, num_tokens) |
| 501 | num_hash_blocks = num_tokens // self.hash_block_size |
| 502 | block_hash_with_group_id = make_block_hash_with_group_id( |
| 503 | block_hash, kv_cache_group_id |
| 504 | ) |
| 505 | already_cached = block.block_hash == block_hash_with_group_id or ( |
| 506 | self.cached_block_hash_to_block.contain( |
| 507 | block_hash_with_group_id, block.block_id |
| 508 | ) |
| 509 | ) |
| 510 | if replace_existing_hashes: |
| 511 | removed_hashes = self._remove_cached_block_hashes(block) |
| 512 | self._emit_block_removed_events(removed_hashes) |
| 513 | already_cached = False |
| 514 | elif ( |
| 515 | not already_cached |
| 516 | and block.block_hash is not None |
| 517 | and block.block_hash_num_tokens is not None |
| 518 | and block.block_hash_num_tokens < num_hash_blocks * self.hash_block_size |
| 519 | ): |
| 520 | removed_hashes = self._remove_cached_block_hashes(block) |
| 521 | self._emit_block_removed_events(removed_hashes) |
| 522 | self._insert_block_hash( |
| 523 | block_hash_with_group_id, |
| 524 | block, |
| 525 | num_tokens=num_hash_blocks * self.hash_block_size, |
| 526 | ) |
| 527 | if self.enable_kv_cache_events and not already_cached: |
| 528 | parent_hash, block_start = self._get_partial_block_parent_hash_and_start( |
| 529 | request, num_tokens |
| 530 | ) |
| 531 | parent_block_hash = ( |
| 532 | maybe_convert_block_hash(parent_hash) |
| 533 | if parent_hash is not None |
| 534 | else None |
| 535 | ) |
| 536 | block_end = num_tokens |
| 537 | curr_mm_idx = 0 |
| 538 | mm_features = request.mm_features |
| 539 | if block_start > 0 and mm_features: |
| 540 | last_mm_pos = mm_features[-1].mm_position |
| 541 | if last_mm_pos.offset + last_mm_pos.length > block_start: |
| 542 | curr_mm_idx, _ = get_mm_features_in_window( |
| 543 | mm_features, block_start, block_end |
| 544 | ) |
| 545 | extra_keys, _ = generate_block_hash_extra_keys( |
| 546 | request, block_start, block_end, curr_mm_idx |
| 547 | ) |
| 548 | self.kv_event_queue.append( |
| 549 | BlockStored( |
| 550 | block_hashes=[maybe_convert_block_hash(block_hash)], |
| 551 | parent_block_hash=parent_block_hash, |
| 552 | token_ids=request.all_token_ids[block_start:block_end], |
| 553 | block_size=block_end - block_start, |
| 554 | lora_id=request.lora_request.adapter_id |
| 555 | if request.lora_request |
| 556 | else None, |
| 557 | medium=self.medium, |
| 558 | lora_name=request.lora_request.name |
| 559 | if request.lora_request |
| 560 | else None, |
| 561 | extra_keys=to_event_extra_keys([extra_keys]), |
| 562 | group_idx=kv_cache_group_id, |
| 563 | session_id=request.session_id, |
| 564 | ) |
| 565 | ) |
| 566 | return block_hash_with_group_id |
| 567 | |
| 568 | def _get_partial_block_hash( |
| 569 | self, |
| 570 | request: Request, |
| 571 | num_tokens: int, |
| 572 | ) -> BlockHash: |
| 573 | assert num_tokens % self.hash_block_size == 0 |
| 574 | num_hash_blocks = num_tokens // self.hash_block_size |
| 575 | assert 0 < num_hash_blocks <= len(request.block_hashes) |
| 576 | |
| 577 | # Each hash_block_size hash chains over its full prefix, so the partial |
| 578 | # entry for any group block size is the hash at that prefix boundary. |
| 579 | return request.block_hashes[num_hash_blocks - 1] |
| 580 | |
| 581 | def _get_partial_block_parent_hash_and_start( |
| 582 | self, |
| 583 | request: Request, |
| 584 | num_tokens: int, |
| 585 | ) -> tuple[BlockHash | None, int]: |
| 586 | num_hash_blocks = num_tokens // self.hash_block_size |
| 587 | parent_hash = ( |
| 588 | request.block_hashes[num_hash_blocks - 2] if num_hash_blocks > 1 else None |
| 589 | ) |
| 590 | block_start = (num_hash_blocks - 1) * self.hash_block_size |
| 591 | return parent_hash, block_start |
| 592 | |
| 593 | def _remove_cached_block_hashes( |
| 594 | self, |
| 595 | block: KVCacheBlock, |
| 596 | ) -> list[BlockHashWithGroupId]: |
| 597 | block_hashes: list[BlockHashWithGroupId] = [] |
| 598 | if block.block_hash is not None: |
| 599 | block_hashes.append(block.block_hash) |
| 600 | block_hashes.extend(self.cached_block_hashes_by_block.pop(block.block_id, ())) |
| 601 | if not block_hashes: |
| 602 | return [] |
| 603 | |
| 604 | removed_hashes: list[BlockHashWithGroupId] = [] |
| 605 | for block_hash in block_hashes: |
| 606 | if ( |
| 607 | self.cached_block_hash_to_block.pop(block_hash, block.block_id) |
| 608 | is not None |
| 609 | ): |
| 610 | removed_hashes.append(block_hash) |
| 611 | block.reset_hash() |
| 612 | return removed_hashes |
| 613 | |
| 614 | def _emit_block_removed_events( |
| 615 | self, |
| 616 | block_hashes: list[BlockHashWithGroupId], |
| 617 | ) -> None: |
| 618 | if not self.enable_kv_cache_events: |
| 619 | return |
| 620 | for block_hash in block_hashes: |
| 621 | self.kv_event_queue.append( |
| 622 | BlockRemoved( |
| 623 | block_hashes=[maybe_convert_block_hash(get_block_hash(block_hash))], |
| 624 | medium=self.medium, |
| 625 | group_idx=get_group_id(block_hash), |
| 626 | ) |
| 627 | ) |
| 628 | |
| 629 | def _insert_block_hash( |
| 630 | self, |
| 631 | block_hash_with_group_id: BlockHashWithGroupId, |
| 632 | block: KVCacheBlock, |
| 633 | num_tokens: int | None, |
| 634 | ) -> None: |
| 635 | if block.block_hash == block_hash_with_group_id: |
| 636 | return |
| 637 | |
| 638 | if self.cached_block_hash_to_block.contain( |
| 639 | block_hash_with_group_id, block.block_id |
| 640 | ): |
| 641 | return |
| 642 | |
| 643 | if block.block_hash is None: |
| 644 | block.set_block_hash(block_hash_with_group_id, num_tokens=num_tokens) |
| 645 | else: |
| 646 | self.cached_block_hashes_by_block.setdefault(block.block_id, set()).add( |
| 647 | block_hash_with_group_id |
| 648 | ) |
| 649 | self.cached_block_hash_to_block.insert(block_hash_with_group_id, block) |
| 650 | |
| 651 | def move_block_hashes( |
| 652 | self, |
| 653 | src_block: KVCacheBlock, |
| 654 | dst_block: KVCacheBlock, |
| 655 | ) -> None: |
| 656 | """Re-point ``src_block``'s prefix-cache entries to ``dst_block``. |
| 657 | |
| 658 | Used when the request owning ``src_block`` keeps writing into it |
| 659 | : the prefix cache holds a private copy (``dst_block``) |
| 660 | under the same hashes instead. Entries stay live; no events emitted. |
| 661 | """ |
| 662 | assert dst_block.block_hash is None |
| 663 | assert dst_block.block_id not in self.cached_block_hashes_by_block |
| 664 | num_tokens = src_block.block_hash_num_tokens |
| 665 | for block_hash in self._remove_cached_block_hashes(src_block): |
| 666 | # `num_tokens` only applies to the first (primary) insertion. |
| 667 | self._insert_block_hash(block_hash, dst_block, num_tokens=num_tokens) |
| 668 | |
| 669 | def get_new_blocks(self, num_blocks: int) -> list[KVCacheBlock]: |
| 670 | """Get new blocks from the free block pool. |
| 671 | |
| 672 | Note that we do not check block cache in this function. |
| 673 | |
| 674 | Args: |
| 675 | num_blocks: The number of blocks to allocate. |
| 676 | |
| 677 | Returns: |
| 678 | A list of new block. |
| 679 | |
| 680 | """ |
| 681 | if num_blocks > self.get_num_free_blocks(): |
| 682 | raise ValueError(f"Cannot get {num_blocks} free blocks from the pool") |
| 683 | |
| 684 | ret: list[KVCacheBlock] = self.free_block_queue.popleft_n(num_blocks) |
| 685 | |
| 686 | if self._reuse_watchers: |
| 687 | self._notify_reuse(ret) |
| 688 | |
| 689 | # In order to only iterate the list once, we duplicated code a bit |
| 690 | if self.enable_caching: |
| 691 | for block in ret: |
| 692 | self._maybe_evict_cached_block(block) |
| 693 | assert block.ref_cnt == 0 |
| 694 | block.ref_cnt += 1 |
| 695 | if self.metrics_collector: |
| 696 | self.metrics_collector.on_block_allocated(block) |
| 697 | else: |
| 698 | for block in ret: |
| 699 | assert block.ref_cnt == 0 |
| 700 | block.ref_cnt += 1 |
| 701 | if self.metrics_collector: |
| 702 | self.metrics_collector.on_block_allocated(block) |
| 703 | return ret |
| 704 | |
| 705 | def _notify_reuse(self, blocks: Iterable[KVCacheBlock]) -> None: |
| 706 | for block in blocks: |
| 707 | watcher = self._reuse_watchers.pop(block.block_id, None) |
| 708 | if watcher is not None: |
| 709 | watcher(block) |
| 710 | |
| 711 | def unpin_blocks( |
| 712 | self, |
| 713 | blocks: Iterable[KVCacheBlock], |
| 714 | on_reuse: Callable[[KVCacheBlock], None], |
| 715 | ) -> None: |
| 716 | """Release references to blocks that stay readable until reused. |
| 717 | |
| 718 | The blocks become last-resort eviction candidates regardless of prefix |
| 719 | caching: they join the tail of the free queue and count as free. |
| 720 | ``on_reuse`` fires when ``get_new_blocks`` hands a block out (or the |
| 721 | cache is reset), so the caller can stop reading it. |
| 722 | """ |
| 723 | released: list[KVCacheBlock] = [] |
| 724 | for block in blocks: |
| 725 | assert block.ref_cnt > 0 and not block.is_null |
| 726 | block.ref_cnt -= 1 |
| 727 | self._reuse_watchers[block.block_id] = on_reuse |
| 728 | if block.ref_cnt == 0: |
| 729 | released.append(block) |
| 730 | self.free_block_queue.append_n(released) |
| 731 | |
| 732 | def _maybe_evict_cached_block(self, block: KVCacheBlock) -> bool: |
| 733 | """If a block is cached in `cached_block_hash_to_block`, we reset its hash |
| 734 | metadata and evict it from the cache. |
| 735 | |
| 736 | Args: |
| 737 | block: The block to evict. |
| 738 | |
| 739 | Returns: |
| 740 | True if the block is evicted, False otherwise. |
| 741 | |
| 742 | """ |
| 743 | # Clean up metrics tracking first to prevent leaks |
| 744 | if self.metrics_collector: |
| 745 | self.metrics_collector.on_block_evicted(block) |
| 746 | |
| 747 | evicted_hashes = self._remove_cached_block_hashes(block) |
| 748 | if not evicted_hashes: |
| 749 | # The block doesn't have hash, eviction is not needed |
| 750 | return False |
| 751 | |
| 752 | self._emit_block_removed_events(evicted_hashes) |
| 753 | return True |
| 754 | |
| 755 | def touch(self, blocks: Sequence[KVCacheBlock]) -> None: |
| 756 | """Touch a block increases its reference count by 1, and may remove |
| 757 | the block from the free queue. This is used when a block is hit by |
| 758 | another request with the same prefix. |
| 759 | |
| 760 | Args: |
| 761 | blocks: A list of blocks to touch. |
| 762 | |
| 763 | """ |
| 764 | for block in blocks: |
| 765 | # ref_cnt=0 means this block is in the free list (i.e. eviction |
| 766 | # candidate), so remove it. |
| 767 | if block.ref_cnt == 0 and not block.is_null: |
| 768 | self.free_block_queue.remove(block) |
| 769 | block.ref_cnt += 1 |
| 770 | if self.metrics_collector: |
| 771 | self.metrics_collector.on_block_accessed(block) |
| 772 | |
| 773 | def is_block_writable(self, block: KVCacheBlock) -> bool: |
| 774 | """Return whether a block can be mutated by its sole owner.""" |
| 775 | return not block.is_null and block.ref_cnt == 1 and block.block_hash is None |
| 776 | |
| 777 | def free_blocks(self, ordered_blocks: Iterable[KVCacheBlock]) -> None: |
| 778 | """Free a list of blocks. The blocks should be ordered by their |
| 779 | eviction priority, where the first block will be evicted first. |
| 780 | |
| 781 | Args: |
| 782 | ordered_blocks: A list of blocks to free ordered by their eviction |
| 783 | priority. |
| 784 | |
| 785 | """ |
| 786 | # Identify blocks with hash (LRU cache) and without it (never match APC) |
| 787 | blocks_to_evict_last = [] |
| 788 | blocks_to_evict_first = [] |
| 789 | other_pools: dict[BlockPool, list[KVCacheBlock]] = {} |
| 790 | for block in ordered_blocks: |
| 791 | if block.pool is not None and block.pool is not self: |
| 792 | other_pools.setdefault(block.pool, []).append(block) |
| 793 | continue |
| 794 | block.ref_cnt -= 1 |
| 795 | if block.ref_cnt == 0 and not block.is_null: |
| 796 | if block.block_hash is None or not self.enable_caching: |
| 797 | # LIFO reuse of non-cached blocks for better GPU locality. |
| 798 | blocks_to_evict_first.append(block) |
| 799 | else: |
| 800 | # FIFO reuse of cached blocks for LRU eviction behavior. |
| 801 | blocks_to_evict_last.append(block) |
| 802 | |
| 803 | # Blocks to reuse first are prepended to the front of the free queue. |
| 804 | self.free_block_queue.prepend_n(blocks_to_evict_first) |
| 805 | # Blocks to reuse last are appended to the end of the free queue. |
| 806 | self.free_block_queue.append_n(blocks_to_evict_last) |
| 807 | for pool, blocks in other_pools.items(): |
| 808 | pool.free_blocks(blocks) |
| 809 | |
| 810 | def evict_blocks(self, block_ids: set[int]) -> None: |
| 811 | """Evict blocks from the prefix cache by their block IDs. |
| 812 | |
| 813 | only evicts blocks that are currently cached (have a hash). blocks |
| 814 | with ref_cnt > 0 are not freed from the block pool, only evicted |
| 815 | from the prefix cache hash table. |
| 816 | |
| 817 | Args: |
| 818 | block_ids: Set of block IDs to evict from cache. |
| 819 | |
| 820 | """ |
| 821 | for block_id in block_ids: |
| 822 | assert block_id < len(self.blocks), ( |
| 823 | f"Invalid block_id {block_id} >= {len(self.blocks)}. " |
| 824 | f"This indicates a bug in the KV connector - workers should " |
| 825 | f"only report block IDs that were allocated by the scheduler." |
| 826 | ) |
| 827 | block = self.blocks[block_id] |
| 828 | self._maybe_evict_cached_block(block) |
| 829 | |
| 830 | def reset_prefix_cache(self) -> bool: |
| 831 | """Reset prefix cache. This function may be used in RLHF |
| 832 | flows to invalid prefix caching after the weights are updated, |
| 833 | or used for resetting prefix caching status for benchmarking. |
| 834 | |
| 835 | Returns: |
| 836 | bool: True if the prefix cache is successfully reset, |
| 837 | False otherwise. |
| 838 | |
| 839 | """ |
| 840 | num_used_blocks = self.num_gpu_blocks - self.get_num_free_blocks() |
| 841 | if num_used_blocks != 1: # The null block is always marked as used |
| 842 | logger.warning( |
| 843 | "Failed to reset prefix cache because some " |
| 844 | "blocks (%d) are not freed yet", |
| 845 | num_used_blocks - 1, |
| 846 | ) |
| 847 | return False |
| 848 | |
| 849 | # Remove all hashes so that no new blocks will hit. |
| 850 | self.cached_block_hash_to_block = BlockHashToBlockMap() |
| 851 | self.cached_block_hashes_by_block.clear() |
| 852 | if self._reuse_watchers: |
| 853 | self._notify_reuse( |
| 854 | [self.blocks[block_id] for block_id in list(self._reuse_watchers)] |
| 855 | ) |
| 856 | |
| 857 | # Remove all hashes from all blocks. |
| 858 | for block in self.blocks: |
| 859 | block.reset_hash() |
| 860 | |
| 861 | if self.metrics_collector: |
| 862 | self.metrics_collector.reset() |
| 863 | |
| 864 | logger.info("Successfully reset prefix cache") |
| 865 | |
| 866 | if self.enable_kv_cache_events: |
| 867 | self.kv_event_queue.append(AllBlocksCleared()) |
| 868 | |
| 869 | return True |
| 870 | |
| 871 | def get_num_free_blocks(self) -> int: |
| 872 | """Get the number of free blocks in the pool. |
| 873 | |
| 874 | Returns: |
| 875 | The number of free blocks. |
| 876 | |
| 877 | """ |
| 878 | return self.free_block_queue.num_free_blocks |
| 879 | |
| 880 | def get_usage(self) -> float: |
| 881 | """Get the KV cache usage. |
| 882 | |
| 883 | Returns: |
| 884 | The KV cache usage (between 0.0 and 1.0). |
| 885 | |
| 886 | """ |
| 887 | # Subtract 1 to account for null block. |
| 888 | total_gpu_blocks = self.num_gpu_blocks - 1 |
| 889 | if not total_gpu_blocks: |
| 890 | return 0 |
| 891 | return 1.0 - (self.get_num_free_blocks() / total_gpu_blocks) |
| 892 | |
| 893 | def take_events(self) -> list[KVCacheEvent]: |
| 894 | """Atomically takes all events and clears the queue. |
| 895 | |
| 896 | Returns: |
| 897 | A list of KV cache events. |
| 898 | |
| 899 | """ |
| 900 | if not self.enable_kv_cache_events: |
| 901 | return [] |
| 902 | events = self.kv_event_queue |
| 903 | self.kv_event_queue = [] |
| 904 | return events |