[Bugfix][MooncakeStore] Fix finish-time save crash on hybrid models (#54643)
Signed-off-by: zhewenl <zhewenl@users.noreply.github.com> Co-authored-by: zhewenl <zhewenl@users.noreply.github.com>
Zhewen Li committed
Sep 7, 2026 at 00:18 UTC
34b1e9f7a61e802b870c7d7f78bf90b000b691cc
2 files changed
+33
-1
tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py
+27
@@ -199,6 +199,33 @@ def _add_unfinished_request(
199
)
200
201
202
+def test_pending_load_for_non_chosen_connector_is_dropped():
203
+ """A MultiConnector loser must not turn its proposed load into a save."""
204
+ scheduler = _make_bare_scheduler()
205
+ request = SimpleNamespace(
206
+ request_id="req-0",
207
+ block_hashes=[b"h0", b"h1", b"h2"],
208
+ )
209
+ blocks = SimpleNamespace(get_block_ids=lambda: ([1, 2], [9]))
210
+ scheduler.load_specs["req-0"] = LoadSpec(
211
+ vllm_cached_tokens=0,
212
+ kvpool_cached_tokens=48,
213
+ can_load=False,
214
+ )
215
+
216
+ scheduler.update_state_after_alloc(request, blocks, num_external_tokens=0)
217
+ meta = scheduler.build_connector_meta(_make_pending_load_scheduler_output())
218
+
219
+ # MultiConnector exposes the real allocation to every child, but only the
220
+ # chosen child receives external tokens. The losing store connector must
221
+ # neither retain those blocks for a pending load nor enqueue a save from
222
+ # the rejected speculative LoadSpec.
223
+ assert scheduler._unfinished_requests["req-0"][1] == ()
224
+ assert meta.requests == []
225
+ assert "req-0" not in scheduler.load_specs
226
+ assert "req-0" not in scheduler._request_trackers
227
+
228
+
229
def test_update_state_excludes_nontransfer_groups():
230
"""Store metadata must match the worker's registered cache groups."""
231
scheduler = _make_bare_scheduler()
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py
+6
-1
@@ -372,7 +372,12 @@ class MooncakeStoreScheduler:
372
) in self._unfinished_requests.items():
373
if request_id not in request_ids and request_id not in cached_reqs.req_ids:
374
load_spec = self.load_specs.pop(request_id, None)
375
- if not load_spec:
375
+ # A load spec may have been proposed by this connector's
376
+ # lookup but rejected by MultiConnector in favor of another
377
+ # connector. Only the chosen connector may issue the pending
378
+ # load; the normal store path gets its blocks later from
379
+ # SchedulerOutput once the request is actually scheduled.
380
+ if load_spec is None or not load_spec.can_load:
381
continue
382
num_tokens_to_compute = load_spec.kvpool_cached_tokens
383
request_tracker = RequestTracker(