[Perf][KV Offload] Publish async lookup results per request group

tests/v1/kv_offload/tiering/test_async_lookup.py

@@ -79,6 +79,162 @@ def test_multiple_keys_single_step(self):
         assert mgr.lookup(_key(4), ctx) is False
         mgr.shutdown()
 
+    def test_publishes_completed_group_before_next_group_finishes(
+        self, monkeypatch: pytest.MonkeyPatch
+    ):
+        """A completed group is drainable while a later group is blocked."""
+        key_a = _key(1)
+        key_b = _key(2)
+        ctx_a = _ctx("req_a")
+        ctx_b = _ctx("req_b")
+        mgr = InMemoryLookupManager(existing_keys={key_a})
+        group_b_started = threading.Event()
+        release_group_b = threading.Event()
+        batch_lookup = mgr.batch_lookup
+
+        def blocking_lookup(keys, req_context):
+            if req_context.req_id == ctx_b.req_id:
+                group_b_started.set()
+                if not release_group_b.wait(timeout=5):
+                    raise TimeoutError("Test did not release the blocked group")
+            return batch_lookup(keys, req_context)
+
+        monkeypatch.setattr(mgr, "batch_lookup", blocking_lookup)
+        try:
+            assert mgr.lookup(key_a, ctx_a) is None
+            assert mgr.lookup(key_b, ctx_b) is None
+            mgr.flush()
+            assert group_b_started.wait(timeout=5)
+
+            assert mgr.lookup(key_a, ctx_a) is True
+            assert mgr.lookup(key_b, ctx_b) is None
+            assert mgr._lookup_state[key_b].phase is LookupPhase.IN_FLIGHT
+            assert mgr._pending_results.empty()
+
+            release_group_b.set()
+            group_b_results = mgr._pending_results.get(timeout=5)
+            mgr._pending_results.put(group_b_results)
+            mgr.flush()
+            assert mgr.lookup(key_b, ctx_b) is False
+        finally:
+            release_group_b.set()
+            mgr.shutdown()
+
+        assert mgr.batch_lookup_calls == 2
+
+    def test_group_order_blocks_later_group_until_first_finishes(
+        self, monkeypatch: pytest.MonkeyPatch
+    ):
+        """A group behind a blocked group remains unresolved until its turn."""
+        key_a = _key(1)
+        key_b = _key(2)
+        ctx_a = _ctx("req_a")
+        ctx_b = _ctx("req_b")
+        mgr = InMemoryLookupManager(existing_keys={key_a, key_b})
+        group_a_started = threading.Event()
+        group_b_started = threading.Event()
+        release_group_b = threading.Event()
+        batch_lookup = mgr.batch_lookup
+
+        def blocking_lookup(keys, req_context):
+            if req_context.req_id == ctx_a.req_id:
+                group_a_started.set()
+            if req_context.req_id == ctx_b.req_id:
+                group_b_started.set()
+                if not release_group_b.wait(timeout=5):
+                    raise TimeoutError("Test did not release the blocked group")
+            return batch_lookup(keys, req_context)
+
+        monkeypatch.setattr(mgr, "batch_lookup", blocking_lookup)
+        try:
+            assert mgr.lookup(key_b, ctx_b) is None
+            assert mgr.lookup(key_a, ctx_a) is None
+            mgr.flush()
+            assert group_b_started.wait(timeout=5)
+
+            assert not group_a_started.is_set()
+            assert mgr.lookup(key_a, ctx_a) is None
+            assert mgr._lookup_state[key_a].phase is LookupPhase.IN_FLIGHT
+            assert mgr._pending_results.empty()
+
+            release_group_b.set()
+            first_group_results = mgr._pending_results.get(timeout=5)
+            second_group_results = mgr._pending_results.get(timeout=5)
+            assert [result[0] for result in first_group_results] == [key_b]
+            assert [result[0] for result in second_group_results] == [key_a]
+            mgr._pending_results.put(first_group_results)
+            mgr._pending_results.put(second_group_results)
+            mgr.drain_results()
+            assert mgr.lookup(key_b, ctx_b) is True
+            assert mgr.lookup(key_a, ctx_a) is True
+        finally:
+            release_group_b.set()
+            mgr.shutdown()
+
+    def test_lazy_results_publish_only_after_group_is_consumed(
+        self, monkeypatch: pytest.MonkeyPatch
+    ):
+        """A lazy backend result is published only after full consumption."""
+        key_a = _key(1)
+        key_b = _key(2)
+        ctx = _ctx()
+        mgr = InMemoryLookupManager(existing_keys={key_a})
+        iterator_started = threading.Event()
+        release_iterator = threading.Event()
+
+        def lazy_lookup(keys, req_context):
+            def results():
+                iterator_started.set()
+                yield keys[0] in mgr._existing
+                if not release_iterator.wait(timeout=5):
+                    raise TimeoutError("Test did not release the lazy result")
+                yield keys[1] in mgr._existing
+
+            return results()
+
+        monkeypatch.setattr(mgr, "batch_lookup", lazy_lookup)
+        try:
+            assert mgr.lookup(key_a, ctx) is None
+            assert mgr.lookup(key_b, ctx) is None
+            mgr.flush()
+            assert iterator_started.wait(timeout=5)
+            assert mgr._pending_results.empty()
+
+            release_iterator.set()
+            group_results = mgr._pending_results.get(timeout=5)
+            assert group_results == [
+                (key_a, mgr._lookup_state[key_a].generation, True),
+                (key_b, mgr._lookup_state[key_b].generation, False),
+            ]
+            mgr._pending_results.put(group_results)
+            mgr.drain_results()
+            assert mgr.lookup(key_a, ctx) is True
+            assert mgr.lookup(key_b, ctx) is False
+        finally:
+            release_iterator.set()
+            mgr.shutdown()
+
+    def test_backend_exception_publishes_misses(self, monkeypatch: pytest.MonkeyPatch):
+        """A backend call exception retains the existing miss fallback."""
+        key = _key(1)
+        ctx = _ctx()
+        mgr = InMemoryLookupManager(existing_keys={key})
+
+        def raising_lookup(keys, req_context):
+            raise RuntimeError("backend error")
+
+        monkeypatch.setattr(mgr, "batch_lookup", raising_lookup)
+        try:
+            assert mgr.lookup(key, ctx) is None
+            mgr.flush()
+            group_results = mgr._pending_results.get(timeout=5)
+            assert group_results == [(key, mgr._lookup_state[key].generation, False)]
+            mgr._pending_results.put(group_results)
+            mgr.drain_results()
+            assert mgr.lookup(key, ctx) is False
+        finally:
+            mgr.shutdown()
+
     def test_cleanup_removes_entries(self):
         mgr = InMemoryLookupManager(existing_keys={_key(1)})
         ctx = _ctx("req_a")

vllm/v1/kv_offload/tiering/async_lookup.py

@@ -105,7 +105,7 @@ def __init__(
             list[tuple[OffloadKey, ReqContext, int]] | None
         ] = queue.SimpleQueue()
 
-        # Worker → scheduler: completed result batches.
+        # Worker → scheduler: completed request-group result batches.
         # 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, int, bool]]] = (
@@ -262,7 +262,6 @@ def _worker(self) -> None:
             if not batches:
                 continue
 
-            results: list[tuple[OffloadKey, int, bool]] = []
             for req_context, entries in batches.values():
                 keys = [key for key, _ in entries]
                 try:
@@ -276,9 +275,11 @@ def _worker(self) -> None:
                     )
                     hits = (False for _ in keys)
 
+                group_results: list[tuple[OffloadKey, int, bool]] = []
                 for (key, generation), hit in zip(entries, hits):
-                    results.append((key, generation, hit))
+                    group_results.append((key, generation, hit))
 
-            # Post the entire batch as one item — no lock needed.
-            if results:
-                self._pending_results.put(results)
+                # Post each completed request group before processing the next
+                # group. The list is not modified after it is enqueued.
+                if group_results:
+                    self._pending_results.put(group_results)