Count store offers, not lookups, for CPU offload store_threshold (#52227)

Signed-off-by: okorzh-amd <okorzh-amd@users.noreply.github.com>
Co-authored-by: okorzh-amd <okorzh-amd@users.noreply.github.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

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():

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

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.