diff --git a/tests/v1/kv_offload/cpu/test_manager.py b/tests/v1/kv_offload/cpu/test_manager.py index 2548c79727..a2b80418aa 100644 --- a/tests/v1/kv_offload/cpu/test_manager.py +++ b/tests/v1/kv_offload/cpu/test_manager.py @@ -965,6 +965,28 @@ def test_filter_reused_manager(): assert manager.counts.get(to_keys([1])[0]) == 1 +def test_filter_reused_manager_oversized_offer_makes_progress(): + """An offer larger than the tracker capacity must not churn forever.""" + manager = make_cpu_manager( + num_blocks=4, + cache_policy="lru", + store_threshold=2, + max_tracker_size=3, + ) + keys = to_keys([1, 2, 3, 4]) + stored_keys: set[OffloadKey] = set() + + for _ in range(4): + remaining_keys = [key for key in keys if key not in stored_keys] + output = manager.prepare_store(remaining_keys, _EMPTY_REQ_CTX) + assert output is not None + if output.keys_to_store: + manager.complete_store(output.keys_to_store, _EMPTY_REQ_CTX) + stored_keys.update(output.keys_to_store) + + assert stored_keys == set(keys) + + def test_evictable_cache_block_count(): """ Verifies _num_evictable_cache_blocks is maintained correctly through the diff --git a/vllm/v1/kv_offload/cpu/manager.py b/vllm/v1/kv_offload/cpu/manager.py index 484ac8b088..9082c83de4 100644 --- a/vllm/v1/kv_offload/cpu/manager.py +++ b/vllm/v1/kv_offload/cpu/manager.py @@ -103,15 +103,25 @@ class CPUOffloadingManager(OffloadingManager): ) -> CPULoadStoreSpec: return CPULoadStoreSpec([block.block_id for block in blocks]) - def _record_access(self, key: OffloadKey) -> None: - """Count one observation of ``key`` for store admission.""" + def _record_accesses(self, keys: Collection[OffloadKey]) -> None: + """Record an offer without evicting its tracked candidates.""" assert self.counts is not None - if key in self.counts: - self.counts.move_to_end(key) - self.counts[key] += 1 - else: + protected: set[OffloadKey] = set() + for key in keys: + if key in self.counts: + self.counts.move_to_end(key) + self.counts[key] += 1 + protected.add(key) + + num_unprotected = len(self.counts) - len(protected) + for key in keys: + if key in self.counts: + continue if len(self.counts) >= self.max_tracker_size: + if num_unprotected == 0: + continue self.counts.popitem(last=False) + num_unprotected -= 1 self.counts[key] = 1 # --- OffloadingManager interface --- @@ -173,8 +183,7 @@ class CPUOffloadingManager(OffloadingManager): ) -> PrepareStoreOutput | None: if self.counts is not None: num_keys = len(keys) - for key in keys: - self._record_access(key) + self._record_accesses(keys) keys = [k for k in keys if self.counts.get(k, 0) >= self.store_threshold] self.stores_skipped_in_current_batch += num_keys - len(keys) # filter out blocks that are already stored