[CPUOffloadingManager] Maintain evictable list in LRUCachePolicy (#46216)

Signed-off-by: <>
Co-authored-by: Varun Sundar Rabindranath <varun-sundar-rabindranath@h100-01.nemg-001.lab.rdu2.dc.redhat.com>

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}),
     )
 
 

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

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:

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

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]