diff --git a/tests/v1/kv_offload/cpu/test_manager.py b/tests/v1/kv_offload/cpu/test_manager.py index aa4fb829597..d568357224c 100644 --- a/tests/v1/kv_offload/cpu/test_manager.py +++ b/tests/v1/kv_offload/cpu/test_manager.py @@ -294,25 +294,25 @@ def test_cpu_manager(): # prepare store with no space ([2, 3] is being loaded) assert cpu_manager.prepare_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX) is None - # complete load [2, 3] + # complete load [2, 3]. Load changes the eviction list, making 2, 3 recent. cpu_manager.complete_load(to_keys([2, 3]), _EMPTY_REQ_CTX) - # prepare store [6, 7, 8] -> evicts [2, 3, 4] (oldest) + # prepare store [6, 7, 8] -> evicts [4, 5, 2] (oldest) prepare_store_output = cpu_manager.prepare_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX) verify_store_output( prepare_store_output, ExpectedPrepareStoreOutput( keys_to_store=[6, 7, 8], - store_block_ids=[3, 2, 1], - evicted_keys=[2, 3, 4], + store_block_ids=[1, 0, 3], + evicted_keys=[4, 5, 2], ), ) # complete store [6, 7, 8] cpu_manager.complete_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX) - # touch [5, 6, 7] (move to end of LRU order) - cpu_manager.touch(to_keys([5, 6, 7]), _EMPTY_REQ_CTX) + # touch [3, 6, 7] (move to end of LRU order) + cpu_manager.touch(to_keys([3, 6, 7]), _EMPTY_REQ_CTX) # prepare store [7, 9] -> evicts [8] (oldest following previous touch) prepare_store_output = cpu_manager.prepare_store(to_keys([9]), _EMPTY_REQ_CTX) @@ -320,7 +320,7 @@ def test_cpu_manager(): prepare_store_output, ExpectedPrepareStoreOutput( keys_to_store=[9], - store_block_ids=[1], + store_block_ids=[3], evicted_keys=[8], ), ) @@ -335,7 +335,7 @@ def test_cpu_manager(): verify_events( cpu_manager.take_events(), expected_stores=({3, 4, 5}, {6, 7, 8}), - expected_evictions=({2, 3, 4}, {8}), + expected_evictions=({4, 5, 2}, {8}), ) diff --git a/tests/v1/kv_offload/tiering/test_tiering_offloading.py b/tests/v1/kv_offload/tiering/test_tiering_offloading.py index 3caff59c2d6..de37afc9a93 100644 --- a/tests/v1/kv_offload/tiering/test_tiering_offloading.py +++ b/tests/v1/kv_offload/tiering/test_tiering_offloading.py @@ -295,6 +295,8 @@ class TestTieringOffloadingManager: self.manager.prepare_store(blocks, _CTX) self.manager.complete_store(blocks, _CTX, success=True) self._simulate_on_schedule_end() + # for secondary tiers to drain jobs, so primary tier's blocks are evictable. + self._simulate_on_schedule_end() self.secondary_tier1.touch = MagicMock(wraps=self.secondary_tier1.touch) self.secondary_tier2.touch = MagicMock(wraps=self.secondary_tier2.touch) @@ -303,7 +305,7 @@ class TestTieringOffloadingManager: self.manager.touch(blocks, _CTX) # Verify touch was called on primary tier (check LRU order) - primary_keys = list(self.primary_tier._policy.blocks.keys()) + primary_keys = list(self.primary_tier._policy.evictable_blocks.keys()) assert primary_keys[-3:] == list(reversed(blocks)) # Verify touch was propagated to all secondary tiers diff --git a/vllm/v1/kv_offload/cpu/manager.py b/vllm/v1/kv_offload/cpu/manager.py index 7d92844d1f4..b48abecec1b 100644 --- a/vllm/v1/kv_offload/cpu/manager.py +++ b/vllm/v1/kv_offload/cpu/manager.py @@ -140,6 +140,7 @@ class CPUOffloadingManager(OffloadingManager): assert block is not None, f"Block {key!r} not found in cache" assert block.is_ready, f"Block {key!r} is not ready for reading" if block.ref_cnt == 0: + self._policy.mark_non_evictable(key) self._num_evictable_cache_blocks -= 1 # ref_cnt 0 -> 1 assert self._num_evictable_cache_blocks >= 0 block.ref_cnt += 1 @@ -161,6 +162,7 @@ class CPUOffloadingManager(OffloadingManager): block.ref_cnt -= 1 if block.ref_cnt == 0: self._num_evictable_cache_blocks += 1 # ref_cnt 1 -> 0 + self._policy.mark_evictable(key) @override def prepare_store( @@ -248,6 +250,7 @@ class CPUOffloadingManager(OffloadingManager): if block is not None and not block.is_ready: block.ref_cnt = 0 self._num_evictable_cache_blocks += 1 + self._policy.mark_evictable(key) stored_keys.append(key) else: for key in keys: diff --git a/vllm/v1/kv_offload/cpu/policies/base.py b/vllm/v1/kv_offload/cpu/policies/base.py index 0febfe90d61..f898a60b0f6 100644 --- a/vllm/v1/kv_offload/cpu/policies/base.py +++ b/vllm/v1/kv_offload/cpu/policies/base.py @@ -82,3 +82,11 @@ class CachePolicy(ABC): Ghost lists and adaptive state are also reset. """ + + def mark_evictable(self, key: OffloadKey) -> None: + """Called when a block's ref_cnt transitions to 0.""" + return + + def mark_non_evictable(self, key: OffloadKey) -> None: + """Called when a block's ref_cnt transitions from 0.""" + return diff --git a/vllm/v1/kv_offload/cpu/policies/lru.py b/vllm/v1/kv_offload/cpu/policies/lru.py index 75fbc6015e1..47e18f5565f 100644 --- a/vllm/v1/kv_offload/cpu/policies/lru.py +++ b/vllm/v1/kv_offload/cpu/policies/lru.py @@ -10,11 +10,18 @@ from vllm.v1.kv_offload.cpu.policies.base import BlockStatus, CachePolicy class LRUCachePolicy(CachePolicy): - """LRU cache policy backed by a single OrderedDict.""" + """ + LRU Caching policy that keeps a dedicated evictable list for fast eviction. + A use is indicated by, + - First time the key is added (store). + - Load job completion + - touch + """ def __init__(self, cache_capacity: int): - # cache_capacity unused by LRU but accepted for a uniform constructor - self.blocks: OrderedDict[OffloadKey, BlockStatus] = OrderedDict() + # Blocks with ref_cnt 0 (not participating in any loads/stores) ordered in LRU + self.evictable_blocks: OrderedDict[OffloadKey, None] = OrderedDict() + self.blocks: dict[OffloadKey, BlockStatus] = {} @override def get(self, key: OffloadKey) -> BlockStatus | None: @@ -23,19 +30,25 @@ class LRUCachePolicy(CachePolicy): @override def insert(self, key: OffloadKey, block: BlockStatus) -> None: self.blocks[key] = block + if block.ref_cnt == 0: + self.evictable_blocks[key] = None @override def remove(self, key: OffloadKey) -> None: del self.blocks[key] + self.evictable_blocks.pop(key, None) @override def touch(self, keys: Iterable[OffloadKey]) -> None: for key in reversed(list(keys)): - if key in self.blocks: - self.blocks.move_to_end(key) + if key in self.evictable_blocks: + self.evictable_blocks.move_to_end(key) + # active blocks are untouched as they are non-evictable now. They + # will eventually reach the end of evictable_blocks when they finish. @override def clear(self) -> None: + self.evictable_blocks.clear() self.blocks.clear() @override @@ -44,14 +57,33 @@ class LRUCachePolicy(CachePolicy): ) -> list[tuple[OffloadKey, BlockStatus]] | None: if n == 0: return [] + candidates: list[tuple[OffloadKey, BlockStatus]] = [] - for key, block in self.blocks.items(): - if block.ref_cnt == 0 and key not in protected: - candidates.append((key, block)) - if len(candidates) == n: - break + for key, _ in self.evictable_blocks.items(): + if key in protected: + continue + + block = self.blocks[key] + assert block.ref_cnt == 0 + candidates.append((key, block)) + if len(candidates) == n: + break + if len(candidates) < n: return None for key, _ in candidates: + del self.evictable_blocks[key] del self.blocks[key] return candidates + + @override + def mark_evictable(self, key: OffloadKey) -> None: + # blocks can become evictable when, + # store completes - i.e. ref_cnt -1 -> 0 # not in evictable list + # all loads complete - i.e ref_cnt 1 -> 0 # not in evictable list + self.evictable_blocks[key] = None + + @override + def mark_non_evictable(self, key: OffloadKey) -> None: + # key must have been in the evictable list. + del self.evictable_blocks[key]