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, 22:00 UTC
[KV Connector] Handle Mooncake finish after preemption (#43281)
Signed-off-by: Zhewen Li <zhewenli@inferact.ai> Co-authored-by: Zhewen Li <zhewenli@inferact.ai>
Zhewen Li committed
May 25, 2026 at 01:58 UTC
873758c13a64742e2a0247e0f2c62cadf027dd2b
3 files changed
+71
-3
tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py
+57
@@ -43,6 +43,21 @@ def _make_scheduler_output(*, scheduled_spec_tokens: list[int] | None):
43
)
44
45
46
+def _make_preemption_scheduler_output():
47
+ return SimpleNamespace(
48
+ finished_req_ids=set(),
49
+ preempted_req_ids={"req-0"},
50
+ scheduled_new_reqs=[],
51
+ scheduled_cached_reqs=SimpleNamespace(
52
+ req_ids=[],
53
+ new_block_ids=[],
54
+ num_computed_tokens=[],
55
+ ),
56
+ num_scheduled_tokens={},
57
+ scheduled_spec_decode_tokens={},
58
+ )
59
+
60
+
61
def _add_unfinished_request(
62
scheduler: MooncakeStoreScheduler,
63
*,
@@ -112,6 +127,48 @@ def test_cached_request_without_spec_decode_keeps_current_step_save_overlap():
127
assert tracker.num_saved_tokens == 48
128
129
130
+def test_preemption_resets_tracker_before_request_finished():
131
+ scheduler = _make_bare_scheduler()
132
+ _add_unfinished_request(
133
+ scheduler,
134
+ token_ids=list(range(44)),
135
+ block_hashes=[b"h0", b"h1"],
136
+ prefill_end_tokens=48,
137
+ )
138
+
139
+ scheduler.build_connector_meta(_make_preemption_scheduler_output())
140
+
141
+ tracker = scheduler._request_trackers["req-0"]
142
+ assert tracker.token_len == 0
143
+ assert tracker.allocated_block_ids == ()
144
+ assert tracker.num_saved_tokens == 0
145
+ assert tracker.token_ids is None
146
+ assert tracker.prefill_end_tokens == 0
147
+ request = SimpleNamespace(request_id="req-0")
148
+ assert scheduler.request_finished(request, ([0, 1],)) == (False, None)
149
+
150
+
151
+def test_preemption_clears_stale_load_state():
152
+ scheduler = _make_bare_scheduler()
153
+ _make_pending_load_unfinished_request(
154
+ scheduler,
155
+ num_tokens=48,
156
+ block_hashes=[b"h0", b"h1", b"h2"],
157
+ block_ids=([10, 11, 12],),
158
+ )
159
+ scheduler.load_specs["req-0"] = LoadSpec(
160
+ vllm_cached_tokens=0,
161
+ kvpool_cached_tokens=48,
162
+ can_load=True,
163
+ )
164
+
165
+ meta = scheduler.build_connector_meta(_make_preemption_scheduler_output())
166
+
167
+ assert meta.requests == []
168
+ assert "req-0" not in scheduler.load_specs
169
+ assert "req-0" not in scheduler._unfinished_requests
170
+
171
+
172
def _make_pending_load_unfinished_request(
173
scheduler: MooncakeStoreScheduler,
174
*,
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py
+7
@@ -173,6 +173,13 @@ class RequestTracker:
173
# request it includes previously-generated tokens, which are re-prefilled.
174
prefill_end_tokens: int = 0
175
176
+ def reset(self) -> None:
177
+ self.token_len = 0
178
+ self.allocated_block_ids = ()
179
+ self.num_saved_tokens = 0
180
+ self.token_ids = None
181
+ self.prefill_end_tokens = 0
182
+
183
def update(
184
self,
185
new_block_ids: tuple[list[int], ...] | list[int],
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py
+7
-3
@@ -171,7 +171,9 @@ class MooncakeStoreScheduler:
171
preempted_ids = scheduler_output.preempted_req_ids or set()
172
self._preempted_req_ids.update(preempted_ids)
173
for req_id in preempted_ids:
174
- self._request_trackers.pop(req_id, None)
174
+ self.load_specs.pop(req_id, None)
175
+ if request_tracker := self._request_trackers.get(req_id):
176
+ request_tracker.reset()
177
self._unfinished_requests.pop(req_id, None)
178
179
meta = MooncakeStoreConnectorMetadata(
@@ -355,8 +357,10 @@ class MooncakeStoreScheduler:
357
if self.kv_role == "kv_consumer":
358
return False, None
359
tracker = self._request_trackers.get(request.request_id)
358
- assert tracker is not None
359
- if tracker.num_saved_tokens <= 0:
360
+ # Missing tracker can happen when the request is aborted before the
361
+ # connector observes the normal finished lifecycle or is preempted
362
+ # before finishing.
363
+ if tracker is None or tracker.num_saved_tokens <= 0:
364
return False, None
365
total_blocks = sum(len(g) for g in block_ids)
366
delay_free_blocks = total_blocks > 0