[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>

tests/v1/kv_offload/tiering/test_async_lookup.py

@@ -26,10 +26,12 @@ class InMemoryLookupManager(AsyncLookupManager):
         super().__init__(tier_type="test")
         self._existing = existing_keys or set()
         self._results_ready = threading.Event()
+        self.batch_lookup_calls = 0
 
     def batch_lookup(
         self, keys: list[OffloadKey], req_context: ReqContext
     ) -> Iterable[bool]:
+        self.batch_lookup_calls += 1
         results = [k in self._existing for k in keys]
         self._results_ready.set()
         return results
@@ -132,6 +134,15 @@ class TestAsyncLookupManager:
         assert mgr._lookup_queue.empty()
         mgr.shutdown()
 
+    def test_flush_skips_lookup_cleaned_up_before_submit(self):
+        mgr = InMemoryLookupManager()
+        mgr.lookup(_key(1), _ctx("req_a"))
+        mgr.cleanup("req_a")
+        mgr.flush()
+        mgr.shutdown()
+
+        assert mgr.batch_lookup_calls == 0
+
     def test_repeated_lookup_same_key_no_duplicate_batch(self):
         mgr = InMemoryLookupManager(existing_keys={_key(1)})
         ctx = _ctx()

vllm/v1/kv_offload/tiering/async_lookup.py

@@ -158,9 +158,16 @@ class AsyncLookupManager(ABC):
         step's drain_results().  Safe to call with an empty batch (no-op).
         """
         self._need_to_drain = True
-        if self._lookup_batch:
-            self._lookup_queue.put(self._lookup_batch)
-            self._lookup_batch = []
+        batch = self._lookup_batch
+        self._lookup_batch = []
+        batch = [
+            (key, req_context, generation)
+            for key, req_context, generation in batch
+            if (state := self._lookup_state.get(key)) is not None
+            and state.generation == generation
+        ]
+        if batch:
+            self._lookup_queue.put(batch)
 
     def drain_results(self) -> None:
         """Apply pending worker results to _lookup_state.