wehub-resource-sync
829 lines
32 KiB
Python
829 lines
32 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
from collections.abc import Iterable, Sequence
|
|
from typing import Any
|
|
|
|
from vllm.distributed.kv_events import (
|
|
MEDIUM_GPU,
|
|
AllBlocksCleared,
|
|
BlockRemoved,
|
|
BlockStored,
|
|
KVCacheEvent,
|
|
)
|
|
from vllm.logger import init_logger
|
|
from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector
|
|
from vllm.v1.core.kv_cache_utils import (
|
|
BlockHash,
|
|
BlockHashWithGroupId,
|
|
ExternalBlockHash,
|
|
FreeKVCacheBlockQueue,
|
|
KVCacheBlock,
|
|
generate_block_hash_extra_keys,
|
|
get_block_hash,
|
|
get_group_id,
|
|
make_block_hash_with_group_id,
|
|
maybe_convert_block_hash,
|
|
resolve_block_hashes,
|
|
)
|
|
from vllm.v1.request import Request
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class BlockHashToBlockMap:
|
|
"""
|
|
Cache of blocks that are used for prefix caching. It caches blocks
|
|
from hash directly to a block or multiple blocks
|
|
(i.e. {block_hash: KVCacheBlocks})
|
|
- Mostly block_hash maps to a single KVCacheBlock, and KVCacheBlocks
|
|
would simply be a KVCacheBlock.
|
|
- Otherwise, KVCacheBlocks is a dict from {block_id: KVCacheBlock}
|
|
|
|
A cached block is a full block with a block hash that can be used
|
|
for prefix caching.
|
|
The cached block may be used by running requests or in the
|
|
free_block_queue that could potentially be evicted.
|
|
|
|
NOTE #1: We currently don't de-duplicate the blocks in the cache,
|
|
meaning that if a block becomes full and is cached, we don't check
|
|
if there is already an identical block in the cache. This is because
|
|
we want to make sure the allocated block IDs won't change so that
|
|
block tables are append-only.
|
|
NOTE #2: The union type is introduced in order to reduce GC costs
|
|
from the inner dict.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self._cache: dict[
|
|
BlockHashWithGroupId, KVCacheBlock | dict[int, KVCacheBlock]
|
|
] = {}
|
|
|
|
def get_one_block(self, key: BlockHashWithGroupId) -> KVCacheBlock | None:
|
|
"""
|
|
Gets any block with the given block hash key.
|
|
"""
|
|
blocks = self._cache.get(key)
|
|
if blocks is not None:
|
|
if isinstance(blocks, KVCacheBlock):
|
|
return blocks
|
|
if isinstance(blocks, dict):
|
|
return next(iter(blocks.values()))
|
|
self._unexpected_blocks_type(blocks)
|
|
return None
|
|
|
|
def contain(self, key: BlockHashWithGroupId, block_id: int) -> bool:
|
|
"""
|
|
Checks whether the key maps to the given block ID.
|
|
"""
|
|
blocks = self._cache.get(key)
|
|
if blocks is None:
|
|
return False
|
|
if isinstance(blocks, KVCacheBlock):
|
|
return blocks.block_id == block_id
|
|
if isinstance(blocks, dict):
|
|
return block_id in blocks
|
|
self._unexpected_blocks_type(blocks)
|
|
return False
|
|
|
|
def insert(self, key: BlockHashWithGroupId, block: KVCacheBlock) -> None:
|
|
"""
|
|
Inserts the KVCacheBlock to the cache
|
|
"""
|
|
blocks = self._cache.get(key)
|
|
if blocks is None:
|
|
# When key is not found, attach a single block to the key
|
|
self._cache[key] = block
|
|
elif isinstance(blocks, KVCacheBlock):
|
|
# If there's a block with the same key, merge the original block
|
|
# and the new block into a dict
|
|
self._cache[key] = {blocks.block_id: blocks, block.block_id: block}
|
|
elif isinstance(blocks, dict):
|
|
# If it's already a dict, simply insert the block
|
|
blocks[block.block_id] = block
|
|
else:
|
|
self._unexpected_blocks_type(blocks)
|
|
|
|
def pop(self, key: BlockHashWithGroupId, block_id: int) -> KVCacheBlock | None:
|
|
"""
|
|
Checks if block_hash exists and pop block_id from the cache
|
|
"""
|
|
blocks = self._cache.pop(key, None)
|
|
if blocks is None:
|
|
# block_hash not found in the cache
|
|
return None
|
|
# TODO(Jialin): If key is found, block_id should always present
|
|
# in blocks. We currently keep the original behaviour for safety.
|
|
#
|
|
# Will add block_id == blocks.block_id assertion and
|
|
# use del blocks[block_id] instead as followup.
|
|
if isinstance(blocks, KVCacheBlock):
|
|
if blocks.block_id == block_id:
|
|
return blocks
|
|
# If the single block ID doesn't match, we should put the
|
|
# block back (it should happen rarely)
|
|
self._cache[key] = blocks
|
|
return None
|
|
if isinstance(blocks, dict):
|
|
# Try to pop block_id from the block dict, and if dict still
|
|
# contain blocks, put back to the cache.
|
|
block = blocks.pop(block_id, None)
|
|
if len(blocks) > 0:
|
|
self._cache[key] = blocks
|
|
return block
|
|
self._unexpected_blocks_type(blocks)
|
|
return None
|
|
|
|
def __len__(self) -> int:
|
|
return len(self._cache)
|
|
|
|
def _unexpected_blocks_type(self, blocks: Any) -> None:
|
|
raise AssertionError(f"Invalid KV cache block type {type(blocks)}")
|
|
|
|
|
|
class BlockPool:
|
|
"""BlockPool that manages KVCacheBlocks.
|
|
It provides methods to allocate, free and cache the kv cache blocks. The
|
|
free_block_queue stores the free blocks in eviction order to enable
|
|
allocation, free, and cache eviction. The cached_block_hash_to_block
|
|
maps between block hash and cached block to support finding cached blocks
|
|
by their block hash.
|
|
|
|
Args:
|
|
num_gpu_blocks: The number of blocks in the pool.
|
|
enable_caching: Whether to enable prefix caching.
|
|
hash_block_size: The block size of which the block hashes are computed.
|
|
The actual block size usually equals hash_block_size, but in cases
|
|
where different KV cache groups have different block sizes, the
|
|
actual block size can be a multiple of hash_block_size.
|
|
enable_kv_cache_events: Whether to enable kv cache events.
|
|
metrics_collector: Optional metrics collector for tracking block residency.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
num_gpu_blocks: int,
|
|
enable_caching: bool,
|
|
hash_block_size: int,
|
|
enable_kv_cache_events: bool = False,
|
|
metrics_collector: KVCacheMetricsCollector | None = None,
|
|
):
|
|
assert isinstance(num_gpu_blocks, int) and num_gpu_blocks > 0
|
|
self.num_gpu_blocks = num_gpu_blocks
|
|
self.enable_caching = enable_caching
|
|
self.hash_block_size = hash_block_size
|
|
# All kv-cache blocks.
|
|
self.blocks: list[KVCacheBlock] = [
|
|
KVCacheBlock(idx) for idx in range(num_gpu_blocks)
|
|
]
|
|
# Free block queue that constructs and manipulates a doubly linked
|
|
# list of free blocks (including eviction candidates when caching is
|
|
# enabled).
|
|
self.free_block_queue = FreeKVCacheBlockQueue(self.blocks)
|
|
|
|
# Cache for block lookup
|
|
self.cached_block_hash_to_block: BlockHashToBlockMap = BlockHashToBlockMap()
|
|
self.cached_block_hashes_by_block: dict[int, set[BlockHashWithGroupId]] = {}
|
|
|
|
# To represent a placeholder block with block_id=0.
|
|
# The ref_cnt of null_block is not maintained, needs special care to
|
|
# avoid freeing it.
|
|
self.null_block = self.free_block_queue.popleft()
|
|
self.null_block.is_null = True
|
|
|
|
self.enable_kv_cache_events = enable_kv_cache_events
|
|
self.kv_event_queue: list[KVCacheEvent] = []
|
|
|
|
self.metrics_collector = metrics_collector
|
|
|
|
def get_cached_block(
|
|
self, block_hash: BlockHash, kv_cache_group_ids: list[int]
|
|
) -> list[KVCacheBlock] | None:
|
|
"""Get the cached block by the block hash for each group in
|
|
`kv_cache_group_ids`, or None if cache miss for any group.
|
|
If there are duplicated blocks, we return the first block in the cache.
|
|
|
|
Args:
|
|
block_hash: The hash value of the block.
|
|
kv_cache_group_ids: The ids of the KV cache groups.
|
|
|
|
Returns:
|
|
The cached blocks if exists, or None.
|
|
"""
|
|
cached_blocks = []
|
|
for group_id in kv_cache_group_ids:
|
|
block_hash_with_group_id = make_block_hash_with_group_id(
|
|
block_hash, group_id
|
|
)
|
|
block = self.cached_block_hash_to_block.get_one_block(
|
|
block_hash_with_group_id
|
|
)
|
|
if not block:
|
|
return None
|
|
cached_blocks.append(block)
|
|
return cached_blocks
|
|
|
|
def cache_full_blocks(
|
|
self,
|
|
request: Request,
|
|
blocks: list[KVCacheBlock],
|
|
num_cached_blocks: int,
|
|
num_full_blocks: int,
|
|
block_size: int,
|
|
kv_cache_group_id: int,
|
|
block_mask: list[bool] | None = None,
|
|
) -> None:
|
|
"""Cache a list of full blocks for prefix caching.
|
|
This function takes a list of blocks that will have their block hash
|
|
metadata to be updated and cached. Given a request, it updates the
|
|
metadata for each block and caching it in the
|
|
`cached_block_hash_to_block`.
|
|
The block hashes values are computed by the Request object immediately
|
|
when it is created and when new tokens are appended.
|
|
|
|
Args:
|
|
request: The request to cache the blocks.
|
|
blocks: All blocks in the request.
|
|
num_cached_blocks: The number of blocks that are already cached.
|
|
num_full_blocks: The number of blocks that are full and should
|
|
be cached after this function.
|
|
block_size: Number of tokens in each block.
|
|
kv_cache_group_id: The id of the KV cache group.
|
|
block_mask: Optional mask aligned with
|
|
``blocks[num_cached_blocks:num_full_blocks]``. When provided,
|
|
blocks where the mask is False are skipped (treated like null
|
|
blocks). Used by groups whose ``find_longest_cache_hit`` only
|
|
consults a subset of blocks (e.g. SWA tail-window), so blocks
|
|
that can never serve a hit stay out of the prefix-cache hash
|
|
map.
|
|
"""
|
|
if num_cached_blocks >= num_full_blocks:
|
|
return
|
|
new_full_blocks = blocks[num_cached_blocks:num_full_blocks]
|
|
assert block_mask is None or len(block_mask) == len(new_full_blocks)
|
|
block_hashes = resolve_block_hashes(
|
|
request.block_hashes, self.hash_block_size, block_size
|
|
)
|
|
|
|
new_block_hashes = block_hashes[num_cached_blocks:]
|
|
new_hashes: list[ExternalBlockHash] | None = (
|
|
[] if self.enable_kv_cache_events else None
|
|
)
|
|
for i, blk in enumerate(new_full_blocks):
|
|
# Some blocks may be null or masked out when enabling sparse attention
|
|
# like sliding window attention, or Mamba models with prefix-caching
|
|
# in align mode. We skip null blocks here.
|
|
if blk.is_null or (block_mask is not None and not block_mask[i]):
|
|
continue
|
|
block_hash = new_block_hashes[i]
|
|
num_hash_tokens = (num_cached_blocks + i + 1) * block_size
|
|
|
|
# Update and added the full block to the cache.
|
|
block_hash_with_group_id = make_block_hash_with_group_id(
|
|
block_hash, kv_cache_group_id
|
|
)
|
|
if blk.block_hash is not None:
|
|
# The only valid case where a "new full block" already has a
|
|
# hash is partial->full promotion of the same cache block.
|
|
assert (
|
|
blk.block_hash_num_tokens is not None
|
|
and blk.block_hash_num_tokens < num_hash_tokens
|
|
)
|
|
removed_hashes = self._remove_cached_block_hashes(blk)
|
|
self._emit_block_removed_events(removed_hashes)
|
|
self._insert_block_hash(
|
|
block_hash_with_group_id,
|
|
blk,
|
|
num_tokens=num_hash_tokens,
|
|
)
|
|
if new_hashes is not None:
|
|
new_hashes.append(maybe_convert_block_hash(block_hash))
|
|
|
|
if self.enable_kv_cache_events:
|
|
if num_cached_blocks == 0:
|
|
parent_block_hash: ExternalBlockHash | None = None
|
|
else:
|
|
parent_block_hash = maybe_convert_block_hash(
|
|
block_hashes[num_cached_blocks - 1]
|
|
)
|
|
|
|
# Calculate token range for the blocks being cached
|
|
start_token_idx = num_cached_blocks * block_size
|
|
end_token_idx = num_full_blocks * block_size
|
|
|
|
# Generate extra keys for each block individually.
|
|
# Each block may have different extra_keys (e.g., different MM
|
|
# features, or cache_salt only for the first block).
|
|
# Skip null/masked-out blocks to match the length of new_hashes.
|
|
extra_keys_list: list[tuple[Any, ...] | None] = []
|
|
curr_mm_idx = 0
|
|
for i in range(num_cached_blocks, num_full_blocks):
|
|
if blocks[i].is_null:
|
|
continue
|
|
if block_mask is not None and not block_mask[i - num_cached_blocks]:
|
|
continue
|
|
block_start = i * block_size
|
|
block_end = block_start + block_size
|
|
extra_keys, curr_mm_idx = generate_block_hash_extra_keys(
|
|
request, block_start, block_end, curr_mm_idx
|
|
)
|
|
extra_keys_list.append(extra_keys)
|
|
|
|
self.kv_event_queue.append(
|
|
self._build_block_stored_event(
|
|
request,
|
|
block_hashes=new_hashes,
|
|
parent_block_hash=parent_block_hash,
|
|
start_token_idx=start_token_idx,
|
|
end_token_idx=end_token_idx,
|
|
block_size=block_size,
|
|
kv_cache_group_id=kv_cache_group_id,
|
|
extra_keys_list=extra_keys_list,
|
|
)
|
|
)
|
|
|
|
def _build_block_stored_event(
|
|
self,
|
|
request: Request,
|
|
block_hashes: list[ExternalBlockHash] | None,
|
|
parent_block_hash: ExternalBlockHash | None,
|
|
start_token_idx: int,
|
|
end_token_idx: int,
|
|
block_size: int,
|
|
kv_cache_group_id: int,
|
|
extra_keys_list: list[tuple[Any, ...] | None],
|
|
) -> BlockStored:
|
|
"""Build a ``BlockStored`` KV event for ``request``.
|
|
|
|
Shared by ``cache_full_blocks`` (newly cached blocks) and
|
|
``emit_cached_block_events`` (prefix-cache-reused blocks) so both emit
|
|
identical event shapes for downstream consumers.
|
|
"""
|
|
return BlockStored(
|
|
block_hashes=block_hashes,
|
|
parent_block_hash=parent_block_hash,
|
|
token_ids=request.all_token_ids[start_token_idx:end_token_idx],
|
|
block_size=block_size,
|
|
lora_id=request.lora_request.adapter_id if request.lora_request else None,
|
|
medium=MEDIUM_GPU,
|
|
lora_name=request.lora_request.name if request.lora_request else None,
|
|
extra_keys=extra_keys_list if extra_keys_list else None,
|
|
group_idx=kv_cache_group_id,
|
|
)
|
|
|
|
def emit_cached_block_events(
|
|
self,
|
|
request: Request,
|
|
num_cached_blocks: int,
|
|
block_size: int,
|
|
kv_cache_group_id: int,
|
|
) -> None:
|
|
"""Generate BlockStored events for blocks reused from prefix cache.
|
|
|
|
Unlike cache_full_blocks(), this does NOT modify block state —
|
|
the blocks are already cached. It only generates events so that
|
|
external consumers (e.g. gateway) can learn about reused blocks.
|
|
|
|
Args:
|
|
request: The request whose prefix cache blocks were reused.
|
|
num_cached_blocks: Number of blocks that were cache hits.
|
|
block_size: Number of tokens per block.
|
|
kv_cache_group_id: The KV cache group ID.
|
|
"""
|
|
if not self.enable_kv_cache_events or num_cached_blocks == 0:
|
|
return
|
|
|
|
block_hashes = resolve_block_hashes(
|
|
request.block_hashes, self.hash_block_size, block_size
|
|
)
|
|
|
|
# Collect external hashes and extra_keys for cached blocks.
|
|
cached_hashes: list[ExternalBlockHash] = []
|
|
extra_keys_list: list[tuple[Any, ...] | None] = []
|
|
curr_mm_idx = 0
|
|
for i in range(num_cached_blocks):
|
|
block_start = i * block_size
|
|
block_end = block_start + block_size
|
|
cached_hashes.append(maybe_convert_block_hash(block_hashes[i]))
|
|
extra_keys, curr_mm_idx = generate_block_hash_extra_keys(
|
|
request, block_start, block_end, curr_mm_idx
|
|
)
|
|
extra_keys_list.append(extra_keys)
|
|
|
|
if not cached_hashes:
|
|
return
|
|
|
|
# Prefix-cache hits always form a contiguous prefix starting at block 0,
|
|
# so the first (and thus the whole group's) parent block hash is None.
|
|
parent_block_hash: ExternalBlockHash | None = None
|
|
start_token_idx = 0
|
|
end_token_idx = num_cached_blocks * block_size
|
|
|
|
logger.debug(
|
|
"EmitCachedBlock event: block_size=%d, "
|
|
"num_cached_blocks=%d, parent_block_hash=%s, "
|
|
"token_ids_len=%d, group_idx=%s",
|
|
block_size,
|
|
num_cached_blocks,
|
|
parent_block_hash,
|
|
len(request.all_token_ids[start_token_idx:end_token_idx]),
|
|
kv_cache_group_id,
|
|
)
|
|
|
|
self.kv_event_queue.append(
|
|
self._build_block_stored_event(
|
|
request,
|
|
block_hashes=cached_hashes,
|
|
parent_block_hash=parent_block_hash,
|
|
start_token_idx=start_token_idx,
|
|
end_token_idx=end_token_idx,
|
|
block_size=block_size,
|
|
kv_cache_group_id=kv_cache_group_id,
|
|
extra_keys_list=extra_keys_list,
|
|
)
|
|
)
|
|
|
|
def cache_partial_block(
|
|
self,
|
|
request: Request,
|
|
block: KVCacheBlock,
|
|
num_tokens: int,
|
|
kv_cache_group_id: int,
|
|
block_size: int,
|
|
) -> BlockHashWithGroupId | None:
|
|
"""Register a partial prefix-cache entry for an existing block.
|
|
|
|
Prefix-cache keys normally identify full cache blocks. A partial entry
|
|
makes an existing cache block reachable from a fine-grained prefix
|
|
boundary inside that block without allocating or copying a new
|
|
``KVCacheBlock``.
|
|
|
|
The partial entry is lookup metadata owned by ``block``. If ``block``
|
|
has no primary hash, the key becomes its primary hash. If the block
|
|
already has a primary hash, the partial entry is tracked in
|
|
``cached_block_hashes_by_block`` so eviction, reset, and promotion can
|
|
remove every hash key that points to the block.
|
|
|
|
Args:
|
|
request: Request whose token IDs and block hashes define the
|
|
partial entry.
|
|
block: Existing cache block to make reachable from the partial
|
|
prefix boundary.
|
|
num_tokens: Prefix length represented by the partial entry. It
|
|
must be a positive multiple of ``self.hash_block_size`` and
|
|
cannot exceed the request's computed block hashes.
|
|
kv_cache_group_id: KV cache group that owns the partial entry.
|
|
block_size: Cache block size for the owning group. The partial
|
|
entry hash itself is always the prefix-chain hash at
|
|
``num_tokens``; ``block_size`` is used to assert that the
|
|
entry is partial within the owning cache block.
|
|
|
|
Returns:
|
|
The hash key with group ID if a partial entry can be registered;
|
|
otherwise ``None`` for null blocks.
|
|
"""
|
|
if block.is_null:
|
|
return None
|
|
|
|
assert block_size > self.hash_block_size
|
|
assert block_size % self.hash_block_size == 0
|
|
assert num_tokens % block_size != 0
|
|
block_hash = self._get_partial_block_hash(request, num_tokens)
|
|
num_hash_blocks = num_tokens // self.hash_block_size
|
|
block_hash_with_group_id = make_block_hash_with_group_id(
|
|
block_hash, kv_cache_group_id
|
|
)
|
|
already_cached = block.block_hash == block_hash_with_group_id or (
|
|
self.cached_block_hash_to_block.contain(
|
|
block_hash_with_group_id, block.block_id
|
|
)
|
|
)
|
|
if (
|
|
not already_cached
|
|
and block.block_hash is not None
|
|
and block.block_hash_num_tokens is not None
|
|
and block.block_hash_num_tokens < num_hash_blocks * self.hash_block_size
|
|
):
|
|
removed_hashes = self._remove_cached_block_hashes(block)
|
|
self._emit_block_removed_events(removed_hashes)
|
|
self._insert_block_hash(
|
|
block_hash_with_group_id,
|
|
block,
|
|
num_tokens=num_hash_blocks * self.hash_block_size,
|
|
)
|
|
if self.enable_kv_cache_events and not already_cached:
|
|
parent_hash, block_start = self._get_partial_block_parent_hash_and_start(
|
|
request, num_tokens
|
|
)
|
|
parent_block_hash = (
|
|
maybe_convert_block_hash(parent_hash)
|
|
if parent_hash is not None
|
|
else None
|
|
)
|
|
block_end = num_tokens
|
|
curr_mm_idx = -1 if block_start > 0 else 0
|
|
extra_keys, _ = generate_block_hash_extra_keys(
|
|
request, block_start, block_end, curr_mm_idx
|
|
)
|
|
self.kv_event_queue.append(
|
|
BlockStored(
|
|
block_hashes=[maybe_convert_block_hash(block_hash)],
|
|
parent_block_hash=parent_block_hash,
|
|
token_ids=request.all_token_ids[block_start:block_end],
|
|
block_size=block_end - block_start,
|
|
lora_id=request.lora_request.adapter_id
|
|
if request.lora_request
|
|
else None,
|
|
medium=MEDIUM_GPU,
|
|
lora_name=request.lora_request.name
|
|
if request.lora_request
|
|
else None,
|
|
extra_keys=[extra_keys],
|
|
group_idx=kv_cache_group_id,
|
|
)
|
|
)
|
|
return block_hash_with_group_id
|
|
|
|
def _get_partial_block_hash(
|
|
self,
|
|
request: Request,
|
|
num_tokens: int,
|
|
) -> BlockHash:
|
|
assert num_tokens % self.hash_block_size == 0
|
|
num_hash_blocks = num_tokens // self.hash_block_size
|
|
assert 0 < num_hash_blocks <= len(request.block_hashes)
|
|
|
|
# Each hash_block_size hash chains over its full prefix, so the partial
|
|
# entry for any group block size is the hash at that prefix boundary.
|
|
return request.block_hashes[num_hash_blocks - 1]
|
|
|
|
def _get_partial_block_parent_hash_and_start(
|
|
self,
|
|
request: Request,
|
|
num_tokens: int,
|
|
) -> tuple[BlockHash | None, int]:
|
|
num_hash_blocks = num_tokens // self.hash_block_size
|
|
parent_hash = (
|
|
request.block_hashes[num_hash_blocks - 2] if num_hash_blocks > 1 else None
|
|
)
|
|
block_start = (num_hash_blocks - 1) * self.hash_block_size
|
|
return parent_hash, block_start
|
|
|
|
def _remove_cached_block_hashes(
|
|
self,
|
|
block: KVCacheBlock,
|
|
) -> list[BlockHashWithGroupId]:
|
|
block_hashes: list[BlockHashWithGroupId] = []
|
|
if block.block_hash is not None:
|
|
block_hashes.append(block.block_hash)
|
|
block_hashes.extend(self.cached_block_hashes_by_block.pop(block.block_id, ()))
|
|
if not block_hashes:
|
|
return []
|
|
|
|
removed_hashes: list[BlockHashWithGroupId] = []
|
|
for block_hash in block_hashes:
|
|
if (
|
|
self.cached_block_hash_to_block.pop(block_hash, block.block_id)
|
|
is not None
|
|
):
|
|
removed_hashes.append(block_hash)
|
|
block.reset_hash()
|
|
return removed_hashes
|
|
|
|
def _emit_block_removed_events(
|
|
self,
|
|
block_hashes: list[BlockHashWithGroupId],
|
|
) -> None:
|
|
if not self.enable_kv_cache_events:
|
|
return
|
|
for block_hash in block_hashes:
|
|
self.kv_event_queue.append(
|
|
BlockRemoved(
|
|
block_hashes=[maybe_convert_block_hash(get_block_hash(block_hash))],
|
|
medium=MEDIUM_GPU,
|
|
group_idx=get_group_id(block_hash),
|
|
)
|
|
)
|
|
|
|
def _insert_block_hash(
|
|
self,
|
|
block_hash_with_group_id: BlockHashWithGroupId,
|
|
block: KVCacheBlock,
|
|
num_tokens: int | None,
|
|
) -> None:
|
|
if block.block_hash == block_hash_with_group_id:
|
|
return
|
|
|
|
if self.cached_block_hash_to_block.contain(
|
|
block_hash_with_group_id, block.block_id
|
|
):
|
|
return
|
|
|
|
if block.block_hash is None:
|
|
block.set_block_hash(block_hash_with_group_id, num_tokens=num_tokens)
|
|
else:
|
|
self.cached_block_hashes_by_block.setdefault(block.block_id, set()).add(
|
|
block_hash_with_group_id
|
|
)
|
|
self.cached_block_hash_to_block.insert(block_hash_with_group_id, block)
|
|
|
|
def move_block_hashes(
|
|
self,
|
|
src_block: KVCacheBlock,
|
|
dst_block: KVCacheBlock,
|
|
) -> None:
|
|
"""Re-point ``src_block``'s prefix-cache entries to ``dst_block``.
|
|
|
|
Used when the request owning ``src_block`` keeps writing into it
|
|
: the prefix cache holds a private copy (``dst_block``)
|
|
under the same hashes instead. Entries stay live; no events emitted.
|
|
"""
|
|
assert dst_block.block_hash is None
|
|
assert dst_block.block_id not in self.cached_block_hashes_by_block
|
|
num_tokens = src_block.block_hash_num_tokens
|
|
for block_hash in self._remove_cached_block_hashes(src_block):
|
|
# `num_tokens` only applies to the first (primary) insertion.
|
|
self._insert_block_hash(block_hash, dst_block, num_tokens=num_tokens)
|
|
|
|
def get_new_blocks(self, num_blocks: int) -> list[KVCacheBlock]:
|
|
"""Get new blocks from the free block pool.
|
|
|
|
Note that we do not check block cache in this function.
|
|
|
|
Args:
|
|
num_blocks: The number of blocks to allocate.
|
|
|
|
Returns:
|
|
A list of new block.
|
|
"""
|
|
if num_blocks > self.get_num_free_blocks():
|
|
raise ValueError(f"Cannot get {num_blocks} free blocks from the pool")
|
|
|
|
ret: list[KVCacheBlock] = self.free_block_queue.popleft_n(num_blocks)
|
|
|
|
# In order to only iterate the list once, we duplicated code a bit
|
|
if self.enable_caching:
|
|
for block in ret:
|
|
self._maybe_evict_cached_block(block)
|
|
assert block.ref_cnt == 0
|
|
block.ref_cnt += 1
|
|
if self.metrics_collector:
|
|
self.metrics_collector.on_block_allocated(block)
|
|
else:
|
|
for block in ret:
|
|
assert block.ref_cnt == 0
|
|
block.ref_cnt += 1
|
|
if self.metrics_collector:
|
|
self.metrics_collector.on_block_allocated(block)
|
|
return ret
|
|
|
|
def _maybe_evict_cached_block(self, block: KVCacheBlock) -> bool:
|
|
"""
|
|
If a block is cached in `cached_block_hash_to_block`, we reset its hash
|
|
metadata and evict it from the cache.
|
|
|
|
Args:
|
|
block: The block to evict.
|
|
|
|
Returns:
|
|
True if the block is evicted, False otherwise.
|
|
"""
|
|
# Clean up metrics tracking first to prevent leaks
|
|
if self.metrics_collector:
|
|
self.metrics_collector.on_block_evicted(block)
|
|
|
|
evicted_hashes = self._remove_cached_block_hashes(block)
|
|
if not evicted_hashes:
|
|
# The block doesn't have hash, eviction is not needed
|
|
return False
|
|
|
|
self._emit_block_removed_events(evicted_hashes)
|
|
return True
|
|
|
|
def touch(self, blocks: Sequence[KVCacheBlock]) -> None:
|
|
"""Touch a block increases its reference count by 1, and may remove
|
|
the block from the free queue. This is used when a block is hit by
|
|
another request with the same prefix.
|
|
|
|
Args:
|
|
blocks: A list of blocks to touch.
|
|
"""
|
|
for block in blocks:
|
|
# ref_cnt=0 means this block is in the free list (i.e. eviction
|
|
# candidate), so remove it.
|
|
if block.ref_cnt == 0 and not block.is_null:
|
|
self.free_block_queue.remove(block)
|
|
block.ref_cnt += 1
|
|
if self.metrics_collector:
|
|
self.metrics_collector.on_block_accessed(block)
|
|
|
|
def free_blocks(self, ordered_blocks: Iterable[KVCacheBlock]) -> None:
|
|
"""Free a list of blocks. The blocks should be ordered by their
|
|
eviction priority, where the first block will be evicted first.
|
|
|
|
Args:
|
|
ordered_blocks: A list of blocks to free ordered by their eviction
|
|
priority.
|
|
"""
|
|
# Identify blocks with hash (LRU cache) and without it (will never match in APC)
|
|
blocks_with_hash = []
|
|
blocks_without_hash = []
|
|
for block in ordered_blocks:
|
|
block.ref_cnt -= 1
|
|
if block.ref_cnt == 0 and not block.is_null:
|
|
if block.block_hash is None:
|
|
blocks_without_hash.append(block)
|
|
else:
|
|
blocks_with_hash.append(block)
|
|
|
|
# Blocks without hash always get evicted first - prepend them last to the tail
|
|
self.free_block_queue.prepend_n(blocks_without_hash)
|
|
self.free_block_queue.append_n(blocks_with_hash)
|
|
|
|
def evict_blocks(self, block_ids: set[int]) -> None:
|
|
"""evict blocks from the prefix cache by their block IDs.
|
|
|
|
only evicts blocks that are currently cached (have a hash). blocks
|
|
with ref_cnt > 0 are not freed from the block pool, only evicted
|
|
from the prefix cache hash table.
|
|
|
|
Args:
|
|
block_ids: Set of block IDs to evict from cache.
|
|
"""
|
|
for block_id in block_ids:
|
|
assert block_id < len(self.blocks), (
|
|
f"Invalid block_id {block_id} >= {len(self.blocks)}. "
|
|
f"This indicates a bug in the KV connector - workers should "
|
|
f"only report block IDs that were allocated by the scheduler."
|
|
)
|
|
block = self.blocks[block_id]
|
|
self._maybe_evict_cached_block(block)
|
|
|
|
def reset_prefix_cache(self) -> bool:
|
|
"""Reset prefix cache. This function may be used in RLHF
|
|
flows to invalid 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.
|
|
"""
|
|
num_used_blocks = self.num_gpu_blocks - self.get_num_free_blocks()
|
|
if num_used_blocks != 1: # The null block is always marked as used
|
|
logger.warning(
|
|
"Failed to reset prefix cache because some "
|
|
"blocks (%d) are not freed yet",
|
|
num_used_blocks - 1,
|
|
)
|
|
return False
|
|
|
|
# Remove all hashes so that no new blocks will hit.
|
|
self.cached_block_hash_to_block = BlockHashToBlockMap()
|
|
self.cached_block_hashes_by_block.clear()
|
|
|
|
# Remove all hashes from all blocks.
|
|
for block in self.blocks:
|
|
block.reset_hash()
|
|
|
|
if self.metrics_collector:
|
|
self.metrics_collector.reset()
|
|
|
|
logger.info("Successfully reset prefix cache")
|
|
|
|
if self.enable_kv_cache_events:
|
|
self.kv_event_queue.append(AllBlocksCleared())
|
|
|
|
return True
|
|
|
|
def get_num_free_blocks(self) -> int:
|
|
"""Get the number of free blocks in the pool.
|
|
|
|
Returns:
|
|
The number of free blocks.
|
|
"""
|
|
return self.free_block_queue.num_free_blocks
|
|
|
|
def get_usage(self) -> float:
|
|
"""Get the KV cache usage.
|
|
|
|
Returns:
|
|
The KV cache usage (between 0.0 and 1.0).
|
|
"""
|
|
|
|
# Subtract 1 to account for null block.
|
|
total_gpu_blocks = self.num_gpu_blocks - 1
|
|
if not total_gpu_blocks:
|
|
return 0
|
|
return 1.0 - (self.get_num_free_blocks() / total_gpu_blocks)
|
|
|
|
def take_events(self) -> list[KVCacheEvent]:
|
|
"""Atomically takes all events and clears the queue.
|
|
|
|
Returns:
|
|
A list of KV cache events.
|
|
"""
|
|
if not self.enable_kv_cache_events:
|
|
return []
|
|
events = self.kv_event_queue
|
|
self.kv_event_queue = []
|
|
return events
|