[Bugfix][KV Offloading] Offload last block at request finish and prevent reuse race (#48596)
Signed-off-by: Alex <alex.tech.lab@outlook.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>
AlexHuang committed
Jul 17, 2026 at 21:50 UTC
f38f3d11fb7b180c4d66f4af82cf770f335391cc
3 files changed
+116
-31
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+51
-2
@@ -64,6 +64,47 @@ def test_scheduler_reports_allocation_failure(request_runner):
64
assert reduced[_ConnectorMetricName.ALLOCATION_FAILURE] == 1
65
66
67
+@pytest.mark.parametrize("async_scheduling", [True, False])
68
+@pytest.mark.parametrize("prompt_offset", [-1, -2])
69
+def test_last_block_offloaded_at_request_finish(
70
+ request_runner, async_scheduling: bool, prompt_offset: int
71
+):
72
+ """EOS fills the last block at request finish — verify req_status is kept alive.
73
+
74
+ prompt = block_size + prompt_offset tokens → not a full block at schedule time,
75
+ so _build_store_jobs creates no store job. After EOS, request_finished
76
+ keeps req_status alive so _build_store_jobs can process it on the next step.
77
+
78
+ prompt_offset=-1: EOS fills the block → store job created on next step.
79
+ prompt_offset=-2: block remains partial → no store job, cleanup in
80
+ _build_store_jobs deletes req_status.
81
+ """
82
+ block_size = 4
83
+ runner = request_runner(
84
+ block_size=block_size,
85
+ num_gpu_blocks=10,
86
+ async_scheduling=async_scheduling,
87
+ )
88
+ # prompt = block_size + prompt_offset tokens
89
+ runner.new_request(token_ids=[0] * (block_size + prompt_offset))
90
+ runner.manager.prepare_store.side_effect = lambda keys, req_context: (
91
+ generate_store_output(list(keys))
92
+ )
93
+
94
+ # Run with one step (EOS)
95
+ runner.run(
96
+ decoded_tokens=[EOS_TOKEN_ID],
97
+ )
98
+
99
+ cs = runner.connector_scheduler
100
+ # Verify req_status is kept alive for _build_store_jobs to process
101
+ # regardless of whether there are storable blocks
102
+ assert "0" in cs._req_status, (
103
+ "req_status was deleted but should be kept alive "
104
+ "for _build_store_jobs to process finished_req_ids."
105
+ )
106
+
107
+
108
def test_scheduler_reports_lookup_sync_delay(request_runner):
109
runner = request_runner(
110
block_size=4,
@@ -560,11 +601,19 @@ def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bo
601
# Group 1 (sliding window, window=2): only the last 2 blocks
602
# are within the window → loads blocks 1,2
603
expected_loaded=((0, 0), (0, 1), (0, 2), (1, 1), (1, 2)),
604
+ # The deferred store from the previous request's last block
605
+ # completes during this step, and its blocks are flushed because
606
+ # they were reallocated to the new request.
607
+ # Only block 1 (sliding window group) is stored — block 0's
608
+ # deferred store is flushed because it was reallocated.
609
+ expected_stored=((0, 1),),
610
+ expected_flushed=((0, 1),),
611
)
612
565
- # one touch in get_num_new_matched_tokens x 2 groups
613
+ # 4 touch calls: 2 from get_num_new_matched_tokens (2 groups)
614
+ # + 2 from _get_reqs_to_store (2 groups)
615
touch_calls = runner.manager.touch.call_args_list
567
- assert len(touch_calls) == 2
616
+ assert len(touch_calls) == 4
617
# full attention group touched all 3 blocks
618
assert len(touch_calls[0].args[0]) == 3
619
# sliding window group touched just the last 2 blocks
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+51
-29
@@ -3,7 +3,7 @@
3
import time
4
from collections.abc import Iterable, Sequence
5
from dataclasses import dataclass, field
6
-from itertools import islice
6
+from itertools import chain, islice
7
from typing import Any, NamedTuple
8
9
from vllm.config import VllmConfig
@@ -444,6 +444,24 @@ class OffloadingConnectorScheduler:
444
if not pending:
445
del self._block_id_to_pending_jobs[bid]
446
447
+ def _calc_num_offloadable_tokens(
448
+ self, req_status: RequestOffloadState, num_computed_tokens: int
449
+ ) -> int:
450
+ num = min(num_computed_tokens, req_status.req.num_tokens)
451
+ max_offload_tokens = req_status.max_offload_tokens
452
+ if max_offload_tokens is not None:
453
+ num = min(num, max_offload_tokens)
454
+ if self.config.offload_prompt_only:
455
+ num = min(num, req_status.req.num_prompt_tokens)
456
+ return num
457
+
458
+ def _maybe_cleanup_finished_req(
459
+ self, req_id: str, req_status: RequestOffloadState
460
+ ) -> None:
461
+ """Clean up req_status if finished and no in-flight jobs."""
462
+ if req_status.req.is_finished() and not req_status.transfer_jobs:
463
+ del self._req_status[req_id]
464
+
465
def _maximal_prefix_lookup(
466
self, keys: Iterable[OffloadKey], req_context: ReqContext
467
) -> int | None:
@@ -911,28 +929,24 @@ class OffloadingConnectorScheduler:
929
) -> dict[int, TransferJob]:
930
blocks_per_chunk = self.config.blocks_per_chunk
931
store_jobs: dict[int, TransferJob] = {}
914
- for req_id in scheduler_output.num_scheduled_tokens:
932
+ for req_id in chain(
933
+ scheduler_output.num_scheduled_tokens,
934
+ scheduler_output.finished_req_ids or (),
935
+ ):
936
req_status = self._req_status.get(req_id)
937
if req_status is None:
938
continue
939
req = req_status.req
940
920
- num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id]
921
- num_tokens_after_batch = req.num_computed_tokens + num_scheduled_tokens
922
- # with async scheduling, some tokens may be missing
923
- num_offloadable_tokens = min(num_tokens_after_batch, req.num_tokens)
924
- max_offload_tokens = req_status.max_offload_tokens
925
- if max_offload_tokens is not None:
926
- num_offloadable_tokens = min(num_offloadable_tokens, max_offload_tokens)
927
-
928
- # Skip decode-phase chunks: clamp to the prompt length so only
929
- # prefill chunks become eligible for store. next_stored_chunk_idx
930
- # never advances past this boundary, so decode chunks are never
931
- # queued in this or any later step.
932
- if self.config.offload_prompt_only:
933
- num_offloadable_tokens = min(
934
- num_offloadable_tokens, req.num_prompt_tokens
935
- )
941
+ if req.is_finished():
942
+ num_tokens_after_batch = req.num_tokens
943
+ else:
944
+ num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id]
945
+ num_tokens_after_batch = req.num_computed_tokens + num_scheduled_tokens
946
+
947
+ num_offloadable_tokens = self._calc_num_offloadable_tokens(
948
+ req_status, num_tokens_after_batch
949
+ )
950
951
# Filter out chunks skipped due to sliding window attention / SSM
952
# or unreachable by the load path's alignment constraints.
@@ -983,6 +997,7 @@ class OffloadingConnectorScheduler:
997
998
if not new_offload_keys:
999
req_status.advance_stored_idx(num_offloadable_tokens)
1000
+ self._maybe_cleanup_finished_req(req_id, req_status)
1001
continue
1002
1003
store_output = self.manager.prepare_store(
@@ -993,10 +1008,12 @@ class OffloadingConnectorScheduler:
1008
_ConnectorMetricName.ALLOCATION_FAILURE
1009
)
1010
logger.warning("Request %s: cannot store chunks", req_id)
1011
+ self._maybe_cleanup_finished_req(req_id, req_status)
1012
continue
1013
1014
if not store_output.keys_to_store:
1015
req_status.advance_stored_idx(num_offloadable_tokens)
1016
+ self._maybe_cleanup_finished_req(req_id, req_status)
1017
continue
1018
1019
self._touch(req_status)
@@ -1093,6 +1110,13 @@ class OffloadingConnectorScheduler:
1110
job_id,
1111
)
1112
1113
+ if req.is_finished():
1114
+ # Register non-sliding-window blocks for flush detection.
1115
+ for bid in non_sliding_window_block_ids:
1116
+ self._block_id_to_pending_jobs.setdefault(bid, set()).add(job_id)
1117
+ if bid in self._current_batch_allocated_block_ids:
1118
+ self._current_batch_jobs_to_flush.add(job_id)
1119
+
1120
return store_jobs
1121
1122
def build_connector_meta(
@@ -1255,8 +1279,6 @@ class OffloadingConnectorScheduler:
1279
Optional KVTransferParams to be included in the request outputs
1280
returned by the engine.
1281
"""
1258
- # TODO(orozery): possibly kickoff offload for last block
1259
- # which may have been deferred due to async scheduling
1282
req_status = self._req_status.get(request.request_id)
1283
1284
if req_status is None:
@@ -1269,20 +1291,20 @@ class OffloadingConnectorScheduler:
1291
1292
self.manager.on_request_finished(req_status.req_context)
1293
self._maybe_observe_lookup_async_delay(req_status)
1272
- if not req_status.transfer_jobs:
1273
- # No in-flight jobs: no later complete_store()/complete_load() calls
1274
- # need this request's state.
1275
- del self._req_status[request.request_id]
1276
- return False, None
1294
1278
- # In-flight jobs remain after the request stopped. Their completion may
1279
- # still call manager.complete_store()/complete_load(), so keep req_status.
1280
- # Pending stores outlive the request's block ownership; register them so
1281
- # future reuse of those blocks triggers a flush.
1295
+ # Update offload keys with final block hash so _build_store_jobs can
1296
+ # create store jobs for the last block(s) on the next schedule step.
1297
+ req_status.update_offload_keys()
1298
+
1299
+ # Keep req_status alive: _build_store_jobs will process finished_req_ids
1300
+ # on the next step and handle cleanup after creating store jobs.
1301
+ # Register non_sliding_window_block_ids so future block reuse triggers
1302
+ # a flush via _block_id_to_pending_jobs.
1303
for job_id in req_status.transfer_jobs:
1304
job_status = self._jobs[job_id]
1305
for bid in job_status.non_sliding_window_block_ids or ():
1306
self._block_id_to_pending_jobs.setdefault(bid, set()).add(job_id)
1307
+
1308
return False, None
1309
1310
def take_events(self) -> Iterable[KVCacheEvent]:
vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py
+14
@@ -280,7 +280,21 @@ class OffloadingConnectorWorker:
280
281
def handle_preemptions(self, kv_connector_metadata: OffloadingConnectorMetadata):
282
assert self.worker is not None
283
+
284
+ # Pop jobs_to_flush from store_jobs into _unsubmitted_store_jobs
285
+ # so the existing submission loop below submits them before wait().
286
+ if kv_connector_metadata.jobs_to_flush:
287
+ for job_id in kv_connector_metadata.jobs_to_flush:
288
+ entry = kv_connector_metadata.store_jobs.pop(job_id, None)
289
+ if entry is not None:
290
+ assert isinstance(entry.src_spec, GPULoadStoreSpec)
291
+ self._unsubmitted_store_jobs.append(
292
+ (job_id, entry.src_spec, entry.dst_spec)
293
+ )
294
+
295
+ # Submit deferred stores from previous step (and jobs_to_flush above).
296
for job_id, src_spec, dst_spec in self._unsubmitted_store_jobs:
297
+ assert isinstance(src_spec, GPULoadStoreSpec)
298
success = self.worker.submit_store(job_id, src_spec, dst_spec)
299
assert success
300
self._unsubmitted_store_jobs.clear()