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>
Oxana Korzh committed
Aug 27, 2026 at 21:21 UTC
e6bfe03ad73a3330cb427885aa90d97a12e1c704
3 files changed
+35
-33
tests/v1/kv_offload/cpu/test_manager.py
+18
-22
@@ -932,41 +932,37 @@ def test_filter_reused_manager():
932
max_tracker_size=3,
933
)
934
935
- # Lookup [1, 2] -> 1st time, added to tracker but not eligible for store yet
935
+ # lookup() does not count towards store admission
936
assert manager.lookup(to_key(1), _EMPTY_REQ_CTX) is LookupResult.MISS
937
assert manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.MISS
938
+ assert not manager.counts
939
939
- # prepare store [1, 2] -> should be filtered
940
+ # 1st offer of [1, 2] -> tracked at count 1, not eligible yet
941
prepare_store_output = manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
942
assert prepare_store_output is not None
943
assert prepare_store_output.keys_to_store == []
944
+ assert manager.counts[to_keys([1])[0]] == 1
945
+ assert manager.counts[to_keys([2])[0]] == 1
946
944
- # Lookup [1] -> 2nd time, eligible now
945
- assert manager.lookup(to_key(1), _EMPTY_REQ_CTX) is LookupResult.MISS
946
-
947
- # prepare store [1, 2] -> [1] should be eligible, [2] should be filtered
947
+ # 2nd offer -> the whole repeated prefix becomes eligible at once
948
prepare_store_output = manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
949
assert prepare_store_output is not None
950
- assert prepare_store_output.keys_to_store == to_keys([1])
951
-
952
- # Lookup [3, 4] -> 1st time
953
- # (evicts [2] from tracker since max_size is 3 and tracker has [1])
954
- assert manager.lookup(to_key(3), _EMPTY_REQ_CTX) is LookupResult.MISS
955
- assert manager.lookup(to_key(4), _EMPTY_REQ_CTX) is LookupResult.MISS
956
- # Verify [2] was evicted from the tracker (tracker now has: [1], [3], [4])
957
- assert to_keys([2])[0] not in manager.counts
958
-
959
- # Lookup [2] again -> (this adds [2] back to the tracker as 1st time)
960
- assert manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.MISS
961
- # Verify [2] was re-added with count=1 (not eligible yet)
962
- assert manager.counts.get(to_keys([2])[0]) == 1
950
+ assert prepare_store_output.keys_to_store == to_keys([1, 2])
951
+ manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
952
964
- # prepare store [2] -> should still be filtered out since count was reset
965
- prepare_store_output = manager.prepare_store(to_keys([2]), _EMPTY_REQ_CTX)
953
+ # Offer [3, 4] -> 1st time, evicting the tracker's LRU entry [1]
954
+ prepare_store_output = manager.prepare_store(to_keys([3, 4]), _EMPTY_REQ_CTX)
955
assert prepare_store_output is not None
956
assert prepare_store_output.keys_to_store == []
957
+ assert to_keys([1])[0] not in manager.counts
958
+ assert manager.counts[to_keys([3])[0]] == 1
959
+ assert manager.counts[to_keys([4])[0]] == 1
960
969
- manager.complete_store(to_keys([1]), _EMPTY_REQ_CTX)
961
+ # [1] re-enters the tracker at count 1, so it is filtered again
962
+ prepare_store_output = manager.prepare_store(to_keys([1]), _EMPTY_REQ_CTX)
963
+ assert prepare_store_output is not None
964
+ assert prepare_store_output.keys_to_store == []
965
+ assert manager.counts.get(to_keys([1])[0]) == 1
966
967
968
def test_evictable_cache_block_count():
vllm/v1/kv_offload/cpu/manager.py
+13
-8
@@ -103,6 +103,17 @@ 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."""
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:
113
+ if len(self.counts) >= self.max_tracker_size:
114
+ self.counts.popitem(last=False)
115
+ self.counts[key] = 1
116
+
117
# --- OffloadingManager interface ---
118
119
@override
@@ -111,14 +122,6 @@ class CPUOffloadingManager(OffloadingManager):
122
123
@override
124
def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult:
114
- if self.counts is not None:
115
- if key in self.counts:
116
- self.counts.move_to_end(key)
117
- self.counts[key] += 1
118
- else:
119
- if len(self.counts) >= self.max_tracker_size:
120
- self.counts.popitem(last=False)
121
- self.counts[key] = 1
125
block = self._policy.get(key)
126
if block is None:
127
return LookupResult.MISS
@@ -170,6 +173,8 @@ class CPUOffloadingManager(OffloadingManager):
173
) -> PrepareStoreOutput | None:
174
if self.counts is not None:
175
num_keys = len(keys)
176
+ for key in keys:
177
+ self._record_access(key)
178
keys = [k for k in keys if self.counts.get(k, 0) >= self.store_threshold]
179
self.stores_skipped_in_current_batch += num_keys - len(keys)
180
# filter out blocks that are already stored
vllm/v1/kv_offload/cpu/spec.py
+4
-3
@@ -123,9 +123,10 @@ class CPUOffloadingSpec(OffloadingSpec):
123
@override
124
def get_manager(self) -> OffloadingManager:
125
if not self._manager:
126
- # store_threshold: how many times a block must appear in lookup()
127
- # before it is eligible for CPU offloading. Values < 2 disable
128
- # filtering (a threshold of 1 equals no filter; 0 is the default).
126
+ # store_threshold: how many times a block must be offered for
127
+ # storage before it is eligible for CPU offloading. Values < 2
128
+ # disable filtering (a threshold of 1 equals no filter; 0 is the
129
+ # default).
130
store_threshold = int(self.extra_config.get("store_threshold", 0))
131
132
# Maximum entries in the internal tracker's LRU table.