diff --git a/tests/v1/kv_offload/cpu/test_manager.py b/tests/v1/kv_offload/cpu/test_manager.py index 6b33041791..2548c79727 100644 --- a/tests/v1/kv_offload/cpu/test_manager.py +++ b/tests/v1/kv_offload/cpu/test_manager.py @@ -932,41 +932,37 @@ def test_filter_reused_manager(): max_tracker_size=3, ) - # Lookup [1, 2] -> 1st time, added to tracker but not eligible for store yet + # lookup() does not count towards store admission assert manager.lookup(to_key(1), _EMPTY_REQ_CTX) is LookupResult.MISS assert manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.MISS + assert not manager.counts - # prepare store [1, 2] -> should be filtered + # 1st offer of [1, 2] -> tracked at count 1, not eligible yet prepare_store_output = manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) assert prepare_store_output is not None assert prepare_store_output.keys_to_store == [] + assert manager.counts[to_keys([1])[0]] == 1 + assert manager.counts[to_keys([2])[0]] == 1 - # Lookup [1] -> 2nd time, eligible now - assert manager.lookup(to_key(1), _EMPTY_REQ_CTX) is LookupResult.MISS - - # prepare store [1, 2] -> [1] should be eligible, [2] should be filtered + # 2nd offer -> the whole repeated prefix becomes eligible at once prepare_store_output = manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) assert prepare_store_output is not None - assert prepare_store_output.keys_to_store == to_keys([1]) + assert prepare_store_output.keys_to_store == to_keys([1, 2]) + manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX) - # Lookup [3, 4] -> 1st time - # (evicts [2] from tracker since max_size is 3 and tracker has [1]) - assert manager.lookup(to_key(3), _EMPTY_REQ_CTX) is LookupResult.MISS - assert manager.lookup(to_key(4), _EMPTY_REQ_CTX) is LookupResult.MISS - # Verify [2] was evicted from the tracker (tracker now has: [1], [3], [4]) - assert to_keys([2])[0] not in manager.counts - - # Lookup [2] again -> (this adds [2] back to the tracker as 1st time) - assert manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.MISS - # Verify [2] was re-added with count=1 (not eligible yet) - assert manager.counts.get(to_keys([2])[0]) == 1 - - # prepare store [2] -> should still be filtered out since count was reset - prepare_store_output = manager.prepare_store(to_keys([2]), _EMPTY_REQ_CTX) + # Offer [3, 4] -> 1st time, evicting the tracker's LRU entry [1] + prepare_store_output = manager.prepare_store(to_keys([3, 4]), _EMPTY_REQ_CTX) assert prepare_store_output is not None assert prepare_store_output.keys_to_store == [] + assert to_keys([1])[0] not in manager.counts + assert manager.counts[to_keys([3])[0]] == 1 + assert manager.counts[to_keys([4])[0]] == 1 - manager.complete_store(to_keys([1]), _EMPTY_REQ_CTX) + # [1] re-enters the tracker at count 1, so it is filtered again + prepare_store_output = manager.prepare_store(to_keys([1]), _EMPTY_REQ_CTX) + assert prepare_store_output is not None + assert prepare_store_output.keys_to_store == [] + assert manager.counts.get(to_keys([1])[0]) == 1 def test_evictable_cache_block_count(): diff --git a/vllm/v1/kv_offload/cpu/manager.py b/vllm/v1/kv_offload/cpu/manager.py index 0eef5cf7e7..484ac8b088 100644 --- a/vllm/v1/kv_offload/cpu/manager.py +++ b/vllm/v1/kv_offload/cpu/manager.py @@ -103,6 +103,17 @@ 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.""" + assert self.counts is not None + if key in self.counts: + self.counts.move_to_end(key) + self.counts[key] += 1 + else: + if len(self.counts) >= self.max_tracker_size: + self.counts.popitem(last=False) + self.counts[key] = 1 + # --- OffloadingManager interface --- @override @@ -111,14 +122,6 @@ class CPUOffloadingManager(OffloadingManager): @override def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: - if self.counts is not None: - if key in self.counts: - self.counts.move_to_end(key) - self.counts[key] += 1 - else: - if len(self.counts) >= self.max_tracker_size: - self.counts.popitem(last=False) - self.counts[key] = 1 block = self._policy.get(key) if block is None: return LookupResult.MISS @@ -170,6 +173,8 @@ class CPUOffloadingManager(OffloadingManager): ) -> PrepareStoreOutput | None: if self.counts is not None: num_keys = len(keys) + for key in keys: + self._record_access(key) 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 diff --git a/vllm/v1/kv_offload/cpu/spec.py b/vllm/v1/kv_offload/cpu/spec.py index 91d20c10b1..6915a821d5 100644 --- a/vllm/v1/kv_offload/cpu/spec.py +++ b/vllm/v1/kv_offload/cpu/spec.py @@ -123,9 +123,10 @@ class CPUOffloadingSpec(OffloadingSpec): @override def get_manager(self) -> OffloadingManager: if not self._manager: - # store_threshold: how many times a block must appear in lookup() - # before it is eligible for CPU offloading. Values < 2 disable - # filtering (a threshold of 1 equals no filter; 0 is the default). + # store_threshold: how many times a block must be offered for + # storage before it is eligible for CPU offloading. Values < 2 + # disable filtering (a threshold of 1 equals no filter; 0 is the + # default). store_threshold = int(self.extra_config.get("store_threshold", 0)) # Maximum entries in the internal tracker's LRU table.