A high-throughput and memory-efficient inference and serving engine for LLMs
[CPUOffloading] Guard CPU eviction check (#45757)
Signed-off-by: Varun Sundar Rabindranath <varun-sundar-rabindranath@h100-01.nemg-001.lab.rdu2.dc.redhat.com> Co-authored-by: Varun Sundar Rabindranath <varun-sundar-rabindranath@h100-01.nemg-001.lab.rdu2.dc.redhat.com>
Varun Sundar Rabindranath committed
Jun 17, 2026 at 22:34 UTC
ed938ad7db9c28e3725058037c41285d8f46869e
2 files changed
+114
tests/v1/kv_offload/cpu/test_manager.py
+93
@@ -689,3 +689,96 @@ def test_filter_reused_manager():
689
assert prepare_store_output.keys_to_store == []
690
691
manager.complete_store(to_keys([1]), _EMPTY_REQ_CTX)
692
+
693
+
694
+def test_evictable_cache_block_count():
695
+ """
696
+ Verifies _num_evictable_cache_blocks is maintained correctly through the
697
+ full store/load lifecycle, eviction, failed stores, concurrent loads,
698
+ reset_cache, and the early-exit fast path in prepare_store.
699
+ """
700
+ manager = make_cpu_manager(num_blocks=4, cache_policy="lru")
701
+
702
+ # Initially no blocks allocated.
703
+ assert manager._num_evictable_cache_blocks == 0
704
+
705
+ # Initial cache state [x, x, x, x]
706
+
707
+ # We get 3 blocks from the cache.
708
+ manager.prepare_store(to_keys([1, 2, 3]), _EMPTY_REQ_CTX)
709
+ # cache state [1', 2', 3', x] <- 1', 2', 3' are actively being used.
710
+ assert manager._num_evictable_cache_blocks == 0
711
+
712
+ # Completing stores makes them idle.
713
+ manager.complete_store(to_keys([1, 2, 3]), _EMPTY_REQ_CTX)
714
+ # cache state [1, 2, 3, x] <- 1, 2, 3 blocks are idle.
715
+ assert manager._num_evictable_cache_blocks == 3
716
+
717
+ # prepare_load pins a block: idle count decrements once even if the
718
+ # same block is loaded by two concurrent callers.
719
+ manager.prepare_load(to_keys([1]), _EMPTY_REQ_CTX)
720
+ # cache state [1', 2, 3, x] <- 2, 3 blocks are idle.
721
+ assert manager._num_evictable_cache_blocks == 2
722
+ manager.prepare_load(to_keys([1]), _EMPTY_REQ_CTX) # 2nd concurrent load
723
+ # cache state [1', 2, 3, x] <- 2, 3 blocks are idle.
724
+ assert manager._num_evictable_cache_blocks == 2 # no double-decrement
725
+
726
+ # First complete_load does not restore idle (ref_cnt still 1).
727
+ manager.complete_load(to_keys([1]), _EMPTY_REQ_CTX)
728
+ # cache state [1', 2, 3, x] <- 2, 3 blocks are idle.
729
+ assert manager._num_evictable_cache_blocks == 2
730
+ # Second complete_load drops ref_cnt to 0 -> block becomes idle again.
731
+ manager.complete_load(to_keys([1]), _EMPTY_REQ_CTX)
732
+ # cache state [1, 2, 3, x] <- 1, 2, 3 blocks are idle.
733
+ assert manager._num_evictable_cache_blocks == 3
734
+
735
+ # Eviction decrements idle count.
736
+ # Cache has 3 stored blocks and 1 free slot. Storing 3 new keys needs 2 eviction.
737
+ manager.prepare_store(to_keys([4, 5, 6]), _EMPTY_REQ_CTX)
738
+ # cache state [1, 4', 5', 6'] <- block 1 is idle
739
+ assert manager._num_evictable_cache_blocks == 1
740
+
741
+ # Failed store does not increment idle count (block discarded from cache).
742
+ manager.complete_store(to_keys([4, 5, 6]), _EMPTY_REQ_CTX, success=False)
743
+ # cache state [1, x, x, x] <- block 1 is idle. Other returned to cache.
744
+ assert manager._num_evictable_cache_blocks == 1
745
+
746
+ # reset_cache zeroes the count unconditionally.
747
+ manager.reset_cache()
748
+ # cache state [x, x, x, x]
749
+ assert manager._num_evictable_cache_blocks == 0
750
+
751
+ # setup 3 blocks with loads so idle count drops to 0.
752
+ manager.prepare_store(to_keys([10, 11, 12]), _EMPTY_REQ_CTX)
753
+ manager.complete_store(to_keys([10, 11, 12]), _EMPTY_REQ_CTX)
754
+ manager.prepare_load(to_keys([10, 11, 12]), _EMPTY_REQ_CTX)
755
+ # cache state [10', 11', 12', x]
756
+ assert manager._num_evictable_cache_blocks == 0
757
+
758
+ # prepare_store requiring eviction must return None immediately (fast exit).
759
+ # Spy on policy.evict to confirm the fast path short-circuits before calling it.
760
+ evict_called = False
761
+ original_evict = manager._policy.evict
762
+
763
+ def spy_evict(*args, **kwargs):
764
+ nonlocal evict_called
765
+ evict_called = True
766
+ return original_evict(*args, **kwargs)
767
+
768
+ manager._policy.evict = spy_evict # type: ignore[method-assign]
769
+ # cache state [10', 11', 12', x] <- cannot evict anything
770
+ assert manager.prepare_store(to_keys([14, 15]), _EMPTY_REQ_CTX) is None
771
+ assert not evict_called, (
772
+ "_num_evictable_cache_blocks==0 should short-circuit before evict()"
773
+ )
774
+
775
+ # After releasing the loads, eviction becomes possible again.
776
+ manager.complete_load(to_keys([10, 11, 12]), _EMPTY_REQ_CTX)
777
+ # cache state [10, 11, 12, x] <- 10, 11, 12 are idle
778
+ assert manager._num_evictable_cache_blocks == 3
779
+ assert manager.prepare_store(to_keys([14, 15]), _EMPTY_REQ_CTX) is not None
780
+ # cache state [10, 11, 14', 15'] <- 10, 11 are idle
781
+ assert manager._num_evictable_cache_blocks == 2
782
+ manager.complete_store(to_keys([14, 15]), _EMPTY_REQ_CTX)
783
+ # cache state [10, 11, 14, 15] <- all blocks idle
784
+ assert manager._num_evictable_cache_blocks == 4
vllm/v1/kv_offload/cpu/manager.py
+21
@@ -59,6 +59,9 @@ class CPUOffloadingManager(OffloadingManager):
59
f"Supported: {list(_CACHE_POLICIES)}"
60
)
61
self._policy: CachePolicy = policy_cls(cache_capacity=num_blocks)
62
+ # Track the number of blocks in the cache that are evictable. i.e. ref_cnt 0.
63
+ self._num_evictable_cache_blocks: int = 0
64
+
65
self.store_threshold: int = store_threshold
66
self.max_tracker_size: int = max_tracker_size
67
self.stores_skipped_in_current_batch: int = 0
@@ -133,6 +136,9 @@ class CPUOffloadingManager(OffloadingManager):
136
block = self._policy.get(key)
137
assert block is not None, f"Block {key!r} not found in cache"
138
assert block.is_ready, f"Block {key!r} is not ready for reading"
139
+ if block.ref_cnt == 0:
140
+ self._num_evictable_cache_blocks -= 1 # ref_cnt 0 -> 1
141
+ assert self._num_evictable_cache_blocks >= 0
142
block.ref_cnt += 1
143
blocks.append(block)
144
return self._get_load_store_spec(keys, blocks)
@@ -150,6 +156,8 @@ class CPUOffloadingManager(OffloadingManager):
156
assert block is not None, f"Block {key!r} not found"
157
assert block.ref_cnt > 0, f"Block {key!r} ref_cnt is already 0"
158
block.ref_cnt -= 1
159
+ if block.ref_cnt == 0:
160
+ self._num_evictable_cache_blocks += 1 # ref_cnt 1 -> 0
161
162
@override
163
def prepare_store(
@@ -175,12 +183,23 @@ class CPUOffloadingManager(OffloadingManager):
183
184
to_evict: list[OffloadKey] = []
185
if num_blocks_to_evict > 0:
186
+ if num_blocks_to_evict > self._num_evictable_cache_blocks:
187
+ # Eviction will fail.
188
+ return None
189
+ # There is a still a chance for eviction failure as some of the
190
+ # idle blocks might be in the protected list.
191
+
192
# Blocks from the original input are excluded from eviction candidates:
193
# a block that was already stored must remain in the cache after this call.
194
protected = set(keys)
195
evicted = self._policy.evict(num_blocks_to_evict, protected)
196
if evicted is None:
197
return None
198
+
199
+ # cache-policy removes only idle blocks.
200
+ self._num_evictable_cache_blocks -= len(evicted)
201
+ assert self._num_evictable_cache_blocks >= 0
202
+
203
for key, block in evicted:
204
self._free_block(block)
205
to_evict.append(key)
@@ -225,6 +244,7 @@ class CPUOffloadingManager(OffloadingManager):
244
block = self._policy.get(key)
245
if block is not None and not block.is_ready:
246
block.ref_cnt = 0
247
+ self._num_evictable_cache_blocks += 1
248
stored_keys.append(key)
249
else:
250
for key in keys:
@@ -250,6 +270,7 @@ class CPUOffloadingManager(OffloadingManager):
270
# flushes in-flight load job IDs to the workers before any new stores
271
# can begin, preventing a cross-direction data race on reused offload block IDs.
272
self._policy.clear()
273
+ self._num_evictable_cache_blocks = 0
274
275
self._free_list.clear()
276
self._num_allocated_blocks = 0