A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 7, 06:30 UTC
[Bugfix][CPUOffloadingManager] Prevent eviction of already-stored blocks in LRU/ARC `prepare_store()` (#35846)
Signed-off-by: Ronen Schaffer <ronen.schaffer@ibm.com>
Ronen Schaffer committed
Mar 4, 2026 at 14:25 UTC
bb6888b8b1a03e683ff4ed5f1fb6df5a0582fd6f
3 files changed
+67
-5
tests/v1/kv_offload/test_cpu_manager.py
+49
@@ -4,6 +4,7 @@ from collections.abc import Iterable
4
from dataclasses import dataclass
5
6
import numpy as np
7
+import pytest
8
9
from vllm.v1.core.kv_cache_utils import BlockHash
10
from vllm.v1.kv_offload.abstract import (
@@ -78,6 +79,54 @@ def verify_events(
79
assert tuple(stores) == to_hash_sets(expected_stores)
80
81
82
+@pytest.mark.parametrize("manager_class", [LRUOffloadingManager, ARCOffloadingManager])
83
+def test_already_stored_block_not_evicted_during_prepare_store(manager_class):
84
+ """
85
+ Regression test: a block that is already stored must not be evicted
86
+ by prepare_store() when it needs to make room for new blocks.
87
+ Applies to both LRUOffloadingManager and ARCOffloadingManager.
88
+
89
+ Scenario:
90
+ - Store blocks [1, 2] and complete.
91
+ - touch([1]) makes block 2 the LRU candidate.
92
+ - prepare_store([2, 3, 4, 5]):
93
+ * block 2 is filtered out as "already stored"
94
+ * but without the fix, block 2 would be evicted as the LRU
95
+ candidate to make room for [3, 4, 5]
96
+ - After complete_store([2, 3, 4, 5]), block 2 must still be present.
97
+ """
98
+ block_size = 256
99
+ cpu_backend = CPUBackend(block_size=block_size, num_blocks=4)
100
+ manager = manager_class(cpu_backend, enable_events=True)
101
+
102
+ # store [1, 2] and complete
103
+ manager.prepare_store(to_hashes([1, 2]))
104
+ manager.complete_store(to_hashes([1, 2]))
105
+
106
+ # touch [1] to make block 2 the LRU candidate
107
+ manager.touch(to_hashes([1]))
108
+
109
+ # prepare_store([2, 3, 4, 5]):
110
+ # - block 2 is already stored → filtered out of block_hashes_to_store
111
+ # - block 2 must NOT be evicted even though it is the LRU candidate
112
+ # - block 1 (ID 0) is evicted instead; new blocks [3,4,5] get IDs 2,3,0
113
+ prepare_store_output = manager.prepare_store(to_hashes([2, 3, 4, 5]))
114
+ verify_store_output(
115
+ prepare_store_output,
116
+ ExpectedPrepareStoreOutput(
117
+ block_hashes_to_store=[3, 4, 5],
118
+ store_block_ids=[2, 3, 0],
119
+ block_hashes_evicted=[1], # block 1 evicted, not block 2
120
+ ),
121
+ )
122
+
123
+ # complete_store must not silently drop block 2
124
+ manager.complete_store(to_hashes([2, 3, 4, 5]))
125
+
126
+ # block 2 must still be present in the cache
127
+ assert manager.lookup(to_hashes([2])) == 1
128
+
129
+
130
def test_cpu_manager():
131
"""
132
Tests LRUOffloadingManager with a CPUBackend.
vllm/v1/kv_offload/arc_manager.py
+9
-3
@@ -123,8 +123,10 @@ class ARCOffloadingManager(OffloadingManager):
123
def prepare_store(
124
self, block_hashes: Iterable[BlockHash]
125
) -> PrepareStoreOutput | None:
126
+ block_hashes_list = list(block_hashes)
127
+
128
block_hashes_to_store = []
127
- for block_hash in block_hashes:
129
+ for block_hash in block_hashes_list:
130
if block_hash not in self.t1 and block_hash not in self.t2:
131
block_hashes_to_store.append(block_hash)
132
@@ -140,12 +142,16 @@ class ARCOffloadingManager(OffloadingManager):
142
)
143
144
to_evict = []
145
+ if num_blocks_to_evict > 0:
146
+ # Blocks from the original input are excluded from eviction candidates:
147
+ # a block that was already stored must remain in the cache after this call.
148
+ protected = set(block_hashes_list)
149
while num_blocks_to_evict > 0:
150
block_to_evict = None
151
if len(self.t1) >= int(self.target_t1_size):
152
# try to evict the least recently used (oldest) block from T1
153
for block_hash, block in self.t1.items():
148
- if block.ref_cnt == 0:
154
+ if block.ref_cnt == 0 and block_hash not in protected:
155
block_to_evict = (block_hash, block)
156
eviction_t = self.t1
157
eviction_b = self.b1
@@ -153,7 +159,7 @@ class ARCOffloadingManager(OffloadingManager):
159
if not block_to_evict:
160
# try to evict the least recently used (oldest) block from T2
161
for block_hash, block in self.t2.items():
156
- if block.ref_cnt == 0:
162
+ if block.ref_cnt == 0 and block_hash not in protected:
163
block_to_evict = (block_hash, block)
164
eviction_t = self.t2
165
eviction_b = self.b2
vllm/v1/kv_offload/lru_manager.py
+9
-2
@@ -57,9 +57,13 @@ class LRUOffloadingManager(OffloadingManager):
57
def prepare_store(
58
self, block_hashes: Iterable[BlockHash]
59
) -> PrepareStoreOutput | None:
60
+ block_hashes_list = list(block_hashes)
61
+
62
# filter out blocks that are already stored
63
block_hashes_to_store = [
62
- block_hash for block_hash in block_hashes if block_hash not in self.blocks
64
+ block_hash
65
+ for block_hash in block_hashes_list
66
+ if block_hash not in self.blocks
67
]
68
69
num_blocks_to_evict = (
@@ -69,8 +73,11 @@ class LRUOffloadingManager(OffloadingManager):
73
# build list of blocks to evict
74
to_evict = []
75
if num_blocks_to_evict > 0:
76
+ # Blocks from the original input are excluded from eviction candidates:
77
+ # a block that was already stored must remain in the cache after this call.
78
+ protected = set(block_hashes_list)
79
for block_hash, block in self.blocks.items():
73
- if block.ref_cnt == 0:
80
+ if block.ref_cnt == 0 and block_hash not in protected:
81
to_evict.append(block_hash)
82
num_blocks_to_evict -= 1
83
if num_blocks_to_evict == 0: