完整部署快照
This commit is contained in:
root
272
vllm/core/evictor_v2.py
Normal file
272
vllm/core/evictor_v2.py
Normal file
@@ -0,0 +1,272 @@
|
||||
import enum
|
||||
import heapq
|
||||
import os
|
||||
from abc import ABC, abstractmethod
|
||||
from collections.abc import Mapping
|
||||
from typing import Dict, List, OrderedDict, Tuple
|
||||
|
||||
|
||||
ContentHash = bytes
|
||||
|
||||
|
||||
class EvictionPolicy(enum.Enum):
|
||||
"""Enum for eviction policy used by make_evictor to instantiate the correct
|
||||
Evictor subclass.
|
||||
"""
|
||||
LRU = enum.auto()
|
||||
FREQUENCY_AWARE = 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, ContentHash]:
|
||||
"""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: ContentHash,
|
||||
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: ContentHash, 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
|
||||
"""
|
||||
|
||||
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, ContentHash]:
|
||||
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: ContentHash,
|
||||
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)
|
||||
|
||||
|
||||
class FrequencyAwareEvictor(Evictor):
|
||||
"""Evict the least frequently reused logical prefix content first.
|
||||
|
||||
Content frequency survives physical block reuse. Heap entries carry a
|
||||
generation and are lazily invalidated so eviction remains O(log N) without
|
||||
allowing stale entries to grow without bound.
|
||||
"""
|
||||
|
||||
_COMPACTION_FACTOR = 2
|
||||
_COMPACTION_SLACK = 1
|
||||
|
||||
def __init__(self):
|
||||
self.free_table: Dict[int, BlockMetaData] = {}
|
||||
self.frequency_by_hash: Dict[ContentHash, int] = {}
|
||||
self._heap: List[Tuple[int, float, int, int, int]] = []
|
||||
self._generations: Dict[int, int] = {}
|
||||
self._next_generation = 0
|
||||
|
||||
@staticmethod
|
||||
def _validate_content_hash(content_hash: ContentHash) -> None:
|
||||
if not isinstance(content_hash, bytes) or len(content_hash) != 32:
|
||||
raise ValueError(
|
||||
"frequency-aware eviction requires a 32-byte content hash")
|
||||
|
||||
def __contains__(self, block_id: int) -> bool:
|
||||
return block_id in self.free_table
|
||||
|
||||
def _heap_key(self, block_id: int, block: BlockMetaData,
|
||||
generation: int) -> Tuple[int, float, int, int, int]:
|
||||
return (
|
||||
self.frequency_by_hash[block.content_hash],
|
||||
block.last_accessed,
|
||||
-block.num_hashed_tokens,
|
||||
block_id,
|
||||
generation,
|
||||
)
|
||||
|
||||
def _push(self, block_id: int) -> None:
|
||||
self._next_generation += 1
|
||||
generation = self._next_generation
|
||||
self._generations[block_id] = generation
|
||||
heapq.heappush(
|
||||
self._heap,
|
||||
self._heap_key(
|
||||
block_id, self.free_table[block_id], generation),
|
||||
)
|
||||
|
||||
def _compact_if_needed(self) -> None:
|
||||
limit = (
|
||||
self._COMPACTION_FACTOR * len(self.free_table)
|
||||
+ self._COMPACTION_SLACK
|
||||
)
|
||||
if len(self._heap) <= limit:
|
||||
return
|
||||
self._heap = [
|
||||
self._heap_key(block_id, block, self._generations[block_id])
|
||||
for block_id, block in self.free_table.items()
|
||||
]
|
||||
heapq.heapify(self._heap)
|
||||
|
||||
def evict(self) -> Tuple[int, ContentHash]:
|
||||
if not self.free_table:
|
||||
raise ValueError("No usable cache memory left")
|
||||
|
||||
while self._heap:
|
||||
entry = heapq.heappop(self._heap)
|
||||
frequency, _, _, block_id, generation = entry
|
||||
block = self.free_table.get(block_id)
|
||||
if (
|
||||
block is None
|
||||
or self._generations.get(block_id) != generation
|
||||
):
|
||||
continue
|
||||
if self.frequency_by_hash[block.content_hash] != frequency:
|
||||
heapq.heappush(
|
||||
self._heap,
|
||||
self._heap_key(block_id, block, generation),
|
||||
)
|
||||
continue
|
||||
|
||||
block = self.free_table.pop(block_id)
|
||||
self._generations.pop(block_id)
|
||||
self._compact_if_needed()
|
||||
return block_id, block.content_hash
|
||||
|
||||
raise RuntimeError("Evictor heap has no usable entry")
|
||||
|
||||
def add(self, block_id: int, content_hash: ContentHash,
|
||||
num_hashed_tokens: int, last_accessed: float):
|
||||
self._validate_content_hash(content_hash)
|
||||
self.frequency_by_hash[content_hash] = (
|
||||
self.frequency_by_hash.get(content_hash, 0) + 1)
|
||||
self.free_table[block_id] = BlockMetaData(
|
||||
content_hash, num_hashed_tokens, last_accessed)
|
||||
self._push(block_id)
|
||||
self._compact_if_needed()
|
||||
|
||||
def update(self, block_id: int, last_accessed: float):
|
||||
self.free_table[block_id].last_accessed = last_accessed
|
||||
self._push(block_id)
|
||||
self._compact_if_needed()
|
||||
|
||||
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)
|
||||
self._generations.pop(block_id)
|
||||
self._compact_if_needed()
|
||||
|
||||
@property
|
||||
def num_blocks(self) -> int:
|
||||
return len(self.free_table)
|
||||
|
||||
|
||||
def eviction_policy_from_env(
|
||||
environ: Mapping[str, str] | None = None,
|
||||
) -> EvictionPolicy:
|
||||
source = os.environ if environ is None else environ
|
||||
value = source.get("BI100_KV_EVICTION_POLICY", "lru").strip().lower()
|
||||
policies = {
|
||||
"lru": EvictionPolicy.LRU,
|
||||
"frequency": EvictionPolicy.FREQUENCY_AWARE,
|
||||
}
|
||||
if value not in policies:
|
||||
raise ValueError(
|
||||
"BI100_KV_EVICTION_POLICY must be one of: frequency, lru")
|
||||
return policies[value]
|
||||
|
||||
|
||||
def make_evictor(eviction_policy: EvictionPolicy) -> Evictor:
|
||||
if eviction_policy == EvictionPolicy.LRU:
|
||||
return LRUEvictor()
|
||||
elif eviction_policy == EvictionPolicy.FREQUENCY_AWARE:
|
||||
return FrequencyAwareEvictor()
|
||||
else:
|
||||
raise ValueError(f"Unknown cache eviction policy: {eviction_policy}")
|
||||
Reference in New Issue
Block a user