[Bugfix][KV Offload] Reuse in-flight async lookup probes (#55823)
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 14, 2026 at 20:02 UTC
c612e2bff0d045a5b2304371793985c330fba436
2 files changed
+112
-21
tests/v1/kv_offload/tiering/test_async_lookup.py
+71
-5
@@ -8,7 +8,7 @@ from collections.abc import Iterable
8
import pytest
9
10
from vllm.v1.kv_offload.base import OffloadKey, ReqContext, make_offload_key
11
-from vllm.v1.kv_offload.tiering.async_lookup import AsyncLookupManager
11
+from vllm.v1.kv_offload.tiering.async_lookup import AsyncLookupManager, LookupPhase
12
13
14
def _key(i: int) -> OffloadKey:
@@ -46,10 +46,13 @@ class TestAsyncLookupManager:
46
def test_found_key_returns_true(self):
47
mgr = InMemoryLookupManager(existing_keys={_key(1)})
48
assert mgr.lookup(_key(1), _ctx()) is None
49
+ assert mgr._lookup_state[_key(1)].phase is LookupPhase.PENDING
50
mgr.flush()
51
+ assert mgr._lookup_state[_key(1)].phase is LookupPhase.IN_FLIGHT
52
mgr._results_ready.wait()
53
mgr._results_ready.clear()
54
assert mgr.lookup(_key(1), _ctx()) is True
55
+ assert mgr._lookup_state[_key(1)].phase is LookupPhase.RESOLVED
56
mgr.shutdown()
57
58
def test_not_found_key_returns_false(self):
@@ -106,6 +109,68 @@ class TestAsyncLookupManager:
109
assert _key(1) not in mgr._lookup_state
110
mgr.shutdown()
111
112
+ def test_cleanup_reuses_in_flight_probe(self, monkeypatch: pytest.MonkeyPatch):
113
+ """A replacement request shares the in-flight probe and its verdict."""
114
+ key = _key(1)
115
+ mgr = InMemoryLookupManager(existing_keys={key})
116
+ ctx_b = _ctx("req_b")
117
+ probe_started = threading.Event()
118
+ release_probe = threading.Event()
119
+ batch_lookup = mgr.batch_lookup
120
+
121
+ def blocking_lookup(keys, req_context):
122
+ probe_started.set()
123
+ if not release_probe.wait(timeout=5):
124
+ raise TimeoutError("Test did not release the backend probe")
125
+ return batch_lookup(keys, req_context)
126
+
127
+ monkeypatch.setattr(mgr, "batch_lookup", blocking_lookup)
128
+ try:
129
+ assert mgr.lookup(key, _ctx("req_a")) is None
130
+ mgr.flush()
131
+ assert probe_started.wait(timeout=5)
132
+ assert mgr._lookup_state[key].phase is LookupPhase.IN_FLIGHT
133
+ mgr.cleanup("req_a")
134
+ assert mgr.lookup(key, ctx_b) is None
135
+ assert mgr._lookup_state[key].phase is LookupPhase.IN_FLIGHT
136
+ mgr.flush()
137
+
138
+ release_probe.set()
139
+ batch = mgr._pending_results.get(timeout=5)
140
+ mgr._pending_results.put(batch)
141
+ replacement_result = mgr.lookup(key, ctx_b)
142
+ finally:
143
+ release_probe.set()
144
+ mgr.shutdown()
145
+
146
+ assert mgr.batch_lookup_calls == 1
147
+ assert replacement_result is True
148
+
149
+ @pytest.mark.parametrize(
150
+ "reclaim_at_shutdown", [False, True], ids=["flush", "shutdown"]
151
+ )
152
+ def test_unclaimed_probe_reclaimed_without_lookup(self, reclaim_at_shutdown: bool):
153
+ """A completed orphan is released by flush or shutdown without a lookup."""
154
+ key = _key(1)
155
+ mgr = InMemoryLookupManager(existing_keys={key})
156
+ try:
157
+ mgr.lookup(key, _ctx("req_a"))
158
+ mgr.flush()
159
+ mgr.cleanup("req_a")
160
+ assert key in mgr._lookup_state
161
+ assert not mgr._req_keys
162
+
163
+ if not reclaim_at_shutdown:
164
+ batch = mgr._pending_results.get(timeout=5)
165
+ mgr._pending_results.put(batch)
166
+ mgr.flush()
167
+ assert key not in mgr._lookup_state
168
+ finally:
169
+ mgr.shutdown()
170
+
171
+ assert not mgr._lookup_state
172
+ assert mgr._pending_results.empty()
173
+
174
def test_stale_result_ignored_after_cleanup_and_key_reuse(self):
175
key = _key(1)
176
mgr = InMemoryLookupManager()
@@ -119,11 +184,12 @@ class TestAsyncLookupManager:
184
generation = mgr._lookup_state[key].generation
185
assert generation != stale_generation
186
187
+ mgr.flush()
188
+ mgr._results_ready.wait()
189
+ mgr._results_ready.clear()
190
+ current_result = mgr._pending_results.get(timeout=5)
191
mgr._pending_results.put([(key, stale_generation, True)])
123
- mgr.drain_results()
124
- assert mgr.lookup(key, ctx_b) is None
125
-
126
- mgr._pending_results.put([(key, generation, False)])
192
+ mgr._pending_results.put(current_result)
193
mgr.drain_results()
194
assert mgr.lookup(key, ctx_b) is False
195
mgr.shutdown()
vllm/v1/kv_offload/tiering/async_lookup.py
+41
-16
@@ -27,8 +27,10 @@ lookup() accumulates new keys in _lookup_batch without touching the queue.
27
flush() is called once per step from the tier's on_schedule_end(), posting
28
the entire batch as a single queue item so the background thread sees one
29
batch per step.
30
-drain_results() is called before any lookup() calls in the same step, so
31
-lookup() is a pure OrderedDict operation.
30
+Results are drained on the first lookup after each flush, at flush(), and
31
+after worker shutdown. In-flight lookups with no remaining request references
32
+are retained until their results are drained, allowing new requests to share
33
+the same probe.
34
"""
35
36
import queue
@@ -36,6 +38,7 @@ import threading
38
from abc import ABC, abstractmethod
39
from collections.abc import Collection, Iterable
40
from dataclasses import dataclass, field
41
+from enum import Enum, auto
42
43
from vllm.logger import init_logger
44
from vllm.v1.kv_offload.base import OffloadKey, ReqContext
@@ -43,10 +46,21 @@ from vllm.v1.kv_offload.base import OffloadKey, ReqContext
46
logger = init_logger(__name__)
47
48
49
+class LookupPhase(Enum):
50
+ """Lifecycle phase of a lookup probe."""
51
+
52
+ PENDING = auto() # Accumulated in _lookup_batch, but not yet submitted.
53
+ IN_FLIGHT = auto() # Submitted to the worker, but not yet resolved.
54
+ RESOLVED = auto() # A final verdict has been applied to the state.
55
+
56
+
57
@dataclass(slots=True)
58
class LookupState:
59
generation: int
49
- result: bool | None = None # True (found), False (not found), None
60
+ phase: LookupPhase = LookupPhase.PENDING
61
+ # None while pending/in flight; True if the key exists; False if absent or
62
+ # explicitly marked missing after a failed load.
63
+ result: bool | None = None
64
request_ids: set[str] = field(default_factory=set) # requests asking for the lookup
65
66
@@ -155,24 +169,28 @@ class AsyncLookupManager(ABC):
169
Called once per step from on_schedule_end() after all lookup() calls
170
are done. The worker receives the full batch and processes it during
171
the model-execution window, maximising time available before the next
158
- step's drain_results(). Safe to call with an empty batch (no-op).
172
+ step's drain_results(). Also drains completed lookups when there
173
+ are no new keys to submit.
174
"""
175
+ self.drain_results()
176
self._need_to_drain = True
177
batch = self._lookup_batch
178
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)
179
+ in_flight_batch = []
180
+ for key, req_context, generation in batch:
181
+ state = self._lookup_state.get(key)
182
+ if state is None or state.generation != generation:
183
+ continue
184
+ assert state.phase is LookupPhase.PENDING
185
+ state.phase = LookupPhase.IN_FLIGHT
186
+ in_flight_batch.append((key, req_context, generation))
187
+ if in_flight_batch:
188
+ self._lookup_queue.put(in_flight_batch)
189
190
def drain_results(self) -> None:
191
"""Apply pending worker results to _lookup_state.
192
175
- Called from lookup() before checking state.
193
+ Called from lookup(), flush(), and shutdown() on the scheduler thread.
194
"""
195
while True:
196
try:
@@ -183,6 +201,10 @@ class AsyncLookupManager(ABC):
201
state = self._lookup_state.get(key)
202
if state is None or state.generation != generation:
203
continue
204
+ if not state.request_ids:
205
+ del self._lookup_state[key]
206
+ continue
207
+ assert state.phase is LookupPhase.IN_FLIGHT
208
# Each lookup generation is enqueued exactly once. A matching
209
# generation must not receive a second result; stale
210
# generations were discarded above.
@@ -192,6 +214,7 @@ class AsyncLookupManager(ABC):
214
"failed-load livelock"
215
)
216
state.result = result
217
+ state.phase = LookupPhase.RESOLVED
218
219
def mark_miss(self, keys: Collection[OffloadKey]) -> None:
220
"""Force the cached verdict for ``keys`` to False after a failed load, so
@@ -201,9 +224,10 @@ class AsyncLookupManager(ABC):
224
state = self._lookup_state.get(key)
225
if state is not None:
226
state.result = False
227
+ state.phase = LookupPhase.RESOLVED
228
229
def cleanup(self, req_id: str) -> None:
206
- """Remove entries no longer needed by any active request.
230
+ """Release request references, retaining in-flight lookups.
231
232
Called from the tier's on_request_finished(). Uses the reverse
233
index to visit only keys associated with this request.
@@ -211,13 +235,14 @@ class AsyncLookupManager(ABC):
235
for key in self._req_keys.pop(req_id, ()):
236
state = self._lookup_state[key]
237
state.request_ids.discard(req_id)
214
- if not state.request_ids:
238
+ if not state.request_ids and state.phase is not LookupPhase.IN_FLIGHT:
239
del self._lookup_state[key]
240
241
def shutdown(self) -> None:
218
- """Stop the worker thread."""
242
+ """Stop the worker thread and drain completed lookups."""
243
self._lookup_queue.put(None) # unblock _worker from _lookup_queue.get()
244
self._thread.join()
245
+ self.drain_results()
246
247
# ------------------------------------------------------------------
248
# Internal helpers