@tduncan / Nvidia-Optimization-Stack-VLLM / commits / d30b1ecd1b

[Bugfix][KV Offloading] Defer request finalization until final store (#49671)

Signed-off-by: Rui Yin <2260891073@qq.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>

Rui Yin committed Jul 26, 2026 at 04:58 UTC d30b1ecd1bdf7c3d92f3b444c4538efd8fbb40ac

3 files changed +111 -58

tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py

+74 -41

@@ -68,7 +68,9 @@ def test_scheduler_reports_allocation_failure(request_runner):

68 runner.run(decoded_tokens=[EOS_TOKEN_ID])

70 reduced = _reduce_kv_connector_stats(runner)

71 - assert reduced[_ConnectorMetricName.ALLOCATION_FAILURE] == 1

71 + # Two attempts: once while running (block becomes full during prefill),

72 + # once from finished_req_ids on the next step.

73 + assert reduced[_ConnectorMetricName.ALLOCATION_FAILURE] == 2

76 @pytest.mark.parametrize("async_scheduling", [True, False])

@@ -76,7 +78,7 @@ def test_scheduler_reports_allocation_failure(request_runner):

78 def test_last_block_offloaded_at_request_finish(

79 request_runner, async_scheduling: bool, prompt_offset: int

80 ):

79 - """EOS fills the last block at request finish — verify req_status is kept alive.

81 + """EOS fills the last block at request finish - verify the final block is stored.

83 prompt = block_size + prompt_offset tokens → not a full block at schedule time,

84 so _build_store_jobs creates no store job. After EOS, request_finished

@@ -98,18 +100,16 @@ def test_last_block_offloaded_at_request_finish(

100 generate_store_output(list(keys))

101 )

102

101 - # Run with one step (EOS)

102 - runner.run(

103 - decoded_tokens=[EOS_TOKEN_ID],

104 - )

103 + if prompt_offset == -1:

104 + # EOS fills the block, so a store job is created for block 0.

105 + runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_stored=(0,))

106 + else:

107 + # Block remains partial, so no store job is created.

108 + runner.run(decoded_tokens=[EOS_TOKEN_ID])

109

110 cs = runner.connector_scheduler

107 - # Verify req_status is kept alive for _build_store_jobs to process

108 - # regardless of whether there are storable blocks

109 - assert "0" in cs._req_status, (

110 - "req_status was deleted but should be kept alive "

111 - "for _build_store_jobs to process finished_req_ids."

112 - )

111 + # After the full run completes, req_status is cleaned up.

112 + assert "0" not in cs._req_status

113

114

115 @pytest.mark.parametrize("async_scheduling", [True, False])

@@ -569,14 +569,14 @@ def test_request_preemption(request_runner, async_scheduling: bool):

569

570

571 @pytest.mark.parametrize("async_scheduling", [True, False])

572 -def test_on_request_finished_is_not_deferred_until_store_completion(

572 +def test_on_request_finished_not_deferred_until_store_completion(

573 request_runner, async_scheduling: bool

574 ):

575 - """on_request_finished fires when no more stores will be submitted.

575 + """on_request_finished fires after the last prepare_store is submitted.

576

577 - A request can finish while its GPU->primary store is still in flight. The

578 - manager-level hook should not wait for that completion; complete_store may

579 - still arrive afterward for already-submitted transfer jobs.

577 + The manager contract guarantees no more submit-side calls (prepare_store)

578 + after on_request_finished. However, complete_store callbacks for

579 + already-submitted transfers may still arrive afterward.

580 """

581 block_size = 4

582 blocks_per_chunk = 3

@@ -613,8 +613,9 @@ def test_on_request_finished_is_not_deferred_until_store_completion(

613 complete_transfers=False,

614 )

615

616 - # Finish the request while its stores are still in flight. The hook should

617 - # fire immediately even though no complete_store has arrived yet.

616 + # Finish the request while its stores are still in flight. The hook fires

617 + # once the last prepare_store is issued (on the next schedule step), even

618 + # though complete_store has not yet been called.

619 runner.run(

620 decoded_tokens=[EOS_TOKEN_ID],

621 complete_transfers=False,

@@ -624,8 +625,7 @@ def test_on_request_finished_is_not_deferred_until_store_completion(

625

626 assert calls == [("on_request_finished", req_id)], calls

627

627 - # Drain the stores afterward. The already-submitted complete_store calls

628 - # are allowed to arrive after on_request_finished.

628 + # Drain the stores afterward. complete_store is allowed after the hook.

629 runner.run(

630 decoded_tokens=[],

631 complete_transfers=True,

@@ -638,11 +638,50 @@ def test_on_request_finished_is_not_deferred_until_store_completion(

638 finished_idx = calls.index(("on_request_finished", req_id))

639 store_indices = [i for i, c in enumerate(calls) if c == ("complete_store", req_id)]

640

641 - # The request-level hook no longer waits for already-submitted transfers.

641 + # complete_store arrives after on_request_finished, as allowed by the contract.

642 assert store_indices, calls

643 assert finished_idx < min(store_indices), calls

644

645

646 +@pytest.mark.parametrize("async_scheduling", [True, False])

647 +def test_on_request_finished_fires_after_final_block_store(

648 + request_runner, async_scheduling: bool

649 +):

650 + """on_request_finished fires after the final-block prepare_store at EOS.

651 +

652 + When EOS fills a partial block, request_finished() keeps req_status alive

653 + so _build_store_jobs can create a store job for it on the next step.

654 + """

655 + block_size = 4

656 + runner = request_runner(

657 + block_size=block_size,

658 + num_gpu_blocks=10,

659 + async_scheduling=async_scheduling,

660 + )

661 +

662 + calls: list[tuple[str, str]] = []

663 + runner.manager.on_request_finished.side_effect = lambda req_context: calls.append(

664 + ("on_request_finished", req_context.req_id)

665 + )

666 +

667 + def prepare_store(keys, req_context):

668 + calls.append(("prepare_store", req_context.req_id))

669 + return generate_store_output(keys)

670 +

671 + runner.manager.prepare_store.side_effect = prepare_store

672 +

673 + runner.new_request(token_ids=[0] * (block_size - 1))

674 + runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_stored=(0,))

675 +

676 + req_id = str(runner.req_id)

677 + assert calls.count(("on_request_finished", req_id)) == 1, calls

678 +

679 + finished_idx = calls.index(("on_request_finished", req_id))

680 + prepare_indices = [i for i, c in enumerate(calls) if c == ("prepare_store", req_id)]

681 + assert prepare_indices, calls

682 + assert finished_idx > max(prepare_indices), calls

683 +

684 +

685 @pytest.mark.parametrize("async_scheduling", [True, False])

686 def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling: bool):

687 block_size = 4

@@ -831,7 +870,10 @@ def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bo

870 touch_calls = runner.manager.touch.call_args_list

871 assert len(touch_calls) == 6

872

834 - runner.run(decoded_tokens=[EOS_TOKEN_ID])

873 + # EOS fills the 7th block (offset 6). The extra schedule step processes

874 + # finished_req_ids and stores block 6 for both groups before the request's

875 + # GPU blocks are freed.

876 + runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_stored=(6,))

877

878 runner.scheduler.reset_prefix_cache()

879

@@ -844,19 +886,11 @@ def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bo

886 # Group 1 (sliding window, window=2): only the last 2 blocks

887 # are within the window → loads blocks 1,2

888 expected_loaded=((0, 0), (0, 1), (0, 2), (1, 1), (1, 2)),

847 - # The deferred store from the previous request's last block

848 - # completes during this step, and its blocks are flushed because

849 - # they were reallocated to the new request.

850 - # Only block 1 (sliding window group) is stored — block 0's

851 - # deferred store is flushed because it was reallocated.

852 - expected_stored=((0, 1),),

853 - expected_flushed=((0, 1),),

889 )

890

856 - # 4 touch calls: 2 from get_num_new_matched_tokens (2 groups)

857 - # + 2 from _get_reqs_to_store (2 groups)

891 + # 2 touch calls from get_num_new_matched_tokens (2 groups)

892 touch_calls = runner.manager.touch.call_args_list

859 - assert len(touch_calls) == 4

893 + assert len(touch_calls) == 2

894 # full attention group touched all 3 blocks

895 assert len(touch_calls[0].args[0]) == 3

896 # sliding window group touched just the last 2 blocks

@@ -1733,8 +1767,8 @@ def test_reset_cache(request_runner, async_scheduling: bool):

1767 def test_reset_cache_finalizes_finished_request_with_pending_store(

1768 request_runner, async_scheduling: bool

1769 ):

1736 - """reset_cache drops a finished request whose in-flight stores it discards

1737 - without calling on_request_finished twice.

1770 + """reset_cache fires on_request_finished for a finished request whose

1771 + in-flight stores it discards, exactly once.

1772 """

1773 block_size = 4

1774 blocks_per_chunk = 3

@@ -1770,16 +1804,15 @@ def test_reset_cache_finalizes_finished_request_with_pending_store(

1804 assert req_status.transfer_jobs, "expected an in-flight store before finish"

1805 assert any(job.is_store for job in cs._jobs.values())

1806

1773 - # Finish the request while its store is still in flight. request_finished

1774 - # fires the hook eagerly, but the entry stays tracked so later completions

1775 - # can still call complete_store().

1807 + # Finish the request while its store is still in flight. The manager hook

1808 + # is deferred because the final store decision has not happened yet.

1809 req_status.req.status = RequestStatus.FINISHED_STOPPED

1810 cs.request_finished(req_status.req)

1778 - assert finalized == [req_id]

1811 + assert finalized == []

1812 assert req_id in cs._req_status

1813

1781 - # reset_cache discards the in-flight store and drops the state without a

1782 - # duplicate on_request_finished call.

1814 + # reset_cache discards both the in-flight and not-yet-prepared final stores,

1815 + # so it issues the deferred notification before dropping the state.

1816 cs.reset_cache()

1817 assert finalized == [req_id]

1818 assert req_id not in cs._req_status

tests/v1/kv_connector/unit/offloading_connector/utils.py

+21 -5

@@ -482,8 +482,13 @@ class RequestRunner:

482 # Strict-always-False frees the request immediately on EOS, but

483 # the worker may still have a deferred store queued. In production

484 # the next request's step drains it; in single-request tests we

485 - # must keep stepping until the scheduler sees no in-flight jobs.

486 - if not self.scheduler.requests and not self.connector_scheduler._jobs:

485 + # must keep stepping until the scheduler sees no in-flight jobs

486 + # and no pending finished_req_ids awaiting build_connector_meta.

487 + if (

488 + not self.scheduler.requests

489 + and not self.connector_scheduler._jobs

490 + and not self.scheduler.finished_req_ids

491 + ):

492 break

493

494 scheduler_output = self.scheduler.schedule()

@@ -544,19 +549,30 @@ class RequestRunner:

549 if (

550 prev_token_id == EOS_TOKEN_ID

551 and prev_token_id != token_id

547 - and (self.scheduler.requests or self.connector_scheduler._jobs)

552 + and (

553 + self.scheduler.requests

554 + or self.connector_scheduler._jobs

555 + or self.scheduler.finished_req_ids

556 + )

557 ):

558 # continue for one more step to allow offloading to kick off

559 continue

560

561 if token_id is None:

562 if self.async_scheduling:

554 - # sample last token

563 + # Flush the previous step's output.

564 engine_outputs = self.scheduler.update_from_output(

565 prev_scheduler_output, prev_model_runner_output

566 )

567 self._record_kv_connector_stats(engine_outputs)

559 - break

568 + prev_model_runner_output = None

569 + if self.scheduler.requests:

570 + # Request still running, just exhausted decoded_tokens.

571 + break

572 + if not self.scheduler.finished_req_ids and (

573 + not complete_transfers or not self.connector_scheduler._jobs

574 + ):

575 + break

576

577 self._parse_transfers()

578

vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py

+16 -12

@@ -278,6 +278,8 @@ class RequestOffloadState:

278 # time.monotonic() of this request's first deferred offload lookup;

279 # None once consumed (observed) or while no lookup is pending.

280 deferred_lookup_start_time: float | None = None

281 + # True once on_request_finished has been signaled to the manager.

282 + finished_signaled: bool = False

283

284 def __post_init__(self) -> None:

285 self.group_states = tuple(

@@ -481,13 +483,6 @@ class OffloadingConnectorScheduler:

483 num = min(num, req_status.req.num_prompt_tokens)

484 return num

485

484 - def _maybe_cleanup_finished_req(

485 - self, req_id: str, req_status: RequestOffloadState

486 - ) -> None:

487 - """Clean up req_status if finished and no in-flight jobs."""

488 - if req_status.req.is_finished() and not req_status.transfer_jobs:

489 - del self._req_status[req_id]

490 -

486 def _maximal_prefix_lookup(

487 self,

488 keys: Iterable[OffloadKey],

@@ -1039,7 +1034,6 @@ class OffloadingConnectorScheduler:

1034

1035 if not new_offload_keys:

1036 req_status.advance_stored_idx(num_offloadable_tokens)

1042 - self._maybe_cleanup_finished_req(req_id, req_status)

1037 continue

1038

1039 store_output = self.manager.prepare_store(

@@ -1050,12 +1044,10 @@ class OffloadingConnectorScheduler:

1044 _ConnectorMetricName.ALLOCATION_FAILURE

1045 )

1046 logger.warning("Request %s: cannot store chunks", req_id)

1053 - self._maybe_cleanup_finished_req(req_id, req_status)

1047 continue

1048

1049 if not store_output.keys_to_store:

1050 req_status.advance_stored_idx(num_offloadable_tokens)

1058 - self._maybe_cleanup_finished_req(req_id, req_status)

1051 continue

1052

1053 self._touch(req_status)

@@ -1199,6 +1191,17 @@ class OffloadingConnectorScheduler:

1191 store_jobs=self._build_store_jobs(scheduler_output),

1192 jobs_to_flush=self._current_batch_jobs_to_flush,

1193 )

1194 +

1195 + # All prepare_store calls for finished requests have been issued.

1196 + # Signal on_request_finished and clean up state where possible.

1197 + for req_id in scheduler_output.finished_req_ids or ():

1198 + req_status = self._req_status.get(req_id)

1199 + if req_status is None:

1200 + continue

1201 + req_status.finished_signaled = True

1202 + self.manager.on_request_finished(req_status.req_context)

1203 + if not req_status.transfer_jobs:

1204 + del self._req_status[req_id]

1205 self._current_batch_load_jobs = {}

1206 self._current_batch_jobs_to_flush = set()

1207 self._current_batch_allocated_block_ids = set()

@@ -1289,7 +1292,7 @@ class OffloadingConnectorScheduler:

1292

1293 del self._jobs[job_id]

1294 req_status.transfer_jobs.remove(job_id)

1292 - if not req_status.transfer_jobs and req_status.req.is_finished():

1295 + if req_status.finished_signaled and not req_status.transfer_jobs:

1296 del self._req_status[job_status.req_id]

1297

1298 def get_stats(self) -> OffloadingConnectorStats | None:

@@ -1331,7 +1334,6 @@ class OffloadingConnectorScheduler:

1334 self.manager.on_request_finished(req_context)

1335 return False, None

1336

1334 - self.manager.on_request_finished(req_status.req_context)

1337 self._maybe_observe_lookup_async_delay(req_status)

1338

1339 # Update offload keys with final block hash so _build_store_jobs can

@@ -1375,6 +1377,8 @@ class OffloadingConnectorScheduler:

1377

1378 for req_id, status in list(self._req_status.items()):

1379 if status.req.is_finished():

1380 + if not status.finished_signaled:

1381 + self.manager.on_request_finished(status.req_context)

1382 del self._req_status[req_id]

1383

1384 # Reset offloading manager cache