diff --git a/tests/v1/kv_offload/tiering/test_async_lookup.py b/tests/v1/kv_offload/tiering/test_async_lookup.py index 05f3f96c6e..ff5cbcd5dd 100644 --- a/tests/v1/kv_offload/tiering/test_async_lookup.py +++ b/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() diff --git a/vllm/v1/kv_offload/tiering/async_lookup.py b/vllm/v1/kv_offload/tiering/async_lookup.py index 2baa1e24b8..f111f233b6 100644 --- a/vllm/v1/kv_offload/tiering/async_lookup.py +++ b/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.