wehub-resource-sync
1733 lines
58 KiB
Python
1733 lines
58 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# Adapted from vllm-mlx (https://github.com/vllm-project/vllm-mlx).
|
|
"""
|
|
Paged KV Cache Manager for oMLX.
|
|
|
|
This module implements block-based paged KV cache management following vLLM's
|
|
architecture (vllm/v1/core/block_pool.py), adapted for MLX on Apple Silicon.
|
|
|
|
Key components:
|
|
- KVCacheBlock: Metadata for each cache block with doubly linked list pointers
|
|
- FreeKVCacheBlockQueue: O(1) doubly linked list for LRU block allocation
|
|
- BlockHashToBlockMap: Hash-to-block cache for prefix caching
|
|
- PagedCacheManager: Main manager with block allocation, prefix caching, and COW
|
|
|
|
Features:
|
|
- Block-based allocation (configurable tokens per block)
|
|
- Reference counting for shared blocks
|
|
- Copy-on-Write (COW) for efficient prefix sharing
|
|
- LRU eviction using doubly linked list (O(1) operations)
|
|
- Chain hashing for prefix caching (hash depends on parent block)
|
|
|
|
Reference: vLLM v1 - vllm/v1/core/block_pool.py, vllm/v1/core/kv_cache_utils.py
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import logging
|
|
import threading
|
|
import time
|
|
from collections.abc import Iterable
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Dict, List, NewType, Optional, Tuple
|
|
|
|
from .interface import CacheManager
|
|
from .stats import BaseCacheStats, PagedCacheStats
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Type alias for block hash (content-based hash for prefix caching)
|
|
BlockHash = NewType("BlockHash", bytes)
|
|
|
|
|
|
def resolve_block_extra_keys(
|
|
block_end: int,
|
|
extra_keys: Optional[Tuple[Any, ...]] = None,
|
|
extra_key_token_start: Optional[int] = None,
|
|
extra_key_ranges: Optional[List[Tuple[int, Tuple[Any, ...]]]] = None,
|
|
) -> Optional[Tuple[Any, ...]]:
|
|
"""Resolve which cache key salt applies to a block ending at ``block_end``.
|
|
|
|
``extra_key_ranges`` takes precedence over ``extra_keys`` and is intended
|
|
for segmented VLM cache keying.
|
|
|
|
``extra_key_ranges`` must be sorted ascending by start position. The
|
|
function uses early break on the first non-matching entry, so unsorted
|
|
input will produce incorrect results.
|
|
"""
|
|
if extra_key_ranges:
|
|
assert all(
|
|
extra_key_ranges[i][0] <= extra_key_ranges[i + 1][0]
|
|
for i in range(len(extra_key_ranges) - 1)
|
|
), "extra_key_ranges must be sorted ascending by start position"
|
|
selected = None
|
|
for start, keys in extra_key_ranges:
|
|
if block_end > start:
|
|
selected = keys
|
|
else:
|
|
break
|
|
return selected
|
|
if extra_keys is not None and (
|
|
extra_key_token_start is None or block_end > extra_key_token_start
|
|
):
|
|
return extra_keys
|
|
return None
|
|
|
|
|
|
def compute_block_hash(
|
|
parent_hash: Optional[BlockHash],
|
|
token_ids: List[int],
|
|
extra_keys: Optional[Tuple[Any, ...]] = None,
|
|
model_name: Optional[str] = None,
|
|
) -> BlockHash:
|
|
"""
|
|
Compute hash for a block based on its content and parent block.
|
|
|
|
This enables prefix caching by creating a chain of hashes where
|
|
each block's hash depends on all previous blocks (similar to vLLM).
|
|
|
|
Args:
|
|
parent_hash: Hash of the previous block, or None for first block
|
|
token_ids: Token IDs in this block
|
|
extra_keys: Additional keys (e.g., LoRA, multimodal)
|
|
model_name: Model name for cache isolation between different models
|
|
|
|
Returns:
|
|
Content-based hash for this block
|
|
"""
|
|
hasher = hashlib.sha256()
|
|
|
|
# Include model name first to isolate caches between different models
|
|
if model_name:
|
|
hasher.update(model_name.encode("utf-8"))
|
|
|
|
# Include parent hash for chain
|
|
if parent_hash:
|
|
hasher.update(parent_hash)
|
|
else:
|
|
# Use fixed seed for reproducibility
|
|
hasher.update(b"omlx-root")
|
|
|
|
# Include token content
|
|
hasher.update(bytes(str(tuple(token_ids)), "utf-8"))
|
|
|
|
# Include extra keys if present
|
|
if extra_keys:
|
|
hasher.update(bytes(str(extra_keys), "utf-8"))
|
|
|
|
return BlockHash(hasher.digest())
|
|
|
|
|
|
# =============================================================================
|
|
# KVCacheBlock - Following vLLM's design
|
|
# =============================================================================
|
|
|
|
@dataclass
|
|
class CacheBlock:
|
|
"""
|
|
KV cache block metadata following vLLM's design.
|
|
|
|
Each block represents a fixed number of tokens (block_size) worth
|
|
of KV cache data. Blocks can be shared across requests via
|
|
reference counting for prefix caching.
|
|
|
|
NOTE: In paged SSD-only mode, blocks do NOT store cache_data in GPU memory.
|
|
All KV cache data is stored on paged SSD via PagedSSDCacheManager, and only
|
|
loaded when needed for inference via BatchGenerator.
|
|
|
|
Attributes:
|
|
block_id: Physical block index (0 to num_blocks - 1)
|
|
ref_count: Reference count for sharing (0 = can be evicted)
|
|
block_hash: Content hash for prefix caching and paged SSD storage key
|
|
prev_free_block: Previous block in free list (doubly linked)
|
|
next_free_block: Next block in free list (doubly linked)
|
|
is_null: True if this is the null/placeholder block
|
|
token_count: Number of tokens stored in this block
|
|
"""
|
|
|
|
block_id: int
|
|
ref_count: int = 0
|
|
block_hash: Optional[BlockHash] = None
|
|
|
|
# Doubly linked list pointers for FreeKVCacheBlockQueue
|
|
prev_free_block: Optional["CacheBlock"] = None
|
|
next_free_block: Optional["CacheBlock"] = None
|
|
|
|
# Special flags
|
|
is_null: bool = False
|
|
|
|
# Metadata
|
|
token_count: int = 0
|
|
last_access: float = field(default_factory=time.time)
|
|
|
|
def is_full(self, block_size: int) -> bool:
|
|
"""Check if block is at capacity."""
|
|
return self.token_count >= block_size
|
|
|
|
def is_shared(self) -> bool:
|
|
"""Check if block is shared (ref_count > 1)."""
|
|
return self.ref_count > 1
|
|
|
|
def reset_hash(self) -> None:
|
|
"""Reset block hash when evicted from cache."""
|
|
self.block_hash = None
|
|
|
|
def touch(self) -> None:
|
|
"""Update last access time."""
|
|
self.last_access = time.time()
|
|
|
|
def __repr__(self) -> str:
|
|
prev_id = self.prev_free_block.block_id if self.prev_free_block else None
|
|
next_id = self.next_free_block.block_id if self.next_free_block else None
|
|
hash_str = f", hash={self.block_hash.hex()[:8]}..." if self.block_hash else ""
|
|
return (
|
|
f"CacheBlock(id={self.block_id}, ref={self.ref_count}, "
|
|
f"tokens={self.token_count}, prev={prev_id}, next={next_id}{hash_str})"
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# FreeKVCacheBlockQueue - O(1) Doubly Linked List (vLLM style)
|
|
# =============================================================================
|
|
|
|
class FreeKVCacheBlockQueue:
|
|
"""
|
|
Doubly linked list of free blocks following vLLM's design.
|
|
|
|
Provides O(1) operations for:
|
|
- popleft(): Allocate block from front (LRU order)
|
|
- remove(): Remove block from middle (when touched by cache hit)
|
|
- append(): Return block to end (when freed)
|
|
|
|
The queue maintains LRU eviction order:
|
|
- Front = least recently used (evict first)
|
|
- Back = most recently used (evict last)
|
|
|
|
Uses fake head/tail sentinels to simplify edge cases.
|
|
"""
|
|
|
|
def __init__(self, blocks: List[CacheBlock]) -> None:
|
|
"""
|
|
Initialize queue with all blocks as free.
|
|
|
|
Args:
|
|
blocks: List of all CacheBlock objects
|
|
"""
|
|
self.num_free_blocks = len(blocks)
|
|
|
|
# Initialize doubly linked list
|
|
for i in range(len(blocks)):
|
|
if i > 0:
|
|
blocks[i].prev_free_block = blocks[i - 1]
|
|
if i < len(blocks) - 1:
|
|
blocks[i].next_free_block = blocks[i + 1]
|
|
|
|
# Create sentinel nodes (never popped)
|
|
self.fake_head = CacheBlock(block_id=-1)
|
|
self.fake_tail = CacheBlock(block_id=-2)
|
|
|
|
if blocks:
|
|
self.fake_head.next_free_block = blocks[0]
|
|
blocks[0].prev_free_block = self.fake_head
|
|
self.fake_tail.prev_free_block = blocks[-1]
|
|
blocks[-1].next_free_block = self.fake_tail
|
|
else:
|
|
self.fake_head.next_free_block = self.fake_tail
|
|
self.fake_tail.prev_free_block = self.fake_head
|
|
|
|
def popleft(self) -> CacheBlock:
|
|
"""
|
|
Pop and return the first (LRU) free block.
|
|
|
|
Raises:
|
|
ValueError: If no free blocks available
|
|
"""
|
|
if self.fake_head.next_free_block is self.fake_tail:
|
|
raise ValueError("No free blocks available")
|
|
|
|
block = self.fake_head.next_free_block
|
|
assert block is not None
|
|
|
|
# Remove from list
|
|
self.fake_head.next_free_block = block.next_free_block
|
|
if block.next_free_block:
|
|
block.next_free_block.prev_free_block = self.fake_head
|
|
|
|
block.prev_free_block = None
|
|
block.next_free_block = None
|
|
self.num_free_blocks -= 1
|
|
|
|
return block
|
|
|
|
def popleft_n(self, n: int) -> List[CacheBlock]:
|
|
"""
|
|
Pop n blocks from the front.
|
|
|
|
Args:
|
|
n: Number of blocks to allocate
|
|
|
|
Returns:
|
|
List of n free blocks
|
|
|
|
Raises:
|
|
AssertionError: If not enough free blocks
|
|
"""
|
|
if n == 0:
|
|
return []
|
|
|
|
assert self.num_free_blocks >= n, f"Need {n} blocks, have {self.num_free_blocks}"
|
|
|
|
result = []
|
|
curr = self.fake_head.next_free_block
|
|
|
|
for _ in range(n):
|
|
assert curr is not None and curr is not self.fake_tail
|
|
result.append(curr)
|
|
last = curr
|
|
curr = curr.next_free_block
|
|
# Clear pointers
|
|
last.prev_free_block = None
|
|
last.next_free_block = None
|
|
|
|
# Reconnect list
|
|
self.fake_head.next_free_block = curr
|
|
if curr:
|
|
curr.prev_free_block = self.fake_head
|
|
|
|
self.num_free_blocks -= n
|
|
return result
|
|
|
|
def remove(self, block: CacheBlock) -> None:
|
|
"""
|
|
Remove a block from the middle of the queue.
|
|
|
|
Used when a free block is "touched" (reused by prefix cache hit).
|
|
|
|
Args:
|
|
block: Block to remove
|
|
|
|
Raises:
|
|
RuntimeError: If block not in queue
|
|
"""
|
|
if block.prev_free_block is None or block.next_free_block is None:
|
|
raise RuntimeError(f"Block {block.block_id} not in free queue")
|
|
|
|
# Unlink
|
|
block.prev_free_block.next_free_block = block.next_free_block
|
|
block.next_free_block.prev_free_block = block.prev_free_block
|
|
block.prev_free_block = None
|
|
block.next_free_block = None
|
|
|
|
self.num_free_blocks -= 1
|
|
|
|
def append(self, block: CacheBlock) -> None:
|
|
"""
|
|
Append a block to the end (MRU position).
|
|
|
|
Args:
|
|
block: Block to append
|
|
"""
|
|
last = self.fake_tail.prev_free_block
|
|
assert last is not None
|
|
|
|
last.next_free_block = block
|
|
block.prev_free_block = last
|
|
block.next_free_block = self.fake_tail
|
|
self.fake_tail.prev_free_block = block
|
|
|
|
self.num_free_blocks += 1
|
|
|
|
def append_n(self, blocks: List[CacheBlock]) -> None:
|
|
"""
|
|
Append multiple blocks to the end.
|
|
|
|
Args:
|
|
blocks: Blocks to append (in order)
|
|
"""
|
|
if not blocks:
|
|
return
|
|
|
|
last = self.fake_tail.prev_free_block
|
|
assert last is not None
|
|
|
|
for block in blocks:
|
|
block.prev_free_block = last
|
|
last.next_free_block = block
|
|
last = block
|
|
|
|
last.next_free_block = self.fake_tail
|
|
self.fake_tail.prev_free_block = last
|
|
|
|
self.num_free_blocks += len(blocks)
|
|
|
|
def get_all_free_blocks(self) -> List[CacheBlock]:
|
|
"""Get all free blocks (for testing)."""
|
|
result = []
|
|
curr = self.fake_head.next_free_block
|
|
while curr and curr is not self.fake_tail:
|
|
result.append(curr)
|
|
curr = curr.next_free_block
|
|
return result
|
|
|
|
|
|
# =============================================================================
|
|
# BlockHashToBlockMap - Hash-based prefix cache (vLLM style)
|
|
# =============================================================================
|
|
|
|
class BlockHashToBlockMap:
|
|
"""
|
|
Cache mapping block hashes to blocks for prefix caching.
|
|
|
|
Follows vLLM's design where the same hash can map to multiple
|
|
blocks (for different KV cache groups in hybrid models).
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
self._cache: Dict[BlockHash, CacheBlock | Dict[int, CacheBlock]] = {}
|
|
|
|
def get_block(self, block_hash: BlockHash) -> Optional[CacheBlock]:
|
|
"""Get any block with the given hash."""
|
|
blocks = self._cache.get(block_hash)
|
|
if blocks is None:
|
|
return None
|
|
if isinstance(blocks, CacheBlock):
|
|
return blocks
|
|
if isinstance(blocks, dict):
|
|
return next(iter(blocks.values()))
|
|
return None
|
|
|
|
def insert(self, block_hash: BlockHash, block: CacheBlock) -> None:
|
|
"""Insert a block into the cache."""
|
|
existing = self._cache.get(block_hash)
|
|
if existing is None:
|
|
self._cache[block_hash] = block
|
|
elif isinstance(existing, CacheBlock):
|
|
self._cache[block_hash] = {
|
|
existing.block_id: existing,
|
|
block.block_id: block,
|
|
}
|
|
elif isinstance(existing, dict):
|
|
existing[block.block_id] = block
|
|
|
|
def pop(self, block_hash: BlockHash, block_id: int) -> Optional[CacheBlock]:
|
|
"""Remove and return a specific block from the cache."""
|
|
blocks = self._cache.pop(block_hash, None)
|
|
if blocks is None:
|
|
return None
|
|
|
|
if isinstance(blocks, CacheBlock):
|
|
if blocks.block_id == block_id:
|
|
return blocks
|
|
# Wrong block ID, put it back
|
|
self._cache[block_hash] = blocks
|
|
return None
|
|
|
|
if isinstance(blocks, dict):
|
|
block = blocks.pop(block_id, None)
|
|
if blocks: # Still has other blocks
|
|
self._cache[block_hash] = blocks
|
|
return block
|
|
|
|
return None
|
|
|
|
def __len__(self) -> int:
|
|
return len(self._cache)
|
|
|
|
def clear(self) -> None:
|
|
self._cache.clear()
|
|
|
|
|
|
# =============================================================================
|
|
# BlockTable - Per-request block mapping
|
|
# =============================================================================
|
|
|
|
@dataclass
|
|
class BlockTable:
|
|
"""
|
|
Per-request block table mapping logical to physical blocks.
|
|
|
|
Similar to vLLM's block table, this maps a request's token positions
|
|
to physical cache blocks.
|
|
|
|
Attributes:
|
|
request_id: Unique request identifier
|
|
block_ids: List of physical block IDs
|
|
num_tokens: Total number of cached tokens
|
|
"""
|
|
|
|
request_id: str
|
|
block_ids: List[int] = field(default_factory=list)
|
|
num_tokens: int = 0
|
|
|
|
def add_block(self, block_id: int, num_tokens: int) -> None:
|
|
"""Add a block to the table."""
|
|
self.block_ids.append(block_id)
|
|
self.num_tokens += num_tokens
|
|
|
|
def __len__(self) -> int:
|
|
return len(self.block_ids)
|
|
|
|
def copy(self, new_request_id: str) -> "BlockTable":
|
|
"""Create a copy with new request ID."""
|
|
return BlockTable(
|
|
request_id=new_request_id,
|
|
block_ids=self.block_ids.copy(),
|
|
num_tokens=self.num_tokens,
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# PagedCacheManager - Main manager (vLLM BlockPool style)
|
|
# =============================================================================
|
|
|
|
class PagedCacheManager(CacheManager):
|
|
"""
|
|
Paged KV cache manager following vLLM's BlockPool architecture.
|
|
|
|
Features:
|
|
- Block allocation/deallocation with reference counting
|
|
- Prefix sharing via chain-based hash deduplication
|
|
- Copy-on-Write for efficient forking
|
|
- O(1) LRU eviction using doubly linked list
|
|
|
|
Implements the CacheManager ABC interface for consistency with other
|
|
cache implementations in oMLX.
|
|
|
|
Args:
|
|
block_size: Number of tokens per block (default: 64)
|
|
max_blocks: Maximum number of blocks to allocate (default: 1000)
|
|
enable_caching: Whether to enable prefix caching (default: True)
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
block_size: int = 64,
|
|
max_blocks: int = 1000,
|
|
enable_caching: bool = True,
|
|
model_name: str = "",
|
|
initial_blocks: int = 256,
|
|
):
|
|
self.block_size = block_size
|
|
self.max_blocks = max_blocks
|
|
self.enable_caching = enable_caching
|
|
self.model_name = model_name
|
|
self.initial_blocks = initial_blocks
|
|
|
|
# Warn if model_name is not set (cache isolation may not work)
|
|
if not model_name:
|
|
logger.warning(
|
|
"PagedCacheManager initialized without model_name. "
|
|
"Cache isolation between models may not work correctly."
|
|
)
|
|
|
|
# Dynamic allocation tracking
|
|
# Only create initial blocks; grow dynamically as needed up to max_blocks
|
|
initial_count = min(initial_blocks, max_blocks)
|
|
self._current_allocated_count = initial_count
|
|
|
|
# Create only initial blocks (memory optimization)
|
|
self.blocks: List[CacheBlock] = [
|
|
CacheBlock(block_id=i) for i in range(initial_count)
|
|
]
|
|
|
|
# Free block queue (doubly linked list for O(1) LRU)
|
|
self.free_block_queue = FreeKVCacheBlockQueue(self.blocks)
|
|
|
|
# Hash-to-block cache for prefix caching
|
|
self.cached_block_hash_to_block = BlockHashToBlockMap()
|
|
|
|
# Request to block table mapping
|
|
self.request_tables: Dict[str, BlockTable] = {}
|
|
|
|
# Allocated blocks (for fast lookup)
|
|
self.allocated_blocks: Dict[int, CacheBlock] = {}
|
|
|
|
# Reserve null block (block 0) - never freed
|
|
self.null_block = self.free_block_queue.popleft()
|
|
self.null_block.is_null = True
|
|
self.null_block.ref_count = 1
|
|
self.allocated_blocks[self.null_block.block_id] = self.null_block
|
|
|
|
# Statistics - track actual created blocks, not max
|
|
self.stats = PagedCacheStats(
|
|
total_blocks=initial_count,
|
|
allocated_blocks=1, # null block
|
|
free_blocks=initial_count - 1,
|
|
)
|
|
|
|
# Thread safety
|
|
self._lock = threading.RLock()
|
|
|
|
# paged SSD cache manager for storage (set via set_paged_ssd_cache_manager)
|
|
self._paged_ssd_cache_manager: Optional[Any] = None
|
|
|
|
logger.info(
|
|
f"PagedCacheManager initialized: block_size={block_size}, "
|
|
f"initial_blocks={initial_count}, max_blocks={max_blocks}, "
|
|
f"max_tokens={block_size * max_blocks}"
|
|
)
|
|
|
|
def set_paged_ssd_cache_manager(self, paged_ssd_cache_manager: Any) -> None:
|
|
"""
|
|
Set paged SSD cache manager for tiered storage.
|
|
|
|
When set, evicted blocks will be saved to paged SSD before clearing
|
|
GPU memory, allowing them to be restored later.
|
|
|
|
Args:
|
|
paged_ssd_cache_manager: PagedSSDCacheManager instance
|
|
"""
|
|
self._paged_ssd_cache_manager = paged_ssd_cache_manager
|
|
logger.info("paged SSD cache manager connected to PagedCacheManager")
|
|
|
|
# =========================================================================
|
|
# Dynamic Block Pool Growth (Elastic KV Cache)
|
|
# =========================================================================
|
|
|
|
def _grow_blocks(self, additional_blocks: int) -> int:
|
|
"""
|
|
Dynamically expand the block pool.
|
|
|
|
This method is called when more blocks are needed but the current
|
|
pool is exhausted. It creates new blocks up to max_blocks limit.
|
|
|
|
Args:
|
|
additional_blocks: Number of blocks to add.
|
|
|
|
Returns:
|
|
Number of blocks actually created.
|
|
"""
|
|
with self._lock:
|
|
available = self.max_blocks - self._current_allocated_count
|
|
to_create = min(additional_blocks, available)
|
|
|
|
if to_create <= 0:
|
|
return 0
|
|
|
|
start_id = self._current_allocated_count
|
|
new_blocks = [
|
|
CacheBlock(block_id=i)
|
|
for i in range(start_id, start_id + to_create)
|
|
]
|
|
|
|
self.blocks.extend(new_blocks)
|
|
self.free_block_queue.append_n(new_blocks)
|
|
self._current_allocated_count += to_create
|
|
|
|
self.stats.total_blocks = self._current_allocated_count
|
|
self.stats.free_blocks = self.free_block_queue.num_free_blocks
|
|
|
|
logger.info(
|
|
f"Block pool grown: +{to_create}, "
|
|
f"total={self._current_allocated_count}/{self.max_blocks}"
|
|
)
|
|
return to_create
|
|
|
|
# =========================================================================
|
|
# Block Allocation (vLLM style)
|
|
# =========================================================================
|
|
|
|
def allocate_block(self) -> Optional[CacheBlock]:
|
|
"""
|
|
Allocate a new cache block.
|
|
|
|
Returns:
|
|
CacheBlock if available, None if out of memory.
|
|
"""
|
|
with self._lock:
|
|
if self.free_block_queue.num_free_blocks == 0:
|
|
# Try to grow the block pool dynamically
|
|
grown = self._grow_blocks(min(256, self.max_blocks - self._current_allocated_count))
|
|
if grown == 0:
|
|
logger.warning("Out of cache blocks (max reached)")
|
|
return None
|
|
|
|
block = self.free_block_queue.popleft()
|
|
|
|
# Evict from hash cache if needed
|
|
if self.enable_caching:
|
|
self._maybe_evict_cached_block(block)
|
|
|
|
block.ref_count = 1
|
|
block.touch()
|
|
self.allocated_blocks[block.block_id] = block
|
|
|
|
self.stats.allocated_blocks += 1
|
|
self.stats.free_blocks -= 1
|
|
|
|
return block
|
|
|
|
def get_new_blocks(self, num_blocks: int) -> List[CacheBlock]:
|
|
"""
|
|
Allocate multiple blocks at once (vLLM style).
|
|
|
|
Args:
|
|
num_blocks: Number of blocks to allocate
|
|
|
|
Returns:
|
|
List of allocated blocks
|
|
|
|
Raises:
|
|
ValueError: If not enough free blocks even after dynamic growth
|
|
"""
|
|
with self._lock:
|
|
if num_blocks > self.free_block_queue.num_free_blocks:
|
|
# Try to grow the block pool dynamically
|
|
needed = num_blocks - self.free_block_queue.num_free_blocks
|
|
self._grow_blocks(needed + 128) # Extra buffer for future allocations
|
|
|
|
if num_blocks > self.free_block_queue.num_free_blocks:
|
|
raise ValueError(
|
|
f"Cannot allocate {num_blocks} blocks, "
|
|
f"only {self.free_block_queue.num_free_blocks} available "
|
|
f"(max={self.max_blocks})"
|
|
)
|
|
|
|
blocks = self.free_block_queue.popleft_n(num_blocks)
|
|
|
|
for block in blocks:
|
|
if self.enable_caching:
|
|
self._maybe_evict_cached_block(block)
|
|
|
|
block.ref_count = 1
|
|
block.touch()
|
|
self.allocated_blocks[block.block_id] = block
|
|
|
|
self.stats.allocated_blocks += num_blocks
|
|
self.stats.free_blocks -= num_blocks
|
|
|
|
return blocks
|
|
|
|
def _maybe_evict_cached_block(self, block: CacheBlock) -> bool:
|
|
"""
|
|
Evict a block from the hash cache if present.
|
|
|
|
In paged SSD-only mode, block data is always on paged SSD, so this just
|
|
removes the block from the cache index.
|
|
|
|
Args:
|
|
block: Block to evict
|
|
|
|
Returns:
|
|
True if block was evicted from cache
|
|
"""
|
|
if block.block_hash is None:
|
|
return False
|
|
|
|
evicted = self.cached_block_hash_to_block.pop(
|
|
block.block_hash, block.block_id
|
|
)
|
|
|
|
if evicted:
|
|
block.reset_hash()
|
|
self.stats.evictions += 1
|
|
return True
|
|
|
|
return False
|
|
|
|
def free_block(self, block_id: int) -> bool:
|
|
"""
|
|
Free a cache block (decrements ref_count, frees if 0).
|
|
|
|
Returns:
|
|
True if block was freed, False if still referenced.
|
|
"""
|
|
with self._lock:
|
|
if block_id not in self.allocated_blocks:
|
|
logger.warning(f"Attempted to free unknown block: {block_id}")
|
|
return False
|
|
|
|
block = self.allocated_blocks[block_id]
|
|
if block.is_null:
|
|
return False # Never free null block
|
|
|
|
block.ref_count -= 1
|
|
|
|
if block.ref_count <= 0:
|
|
# Remove from hash cache
|
|
if block.block_hash is not None:
|
|
self.cached_block_hash_to_block.pop(block.block_hash, block.block_id)
|
|
|
|
# Remove from allocated
|
|
del self.allocated_blocks[block_id]
|
|
|
|
# Add to free queue (back = MRU)
|
|
self.free_block_queue.append(block)
|
|
|
|
self.stats.allocated_blocks -= 1
|
|
self.stats.free_blocks += 1
|
|
self.stats.total_tokens_cached -= block.token_count
|
|
|
|
return True
|
|
|
|
return False
|
|
|
|
def free_blocks(self, blocks: Iterable[CacheBlock]) -> None:
|
|
"""
|
|
Free multiple blocks (vLLM style).
|
|
|
|
Blocks with ref_count=0 are added to the free queue.
|
|
|
|
Args:
|
|
blocks: Blocks to free (in eviction order)
|
|
"""
|
|
with self._lock:
|
|
blocks_list = list(blocks)
|
|
to_free = []
|
|
|
|
for block in blocks_list:
|
|
if block.is_null:
|
|
continue
|
|
|
|
block.ref_count -= 1
|
|
|
|
if block.ref_count <= 0:
|
|
# Remove from hash cache
|
|
if block.block_hash is not None:
|
|
self.cached_block_hash_to_block.pop(block.block_hash, block.block_id)
|
|
|
|
del self.allocated_blocks[block.block_id]
|
|
to_free.append(block)
|
|
self.stats.allocated_blocks -= 1
|
|
self.stats.free_blocks += 1
|
|
self.stats.total_tokens_cached -= block.token_count
|
|
|
|
# Add to free queue (back = MRU, evicted last)
|
|
self.free_block_queue.append_n(to_free)
|
|
|
|
def touch(self, blocks: Iterable[CacheBlock]) -> None:
|
|
"""
|
|
Touch blocks to prevent eviction (cache hit, vLLM style).
|
|
|
|
Increments ref_count and removes from free queue if needed.
|
|
|
|
Args:
|
|
blocks: Blocks to touch
|
|
"""
|
|
with self._lock:
|
|
for block in blocks:
|
|
if block.ref_count == 0 and not block.is_null:
|
|
# Block is in free queue, remove it
|
|
try:
|
|
self.free_block_queue.remove(block)
|
|
self.stats.free_blocks -= 1
|
|
self.stats.allocated_blocks += 1
|
|
self.allocated_blocks[block.block_id] = block
|
|
except RuntimeError:
|
|
pass # Block not in queue
|
|
|
|
block.ref_count += 1
|
|
block.touch()
|
|
|
|
# =========================================================================
|
|
# Reference Counting
|
|
# =========================================================================
|
|
|
|
def increment_ref(self, block_id: int) -> bool:
|
|
"""Increment reference count for a block."""
|
|
with self._lock:
|
|
if block_id not in self.allocated_blocks:
|
|
return False
|
|
|
|
block = self.allocated_blocks[block_id]
|
|
block.ref_count += 1
|
|
block.touch()
|
|
|
|
if block.ref_count == 2:
|
|
self.stats.shared_blocks += 1
|
|
|
|
return True
|
|
|
|
def decrement_ref(self, block_id: int) -> bool:
|
|
"""Decrement reference count (alias for free_block)."""
|
|
return self.free_block(block_id)
|
|
|
|
def release_for_eviction(self, block_ids: List[int]) -> int:
|
|
"""
|
|
Release blocks for eviction without removing from allocated_blocks.
|
|
|
|
Decrements ref_count so blocks become evictable, but keeps them in
|
|
allocated_blocks so they can still be found by cache lookups.
|
|
|
|
Args:
|
|
block_ids: List of block IDs to release.
|
|
|
|
Returns:
|
|
Number of blocks released.
|
|
"""
|
|
released = 0
|
|
with self._lock:
|
|
for block_id in block_ids:
|
|
block = self.allocated_blocks.get(block_id)
|
|
if block is None or block.is_null:
|
|
continue
|
|
|
|
if block.ref_count > 0:
|
|
block.ref_count -= 1
|
|
released += 1
|
|
|
|
if block.ref_count == 1:
|
|
self.stats.shared_blocks = max(0, self.stats.shared_blocks - 1)
|
|
|
|
return released
|
|
|
|
# =========================================================================
|
|
# Prefix Caching (vLLM chain-hash style)
|
|
# =========================================================================
|
|
|
|
def get_cached_block(self, block_hash: BlockHash) -> Optional[CacheBlock]:
|
|
"""
|
|
Get a cached block by its hash (vLLM style).
|
|
|
|
Args:
|
|
block_hash: Content hash of the block
|
|
|
|
Returns:
|
|
Cached block if found, None otherwise
|
|
"""
|
|
if not self.enable_caching:
|
|
return None
|
|
|
|
with self._lock:
|
|
block = self.cached_block_hash_to_block.get_block(block_hash)
|
|
if block:
|
|
self.stats.hits += 1
|
|
else:
|
|
self.stats.misses += 1
|
|
return block
|
|
|
|
def cache_full_blocks(
|
|
self,
|
|
blocks: List[CacheBlock],
|
|
token_ids: List[int],
|
|
num_cached_blocks: int,
|
|
num_full_blocks: int,
|
|
extra_keys: Optional[Tuple[Any, ...]] = None,
|
|
) -> None:
|
|
"""
|
|
Cache full blocks for prefix caching (vLLM style).
|
|
|
|
Computes chain hashes and adds blocks to the cache.
|
|
|
|
Args:
|
|
blocks: All blocks for the request
|
|
token_ids: All token IDs for the request
|
|
num_cached_blocks: Number of blocks already cached
|
|
num_full_blocks: Number of full blocks to cache
|
|
extra_keys: Additional keys for hash (e.g., VLM image hash)
|
|
"""
|
|
if not self.enable_caching:
|
|
return
|
|
|
|
if num_cached_blocks >= num_full_blocks:
|
|
return
|
|
|
|
with self._lock:
|
|
# Get parent hash from last cached block
|
|
parent_hash = None
|
|
if num_cached_blocks > 0:
|
|
parent_hash = blocks[num_cached_blocks - 1].block_hash
|
|
|
|
for i in range(num_cached_blocks, num_full_blocks):
|
|
block = blocks[i]
|
|
if block.block_hash is not None:
|
|
parent_hash = block.block_hash
|
|
continue # Already cached
|
|
|
|
# Get tokens for this block
|
|
start = i * self.block_size
|
|
end = start + self.block_size
|
|
block_tokens = token_ids[start:end]
|
|
|
|
# Compute chain hash
|
|
block_hash = compute_block_hash(
|
|
parent_hash, block_tokens,
|
|
extra_keys=extra_keys, model_name=self.model_name,
|
|
)
|
|
block.block_hash = block_hash
|
|
block.token_count = len(block_tokens)
|
|
|
|
# Add to cache
|
|
self.cached_block_hash_to_block.insert(block_hash, block)
|
|
|
|
parent_hash = block_hash
|
|
|
|
def get_computed_blocks(
|
|
self,
|
|
token_ids: List[int],
|
|
extra_keys: Optional[Tuple[Any, ...]] = None,
|
|
extra_key_token_start: Optional[int] = None,
|
|
extra_key_ranges: Optional[List[Tuple[int, Tuple[Any, ...]]]] = None,
|
|
) -> Tuple[List[CacheBlock], int]:
|
|
"""
|
|
Find cached blocks for a token prefix (vLLM style).
|
|
|
|
Args:
|
|
token_ids: Token IDs to look up
|
|
extra_keys: Additional keys for hash (e.g., VLM image hash)
|
|
|
|
Returns:
|
|
Tuple of (cached_blocks, num_cached_tokens)
|
|
"""
|
|
if not self.enable_caching:
|
|
return [], 0
|
|
|
|
with self._lock:
|
|
cached_blocks = []
|
|
parent_hash = None
|
|
num_cached_tokens = 0
|
|
|
|
num_full_blocks = len(token_ids) // self.block_size
|
|
|
|
for i in range(num_full_blocks):
|
|
start = i * self.block_size
|
|
end = start + self.block_size
|
|
block_tokens = token_ids[start:end]
|
|
block_extra_keys = resolve_block_extra_keys(
|
|
end,
|
|
extra_keys=extra_keys,
|
|
extra_key_token_start=extra_key_token_start,
|
|
extra_key_ranges=extra_key_ranges,
|
|
)
|
|
|
|
# Compute expected hash
|
|
block_hash = compute_block_hash(
|
|
parent_hash, block_tokens,
|
|
extra_keys=block_extra_keys, model_name=self.model_name,
|
|
)
|
|
|
|
# Look up in cache
|
|
cached_block = self.cached_block_hash_to_block.get_block(block_hash)
|
|
|
|
# Lazy restore: if not in memory but exists on SSD, register it
|
|
if cached_block is None and self._paged_ssd_cache_manager is not None:
|
|
if self._paged_ssd_cache_manager.has_block(block_hash):
|
|
# Use standard allocation path so we handle an empty
|
|
# free queue gracefully (grow/evict) and keep stats in sync.
|
|
block = self.allocate_block()
|
|
if block is not None:
|
|
block.block_hash = block_hash
|
|
block.token_count = self.block_size
|
|
# Cold-registered blocks are metadata-only until a
|
|
# request claims them via increment_ref().
|
|
block.ref_count = 0
|
|
self.cached_block_hash_to_block.insert(
|
|
block_hash, block
|
|
)
|
|
cached_block = block
|
|
|
|
if cached_block is None:
|
|
self.stats.misses += 1
|
|
break # Cache miss, stop here
|
|
|
|
cached_blocks.append(cached_block)
|
|
parent_hash = block_hash
|
|
num_cached_tokens += self.block_size
|
|
self.stats.hits += 1
|
|
|
|
return cached_blocks, num_cached_tokens
|
|
|
|
# =========================================================================
|
|
# Legacy hash methods (for backwards compatibility)
|
|
# =========================================================================
|
|
|
|
def find_cached_block(
|
|
self,
|
|
tokens: List[int],
|
|
parent_hash: Optional[BlockHash] = None,
|
|
extra_keys: Optional[Tuple[Any, ...]] = None,
|
|
) -> Optional[CacheBlock]:
|
|
"""
|
|
Find a cached block matching the given tokens using chain hash.
|
|
|
|
Args:
|
|
tokens: Token IDs to look up
|
|
parent_hash: Hash of the parent block (for chain), or None for first block
|
|
extra_keys: Additional keys for hash (e.g., VLM image hash)
|
|
|
|
Returns:
|
|
Cached block if found, None otherwise
|
|
"""
|
|
if not self.enable_caching:
|
|
return None
|
|
|
|
with self._lock:
|
|
block_hash = compute_block_hash(
|
|
parent_hash, tokens, extra_keys=extra_keys,
|
|
model_name=self.model_name,
|
|
)
|
|
block = self.cached_block_hash_to_block.get_block(block_hash)
|
|
if block:
|
|
block.touch()
|
|
self.stats.hits += 1
|
|
return block
|
|
|
|
self.stats.misses += 1
|
|
return None
|
|
|
|
def register_block_hash(
|
|
self,
|
|
block: CacheBlock,
|
|
tokens: List[int],
|
|
parent_hash: Optional[BlockHash] = None,
|
|
extra_keys: Optional[Tuple[Any, ...]] = None,
|
|
) -> None:
|
|
"""
|
|
Register a block's hash for deduplication using chain hash.
|
|
|
|
Args:
|
|
block: Block to register
|
|
tokens: Token IDs in this block
|
|
parent_hash: Hash of the parent block (for chain), or None for first block
|
|
extra_keys: Additional keys for hash (e.g., VLM image hash)
|
|
"""
|
|
if not self.enable_caching:
|
|
return
|
|
|
|
with self._lock:
|
|
block_hash = compute_block_hash(
|
|
parent_hash, tokens, extra_keys=extra_keys,
|
|
model_name=self.model_name,
|
|
)
|
|
block.block_hash = block_hash
|
|
self.cached_block_hash_to_block.insert(block_hash, block)
|
|
|
|
# =========================================================================
|
|
# Block Table Management
|
|
# =========================================================================
|
|
|
|
def create_block_table(self, request_id: str) -> BlockTable:
|
|
"""Create a new block table for a request."""
|
|
with self._lock:
|
|
table = BlockTable(request_id=request_id)
|
|
self.request_tables[request_id] = table
|
|
return table
|
|
|
|
def get_block_table(self, request_id: str) -> Optional[BlockTable]:
|
|
"""Get block table for a request."""
|
|
with self._lock:
|
|
return self.request_tables.get(request_id)
|
|
|
|
def get_or_create_block_table(self, request_id: str) -> BlockTable:
|
|
"""Get or create block table for a request."""
|
|
with self._lock:
|
|
if request_id not in self.request_tables:
|
|
self.request_tables[request_id] = BlockTable(request_id=request_id)
|
|
return self.request_tables[request_id]
|
|
|
|
def delete_block_table(self, request_id: str) -> None:
|
|
"""Delete block table and free associated blocks."""
|
|
with self._lock:
|
|
table = self.request_tables.pop(request_id, None)
|
|
if table:
|
|
for block_id in table.block_ids:
|
|
self.free_block(block_id)
|
|
|
|
def add_block_to_table(
|
|
self,
|
|
table: BlockTable,
|
|
block: CacheBlock,
|
|
tokens_in_block: int,
|
|
) -> None:
|
|
"""Add a block to a block table."""
|
|
with self._lock:
|
|
table.block_ids.append(block.block_id)
|
|
block.token_count = tokens_in_block
|
|
table.num_tokens += tokens_in_block
|
|
self.stats.total_tokens_cached += tokens_in_block
|
|
|
|
# =========================================================================
|
|
# Prefix Sharing & COW
|
|
# =========================================================================
|
|
|
|
def find_shared_prefix(
|
|
self,
|
|
tokens: List[int],
|
|
extra_keys: Optional[Tuple[Any, ...]] = None,
|
|
extra_key_token_start: Optional[int] = None,
|
|
extra_key_ranges: Optional[List[Tuple[int, Tuple[Any, ...]]]] = None,
|
|
) -> Tuple[List[int], List[int]]:
|
|
"""
|
|
Find shared prefix blocks for a token sequence.
|
|
|
|
Uses get_computed_blocks for consistent chain-hash lookup.
|
|
"""
|
|
cached_blocks, num_cached_tokens = self.get_computed_blocks(
|
|
tokens,
|
|
extra_keys=extra_keys,
|
|
extra_key_token_start=extra_key_token_start,
|
|
extra_key_ranges=extra_key_ranges,
|
|
)
|
|
|
|
shared_block_ids = [b.block_id for b in cached_blocks]
|
|
remaining_tokens = tokens[num_cached_tokens:]
|
|
|
|
return shared_block_ids, remaining_tokens
|
|
|
|
def fork_block_table(
|
|
self,
|
|
source_table: BlockTable,
|
|
new_request_id: str,
|
|
) -> BlockTable:
|
|
"""
|
|
Fork a block table for a new request (COW).
|
|
"""
|
|
with self._lock:
|
|
new_table = source_table.copy(new_request_id)
|
|
|
|
for block_id in new_table.block_ids:
|
|
self.increment_ref(block_id)
|
|
|
|
self.request_tables[new_request_id] = new_table
|
|
|
|
logger.debug(
|
|
f"Forked block table: {source_table.request_id} -> {new_request_id}, "
|
|
f"blocks={len(new_table.block_ids)}"
|
|
)
|
|
|
|
return new_table
|
|
|
|
def get_blocks_for_generation(
|
|
self,
|
|
table: BlockTable,
|
|
) -> Tuple[List[CacheBlock], bool]:
|
|
"""
|
|
Get blocks for generation, applying COW if needed.
|
|
"""
|
|
with self._lock:
|
|
blocks = []
|
|
was_copied = False
|
|
|
|
for i, block_id in enumerate(table.block_ids):
|
|
block = self.allocated_blocks.get(block_id)
|
|
if not block:
|
|
continue
|
|
|
|
if block.is_shared():
|
|
new_block = self._cow_copy_block(block)
|
|
if new_block:
|
|
table.block_ids[i] = new_block.block_id
|
|
blocks.append(new_block)
|
|
was_copied = True
|
|
self.stats.cow_copies += 1
|
|
else:
|
|
blocks.append(block)
|
|
else:
|
|
blocks.append(block)
|
|
|
|
block.touch()
|
|
|
|
return blocks, was_copied
|
|
|
|
def _cow_copy_block(self, source_block: CacheBlock) -> Optional[CacheBlock]:
|
|
"""
|
|
Create a copy of a block for COW.
|
|
|
|
In paged SSD-only mode, we don't copy data - we just allocate a new block
|
|
with the same metadata. The actual KV data will be loaded from paged SSD
|
|
when needed.
|
|
"""
|
|
new_block = self.allocate_block()
|
|
if not new_block:
|
|
return None
|
|
|
|
new_block.token_count = source_block.token_count
|
|
new_block.block_hash = source_block.block_hash
|
|
|
|
source_block.ref_count -= 1
|
|
if source_block.ref_count == 1:
|
|
self.stats.shared_blocks -= 1
|
|
|
|
logger.debug(
|
|
f"COW copy: block {source_block.block_id} -> {new_block.block_id}"
|
|
)
|
|
|
|
return new_block
|
|
|
|
# =========================================================================
|
|
# Legacy allocation methods (for backwards compatibility)
|
|
# =========================================================================
|
|
|
|
def allocate_blocks_for_tokens(self, num_tokens: int) -> List[CacheBlock]:
|
|
"""Allocate enough blocks to hold num_tokens."""
|
|
num_blocks_needed = (num_tokens + self.block_size - 1) // self.block_size
|
|
return self.get_new_blocks(num_blocks_needed)
|
|
|
|
# =========================================================================
|
|
# Eviction
|
|
# =========================================================================
|
|
|
|
def evict_lru_blocks(self, num_blocks: int) -> int:
|
|
"""
|
|
Evict least recently used blocks.
|
|
|
|
With the doubly linked list, LRU blocks are already at the front
|
|
of the free queue. We just need to pop from front.
|
|
"""
|
|
with self._lock:
|
|
evicted = 0
|
|
|
|
# Get evictable blocks from free queue (they're already LRU ordered)
|
|
for _ in range(min(num_blocks, self.free_block_queue.num_free_blocks)):
|
|
try:
|
|
block = self.free_block_queue.popleft()
|
|
self._maybe_evict_cached_block(block)
|
|
# Put back at end (now available for allocation)
|
|
self.free_block_queue.append(block)
|
|
evicted += 1
|
|
except ValueError:
|
|
break
|
|
|
|
if evicted > 0:
|
|
logger.info(f"Evicted {evicted} LRU blocks from cache")
|
|
|
|
return evicted
|
|
|
|
def handle_memory_pressure(self, requested_blocks: int) -> bool:
|
|
"""Handle memory pressure by evicting blocks."""
|
|
with self._lock:
|
|
if self.free_block_queue.num_free_blocks >= requested_blocks:
|
|
return True
|
|
|
|
needed = requested_blocks - self.free_block_queue.num_free_blocks
|
|
self.evict_lru_blocks(needed)
|
|
|
|
return self.free_block_queue.num_free_blocks >= requested_blocks
|
|
|
|
# =========================================================================
|
|
# Statistics and Properties
|
|
# =========================================================================
|
|
|
|
@property
|
|
def free_blocks(self) -> int:
|
|
"""Number of free blocks available."""
|
|
return self.free_block_queue.num_free_blocks
|
|
|
|
@property
|
|
def usage(self) -> float:
|
|
"""Cache usage ratio (0.0 to 1.0)."""
|
|
total = self.max_blocks - 1 # Exclude null block
|
|
if total == 0:
|
|
return 0.0
|
|
return 1.0 - (self.free_blocks / total)
|
|
|
|
def get_stats(self) -> PagedCacheStats:
|
|
"""Get current cache statistics."""
|
|
with self._lock:
|
|
self.stats.shared_blocks = sum(
|
|
1 for b in self.allocated_blocks.values() if b.ref_count > 1
|
|
)
|
|
self.stats.free_blocks = self.free_block_queue.num_free_blocks
|
|
return self.stats
|
|
|
|
def get_memory_usage(self) -> Dict[str, Any]:
|
|
"""Get memory usage information."""
|
|
with self._lock:
|
|
stats = self.get_stats()
|
|
return {
|
|
"block_size": self.block_size,
|
|
"max_blocks": self.max_blocks,
|
|
"allocated_blocks": stats.allocated_blocks,
|
|
"free_blocks": stats.free_blocks,
|
|
"shared_blocks": stats.shared_blocks,
|
|
"total_tokens_cached": stats.total_tokens_cached,
|
|
"utilization": stats.allocated_blocks / self.max_blocks,
|
|
"cache_hit_rate": (
|
|
stats.hits / (stats.hits + stats.misses)
|
|
if (stats.hits + stats.misses) > 0 else 0
|
|
),
|
|
}
|
|
|
|
def reset_stats(self) -> None:
|
|
"""Reset statistics counters."""
|
|
with self._lock:
|
|
self.stats.hits = 0
|
|
self.stats.misses = 0
|
|
self.stats.cow_copies = 0
|
|
self.stats.evictions = 0
|
|
|
|
def reset_prefix_cache(self) -> bool:
|
|
"""Reset the prefix cache."""
|
|
with self._lock:
|
|
num_used = self.max_blocks - self.free_block_queue.num_free_blocks
|
|
if num_used > 1: # null_block is always "used"
|
|
logger.warning(f"Cannot reset cache: {num_used - 1} blocks in use")
|
|
return False
|
|
|
|
self.cached_block_hash_to_block.clear()
|
|
|
|
for block in self.blocks:
|
|
block.reset_hash()
|
|
|
|
self.stats.evictions = 0
|
|
self.stats.hits = 0
|
|
self.stats.misses = 0
|
|
|
|
logger.info("Prefix cache reset successfully")
|
|
return True
|
|
|
|
def clear(self) -> int:
|
|
"""
|
|
Clear all cached data and reset to initial block count.
|
|
|
|
Returns:
|
|
Number of entries cleared.
|
|
"""
|
|
with self._lock:
|
|
# Count entries before clearing
|
|
cleared_count = len(self.allocated_blocks) - 1 # Exclude null block
|
|
|
|
# Reset to initial blocks (memory optimization)
|
|
initial_count = min(self.initial_blocks, self.max_blocks)
|
|
self._current_allocated_count = initial_count
|
|
|
|
# Recreate blocks and queue with only initial blocks
|
|
self.blocks = [
|
|
CacheBlock(block_id=i) for i in range(initial_count)
|
|
]
|
|
self.free_block_queue = FreeKVCacheBlockQueue(self.blocks)
|
|
|
|
self.cached_block_hash_to_block.clear()
|
|
self.request_tables.clear()
|
|
self.allocated_blocks.clear()
|
|
|
|
# Reserve null block
|
|
self.null_block = self.free_block_queue.popleft()
|
|
self.null_block.is_null = True
|
|
self.null_block.ref_count = 1
|
|
self.allocated_blocks[self.null_block.block_id] = self.null_block
|
|
|
|
self.stats = PagedCacheStats(
|
|
total_blocks=initial_count,
|
|
allocated_blocks=1,
|
|
free_blocks=initial_count - 1,
|
|
)
|
|
|
|
logger.info(
|
|
f"PagedCacheManager cleared (reset to {initial_count} initial blocks)"
|
|
)
|
|
|
|
return max(0, cleared_count)
|
|
|
|
# =========================================================================
|
|
# SSD Cache Support
|
|
# =========================================================================
|
|
|
|
def get_evictable_blocks(self, count: int) -> List[CacheBlock]:
|
|
"""
|
|
Get LRU blocks that can be evicted (metadata cleared).
|
|
|
|
In paged SSD-only mode, blocks don't store data in GPU memory,
|
|
so this returns blocks that can be freed from the metadata index.
|
|
|
|
Args:
|
|
count: Maximum number of blocks to return.
|
|
|
|
Returns:
|
|
List of evictable blocks in LRU order.
|
|
"""
|
|
with self._lock:
|
|
candidates = []
|
|
|
|
# Iterate through free queue (LRU order)
|
|
current = self.free_block_queue.fake_head.next_free_block
|
|
while (
|
|
current is not None
|
|
and current != self.free_block_queue.fake_tail
|
|
and len(candidates) < count
|
|
):
|
|
# Block must not be null and have ref_count == 0
|
|
if not current.is_null and current.ref_count == 0:
|
|
candidates.append(current)
|
|
current = current.next_free_block
|
|
|
|
# Also check allocated blocks with ref_count == 0 (not in free queue yet)
|
|
if len(candidates) < count:
|
|
# Sort by last_access (LRU)
|
|
remaining = []
|
|
for block in self.allocated_blocks.values():
|
|
if (
|
|
not block.is_null
|
|
and block.ref_count == 0
|
|
and block not in candidates
|
|
):
|
|
remaining.append(block)
|
|
|
|
remaining.sort(key=lambda b: b.last_access)
|
|
candidates.extend(remaining[: count - len(candidates)])
|
|
|
|
return candidates[:count]
|
|
|
|
def mark_block_cold(self, block_id: int) -> bool:
|
|
"""
|
|
Mark a block as evicted (metadata preserved).
|
|
|
|
In paged SSD-only mode, this is a no-op since block data is always on paged SSD.
|
|
Kept for API compatibility.
|
|
|
|
Args:
|
|
block_id: Block ID to mark.
|
|
|
|
Returns:
|
|
True if successful, False if block not found or has data users.
|
|
"""
|
|
with self._lock:
|
|
block = self.blocks[block_id] if block_id < len(self.blocks) else None
|
|
if block is None:
|
|
logger.warning(f"Block {block_id} not found")
|
|
return False
|
|
|
|
if block.ref_count > 0:
|
|
logger.warning(
|
|
f"Cannot mark block {block_id}: ref_count={block.ref_count}"
|
|
)
|
|
return False
|
|
|
|
if block.is_null:
|
|
logger.warning(f"Cannot mark null block")
|
|
return False
|
|
|
|
# In paged SSD-only mode, data is already on paged SSD
|
|
self.stats.evictions += 1
|
|
|
|
logger.debug(
|
|
f"Marked block {block_id} "
|
|
f"(hash={block.block_hash.hex()[:16] if block.block_hash else 'None'}...)"
|
|
)
|
|
return True
|
|
|
|
def evict_block_permanently(self, block_id: int) -> bool:
|
|
"""
|
|
Evict a block permanently (removes from metadata index).
|
|
|
|
This method:
|
|
- Removes from hash cache (block won't be found in cache lookups)
|
|
- Returns block to free queue (can be reallocated)
|
|
- Removes from allocated_blocks
|
|
|
|
Note: In paged SSD-only mode, the data remains on paged SSD and may be deleted
|
|
by PagedSSDCacheManager's LRU eviction if needed.
|
|
|
|
Args:
|
|
block_id: Block ID to evict.
|
|
|
|
Returns:
|
|
True if successful, False if block not found or in use.
|
|
"""
|
|
with self._lock:
|
|
block = self.blocks[block_id] if block_id < len(self.blocks) else None
|
|
if block is None:
|
|
logger.warning(f"Block {block_id} not found for permanent eviction")
|
|
return False
|
|
|
|
if block.ref_count > 0:
|
|
logger.warning(
|
|
f"Cannot permanently evict block {block_id}: ref_count={block.ref_count}"
|
|
)
|
|
return False
|
|
|
|
if block.is_null:
|
|
logger.warning(f"Cannot evict null block")
|
|
return False
|
|
|
|
# Remove from hash cache
|
|
if block.block_hash is not None:
|
|
self.cached_block_hash_to_block.pop(block.block_hash, block.block_id)
|
|
|
|
# Clear metadata
|
|
block.reset_hash()
|
|
block.token_count = 0
|
|
|
|
# Remove from allocated_blocks and add to free queue
|
|
if block_id in self.allocated_blocks:
|
|
del self.allocated_blocks[block_id]
|
|
self.stats.allocated_blocks -= 1
|
|
|
|
self.free_block_queue.append(block)
|
|
self.stats.free_blocks += 1
|
|
self.stats.evictions += 1
|
|
|
|
logger.debug(f"Permanently evicted block {block_id}")
|
|
return True
|
|
|
|
def restore_block(
|
|
self,
|
|
block_id: int,
|
|
cache_data: List[Tuple[Any, Any]],
|
|
) -> bool:
|
|
"""
|
|
Restore block data from cold storage.
|
|
|
|
In paged SSD-only mode, this is a no-op since data is always loaded
|
|
directly from paged SSD when needed. Kept for API compatibility.
|
|
|
|
Args:
|
|
block_id: Block ID to restore.
|
|
cache_data: KV cache data (unused in paged SSD-only mode).
|
|
|
|
Returns:
|
|
True if successful, False if block not found.
|
|
"""
|
|
with self._lock:
|
|
block = self.blocks[block_id] if block_id < len(self.blocks) else None
|
|
if block is None:
|
|
logger.warning(f"Block {block_id} not found for restoration")
|
|
return False
|
|
|
|
block.touch()
|
|
|
|
logger.debug(
|
|
f"Block {block_id} touched "
|
|
f"(hash={block.block_hash.hex()[:16] if block.block_hash else 'None'}...)"
|
|
)
|
|
return True
|
|
|
|
def get_cold_blocks(self) -> List[CacheBlock]:
|
|
"""
|
|
Get all blocks that have data on paged SSD.
|
|
|
|
In paged SSD-only mode, returns all blocks with block_hash set
|
|
(i.e., blocks that have data stored on paged SSD).
|
|
|
|
Returns:
|
|
List of blocks with paged SSD data.
|
|
"""
|
|
with self._lock:
|
|
return [b for b in self.blocks if b.block_hash is not None and not b.is_null]
|
|
|
|
@property
|
|
def cold_block_count(self) -> int:
|
|
"""Number of blocks with data on paged SSD."""
|
|
with self._lock:
|
|
return sum(1 for b in self.blocks if b.block_hash is not None and not b.is_null)
|
|
|
|
def get_ref_count_distribution(self) -> Dict[int, int]:
|
|
"""
|
|
Get distribution of blocks by ref_count.
|
|
|
|
Returns:
|
|
Dict mapping ref_count -> number of blocks with that count.
|
|
Only includes ref_counts that have at least one block.
|
|
"""
|
|
with self._lock:
|
|
distribution: Dict[int, int] = {}
|
|
for block in self.allocated_blocks.values():
|
|
rc = block.ref_count
|
|
distribution[rc] = distribution.get(rc, 0) + 1
|
|
return distribution
|
|
|
|
def get_ref_count_summary(self) -> str:
|
|
"""
|
|
Get a compact string summary of ref_count distribution.
|
|
|
|
Returns:
|
|
String like "rc0=5(ssd=3),rc1=100" showing counts per ref_count.
|
|
In paged SSD-only mode, all blocks have data on paged SSD.
|
|
"""
|
|
dist = self.get_ref_count_distribution()
|
|
if not dist:
|
|
return "rc=none"
|
|
|
|
# Count blocks with paged SSD data (blocks with block_hash)
|
|
ssd_count = 0
|
|
with self._lock:
|
|
for block in self.allocated_blocks.values():
|
|
if block.ref_count == 0 and block.block_hash is not None:
|
|
ssd_count += 1
|
|
|
|
parts = []
|
|
for k, v in sorted(dist.items()):
|
|
if k == 0:
|
|
# Show paged SSD count for rc0
|
|
parts.append(f"rc0={v}(ssd={ssd_count})")
|
|
else:
|
|
parts.append(f"rc{k}={v}")
|
|
return ",".join(parts)
|
|
|
|
# =========================================================================
|
|
# CacheManager ABC Interface Implementation
|
|
# =========================================================================
|
|
|
|
def fetch(self, key: Any) -> Tuple[Optional[Any], bool]:
|
|
"""
|
|
Fetch a cached block by its hash.
|
|
|
|
Args:
|
|
key: BlockHash (bytes) to look up.
|
|
|
|
Returns:
|
|
Tuple of (CacheBlock, True) if found, (None, False) otherwise.
|
|
"""
|
|
if not isinstance(key, bytes):
|
|
return None, False
|
|
|
|
block = self.get_cached_block(BlockHash(key))
|
|
if block is not None:
|
|
return block, True
|
|
return None, False
|
|
|
|
def store(self, key: Any, value: Any) -> bool:
|
|
"""
|
|
Store a block in the cache.
|
|
|
|
For PagedCacheManager, use allocate_block() and register_block_hash()
|
|
for the full workflow. This method provides a simplified interface.
|
|
|
|
Args:
|
|
key: BlockHash (bytes) for the block.
|
|
value: CacheBlock to store.
|
|
|
|
Returns:
|
|
True if stored successfully.
|
|
"""
|
|
if not isinstance(key, bytes) or not isinstance(value, CacheBlock):
|
|
return False
|
|
|
|
with self._lock:
|
|
block_hash = BlockHash(key)
|
|
value.block_hash = block_hash
|
|
self.cached_block_hash_to_block.insert(block_hash, value)
|
|
return True
|
|
|
|
def evict(self, key: Any) -> bool:
|
|
"""
|
|
Evict a specific block from the cache.
|
|
|
|
Args:
|
|
key: BlockHash (bytes) or block_id (int) to evict.
|
|
|
|
Returns:
|
|
True if evicted, False if not found.
|
|
"""
|
|
if isinstance(key, bytes):
|
|
# Evict by block hash
|
|
block = self.cached_block_hash_to_block.get_block(BlockHash(key))
|
|
if block is not None:
|
|
return self.evict_block_permanently(block.block_id)
|
|
return False
|
|
elif isinstance(key, int):
|
|
# Evict by block ID
|
|
return self.evict_block_permanently(key)
|
|
return False
|
|
|
|
@property
|
|
def size(self) -> int:
|
|
"""
|
|
Get the current number of allocated blocks.
|
|
|
|
Returns:
|
|
Number of allocated blocks (excluding null block).
|
|
"""
|
|
return max(0, len(self.allocated_blocks) - 1) # Exclude null block
|
|
|
|
@property
|
|
def max_size(self) -> int:
|
|
"""
|
|
Get the maximum number of blocks.
|
|
|
|
Returns:
|
|
Maximum number of blocks.
|
|
"""
|
|
return self.max_blocks
|