A high-throughput and memory-efficient inference and serving engine for LLMs
[Fix][KV offload] Defer `on_request_finished` until in-flight transfers drain (#45823)
Signed-off-by: Ronen Schaffer <ronen.schaffer@ibm.com>
Ronen Schaffer committed
Jun 18, 2026 at 07:05 UTC
f428718ffe7487dae6d713b6dabb93dd59147349
4 files changed
+174
-7
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+129
@@ -243,6 +243,77 @@ def test_request_preemption(request_runner, async_scheduling: bool):
243
assert runner.connector_scheduler._block_id_to_pending_jobs == {}
244
245
246
+@pytest.mark.parametrize("async_scheduling", [True, False])
247
+def test_no_offload_call_after_on_request_finished(
248
+ request_runner, async_scheduling: bool
249
+):
250
+ """on_request_finished is not issued before a per-request offload
251
+ call.
252
+
253
+ A request can finish while its GPU->primary store is still in flight; the
254
+ later worker completion then drives complete_store. The scheduler defers
255
+ on_request_finished until the request is finished AND has no in-flight
256
+ transfer jobs, so complete_store is observed BEFORE on_request_finished,
257
+ and it is called exactly once.
258
+ """
259
+ block_size = 4
260
+ block_size_factor = 3
261
+ offloaded_block_size = block_size * block_size_factor
262
+
263
+ runner = request_runner(
264
+ block_size=block_size,
265
+ num_gpu_blocks=100,
266
+ async_scheduling=async_scheduling,
267
+ block_size_factor=block_size_factor,
268
+ )
269
+
270
+ # Record the order of per-request connector calls on the (mocked) manager.
271
+ # The external list survives manager.reset_mock() between run() calls.
272
+ calls: list[tuple[str, str]] = []
273
+ runner.manager.on_request_finished.side_effect = lambda req_context: calls.append(
274
+ ("on_request_finished", req_context.req_id)
275
+ )
276
+ runner.manager.complete_store.side_effect = (
277
+ lambda keys, req_context, *args, **kwargs: calls.append(
278
+ ("complete_store", req_context.req_id)
279
+ )
280
+ )
281
+ runner.manager.prepare_store.side_effect = lambda keys, req_context: (
282
+ generate_store_output(keys)
283
+ )
284
+
285
+ # Decode a couple of blocks, keeping every transfer in flight
286
+ # (complete_transfers=False) so no store completes while the request runs.
287
+ runner.new_request(token_ids=[0] * offloaded_block_size * 2)
288
+ runner.run(decoded_tokens=[0], complete_transfers=False)
289
+ runner.run(
290
+ decoded_tokens=[0] * (2 * offloaded_block_size),
291
+ complete_transfers=False,
292
+ )
293
+
294
+ # Finish the request, completing its pending stores. on_request_finished is
295
+ # deferred until the stores drain, so it lands after the last complete_store.
296
+ # 4 offloaded blocks are stored (2 prompt + 2 decode) -> 4 * block_size_factor
297
+ # GPU blocks.
298
+ runner.run(
299
+ decoded_tokens=[EOS_TOKEN_ID],
300
+ expected_stored=tuple(range(4 * block_size_factor)),
301
+ )
302
+
303
+ req_id = str(runner.req_id)
304
+
305
+ # on_request_finished is issued exactly once.
306
+ assert calls.count(("on_request_finished", req_id)) == 1, calls
307
+
308
+ finished_idx = calls.index(("on_request_finished", req_id))
309
+ store_indices = [i for i, c in enumerate(calls) if c == ("complete_store", req_id)]
310
+
311
+ # All of the request's complete_store calls must precede its single
312
+ # on_request_finished.
313
+ assert store_indices, calls
314
+ assert max(store_indices) < finished_idx, calls
315
+
316
+
317
@pytest.mark.parametrize("async_scheduling", [True, False])
318
def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling: bool):
319
block_size = 4
@@ -1149,6 +1220,64 @@ def test_reset_cache(request_runner, async_scheduling: bool):
1220
assert group_state.next_stored_block_idx == 0
1221
1222
1223
+@pytest.mark.parametrize("async_scheduling", [True, False])
1224
+def test_reset_cache_finalizes_finished_request_with_pending_store(
1225
+ request_runner, async_scheduling: bool
1226
+):
1227
+ """reset_cache must finalize a finished request whose in-flight stores it
1228
+ discards: call on_request_finished and drop its _req_status entry.
1229
+
1230
+ Otherwise the deferred hook (which waits for the now-discarded jobs to
1231
+ complete) never fires and the entry leaks.
1232
+ """
1233
+ block_size = 4
1234
+ block_size_factor = 3
1235
+ offloaded_block_size = block_size * block_size_factor
1236
+
1237
+ runner = request_runner(
1238
+ block_size=block_size,
1239
+ num_gpu_blocks=100,
1240
+ async_scheduling=async_scheduling,
1241
+ block_size_factor=block_size_factor,
1242
+ )
1243
+
1244
+ finalized: list[str] = []
1245
+ runner.manager.on_request_finished.side_effect = (
1246
+ lambda req_context: finalized.append(req_context.req_id)
1247
+ )
1248
+ runner.manager.prepare_store.side_effect = (
1249
+ lambda keys, req_context: generate_store_output(keys)
1250
+ )
1251
+
1252
+ # Decode a couple of blocks and keep every transfer in flight, so the
1253
+ # request has pending store jobs.
1254
+ runner.new_request(token_ids=[0] * offloaded_block_size * 2)
1255
+ runner.run(decoded_tokens=[0], complete_transfers=False)
1256
+ runner.run(
1257
+ decoded_tokens=[0] * (2 * offloaded_block_size),
1258
+ complete_transfers=False,
1259
+ )
1260
+
1261
+ cs = runner.connector_scheduler
1262
+ req_id = str(runner.req_id)
1263
+ req_status = cs._req_status[req_id]
1264
+ assert req_status.transfer_jobs, "expected an in-flight store before finish"
1265
+ assert any(job.is_store for job in cs._jobs.values())
1266
+
1267
+ # Finish the request while its store is still in flight. request_finished
1268
+ # takes the defer branch (pending jobs), so on_request_finished is NOT
1269
+ # called yet and the entry stays tracked.
1270
+ req_status.req.status = RequestStatus.FINISHED_STOPPED
1271
+ cs.request_finished(req_status.req)
1272
+ assert finalized == []
1273
+ assert req_id in cs._req_status
1274
+
1275
+ # reset_cache discards the in-flight store; it must finalize the request.
1276
+ cs.reset_cache()
1277
+ assert finalized == [req_id]
1278
+ assert req_id not in cs._req_status
1279
+
1280
+
1281
@pytest.mark.parametrize("async_scheduling", [True, False])
1282
def test_swa_alignment_skip(request_runner, async_scheduling: bool):
1283
"""SWA blocks unreachable by the load path are skipped during store.
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+27
-7
@@ -1115,6 +1115,11 @@ class OffloadingConnectorScheduler:
1115
del self._jobs[job_id]
1116
req_status.transfer_jobs.remove(job_id)
1117
if not req_status.transfer_jobs and req_status.req.is_finished():
1118
+ # Deferred from request_finished: the request's last in-flight
1119
+ # job is now done, so fire the finalize hook here, after the
1120
+ # final complete_store/complete_load above (and any submit_store
1121
+ # the complete_store cascade issued).
1122
+ self.manager.on_request_finished(req_status.req_context)
1123
del self._req_status[job_status.req_id]
1124
1125
def get_stats(self) -> OffloadingConnectorStats | None:
@@ -1148,18 +1153,23 @@ class OffloadingConnectorScheduler:
1153
# which may have been deferred due to async scheduling
1154
req_status = self._req_status.get(request.request_id)
1155
1151
- req_context = (
1152
- req_status.req_context if req_status else _create_req_context(request)
1153
- )
1154
- self.manager.on_request_finished(req_context)
1155
-
1156
if req_status is None:
1157
+ # Untracked request (offloading never started): no in-flight jobs,
1158
+ # nothing was deferred, so finalize immediately.
1159
+ self.manager.on_request_finished(_create_req_context(request))
1160
return False, None
1161
+
1162
if not req_status.transfer_jobs:
1163
+ # No in-flight jobs: all per-request calls are done, finalize now.
1164
+ self.manager.on_request_finished(req_status.req_context)
1165
del self._req_status[request.request_id]
1166
return False, None
1161
- # Pending stores will outlive the request's block ownership.
1162
- # Register them so future block reuse triggers a flush.
1167
+
1168
+ # In-flight jobs remain, so defer on_request_finished to
1169
+ # update_connector_output, which fires it once the last job completes
1170
+ # (after the final complete_store and any cascade submit_store it
1171
+ # issues). These pending stores outlive the request's block ownership;
1172
+ # register them so future reuse of those blocks triggers a flush.
1173
for job_id in req_status.transfer_jobs:
1174
job_status = self._jobs[job_id]
1175
for bid in job_status.non_sliding_window_block_ids or ():
@@ -1198,6 +1208,16 @@ class OffloadingConnectorScheduler:
1208
# Flush all in-flight jobs
1209
self._current_batch_jobs_to_flush.update(self._jobs.keys())
1210
1211
+ # A finished request may still be tracked here with in-flight jobs that
1212
+ # this reset discards, so its deferred on_request_finished() would never
1213
+ # fire (completions are skipped as stale) and its _req_status entry would
1214
+ # leak. Finalize such requests now, before resetting the manager.
1215
+ # list() snapshots because we delete while iterating.
1216
+ for req_id, status in list(self._req_status.items()):
1217
+ if status.req.is_finished():
1218
+ self.manager.on_request_finished(status.req_context)
1219
+ del self._req_status[req_id]
1220
+
1221
# Reset offloading manager cache
1222
self.manager.reset_cache()
1223
vllm/v1/kv_offload/base.py
+11
@@ -266,6 +266,17 @@ class OffloadingManager(ABC):
266
"""
267
Called when a request has finished.
268
269
+ By the time this is called, all per-request offload calls for this
270
+ request (prepare_store/complete_store, prepare_load/complete_load,
271
+ touch, lookup) have already been issued, and none will follow. The
272
+ scheduler defers this call until the request is finished and has no
273
+ in-flight transfer jobs.
274
+
275
+ Note this signals only that no further calls will be made; it does NOT
276
+ imply the data has been persisted. Asynchronous transfers already
277
+ submitted for this request (e.g. CPU->secondary cascades) may still be
278
+ in flight. This is the right place to release per-request bookkeeping.
279
+
280
Args:
281
req_context: per-request context.
282
"""
vllm/v1/kv_offload/tiering/base.py
+7
@@ -188,6 +188,13 @@ class SecondaryTierManager(ABC):
188
"""
189
Called when a request has finished.
190
191
+ By the time this is called, all per-request calls for this request
192
+ (submit_store, submit_load, touch) have already been issued, and none
193
+ will follow. Note this does NOT imply the tier's transfers have
194
+ completed: jobs already submitted may still be in flight and will
195
+ report via get_finished_jobs(). This is the right place to release
196
+ per-request bookkeeping.
197
+
198
Args:
199
req_context: per-request context.
200
"""