[Bugfix][KV Offload] Ignore stale async lookup results (#54872)
Signed-off-by: Alex <jihui.huang@daocloud.io> Co-authored-by: OpenAI Codex <noreply@openai.com>
AlexHuang committed
Sep 2, 2026 at 22:10 UTC
35faf957d65acdfc0b19f7cd4e4a2d50d7d73a1b
2 files changed
+56
-27
tests/v1/kv_offload/tiering/test_async_lookup.py
+24
-1
@@ -104,6 +104,28 @@ class TestAsyncLookupManager:
104
assert _key(1) not in mgr._lookup_state
105
mgr.shutdown()
106
107
+ def test_stale_result_ignored_after_cleanup_and_key_reuse(self):
108
+ key = _key(1)
109
+ mgr = InMemoryLookupManager()
110
+ ctx_a = _ctx("req_a")
111
+ ctx_b = _ctx("req_b")
112
+ assert mgr.lookup(key, ctx_a) is None
113
+ stale_generation = mgr._lookup_state[key].generation
114
+ mgr.cleanup("req_a")
115
+
116
+ assert mgr.lookup(key, ctx_b) is None
117
+ generation = mgr._lookup_state[key].generation
118
+ assert generation != stale_generation
119
+
120
+ mgr._pending_results.put([(key, stale_generation, True)])
121
+ mgr.drain_results()
122
+ assert mgr.lookup(key, ctx_b) is None
123
+
124
+ mgr._pending_results.put([(key, generation, False)])
125
+ mgr.drain_results()
126
+ assert mgr.lookup(key, ctx_b) is False
127
+ mgr.shutdown()
128
+
129
def test_flush_no_queue_post_when_empty(self):
130
mgr = InMemoryLookupManager()
131
mgr.flush()
@@ -220,7 +242,8 @@ class TestAsyncLookupManager:
242
243
# (b) A stray/duplicate result for the now-decided key violates the
244
# enqueue-once invariant and must trip the assert.
223
- mgr._pending_results.put([(_key(1), True)])
245
+ generation = mgr._lookup_state[_key(1)].generation
246
+ mgr._pending_results.put([(_key(1), generation, True)])
247
with pytest.raises(AssertionError):
248
mgr.drain_results()
249
mgr.shutdown()
vllm/v1/kv_offload/tiering/async_lookup.py
+32
-26
@@ -45,6 +45,7 @@ logger = init_logger(__name__)
45
46
@dataclass(slots=True)
47
class LookupState:
48
+ generation: int
49
result: bool | None = None # True (found), False (not found), None
50
request_ids: set[str] = field(default_factory=set) # requests asking for the lookup
51
@@ -78,21 +79,24 @@ class AsyncLookupManager(ABC):
79
self._lookup_state: dict[OffloadKey, LookupState] = {}
80
# req_id → keys looked up by that request (reverse index for cleanup).
81
self._req_keys: dict[str, set[OffloadKey]] = {}
82
+ # Results from an older state must not be applied after cleanup removes
83
+ # and recreates a key.
84
+ self._next_generation = 0
85
82
- # Accumulates (key, req_context) pairs during lookup() calls.
86
+ # Accumulates (key, req_context, generation) triples during lookup().
87
# Flushed as one queue item per step by flush().
84
- self._lookup_batch: list[tuple[OffloadKey, ReqContext]] = []
88
+ self._lookup_batch: list[tuple[OffloadKey, ReqContext, int]] = []
89
90
# Scheduler → worker: one full step's batch per item.
91
# None is used as a shutdown sentinel.
92
self._lookup_queue: queue.SimpleQueue[
89
- list[tuple[OffloadKey, ReqContext]] | None
93
+ list[tuple[OffloadKey, ReqContext, int]] | None
94
] = queue.SimpleQueue()
95
96
# Worker → scheduler: completed result batches.
93
- # Each item is a list of (key, found) pairs.
97
+ # Each item is a list of (key, generation, found) triples.
98
# SimpleQueue is explicitly thread-safe for one writer / one reader.
95
- self._pending_results: queue.SimpleQueue[list[tuple[OffloadKey, bool]]] = (
99
+ self._pending_results: queue.SimpleQueue[list[tuple[OffloadKey, int, bool]]] = (
100
queue.SimpleQueue()
101
)
102
self._need_to_drain: bool = False
@@ -137,9 +141,10 @@ class AsyncLookupManager(ABC):
141
req_id = req_context.req_id
142
state = self._lookup_state.get(key)
143
if state is None:
140
- state = LookupState()
144
+ state = LookupState(generation=self._next_generation)
145
+ self._next_generation += 1
146
self._lookup_state[key] = state
142
- self._lookup_batch.append((key, req_context))
147
+ self._lookup_batch.append((key, req_context, state.generation))
148
state.request_ids.add(req_id)
149
self._req_keys.setdefault(req_id, set()).add(key)
150
return state.result
@@ -167,19 +172,19 @@ class AsyncLookupManager(ABC):
172
batch = self._pending_results.get_nowait()
173
except queue.Empty:
174
break
170
- for key, result in batch:
175
+ for key, generation, result in batch:
176
state = self._lookup_state.get(key)
172
- if state is not None:
173
- # A key is enqueued for probing exactly once, so a decided
174
- # verdict must never receive a second result. Enforcing it
175
- # keeps a late/duplicate result from resurrecting a stale
176
- # True and reopening the failed-load livelock.
177
- assert state.result is None, (
178
- "cached key received a second lookup result; the "
179
- "enqueue-once invariant is broken and could reopen the "
180
- "failed-load livelock"
181
- )
182
- state.result = result
177
+ if state is None or state.generation != generation:
178
+ continue
179
+ # Each lookup generation is enqueued exactly once. A matching
180
+ # generation must not receive a second result; stale
181
+ # generations were discarded above.
182
+ assert state.result is None, (
183
+ "cached key received a second lookup result; the "
184
+ "enqueue-once invariant is broken and could reopen the "
185
+ "failed-load livelock"
186
+ )
187
+ state.result = result
188
189
def mark_miss(self, keys: Collection[OffloadKey]) -> None:
190
"""Force the cached verdict for ``keys`` to False after a failed load, so
@@ -218,18 +223,19 @@ class AsyncLookupManager(ABC):
223
break
224
225
# Group by req_id.
221
- batches: dict[str, tuple[ReqContext, list[OffloadKey]]] = {}
222
- for key, req_context in pending:
226
+ batches: dict[str, tuple[ReqContext, list[tuple[OffloadKey, int]]]] = {}
227
+ for key, req_context, generation in pending:
228
req_id = req_context.req_id
229
if req_id not in batches:
230
batches[req_id] = (req_context, [])
226
- batches[req_id][1].append(key)
231
+ batches[req_id][1].append((key, generation))
232
233
if not batches:
234
continue
235
231
- results: list[tuple[OffloadKey, bool]] = []
232
- for req_context, keys in batches.values():
236
+ results: list[tuple[OffloadKey, int, bool]] = []
237
+ for req_context, entries in batches.values():
238
+ keys = [key for key, _ in entries]
239
try:
240
hits = self.batch_lookup(keys, req_context)
241
except Exception as exc:
@@ -241,8 +247,8 @@ class AsyncLookupManager(ABC):
247
)
248
hits = (False for _ in keys)
249
244
- for key, hit in zip(keys, hits):
245
- results.append((key, hit))
250
+ for (key, generation), hit in zip(entries, hits):
251
+ results.append((key, generation, hit))
252
253
# Post the entire batch as one item — no lock needed.
254
if results: