A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 23:15 UTC
[KV Connector] MooncakeStore: don't co-queue save with load to avoid double delayed-free (#43371)
Signed-off-by: Dao Le <Dao007forever@gmail.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Dao007forever committed
May 22, 2026 at 09:14 UTC
b21f3d56d4a2ab5504b56504e87e0475c6d84eb2
2 files changed
+241
tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py
+235
@@ -4,6 +4,8 @@
4
from types import SimpleNamespace
5
6
from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
7
+ LoadSpec,
8
+ ReqMeta,
9
RequestTracker,
10
)
11
from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.scheduler import (
@@ -109,3 +111,236 @@ def test_cached_request_without_spec_decode_keeps_current_step_save_overlap():
111
tracker = scheduler._request_trackers["req-0"]
112
assert tracker.token_len == 48
113
assert tracker.num_saved_tokens == 48
114
+
115
+
116
+def _make_pending_load_unfinished_request(
117
+ scheduler: MooncakeStoreScheduler,
118
+ *,
119
+ num_tokens: int,
120
+ block_hashes: list[bytes],
121
+ block_ids: tuple[list[int], ...] = ([0, 1, 2],),
122
+) -> None:
123
+ request = SimpleNamespace(
124
+ num_tokens=num_tokens,
125
+ block_hashes=block_hashes,
126
+ num_output_placeholders=0,
127
+ )
128
+ scheduler._unfinished_requests["req-0"] = (request, block_ids)
129
+
130
+
131
+def _make_pending_load_scheduler_output() -> SimpleNamespace:
132
+ """scheduler_output for a step where req-0 is parked on a pending load
133
+ (not in scheduled_new_reqs or scheduled_cached_reqs)."""
134
+ return SimpleNamespace(
135
+ finished_req_ids=set(),
136
+ preempted_req_ids=set(),
137
+ scheduled_new_reqs=[],
138
+ scheduled_cached_reqs=SimpleNamespace(
139
+ req_ids=[],
140
+ new_block_ids=[],
141
+ num_computed_tokens=[],
142
+ ),
143
+ num_scheduled_tokens={},
144
+ scheduled_spec_decode_tokens={},
145
+ )
146
+
147
+
148
+def test_pending_load_does_not_co_queue_save():
149
+ # Regression: a cache-hit request waiting on an async load must not also
150
+ # enqueue a save in the same scheduling step. Co-queuing both produces a
151
+ # recv+send pair for the same req_id, and the scheduler's
152
+ # _update_from_kv_xfer_finished then trips `assert req_id in self.requests`
153
+ # when both completions land for the delay-freed request.
154
+ scheduler = _make_bare_scheduler()
155
+ _make_pending_load_unfinished_request(
156
+ scheduler,
157
+ num_tokens=48,
158
+ block_hashes=[b"h0", b"h1", b"h2"],
159
+ )
160
+ scheduler.load_specs["req-0"] = LoadSpec(
161
+ vllm_cached_tokens=0,
162
+ kvpool_cached_tokens=48,
163
+ can_load=True,
164
+ )
165
+
166
+ meta = scheduler.build_connector_meta(_make_pending_load_scheduler_output())
167
+
168
+ assert len(meta.requests) == 1
169
+ req_meta = meta.requests[0]
170
+ assert req_meta.req_id == "req-0"
171
+ # Save must be off so the worker does not call add_stored_request.
172
+ assert req_meta.can_save is False
173
+ # Load is still issued as planned.
174
+ assert req_meta.load_spec is not None
175
+ assert req_meta.load_spec.can_load is True
176
+ # And the tracker's saved-tokens watermark stays at 0 so request_finished
177
+ # later sees `num_saved_tokens <= 0` and frees immediately rather than
178
+ # waiting for a finished_sending that will never come.
179
+ tracker = scheduler._request_trackers["req-0"]
180
+ assert tracker.num_saved_tokens == 0
181
+
182
+
183
+def _make_resumed_unfinished_request(
184
+ scheduler: MooncakeStoreScheduler,
185
+ *,
186
+ token_ids: list[int],
187
+ block_hashes: list[bytes],
188
+ num_computed_tokens: int,
189
+) -> None:
190
+ request = SimpleNamespace(
191
+ all_token_ids=token_ids,
192
+ block_hashes=block_hashes,
193
+ num_computed_tokens=num_computed_tokens,
194
+ num_output_placeholders=0,
195
+ )
196
+ scheduler._unfinished_requests["req-0"] = (request, ([0, 1],))
197
+
198
+
199
+def _make_resumed_scheduler_output(*, num_scheduled_tokens: int) -> SimpleNamespace:
200
+ return SimpleNamespace(
201
+ finished_req_ids=set(),
202
+ preempted_req_ids=set(),
203
+ scheduled_new_reqs=[],
204
+ scheduled_cached_reqs=SimpleNamespace(
205
+ req_ids=["req-0"],
206
+ new_block_ids=[([2],)],
207
+ num_computed_tokens=[0],
208
+ ),
209
+ num_scheduled_tokens={"req-0": num_scheduled_tokens},
210
+ scheduled_spec_decode_tokens={},
211
+ )
212
+
213
+
214
+def test_resumed_from_preemption_with_load_skips_save():
215
+ # On resume-from-preemption with a cache hit, the same co-queueing race
216
+ # applies: the resumed-from-preemption branch in build_connector_meta also
217
+ # passes load_spec.can_load=True. Skip save in this step; subsequent
218
+ # cached_reqs steps will save new tokens normally.
219
+ scheduler = _make_bare_scheduler()
220
+ scheduler._preempted_req_ids = {"req-0"}
221
+ _make_resumed_unfinished_request(
222
+ scheduler,
223
+ token_ids=list(range(48)),
224
+ block_hashes=[b"h0", b"h1", b"h2"],
225
+ num_computed_tokens=0,
226
+ )
227
+ scheduler.load_specs["req-0"] = LoadSpec(
228
+ vllm_cached_tokens=0,
229
+ kvpool_cached_tokens=48,
230
+ can_load=True,
231
+ )
232
+
233
+ meta = scheduler.build_connector_meta(
234
+ _make_resumed_scheduler_output(num_scheduled_tokens=48)
235
+ )
236
+
237
+ assert len(meta.requests) == 1
238
+ req_meta = meta.requests[0]
239
+ assert req_meta.req_id == "req-0"
240
+ assert req_meta.can_save is False
241
+ assert req_meta.load_spec is not None
242
+ assert req_meta.load_spec.can_load is True
243
+ tracker = scheduler._request_trackers["req-0"]
244
+ assert tracker.num_saved_tokens == 0
245
+
246
+
247
+def test_resumed_from_preemption_without_load_still_saves():
248
+ # No load_spec → behavior is unchanged: save proceeds.
249
+ scheduler = _make_bare_scheduler()
250
+ scheduler._preempted_req_ids = {"req-0"}
251
+ _make_resumed_unfinished_request(
252
+ scheduler,
253
+ token_ids=list(range(48)),
254
+ block_hashes=[b"h0", b"h1", b"h2"],
255
+ num_computed_tokens=0,
256
+ )
257
+
258
+ meta = scheduler.build_connector_meta(
259
+ _make_resumed_scheduler_output(num_scheduled_tokens=48)
260
+ )
261
+
262
+ assert len(meta.requests) == 1
263
+ req_meta = meta.requests[0]
264
+ assert req_meta.req_id == "req-0"
265
+ assert req_meta.can_save is True
266
+ assert req_meta.load_spec is None
267
+ tracker = scheduler._request_trackers["req-0"]
268
+ assert tracker.num_saved_tokens == 48
269
+
270
+
271
+# Focused tests for ReqMeta.from_request_tracker — the centralized guard that
272
+# enforces "a ReqMeta never carries both a save and a load".
273
+
274
+
275
+def test_from_request_tracker_load_overrides_caller_skip_save():
276
+ # Caller asks for skip_save=False, but load_spec.can_load=True. The
277
+ # function must force skip_save=True to avoid producing a ReqMeta the
278
+ # worker would enqueue on both kv_send_thread and kv_recv_thread.
279
+ tracker = RequestTracker(
280
+ req_id="req-0",
281
+ token_len=48,
282
+ allocated_block_ids=([0, 1, 2],),
283
+ num_saved_tokens=0,
284
+ )
285
+ load_spec = LoadSpec(vllm_cached_tokens=0, kvpool_cached_tokens=48, can_load=True)
286
+
287
+ req_meta = ReqMeta.from_request_tracker(
288
+ tracker,
289
+ block_size=16,
290
+ load_spec=load_spec,
291
+ skip_save=False,
292
+ block_hashes=[b"h0", b"h1", b"h2"],
293
+ )
294
+
295
+ assert req_meta is not None
296
+ assert req_meta.can_save is False
297
+ assert req_meta.load_spec is load_spec
298
+ assert tracker.num_saved_tokens == 0
299
+
300
+
301
+def test_from_request_tracker_load_with_can_load_false_still_saves():
302
+ # A LoadSpec with can_load=False (e.g., no external tokens to load after
303
+ # update_state_after_alloc) must not suppress the save.
304
+ tracker = RequestTracker(
305
+ req_id="req-0",
306
+ token_len=48,
307
+ allocated_block_ids=([0, 1, 2],),
308
+ num_saved_tokens=0,
309
+ )
310
+ load_spec = LoadSpec(vllm_cached_tokens=0, kvpool_cached_tokens=48, can_load=False)
311
+
312
+ req_meta = ReqMeta.from_request_tracker(
313
+ tracker,
314
+ block_size=16,
315
+ load_spec=load_spec,
316
+ skip_save=False,
317
+ block_hashes=[b"h0", b"h1", b"h2"],
318
+ )
319
+
320
+ assert req_meta is not None
321
+ assert req_meta.can_save is True
322
+ # from_request_tracker clears load_spec when can_load is False.
323
+ assert req_meta.load_spec is None
324
+ assert tracker.num_saved_tokens == 48
325
+
326
+
327
+def test_from_request_tracker_no_load_saves_normally():
328
+ tracker = RequestTracker(
329
+ req_id="req-0",
330
+ token_len=48,
331
+ allocated_block_ids=([0, 1, 2],),
332
+ num_saved_tokens=0,
333
+ )
334
+
335
+ req_meta = ReqMeta.from_request_tracker(
336
+ tracker,
337
+ block_size=16,
338
+ load_spec=None,
339
+ skip_save=False,
340
+ block_hashes=[b"h0", b"h1", b"h2"],
341
+ )
342
+
343
+ assert req_meta is not None
344
+ assert req_meta.can_save is True
345
+ assert req_meta.load_spec is None
346
+ assert tracker.num_saved_tokens == 48
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py
+6
@@ -236,6 +236,12 @@ class ReqMeta:
236
)
237
238
skip_save = skip_save or num_tokens_to_save < chunk_boundary
239
+ # A ReqMeta must never carry both a save AND a load.
240
+ # The save would also be wasted work — the bytes are being looked up
241
+ # in the store right now. Later cached_reqs steps save new tokens
242
+ # normally.
243
+ if load_spec is not None and load_spec.can_load:
244
+ skip_save = True
245
if skip_save and load_spec is None:
246
return None
247