[Bugfix][KV Offload] Skip cleaned-up async lookup batches (#55075)
Signed-off-by: Alex <jihui.huang@daocloud.io> Co-authored-by: OpenAI Codex <noreply@openai.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>
AlexHuang committed
Sep 6, 2026 at 18:27 UTC
9afb878a55eb08fd064dfaab56afd8447e97101e
2 files changed
+21
-3
tests/v1/kv_offload/tiering/test_async_lookup.py
+11
@@ -26,10 +26,12 @@ class InMemoryLookupManager(AsyncLookupManager):
26
super().__init__(tier_type="test")
27
self._existing = existing_keys or set()
28
self._results_ready = threading.Event()
29
+ self.batch_lookup_calls = 0
30
31
def batch_lookup(
32
self, keys: list[OffloadKey], req_context: ReqContext
33
) -> Iterable[bool]:
34
+ self.batch_lookup_calls += 1
35
results = [k in self._existing for k in keys]
36
self._results_ready.set()
37
return results
@@ -132,6 +134,15 @@ class TestAsyncLookupManager:
134
assert mgr._lookup_queue.empty()
135
mgr.shutdown()
136
137
+ def test_flush_skips_lookup_cleaned_up_before_submit(self):
138
+ mgr = InMemoryLookupManager()
139
+ mgr.lookup(_key(1), _ctx("req_a"))
140
+ mgr.cleanup("req_a")
141
+ mgr.flush()
142
+ mgr.shutdown()
143
+
144
+ assert mgr.batch_lookup_calls == 0
145
+
146
def test_repeated_lookup_same_key_no_duplicate_batch(self):
147
mgr = InMemoryLookupManager(existing_keys={_key(1)})
148
ctx = _ctx()
vllm/v1/kv_offload/tiering/async_lookup.py
+10
-3
@@ -158,9 +158,16 @@ class AsyncLookupManager(ABC):
158
step's drain_results(). Safe to call with an empty batch (no-op).
159
"""
160
self._need_to_drain = True
161
- if self._lookup_batch:
162
- self._lookup_queue.put(self._lookup_batch)
163
- self._lookup_batch = []
161
+ batch = self._lookup_batch
162
+ self._lookup_batch = []
163
+ batch = [
164
+ (key, req_context, generation)
165
+ for key, req_context, generation in batch
166
+ if (state := self._lookup_state.get(key)) is not None
167
+ and state.generation == generation
168
+ ]
169
+ if batch:
170
+ self._lookup_queue.put(batch)
171
172
def drain_results(self) -> None:
173
"""Apply pending worker results to _lookup_state.