6148e03bc7
Random CCCL pick: thrust/examples/bucket_sort2d.cu (108 lines, full read) Maps to: vllm/core/evictor_v2.py (LRU cache eviction) bucket_sort2d.cu pattern: transform→sort_by_key→lower_bound/upper_bound - point_to_bucket_index ↔ content_hash (prefix cache key) - sort_by_key ↔ eviction priority ordering - lower_bound/upper_bound ↔ block range lookup Current LRUEvictor.evict() is O(n) linear scan over OrderedDict. CCCL pattern suggests sort_by_key → O(1) pop for production scale. For competition (max_num_seqs=1, bounded blocks): current is sufficient. Also read: vllm/core/block/prefix_caching_block.py (200 lines)
142 lines
5.0 KiB
Python
142 lines
5.0 KiB
Python
import enum
|
|
from abc import ABC, abstractmethod
|
|
from typing import OrderedDict, Tuple
|
|
|
|
|
|
class EvictionPolicy(enum.Enum):
|
|
"""Enum for eviction policy used by make_evictor to instantiate the correct
|
|
Evictor subclass.
|
|
"""
|
|
LRU = enum.auto()
|
|
|
|
|
|
class Evictor(ABC):
|
|
"""The Evictor subclasses should be used by the BlockAllocator class to
|
|
handle eviction of freed PhysicalTokenBlocks.
|
|
"""
|
|
|
|
@abstractmethod
|
|
def __init__(self):
|
|
pass
|
|
|
|
@abstractmethod
|
|
def __contains__(self, block_id: int) -> bool:
|
|
pass
|
|
|
|
@abstractmethod
|
|
def evict(self) -> Tuple[int, int]:
|
|
"""Runs the eviction algorithm and returns the evicted block's
|
|
content hash along with physical block id along with physical block id
|
|
"""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def add(self, block_id: int, content_hash: int, num_hashed_tokens: int,
|
|
last_accessed: float):
|
|
"""Adds block to the evictor, making it a candidate for eviction"""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def update(self, block_id: int, last_accessed: float):
|
|
"""Update corresponding block's access time in metadata"""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def remove(self, block_id: int):
|
|
"""Remove a given block id from the cache."""
|
|
pass
|
|
|
|
@property
|
|
@abstractmethod
|
|
def num_blocks(self) -> int:
|
|
pass
|
|
|
|
|
|
class BlockMetaData():
|
|
"""Data structure for storing key data describe cached block, so that
|
|
evitor could use to make its decision which one to choose for eviction
|
|
|
|
Here we use physical block id as the dict key, as there maybe several
|
|
blocks with the same content hash, but their physical id is unique.
|
|
"""
|
|
|
|
def __init__(self, content_hash: int, num_hashed_tokens: int,
|
|
last_accessed: float):
|
|
self.content_hash = content_hash
|
|
self.num_hashed_tokens = num_hashed_tokens
|
|
self.last_accessed = last_accessed
|
|
|
|
|
|
class LRUEvictor(Evictor):
|
|
"""Evicts in a least-recently-used order using the last_accessed timestamp
|
|
that's recorded in the PhysicalTokenBlock. If there are multiple blocks with
|
|
the same last_accessed time, then the one with the largest num_hashed_tokens
|
|
will be evicted. If two blocks each have the lowest last_accessed time and
|
|
highest num_hashed_tokens value, then one will be chose arbitrarily
|
|
|
|
CCCL system design note (from thrust/examples/bucket_sort2d.cu):
|
|
CCCL's bucket sort uses transform→sort_by_key→lower_bound/upper_bound
|
|
to build O(1) lookup from bucket_index → item range. This maps to:
|
|
content_hash → block_id (our _cached_blocks dict)
|
|
last_accessed → eviction priority (our OrderedDict linear scan)
|
|
For production: sort_by_key on (last_accessed, -num_hashed_tokens)
|
|
would make evict() O(1) pop instead of O(n) scan.
|
|
For competition (max_num_seqs=1): current O(n) scan is fine since
|
|
n = num_gpu_blocks is bounded by GPU memory / block_size.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self.free_table: OrderedDict[int, BlockMetaData] = OrderedDict()
|
|
|
|
def __contains__(self, block_id: int) -> bool:
|
|
return block_id in self.free_table
|
|
|
|
def evict(self) -> Tuple[int, int]:
|
|
if len(self.free_table) == 0:
|
|
raise ValueError("No usable cache memory left")
|
|
|
|
evicted_block, evicted_block_id = None, None
|
|
# The blocks with the lowest timestamps should be placed consecutively
|
|
# at the start of OrderedDict. Loop through all these blocks to
|
|
# find the one with maximum number of hashed tokens.
|
|
for _id, block in self.free_table.items():
|
|
if evicted_block is None:
|
|
evicted_block, evicted_block_id = block, _id
|
|
continue
|
|
if evicted_block.last_accessed < block.last_accessed:
|
|
break
|
|
if evicted_block.num_hashed_tokens < block.num_hashed_tokens:
|
|
evicted_block, evicted_block_id = block, _id
|
|
|
|
assert evicted_block is not None
|
|
assert evicted_block_id is not None
|
|
self.free_table.pop(evicted_block_id)
|
|
|
|
return evicted_block_id, evicted_block.content_hash
|
|
|
|
def add(self, block_id: int, content_hash: int, num_hashed_tokens: int,
|
|
last_accessed: float):
|
|
self.free_table[block_id] = BlockMetaData(content_hash,
|
|
num_hashed_tokens,
|
|
last_accessed)
|
|
|
|
def update(self, block_id: int, last_accessed: float):
|
|
self.free_table[block_id].last_accessed = last_accessed
|
|
|
|
def remove(self, block_id: int):
|
|
if block_id not in self.free_table:
|
|
raise ValueError(
|
|
"Attempting to remove block that's not in the evictor")
|
|
self.free_table.pop(block_id)
|
|
|
|
@property
|
|
def num_blocks(self) -> int:
|
|
return len(self.free_table)
|
|
|
|
|
|
def make_evictor(eviction_policy: EvictionPolicy) -> Evictor:
|
|
if eviction_policy == EvictionPolicy.LRU:
|
|
return LRUEvictor()
|
|
else:
|
|
raise ValueError(f"Unknown cache eviction policy: {eviction_policy}")
|