[Bugfix][KV Offload] Ensure tracker progress for oversized offers (#54759)
Signed-off-by: Cheng Rui <286040359@qq.com> Co-authored-by: Codex <noreply@openai.com>
Cheng Rui committed
Sep 2, 2026 at 04:09 UTC
2a4e3cc3dbddf4f44d69864209361f1e2a70c79a
2 files changed
+39
-8
tests/v1/kv_offload/cpu/test_manager.py
+22
@@ -965,6 +965,28 @@ def test_filter_reused_manager():
965
assert manager.counts.get(to_keys([1])[0]) == 1
966
967
968
+def test_filter_reused_manager_oversized_offer_makes_progress():
969
+ """An offer larger than the tracker capacity must not churn forever."""
970
+ manager = make_cpu_manager(
971
+ num_blocks=4,
972
+ cache_policy="lru",
973
+ store_threshold=2,
974
+ max_tracker_size=3,
975
+ )
976
+ keys = to_keys([1, 2, 3, 4])
977
+ stored_keys: set[OffloadKey] = set()
978
+
979
+ for _ in range(4):
980
+ remaining_keys = [key for key in keys if key not in stored_keys]
981
+ output = manager.prepare_store(remaining_keys, _EMPTY_REQ_CTX)
982
+ assert output is not None
983
+ if output.keys_to_store:
984
+ manager.complete_store(output.keys_to_store, _EMPTY_REQ_CTX)
985
+ stored_keys.update(output.keys_to_store)
986
+
987
+ assert stored_keys == set(keys)
988
+
989
+
990
def test_evictable_cache_block_count():
991
"""
992
Verifies _num_evictable_cache_blocks is maintained correctly through the
vllm/v1/kv_offload/cpu/manager.py
+17
-8
@@ -103,15 +103,25 @@ class CPUOffloadingManager(OffloadingManager):
103
) -> CPULoadStoreSpec:
104
return CPULoadStoreSpec([block.block_id for block in blocks])
105
106
- def _record_access(self, key: OffloadKey) -> None:
107
- """Count one observation of ``key`` for store admission."""
106
+ def _record_accesses(self, keys: Collection[OffloadKey]) -> None:
107
+ """Record an offer without evicting its tracked candidates."""
108
assert self.counts is not None
109
- if key in self.counts:
110
- self.counts.move_to_end(key)
111
- self.counts[key] += 1
112
- else:
109
+ protected: set[OffloadKey] = set()
110
+ for key in keys:
111
+ if key in self.counts:
112
+ self.counts.move_to_end(key)
113
+ self.counts[key] += 1
114
+ protected.add(key)
115
+
116
+ num_unprotected = len(self.counts) - len(protected)
117
+ for key in keys:
118
+ if key in self.counts:
119
+ continue
120
if len(self.counts) >= self.max_tracker_size:
121
+ if num_unprotected == 0:
122
+ continue
123
self.counts.popitem(last=False)
124
+ num_unprotected -= 1
125
self.counts[key] = 1
126
127
# --- OffloadingManager interface ---
@@ -173,8 +183,7 @@ class CPUOffloadingManager(OffloadingManager):
183
) -> PrepareStoreOutput | None:
184
if self.counts is not None:
185
num_keys = len(keys)
176
- for key in keys:
177
- self._record_access(key)
186
+ self._record_accesses(keys)
187
keys = [k for k in keys if self.counts.get(k, 0) >= self.store_threshold]
188
self.stores_skipped_in_current_batch += num_keys - len(keys)
189
# filter out blocks that are already stored