[Bugfix][KV Offload] Count each key once per request toward store_threshold

tests/v1/kv_offload/cpu/test_manager.py

@@ -1166,6 +1166,64 @@ def test_store_threshold_does_not_hide_ready_reuse():
     assert output.evicted_keys == [head]
 
 
+@pytest.mark.parametrize("store_threshold", [2, 3])
+def test_store_threshold_does_not_count_retries_of_a_failed_offer(
+    store_threshold: int,
+):
+    """After prepare_store() returns None, the scheduler offers the same keys
+    again. A retry must not admit a key that only one request offered, nor
+    count that key's skip again."""
+    manager = make_cpu_manager(num_chunks=2, store_threshold=store_threshold)
+    pinned = to_keys([100, 101])
+    warm, cold = to_keys([1, 2])
+
+    # Other requests offer `pinned` and `warm`, so one more offer admits them.
+    for i in range(store_threshold - 1):
+        manager.prepare_store([*pinned, warm], make_req_context(f"warmup-{i}"))
+    # Fill both chunks with stores that stay in flight.
+    pin_ctx = make_req_context("pin")
+    output = manager.prepare_store(pinned, pin_ctx)
+    assert output is not None
+    assert output.keys_to_store == pinned
+    manager.get_stats()  # drop the warm-up skips
+
+    ctx = make_req_context("req")
+    assert manager.prepare_store([warm, cold], ctx) is None
+    assert manager.prepare_store([warm, cold], ctx) is None
+    # The stores complete, so both chunks become evictable.
+    manager.complete_store(pinned, pin_ctx)
+    output = manager.prepare_store([warm, cold], ctx)
+    assert output is not None
+    assert output.keys_to_store == [warm]
+    assert manager.counts is not None
+    assert manager.counts[cold] == 1
+    stats = manager.get_stats()
+    assert stats is not None
+    assert stats.reduce()[CPUOffloadingMetrics.STORES_SKIPPED] == 1
+
+
+def test_store_threshold_counts_offers_after_a_retry_succeeds():
+    """Once a retry succeeds, a later offer by the same request (for example
+    after preemption, which restarts its stores) counts as before."""
+    manager = make_cpu_manager(num_chunks=2, store_threshold=2)
+    pinned = to_keys([100, 101])
+    warm, cold = to_keys([1, 2])
+    manager.prepare_store([*pinned, warm], make_req_context("warmup"))
+    pin_ctx = make_req_context("pin")
+    assert manager.prepare_store(pinned, pin_ctx) is not None
+
+    ctx = make_req_context("req")
+    assert manager.prepare_store([warm, cold], ctx) is None
+    manager.complete_store(pinned, pin_ctx)
+    output = manager.prepare_store([warm, cold], ctx)
+    assert output is not None
+    assert output.keys_to_store == [warm]
+
+    output = manager.prepare_store([cold], ctx)
+    assert output is not None
+    assert output.keys_to_store == [cold]
+
+
 def test_evictable_cache_chunk_count(monkeypatch):
     """Verifies _num_evictable_cache_chunks is maintained correctly through the
     full store/load lifecycle, eviction, failed stores, concurrent loads,

vllm/v1/kv_offload/cpu/manager.py

@@ -43,6 +43,9 @@ class _RequestCacheAccess:
     inserted_keys: set[OffloadKey] = field(default_factory=set)
     reused_keys: set[OffloadKey] = field(default_factory=set)
     store_miss_keys: set[OffloadKey] = field(default_factory=set)
+    # Keys of offers that failed to allocate, until an offer of them succeeds.
+    # Offering them again is a retry, not another use of the content.
+    failed_offer_keys: set[OffloadKey] = field(default_factory=set)
     finished: bool = False
 
 
@@ -126,14 +129,21 @@ def _get_load_store_spec(
     ) -> CPULoadStoreSpec:
         return CPULoadStoreSpec([chunk.chunk_id for chunk in chunks])
 
-    def _record_accesses(self, keys: Collection[OffloadKey]) -> None:
-        """Record an offer without evicting its tracked candidates."""
+    def _record_accesses(
+        self, keys: Collection[OffloadKey], retried: Collection[OffloadKey]
+    ) -> None:
+        """Record an offer without evicting its tracked candidates.
+
+        Keys in ``retried`` were counted by a failed offer of the same request:
+        they stay tracked but don't count again.
+        """
         assert self.counts is not None
         protected: set[OffloadKey] = set()
         for key in keys:
             if key in self.counts:
                 self.counts.move_to_end(key)
-                self.counts[key] += 1
+                if key not in retried:
+                    self.counts[key] += 1
                 protected.add(key)
 
         num_unprotected = len(self.counts) - len(protected)
@@ -264,8 +274,9 @@ def prepare_store(
         req_context: ReqContext,
     ) -> PrepareStoreOutput | None:
         keys = list(keys)
+        state = self._get_request_cache_access(req_context)
         if self.counts is not None:
-            self._record_accesses(keys)
+            self._record_accesses(keys, state.failed_offer_keys)
         # Partition the original offer before admission filtering. Ready
         # resident chunks are request reuses even if their tracker entry was
         # aged out; the threshold applies only to new store candidates.
@@ -280,17 +291,14 @@ def prepare_store(
                     ready_existing_keys.append(key)
 
         if self.counts is not None:
-            num_store_candidates = len(keys_to_store)
-            keys_to_store = [
-                key
-                for key in keys_to_store
-                if self.counts.get(key, 0) >= self.store_threshold
-            ]
-            self.stores_skipped_in_current_batch += num_store_candidates - len(
-                keys_to_store
-            )
+            admitted_keys: list[OffloadKey] = []
+            for key in keys_to_store:
+                if self.counts.get(key, 0) >= self.store_threshold:
+                    admitted_keys.append(key)
+                elif key not in state.failed_offer_keys:
+                    self.stores_skipped_in_current_batch += 1
+            keys_to_store = admitted_keys
 
-        state = self._get_request_cache_access(req_context)
         new_store_misses = [
             key for key in keys_to_store if key not in state.store_miss_keys
         ]
@@ -311,6 +319,7 @@ def prepare_store(
         )
 
         if not keys_to_store:
+            state.failed_offer_keys.difference_update(keys)
             return PrepareStoreOutput(
                 keys_to_store=[],
                 store_spec=self._get_load_store_spec([], []),
@@ -324,6 +333,8 @@ def prepare_store(
         if num_chunks_to_evict > 0:
             if num_chunks_to_evict > self._num_evictable_cache_chunks:
                 # Eviction will fail.
+                if self.counts is not None:
+                    state.failed_offer_keys.update(keys)
                 return None
             # There is a still a chance for eviction failure as some of the
             # idle chunks might be in the protected list.
@@ -333,6 +344,8 @@ def prepare_store(
             protected = set(keys)
             evicted = self._policy.evict(num_chunks_to_evict, protected)
             if evicted is None:
+                if self.counts is not None:
+                    state.failed_offer_keys.update(keys)
                 return None
 
             # cache-policy removes only idle chunks.
@@ -365,6 +378,7 @@ def prepare_store(
         # build store specs for allocated chunks
         store_spec = self._get_load_store_spec(keys_to_store, chunks)
 
+        state.failed_offer_keys.difference_update(keys)
         return PrepareStoreOutput(
             keys_to_store=keys_to_store,
             store_spec=store_spec,