[Bugfix][KV Offload] Ignore stale async lookup results

tests/v1/kv_offload/tiering/test_async_lookup.py

@@ -104,6 +104,28 @@ def test_cleanup_preserves_shared_entries(self):
         assert _key(1) not in mgr._lookup_state
         mgr.shutdown()
 
+    def test_stale_result_ignored_after_cleanup_and_key_reuse(self):
+        key = _key(1)
+        mgr = InMemoryLookupManager()
+        ctx_a = _ctx("req_a")
+        ctx_b = _ctx("req_b")
+        assert mgr.lookup(key, ctx_a) is None
+        stale_generation = mgr._lookup_state[key].generation
+        mgr.cleanup("req_a")
+
+        assert mgr.lookup(key, ctx_b) is None
+        generation = mgr._lookup_state[key].generation
+        assert generation != stale_generation
+
+        mgr._pending_results.put([(key, stale_generation, True)])
+        mgr.drain_results()
+        assert mgr.lookup(key, ctx_b) is None
+
+        mgr._pending_results.put([(key, generation, False)])
+        mgr.drain_results()
+        assert mgr.lookup(key, ctx_b) is False
+        mgr.shutdown()
+
     def test_flush_no_queue_post_when_empty(self):
         mgr = InMemoryLookupManager()
         mgr.flush()
@@ -220,7 +242,8 @@ def test_enqueue_once_invariant_enforced(self):
 
         # (b) A stray/duplicate result for the now-decided key violates the
         # enqueue-once invariant and must trip the assert.
-        mgr._pending_results.put([(_key(1), True)])
+        generation = mgr._lookup_state[_key(1)].generation
+        mgr._pending_results.put([(_key(1), generation, True)])
         with pytest.raises(AssertionError):
             mgr.drain_results()
         mgr.shutdown()

vllm/v1/kv_offload/tiering/async_lookup.py

@@ -45,6 +45,7 @@
 
 @dataclass(slots=True)
 class LookupState:
+    generation: int
     result: bool | None = None  # True (found), False (not found), None
     request_ids: set[str] = field(default_factory=set)  # requests asking for the lookup
 
@@ -78,21 +79,24 @@ def __init__(
         self._lookup_state: dict[OffloadKey, LookupState] = {}
         # req_id → keys looked up by that request (reverse index for cleanup).
         self._req_keys: dict[str, set[OffloadKey]] = {}
+        # Results from an older state must not be applied after cleanup removes
+        # and recreates a key.
+        self._next_generation = 0
 
-        # Accumulates (key, req_context) pairs during lookup() calls.
+        # Accumulates (key, req_context, generation) triples during lookup().
         # Flushed as one queue item per step by flush().
-        self._lookup_batch: list[tuple[OffloadKey, ReqContext]] = []
+        self._lookup_batch: list[tuple[OffloadKey, ReqContext, int]] = []
 
         # Scheduler → worker: one full step's batch per item.
         # None is used as a shutdown sentinel.
         self._lookup_queue: queue.SimpleQueue[
-            list[tuple[OffloadKey, ReqContext]] | None
+            list[tuple[OffloadKey, ReqContext, int]] | None
         ] = queue.SimpleQueue()
 
         # Worker → scheduler: completed result batches.
-        # Each item is a list of (key, found) pairs.
+        # Each item is a list of (key, generation, found) triples.
         # SimpleQueue is explicitly thread-safe for one writer / one reader.
-        self._pending_results: queue.SimpleQueue[list[tuple[OffloadKey, bool]]] = (
+        self._pending_results: queue.SimpleQueue[list[tuple[OffloadKey, int, bool]]] = (
             queue.SimpleQueue()
         )
         self._need_to_drain: bool = False
@@ -137,9 +141,10 @@ def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None:
         req_id = req_context.req_id
         state = self._lookup_state.get(key)
         if state is None:
-            state = LookupState()
+            state = LookupState(generation=self._next_generation)
+            self._next_generation += 1
             self._lookup_state[key] = state
-            self._lookup_batch.append((key, req_context))
+            self._lookup_batch.append((key, req_context, state.generation))
         state.request_ids.add(req_id)
         self._req_keys.setdefault(req_id, set()).add(key)
         return state.result
@@ -167,19 +172,19 @@ def drain_results(self) -> None:
                 batch = self._pending_results.get_nowait()
             except queue.Empty:
                 break
-            for key, result in batch:
+            for key, generation, result in batch:
                 state = self._lookup_state.get(key)
-                if state is not None:
-                    # A key is enqueued for probing exactly once, so a decided
-                    # verdict must never receive a second result. Enforcing it
-                    # keeps a late/duplicate result from resurrecting a stale
-                    # True and reopening the failed-load livelock.
-                    assert state.result is None, (
-                        "cached key received a second lookup result; the "
-                        "enqueue-once invariant is broken and could reopen the "
-                        "failed-load livelock"
-                    )
-                    state.result = result
+                if state is None or state.generation != generation:
+                    continue
+                # Each lookup generation is enqueued exactly once. A matching
+                # generation must not receive a second result; stale
+                # generations were discarded above.
+                assert state.result is None, (
+                    "cached key received a second lookup result; the "
+                    "enqueue-once invariant is broken and could reopen the "
+                    "failed-load livelock"
+                )
+                state.result = result
 
     def mark_miss(self, keys: Collection[OffloadKey]) -> None:
         """Force the cached verdict for ``keys`` to False after a failed load, so
@@ -218,18 +223,19 @@ def _worker(self) -> None:
                 break
 
             # Group by req_id.
-            batches: dict[str, tuple[ReqContext, list[OffloadKey]]] = {}
-            for key, req_context in pending:
+            batches: dict[str, tuple[ReqContext, list[tuple[OffloadKey, int]]]] = {}
+            for key, req_context, generation in pending:
                 req_id = req_context.req_id
                 if req_id not in batches:
                     batches[req_id] = (req_context, [])
-                batches[req_id][1].append(key)
+                batches[req_id][1].append((key, generation))
 
             if not batches:
                 continue
 
-            results: list[tuple[OffloadKey, bool]] = []
-            for req_context, keys in batches.values():
+            results: list[tuple[OffloadKey, int, bool]] = []
+            for req_context, entries in batches.values():
+                keys = [key for key, _ in entries]
                 try:
                     hits = self.batch_lookup(keys, req_context)
                 except Exception as exc:
@@ -241,8 +247,8 @@ def _worker(self) -> None:
                     )
                     hits = (False for _ in keys)
 
-                for key, hit in zip(keys, hits):
-                    results.append((key, hit))
+                for (key, generation), hit in zip(entries, hits):
+                    results.append((key, generation, hit))
 
             # Post the entire batch as one item — no lock needed.
             if results: