[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>
Varun Sundar Rabindranath committed
Jun 22, 2026 at 02:54 UTC
68567ef2dfa5fbaa7cf9c5ff4eb556c17c5b8bbd
5 files changed
+64
-19
tests/v1/kv_offload/cpu/test_manager.py
+8
-8
@@ -294,25 +294,25 @@ def test_cpu_manager():
294
# prepare store with no space ([2, 3] is being loaded)
295
assert cpu_manager.prepare_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX) is None
296
297
- # complete load [2, 3]
297
+ # complete load [2, 3]. Load changes the eviction list, making 2, 3 recent.
298
cpu_manager.complete_load(to_keys([2, 3]), _EMPTY_REQ_CTX)
299
300
- # prepare store [6, 7, 8] -> evicts [2, 3, 4] (oldest)
300
+ # prepare store [6, 7, 8] -> evicts [4, 5, 2] (oldest)
301
prepare_store_output = cpu_manager.prepare_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX)
302
verify_store_output(
303
prepare_store_output,
304
ExpectedPrepareStoreOutput(
305
keys_to_store=[6, 7, 8],
306
- store_block_ids=[3, 2, 1],
307
- evicted_keys=[2, 3, 4],
306
+ store_block_ids=[1, 0, 3],
307
+ evicted_keys=[4, 5, 2],
308
),
309
)
310
311
# complete store [6, 7, 8]
312
cpu_manager.complete_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX)
313
314
- # touch [5, 6, 7] (move to end of LRU order)
315
- cpu_manager.touch(to_keys([5, 6, 7]), _EMPTY_REQ_CTX)
314
+ # touch [3, 6, 7] (move to end of LRU order)
315
+ cpu_manager.touch(to_keys([3, 6, 7]), _EMPTY_REQ_CTX)
316
317
# prepare store [7, 9] -> evicts [8] (oldest following previous touch)
318
prepare_store_output = cpu_manager.prepare_store(to_keys([9]), _EMPTY_REQ_CTX)
@@ -320,7 +320,7 @@ def test_cpu_manager():
320
prepare_store_output,
321
ExpectedPrepareStoreOutput(
322
keys_to_store=[9],
323
- store_block_ids=[1],
323
+ store_block_ids=[3],
324
evicted_keys=[8],
325
),
326
)
@@ -335,7 +335,7 @@ def test_cpu_manager():
335
verify_events(
336
cpu_manager.take_events(),
337
expected_stores=({3, 4, 5}, {6, 7, 8}),
338
- expected_evictions=({2, 3, 4}, {8}),
338
+ expected_evictions=({4, 5, 2}, {8}),
339
)
340
341
tests/v1/kv_offload/tiering/test_tiering_offloading.py
+3
-1
@@ -295,6 +295,8 @@ class TestTieringOffloadingManager:
295
self.manager.prepare_store(blocks, _CTX)
296
self.manager.complete_store(blocks, _CTX, success=True)
297
self._simulate_on_schedule_end()
298
+ # for secondary tiers to drain jobs, so primary tier's blocks are evictable.
299
+ self._simulate_on_schedule_end()
300
301
self.secondary_tier1.touch = MagicMock(wraps=self.secondary_tier1.touch)
302
self.secondary_tier2.touch = MagicMock(wraps=self.secondary_tier2.touch)
@@ -303,7 +305,7 @@ class TestTieringOffloadingManager:
305
self.manager.touch(blocks, _CTX)
306
307
# Verify touch was called on primary tier (check LRU order)
306
- primary_keys = list(self.primary_tier._policy.blocks.keys())
308
+ primary_keys = list(self.primary_tier._policy.evictable_blocks.keys())
309
assert primary_keys[-3:] == list(reversed(blocks))
310
311
# Verify touch was propagated to all secondary tiers
vllm/v1/kv_offload/cpu/manager.py
+3
@@ -140,6 +140,7 @@ class CPUOffloadingManager(OffloadingManager):
140
assert block is not None, f"Block {key!r} not found in cache"
141
assert block.is_ready, f"Block {key!r} is not ready for reading"
142
if block.ref_cnt == 0:
143
+ self._policy.mark_non_evictable(key)
144
self._num_evictable_cache_blocks -= 1 # ref_cnt 0 -> 1
145
assert self._num_evictable_cache_blocks >= 0
146
block.ref_cnt += 1
@@ -161,6 +162,7 @@ class CPUOffloadingManager(OffloadingManager):
162
block.ref_cnt -= 1
163
if block.ref_cnt == 0:
164
self._num_evictable_cache_blocks += 1 # ref_cnt 1 -> 0
165
+ self._policy.mark_evictable(key)
166
167
@override
168
def prepare_store(
@@ -248,6 +250,7 @@ class CPUOffloadingManager(OffloadingManager):
250
if block is not None and not block.is_ready:
251
block.ref_cnt = 0
252
self._num_evictable_cache_blocks += 1
253
+ self._policy.mark_evictable(key)
254
stored_keys.append(key)
255
else:
256
for key in keys:
vllm/v1/kv_offload/cpu/policies/base.py
+8
@@ -82,3 +82,11 @@ class CachePolicy(ABC):
82
83
Ghost lists and adaptive state are also reset.
84
"""
85
+
86
+ def mark_evictable(self, key: OffloadKey) -> None:
87
+ """Called when a block's ref_cnt transitions to 0."""
88
+ return
89
+
90
+ def mark_non_evictable(self, key: OffloadKey) -> None:
91
+ """Called when a block's ref_cnt transitions from 0."""
92
+ return
vllm/v1/kv_offload/cpu/policies/lru.py
+42
-10
@@ -10,11 +10,18 @@ from vllm.v1.kv_offload.cpu.policies.base import BlockStatus, CachePolicy
10
11
12
class LRUCachePolicy(CachePolicy):
13
- """LRU cache policy backed by a single OrderedDict."""
13
+ """
14
+ LRU Caching policy that keeps a dedicated evictable list for fast eviction.
15
+ A use is indicated by,
16
+ - First time the key is added (store).
17
+ - Load job completion
18
+ - touch
19
+ """
20
21
def __init__(self, cache_capacity: int):
16
- # cache_capacity unused by LRU but accepted for a uniform constructor
17
- self.blocks: OrderedDict[OffloadKey, BlockStatus] = OrderedDict()
22
+ # Blocks with ref_cnt 0 (not participating in any loads/stores) ordered in LRU
23
+ self.evictable_blocks: OrderedDict[OffloadKey, None] = OrderedDict()
24
+ self.blocks: dict[OffloadKey, BlockStatus] = {}
25
26
@override
27
def get(self, key: OffloadKey) -> BlockStatus | None:
@@ -23,19 +30,25 @@ class LRUCachePolicy(CachePolicy):
30
@override
31
def insert(self, key: OffloadKey, block: BlockStatus) -> None:
32
self.blocks[key] = block
33
+ if block.ref_cnt == 0:
34
+ self.evictable_blocks[key] = None
35
36
@override
37
def remove(self, key: OffloadKey) -> None:
38
del self.blocks[key]
39
+ self.evictable_blocks.pop(key, None)
40
41
@override
42
def touch(self, keys: Iterable[OffloadKey]) -> None:
43
for key in reversed(list(keys)):
34
- if key in self.blocks:
35
- self.blocks.move_to_end(key)
44
+ if key in self.evictable_blocks:
45
+ self.evictable_blocks.move_to_end(key)
46
+ # active blocks are untouched as they are non-evictable now. They
47
+ # will eventually reach the end of evictable_blocks when they finish.
48
49
@override
50
def clear(self) -> None:
51
+ self.evictable_blocks.clear()
52
self.blocks.clear()
53
54
@override
@@ -44,14 +57,33 @@ class LRUCachePolicy(CachePolicy):
57
) -> list[tuple[OffloadKey, BlockStatus]] | None:
58
if n == 0:
59
return []
60
+
61
candidates: list[tuple[OffloadKey, BlockStatus]] = []
48
- for key, block in self.blocks.items():
49
- if block.ref_cnt == 0 and key not in protected:
50
- candidates.append((key, block))
51
- if len(candidates) == n:
52
- break
62
+ for key, _ in self.evictable_blocks.items():
63
+ if key in protected:
64
+ continue
65
+
66
+ block = self.blocks[key]
67
+ assert block.ref_cnt == 0
68
+ candidates.append((key, block))
69
+ if len(candidates) == n:
70
+ break
71
+
72
if len(candidates) < n:
73
return None
74
for key, _ in candidates:
75
+ del self.evictable_blocks[key]
76
del self.blocks[key]
77
return candidates
78
+
79
+ @override
80
+ def mark_evictable(self, key: OffloadKey) -> None:
81
+ # blocks can become evictable when,
82
+ # store completes - i.e. ref_cnt -1 -> 0 # not in evictable list
83
+ # all loads complete - i.e ref_cnt 1 -> 0 # not in evictable list
84
+ self.evictable_blocks[key] = None
85
+
86
+ @override
87
+ def mark_non_evictable(self, key: OffloadKey) -> None:
88
+ # key must have been in the evictable list.
89
+ del self.evictable_blocks[key]