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

Fix KV offload request-finished lifecycle contract (#46284)

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

Rui Yin committed Jun 24, 2026 at 20:42 UTC cf9fd6457eb7bc18942bf242365a3aa7d6f49ad8

5 files changed +336 -114

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

+28 -26

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

216

217

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

219 -def test_no_offload_call_after_on_request_finished(

219 +def test_on_request_finished_is_not_deferred_until_store_completion(

220 request_runner, async_scheduling: bool

221 ):

222 - """on_request_finished is not issued before a per-request offload

223 - call.

224 -

225 - A request can finish while its GPU->primary store is still in flight; the

226 - later worker completion then drives complete_store. The scheduler defers

227 - on_request_finished until the request is finished AND has no in-flight

228 - transfer jobs, so complete_store is observed BEFORE on_request_finished,

229 - and it is called exactly once.

222 + """on_request_finished fires when no more stores will be submitted.

223 +

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

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

226 + still arrive afterward for already-submitted transfer jobs.

227 """

228 block_size = 4

229 block_size_factor = 3

@@ -263,27 +260,34 @@ def test_no_offload_call_after_on_request_finished(

260 complete_transfers=False,

261 )

262

266 - # Finish the request, completing its pending stores. on_request_finished is

267 - # deferred until the stores drain, so it lands after the last complete_store.

268 - # 4 offloaded blocks are stored (2 prompt + 2 decode) -> 4 * block_size_factor

269 - # GPU blocks.

263 + # Finish the request while its stores are still in flight. The hook should

264 + # fire immediately even though no complete_store has arrived yet.

265 runner.run(

266 decoded_tokens=[EOS_TOKEN_ID],

272 - expected_stored=tuple(range(4 * block_size_factor)),

267 + complete_transfers=False,

268 )

269

270 req_id = str(runner.req_id)

271

272 + assert calls == [("on_request_finished", req_id)], calls

273 +

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

275 + # are allowed to arrive after on_request_finished.

276 + runner.run(

277 + decoded_tokens=[],

278 + complete_transfers=True,

279 + expected_stored=tuple(range(4 * block_size_factor)),

280 + )

281 +

282 # on_request_finished is issued exactly once.

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

284

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

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

287

283 - # All of the request's complete_store calls must precede its single

284 - # on_request_finished.

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

289 assert store_indices, calls

286 - assert max(store_indices) < finished_idx, calls

290 + assert finished_idx < min(store_indices), calls

291

292

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

@@ -1232,11 +1236,8 @@ def test_reset_cache(request_runner, async_scheduling: bool):

1236 def test_reset_cache_finalizes_finished_request_with_pending_store(

1237 request_runner, async_scheduling: bool

1238 ):

1235 - """reset_cache must finalize a finished request whose in-flight stores it

1236 - discards: call on_request_finished and drop its _req_status entry.

1237 -

1238 - Otherwise the deferred hook (which waits for the now-discarded jobs to

1239 - complete) never fires and the entry leaks.

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

1240 + without calling on_request_finished twice.

1241 """

1242 block_size = 4

1243 block_size_factor = 3

@@ -1273,14 +1274,15 @@ def test_reset_cache_finalizes_finished_request_with_pending_store(

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

1275

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

1276 - # takes the defer branch (pending jobs), so on_request_finished is NOT

1277 - # called yet and the entry stays tracked.

1277 + # fires the hook eagerly, but the entry stays tracked so later completions

1278 + # can still call complete_store().

1279 req_status.req.status = RequestStatus.FINISHED_STOPPED

1280 cs.request_finished(req_status.req)

1280 - assert finalized == []

1281 + assert finalized == [req_id]

1282 assert req_id in cs._req_status

1283

1283 - # reset_cache discards the in-flight store; it must finalize the request.

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

1285 + # duplicate on_request_finished call.

1286 cs.reset_cache()

1287 assert finalized == [req_id]

1288 assert req_id not in cs._req_status

tests/v1/kv_offload/tiering/test_tiering_offloading.py

+203 -19

@@ -235,11 +235,16 @@ class TestTieringOffloadingManager:

235 self.manager.on_schedule_end()

236 list(self.manager.take_events())

237

238 + def _start_request(self, req_context: ReqContext = _CTX):

239 + if req_context.req_id not in self.manager._req_state:

240 + self.manager.on_new_request(req_context)

241 +

242 def test_basic_store_to_primary(self, manager_setup):

243 """Test basic store operation to primary tier."""

244 blocks = to_keys(range(3))

245

246 # Prepare store

247 + self._start_request()

248 result = self.manager.prepare_store(blocks, _CTX)

249 assert result is not None

250 assert len(result.keys_to_store) == 3

@@ -262,6 +267,7 @@ class TestTieringOffloadingManager:

267 )

268

269 # Store to primary

270 + self._start_request()

271 result = self.manager.prepare_store(blocks, _CTX)

272 assert result is not None

273

@@ -285,6 +291,7 @@ class TestTieringOffloadingManager:

291 blocks = to_keys(range(3))

292

293 # Store to primary

294 + self._start_request()

295 result = self.manager.prepare_store(blocks, _CTX)

296 assert result is not None

297 self.manager.complete_store(blocks, _CTX, success=True)

@@ -330,6 +337,7 @@ class TestTieringOffloadingManager:

337 blocks = to_keys(range(3))

338

339 # Store blocks

340 + self._start_request()

341 self.manager.prepare_store(blocks, _CTX)

342 self.manager.complete_store(blocks, _CTX, success=True)

343

@@ -366,6 +374,7 @@ class TestTieringOffloadingManager:

374 blocks = to_keys(range(5))

375

376 # Store first 3 blocks to primary

377 + self._start_request()

378 self.manager.prepare_store(blocks[:3], _CTX)

379 self.manager.complete_store(blocks[:3], _CTX, success=True)

380

@@ -377,6 +386,7 @@ class TestTieringOffloadingManager:

386 # Primary tier has capacity of 5 blocks

387 # First, fill the primary tier

388 blocks = to_keys(range(5))

389 + self._start_request()

390 result = self.manager.prepare_store(blocks, _CTX)

391 assert result is not None

392 assert len(result.keys_to_store) == 5

@@ -399,6 +409,7 @@ class TestTieringOffloadingManager:

409 blocks = to_keys(range(3))

410

411 # Store blocks

412 + self._start_request()

413 self.manager.prepare_store(blocks, _CTX)

414 self.manager.complete_store(blocks, _CTX, success=True)

415 self._simulate_on_schedule_end()

@@ -431,6 +442,7 @@ class TestTieringOffloadingManager:

442 )

443

444 # Prepare store

445 + self._start_request()

446 result = self.manager.prepare_store(blocks, _CTX)

447 assert result is not None

448

@@ -521,6 +533,7 @@ class TestTieringOffloadingManager:

533

534 ctx = ReqContext(req_id="req_ctx", kv_transfer_params={"key": "value"})

535

536 + self._start_request(ctx)

537 self.manager.prepare_store(blocks, ctx)

538 self.manager.complete_store(blocks, ctx, success=True)

539

@@ -528,6 +541,172 @@ class TestTieringOffloadingManager:

541 job_metadata = self.secondary_tier1.submit_store.call_args.args[0]

542 assert job_metadata.req_context is ctx

543

544 + def test_on_request_finished_delays_secondary_until_store_submitted(

545 + self, manager_setup

546 + ):

547 + """Manager hook is eager; secondary hooks wait for cascade submission."""

548 + blocks = to_keys(range(2))

549 + ctx = ReqContext(req_id="req_delayed_secondary")

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

551 +

552 + self.primary_tier.on_request_finished = MagicMock(

553 + side_effect=lambda req_context: calls.append(

554 + ("primary_finish", req_context.req_id)

555 + )

556 + )

557 +

558 + original_submit_store1 = self.secondary_tier1.submit_store

559 + original_submit_store2 = self.secondary_tier2.submit_store

560 +

561 + def submit_store1(job_metadata):

562 + calls.append(("submit_store_1", job_metadata.req_context.req_id))

563 + return original_submit_store1(job_metadata)

564 +

565 + def submit_store2(job_metadata):

566 + calls.append(("submit_store_2", job_metadata.req_context.req_id))

567 + return original_submit_store2(job_metadata)

568 +

569 + self.secondary_tier1.submit_store = MagicMock(side_effect=submit_store1)

570 + self.secondary_tier2.submit_store = MagicMock(side_effect=submit_store2)

571 + self.secondary_tier1.on_request_finished = MagicMock(

572 + side_effect=lambda req_context: calls.append(

573 + ("secondary_finish_1", req_context.req_id)

574 + )

575 + )

576 + self.secondary_tier2.on_request_finished = MagicMock(

577 + side_effect=lambda req_context: calls.append(

578 + ("secondary_finish_2", req_context.req_id)

579 + )

580 + )

581 +

582 + self._start_request(ctx)

583 + self.manager.prepare_store(blocks, ctx)

584 + self.manager.on_request_finished(ctx)

585 +

586 + assert calls == [("primary_finish", ctx.req_id)]

587 + self.secondary_tier1.on_request_finished.assert_not_called()

588 + self.secondary_tier2.on_request_finished.assert_not_called()

589 +

590 + self.manager.complete_store(blocks, ctx, success=True)

591 +

592 + assert calls == [

593 + ("primary_finish", ctx.req_id),

594 + ("submit_store_1", ctx.req_id),

595 + ("submit_store_2", ctx.req_id),

596 + ("secondary_finish_1", ctx.req_id),

597 + ("secondary_finish_2", ctx.req_id),

598 + ]

599 +

600 + def test_failed_store_finalizes_finished_request(self, manager_setup):

601 + """Failed primary stores still unblock secondary finalization."""

602 + blocks = to_keys(range(2))

603 + ctx = ReqContext(req_id="req_failed_store_finalize")

604 +

605 + self.secondary_tier1.submit_store = MagicMock(

606 + wraps=self.secondary_tier1.submit_store

607 + )

608 + self.secondary_tier2.submit_store = MagicMock(

609 + wraps=self.secondary_tier2.submit_store

610 + )

611 + self.secondary_tier1.on_request_finished = MagicMock(

612 + wraps=self.secondary_tier1.on_request_finished

613 + )

614 + self.secondary_tier2.on_request_finished = MagicMock(

615 + wraps=self.secondary_tier2.on_request_finished

616 + )

617 +

618 + self._start_request(ctx)

619 + self.manager.prepare_store(blocks, ctx)

620 + self.manager.on_request_finished(ctx)

621 +

622 + self.secondary_tier1.on_request_finished.assert_not_called()

623 + self.secondary_tier2.on_request_finished.assert_not_called()

624 +

625 + self.manager.complete_store(blocks, ctx, success=False)

626 +

627 + self.secondary_tier1.submit_store.assert_not_called()

628 + self.secondary_tier2.submit_store.assert_not_called()

629 + self.secondary_tier1.on_request_finished.assert_called_once_with(ctx)

630 + self.secondary_tier2.on_request_finished.assert_called_once_with(ctx)

631 + assert ctx.req_id not in self.manager._req_state

632 +

633 + def test_zero_store_request_finalizes_immediately(self, manager_setup):

634 + """Requests with no pending stores finalize secondary tiers immediately."""

635 + ctx = ReqContext(req_id="req_zero_store_finalize")

636 +

637 + self.secondary_tier1.on_request_finished = MagicMock(

638 + wraps=self.secondary_tier1.on_request_finished

639 + )

640 + self.secondary_tier2.on_request_finished = MagicMock(

641 + wraps=self.secondary_tier2.on_request_finished

642 + )

643 +

644 + self._start_request(ctx)

645 + self.manager.on_request_finished(ctx)

646 +

647 + self.secondary_tier1.on_request_finished.assert_called_once_with(ctx)

648 + self.secondary_tier2.on_request_finished.assert_called_once_with(ctx)

649 + assert ctx.req_id not in self.manager._req_state

650 +

651 + def test_reset_cache_finalizes_delayed_secondary_request(self, manager_setup):

652 + """reset_cache abandons pending primary stores and finalizes secondaries."""

653 + blocks = to_keys(range(2))

654 + ctx = ReqContext(req_id="req_reset_finalize_secondary")

655 +

656 + self.secondary_tier1.on_request_finished = MagicMock(

657 + wraps=self.secondary_tier1.on_request_finished

658 + )

659 + self.secondary_tier2.on_request_finished = MagicMock(

660 + wraps=self.secondary_tier2.on_request_finished

661 + )

662 +

663 + self._start_request(ctx)

664 + self.manager.prepare_store(blocks, ctx)

665 + self.manager.on_request_finished(ctx)

666 +

667 + self.secondary_tier1.on_request_finished.assert_not_called()

668 + self.secondary_tier2.on_request_finished.assert_not_called()

669 +

670 + self.manager.reset_cache()

671 +

672 + self.secondary_tier1.on_request_finished.assert_called_once_with(ctx)

673 + self.secondary_tier2.on_request_finished.assert_called_once_with(ctx)

674 + assert self.manager._req_state == {}

675 +

676 + def test_reset_cache_clears_pending_primary_stores_for_active_request(

677 + self, manager_setup

678 + ):

679 + """reset_cache drops active pending stores so resumed requests finalize."""

680 + initial_blocks = to_keys(range(2))

681 + resumed_blocks = to_keys(range(2, 4))

682 + ctx = ReqContext(req_id="req_reset_resume")

683 +

684 + self.secondary_tier1.on_request_finished = MagicMock(

685 + wraps=self.secondary_tier1.on_request_finished

686 + )

687 + self.secondary_tier2.on_request_finished = MagicMock(

688 + wraps=self.secondary_tier2.on_request_finished

689 + )

690 +

691 + self._start_request(ctx)

692 + self.manager.prepare_store(initial_blocks, ctx)

693 + assert self.manager._req_state[ctx.req_id].pending_primary_stores == 1

694 +

695 + self.manager.reset_cache()

696 +

697 + assert ctx.req_id in self.manager._req_state

698 + assert self.manager._req_state[ctx.req_id].pending_primary_stores == 0

699 + self.secondary_tier1.on_request_finished.assert_not_called()

700 + self.secondary_tier2.on_request_finished.assert_not_called()

701 +

702 + self.manager.prepare_store(resumed_blocks, ctx)

703 + self.manager.complete_store(resumed_blocks, ctx, success=True)

704 + self.manager.on_request_finished(ctx)

705 +

706 + self.secondary_tier1.on_request_finished.assert_called_once_with(ctx)

707 + self.secondary_tier2.on_request_finished.assert_called_once_with(ctx)

708 + assert ctx.req_id not in self.manager._req_state

709 +

710 def test_on_new_request_lifecycle(self, manager_setup):

711 """Policy defaults to BLOCK_LEVEL, escalates when a tier requests it,

712 and is cleaned up on on_request_finished."""

@@ -535,23 +714,25 @@ class TestTieringOffloadingManager:

714 ctx = ReqContext(req_id="req_policy_lifecycle")

715 result = self.manager.on_new_request(ctx)

716 assert result.policy == OffloadPolicy.BLOCK_LEVEL

717 + assert self.manager._req_state[ctx.req_id].request_level_tiers is None

718 self.manager.on_request_finished(ctx)

719 + assert ctx.req_id not in self.manager._req_state

720

721 # Escalate: tier1 requests REQUEST_LEVEL

541 - self.secondary_tier1.on_new_request = (

542 - lambda req_context: RequestOffloadingContext(

543 - policy=OffloadPolicy.REQUEST_LEVEL

544 - )

722 + self.secondary_tier1.on_new_request = lambda req_context: (

723 + RequestOffloadingContext(policy=OffloadPolicy.REQUEST_LEVEL)

724 )

725

726 ctx = ReqContext(req_id="req_policy_lifecycle_2")

727 result = self.manager.on_new_request(ctx)

728 assert result.policy == OffloadPolicy.REQUEST_LEVEL

550 - assert ctx.req_id in self.manager._request_level_tiers

729 + assert self.manager._req_state[ctx.req_id].request_level_tiers == {

730 + self.secondary_tier1

731 + }

732

733 # Cleanup

734 self.manager.on_request_finished(ctx)

554 - assert ctx.req_id not in self.manager._request_level_tiers

735 + assert ctx.req_id not in self.manager._req_state

736

737 def test_prepare_store_cascades_existing_blocks_to_request_level_tiers(

738 self, manager_setup

@@ -559,6 +740,7 @@ class TestTieringOffloadingManager:

740 """prepare_store cascades hit blocks to request-level tiers only."""

741 # Store some blocks to primary first

742 existing_blocks = to_keys(range(3))

743 + self._start_request()

744 result = self.manager.prepare_store(existing_blocks, _CTX)

745 assert result is not None

746 self.manager.complete_store(existing_blocks, _CTX, success=True)

@@ -566,10 +748,8 @@ class TestTieringOffloadingManager:

748 self._simulate_on_schedule_end()

749

750 # Make tier1 request-level, tier2 stays block-level

569 - self.secondary_tier1.on_new_request = (

570 - lambda req_context: RequestOffloadingContext(

571 - policy=OffloadPolicy.REQUEST_LEVEL

572 - )

751 + self.secondary_tier1.on_new_request = lambda req_context: (

752 + RequestOffloadingContext(policy=OffloadPolicy.REQUEST_LEVEL)

753 )

754

755 ctx = ReqContext(req_id="req_cascade")

@@ -599,14 +779,15 @@ class TestTieringOffloadingManager:

779 # tier2 (block-level) does not get existing blocks here.

780 self.secondary_tier2.submit_store.assert_not_called()

781

602 - def test_reset_cache_clears_all_state(self, manager_setup):

782 + def test_reset_cache_clears_orchestrator_state(self, manager_setup):

783 """reset_cache wipes every kind of orchestrator state and resets

784 primary tier; pending submissions are dropped without being sent

605 - to the secondary tier."""

785 + to the secondary tier. Active request state is retained."""

786 # Cascade — populates primary blocks and leaves cascade jobs

787 # in _transfer_jobs (the synchronous example tier has already

788 # queued completions); reset_cache's drain loop will pick them up.

789 blocks = to_keys(range(3))

790 + self._start_request()

791 self.manager.prepare_store(blocks, _CTX)

792 self.manager.complete_store(blocks, _CTX, success=True)

793 assert self.manager._transfer_jobs

@@ -619,13 +800,14 @@ class TestTieringOffloadingManager:

800 assert self.manager._pending_load_submissions

801

802 # Request-level tier registration.

622 - self.secondary_tier1.on_new_request = (

623 - lambda req_context: RequestOffloadingContext(

624 - policy=OffloadPolicy.REQUEST_LEVEL

625 - )

803 + self.secondary_tier1.on_new_request = lambda req_context: (

804 + RequestOffloadingContext(policy=OffloadPolicy.REQUEST_LEVEL)

805 )

627 - self.manager.on_new_request(ReqContext(req_id="rl"))

628 - assert self.manager._request_level_tiers

806 + rl_ctx = ReqContext(req_id="rl")

807 + self.manager.on_new_request(rl_ctx)

808 + assert self.manager._req_state[rl_ctx.req_id].request_level_tiers == {

809 + self.secondary_tier1

810 + }

811

812 # Mark this step as already polled (reset_cache must clear it).

813 self.manager._processed_jobs_this_step = True

@@ -640,7 +822,7 @@ class TestTieringOffloadingManager:

822 # Orchestrator state cleared.

823 assert self.manager._transfer_jobs == {}

824 assert self.manager._pending_load_submissions == {}

643 - assert self.manager._request_level_tiers == {}

825 + assert set(self.manager._req_state) == {_CTX.req_id, rl_ctx.req_id}

826 assert self.manager._processed_jobs_this_step is False

827

828 # Primary tier reset to a fresh state.

@@ -668,6 +850,7 @@ class TestTieringOffloadingManager:

850

851 # Drive a cascade so a job lands in _transfer_jobs.

852 blocks = to_keys(range(3))

853 + self._start_request()

854 self.manager.prepare_store(blocks, _CTX)

855 self.manager.complete_store(blocks, _CTX, success=True)

856 assert self.manager._transfer_jobs

@@ -696,6 +879,7 @@ class TestTieringOffloadingWithoutSecondaryTiers:

879 blocks = to_keys(range(3))

880

881 # Should work like a regular OffloadingManager

882 + manager.on_new_request(_CTX)

883 result = manager.prepare_store(blocks, _CTX)

884 assert result is not None

885 manager.complete_store(blocks, _CTX, success=True)

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

+11 -19

@@ -1140,11 +1140,6 @@ class OffloadingConnectorScheduler:

1140 del self._jobs[job_id]

1141 req_status.transfer_jobs.remove(job_id)

1142 if not req_status.transfer_jobs and req_status.req.is_finished():

1143 - # Deferred from request_finished: the request's last in-flight

1144 - # job is now done, so fire the finalize hook here, after the

1145 - # final complete_store/complete_load above (and any submit_store

1146 - # the complete_store cascade issued).

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

1143 del self._req_status[job_status.req_id]

1144

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

@@ -1181,20 +1176,23 @@ class OffloadingConnectorScheduler:

1176 if req_status is None:

1177 # Untracked request (offloading never started): no in-flight jobs,

1178 # nothing was deferred, so finalize immediately.

1184 - self.manager.on_request_finished(_create_req_context(request))

1179 + req_context = _create_req_context(request)

1180 + self.manager.on_new_request(req_context)

1181 + self.manager.on_request_finished(req_context)

1182 return False, None

1183

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

1185 +

1186 if not req_status.transfer_jobs:

1188 - # No in-flight jobs: all per-request calls are done, finalize now.

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

1187 + # No in-flight jobs: no later complete_store()/complete_load() calls

1188 + # need this request's state.

1189 del self._req_status[request.request_id]

1190 return False, None

1191

1193 - # In-flight jobs remain, so defer on_request_finished to

1194 - # update_connector_output, which fires it once the last job completes

1195 - # (after the final complete_store and any cascade submit_store it

1196 - # issues). These pending stores outlive the request's block ownership;

1197 - # register them so future reuse of those blocks triggers a flush.

1192 + # In-flight jobs remain after the request stopped. Their completion may

1193 + # still call manager.complete_store()/complete_load(), so keep req_status.

1194 + # Pending stores outlive the request's block ownership; register them so

1195 + # future reuse of those blocks triggers a flush.

1196 for job_id in req_status.transfer_jobs:

1197 job_status = self._jobs[job_id]

1198 for bid in job_status.non_sliding_window_block_ids or ():

@@ -1225,14 +1223,8 @@ class OffloadingConnectorScheduler:

1223 # Flush all in-flight jobs

1224 self._current_batch_jobs_to_flush.update(self._jobs.keys())

1225

1228 - # A finished request may still be tracked here with in-flight jobs that

1229 - # this reset discards, so its deferred on_request_finished() would never

1230 - # fire (completions are skipped as stale) and its _req_status entry would

1231 - # leak. Finalize such requests now, before resetting the manager.

1232 - # list() snapshots because we delete while iterating.

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

1227 if status.req.is_finished():

1235 - self.manager.on_request_finished(status.req_context)

1228 del self._req_status[req_id]

1229

1230 # Reset offloading manager cache

vllm/v1/kv_offload/base.py

+10 -10

@@ -276,16 +276,16 @@ class OffloadingManager(ABC):

276 """

277 Called when a request has finished.

278

279 - By the time this is called, all per-request offload calls for this

280 - request (prepare_store/complete_store, prepare_load/complete_load,

281 - touch, lookup) have already been issued, and none will follow. The

282 - scheduler defers this call until the request is finished and has no

283 - in-flight transfer jobs.

284 -

285 - Note this signals only that no further calls will be made; it does NOT

286 - imply the data has been persisted. Asynchronous transfers already

287 - submitted for this request (e.g. CPU->secondary cascades) may still be

288 - in flight. This is the right place to release per-request bookkeeping.

279 + By the time this is called, the scheduler will issue no more

280 + submit-side calls for this request, such as prepare_store() and

281 + prepare_load(). Completion callbacks for already-submitted transfers

282 + (complete_store() and complete_load()) may still arrive afterward.

283 +

284 + This hook does NOT imply the data has been persisted. Asynchronous

285 + transfers already submitted for this request may still be in flight.

286 + Managers that cascade to lower tiers should delay those tiers'

287 + on_request_finished() calls until no more lower-tier submit calls can

288 + be issued for this request.

289

290 Args:

291 req_context: per-request context.

vllm/v1/kv_offload/tiering/manager.py

+84 -40

@@ -20,7 +20,6 @@ Key Design Principles:

20 protecting blocks from eviction until complete_read() is called

21 """

23 -from collections import defaultdict

23 from collections.abc import Collection, Iterable, Sequence

24 from dataclasses import dataclass, field

@@ -62,6 +61,14 @@ class PendingPromotion:

61 block_ids: list[int] = field(default_factory=list)

64 +@dataclass(slots=True)

65 +class RequestState:

66 + req_context: ReqContext

67 + pending_primary_stores: int = 0

68 + is_finished: bool = False

69 + request_level_tiers: set[SecondaryTierManager] | None = None

70 +

71 +

72 class CPUPrimaryTierOffloadingManager(CPUOffloadingManager):

73 """CPUOffloadingManager with a primary/secondary transfer interface.

@@ -166,12 +173,10 @@ class TieringOffloadingManager(OffloadingManager):

173 # Reset at the end of each step in on_schedule_end().

174 self._processed_jobs_this_step: bool = False

175

169 - # Per-request set of secondary tiers that requested REQUEST_LEVEL

170 - # policy. Populated in on_new_request(),

171 - # cleaned up in on_request_finished().

172 - self._request_level_tiers: defaultdict[str, set[SecondaryTierManager]] = (

173 - defaultdict(set)

174 - )

176 + # Per-request state for prepared GPU->primary stores and finalization.

177 + # Secondary tiers are finalized only after pending primary stores reach

178 + # complete_store(), since complete_store() can still submit cascades.

179 + self._req_state: dict[str, RequestState] = {}

180

181 def _next_job_id(self) -> JobId:

182 """Generate a unique job ID for async transfer tracking."""

@@ -432,9 +437,13 @@ class TieringOffloadingManager(OffloadingManager):

437 if primary_result is None:

438 return None

439

440 + if primary_result.keys_to_store:

441 + state = self._req_state[req_context.req_id]

442 + state.pending_primary_stores += 1

443 +

444 # Step 3: For request-level tiers, cascade blocks already in primary

436 - request_level_tiers = self._request_level_tiers.get(req_context.req_id)

437 - if request_level_tiers is not None:

445 + request_level_tiers = self._req_state[req_context.req_id].request_level_tiers

446 + if request_level_tiers:

447 keys_to_store_set = set(primary_result.keys_to_store)

448 keys_already_in_primary = tuple(

449 k for k in keys if k not in keys_to_store_set

@@ -508,36 +517,38 @@ class TieringOffloadingManager(OffloadingManager):

517 # Step 1: Complete store in primary tier (makes blocks loadable)

518 self.primary_tier.complete_store(keys, req_context, success)

519

511 - if not success:

512 - # If GPU→Primary transfer failed, don't cascade to secondary tiers

513 - return

514 -

515 - # Step 2: Cascade to ALL secondary tiers

516 - # For each secondary tier, call primary.prepare_read() to get the

517 - # LoadStoreSpec AND to increment ref_cnt (protecting blocks from

518 - # eviction during the async transfer). One prepare_read() call per

519 - # secondary tier.

520 - for tier in self.secondary_tiers:

521 - primary_blocks_spec = self.primary_tier.prepare_read(keys, req_context)

520 + if success:

521 + # Step 2: Cascade to ALL secondary tiers

522 + # For each secondary tier, call primary.prepare_read() to get the

523 + # LoadStoreSpec AND to increment ref_cnt (protecting blocks from

524 + # eviction during the async transfer). One prepare_read() call per

525 + # secondary tier.

526 + for tier in self.secondary_tiers:

527 + primary_blocks_spec = self.primary_tier.prepare_read(keys, req_context)

528

523 - # Submit async store job: primary→secondary

524 - job_id = self._next_job_id()

529 + # Submit async store job: primary→secondary

530 + job_id = self._next_job_id()

531

526 - # Track this store job

527 - assert isinstance(primary_blocks_spec, CPULoadStoreSpec)

528 - job_metadata = JobMetadata(

529 - job_id=job_id,

530 - keys=keys,

531 - block_ids=primary_blocks_spec.block_ids,

532 - is_promotion=False,

533 - req_context=req_context,

534 - )

535 - self._transfer_jobs[job_id] = job_metadata

532 + # Track this store job

533 + assert isinstance(primary_blocks_spec, CPULoadStoreSpec)

534 + job_metadata = JobMetadata(

535 + job_id=job_id,

536 + keys=keys,

537 + block_ids=primary_blocks_spec.block_ids,

538 + is_promotion=False,

539 + req_context=req_context,

540 + )

541 + self._transfer_jobs[job_id] = job_metadata

542

537 - tier.submit_store(job_metadata)

543 + tier.submit_store(job_metadata)

544

545 # Note: The async transfers are now in flight. Their completion is

546 # tracked via get_finished_jobs() / _maybe_process_finished_jobs().

547 + req_id = req_context.req_id

548 + state = self._req_state[req_id]

549 + assert state.pending_primary_stores > 0

550 + state.pending_primary_stores -= 1

551 + self._maybe_finalize_request(req_id)

552

553 @override

554 def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext:

@@ -547,14 +558,18 @@ class TieringOffloadingManager(OffloadingManager):

558 Returns REQUEST_LEVEL if ANY secondary tier wants request-level.

559 Only stores REQUEST_LEVEL tier decisions for use in prepare_store.

560 """

561 + state = RequestState(req_context=req_context)

562 for tier in self.secondary_tiers:

563 tier_ctx = tier.on_new_request(req_context)

564 if tier_ctx.policy == OffloadPolicy.REQUEST_LEVEL:

553 - self._request_level_tiers[req_context.req_id].add(tier)

565 + if state.request_level_tiers is None:

566 + state.request_level_tiers = set()

567 + state.request_level_tiers.add(tier)

568 + self._req_state[req_context.req_id] = state

569

570 policy = (

571 OffloadPolicy.REQUEST_LEVEL

557 - if req_context.req_id in self._request_level_tiers

572 + if state.request_level_tiers

573 else OffloadPolicy.BLOCK_LEVEL

574 )

575 return RequestOffloadingContext(policy=policy)

@@ -562,9 +577,26 @@ class TieringOffloadingManager(OffloadingManager):

577 @override

578 def on_request_finished(self, req_context: ReqContext) -> None:

579 self.primary_tier.on_request_finished(req_context)

580 + state = self._req_state[req_context.req_id]

581 + state.is_finished = True

582 + self._maybe_finalize_request(req_context.req_id)

583 +

584 + def _maybe_finalize_request(self, req_id: str) -> None:

585 + """Finalize secondary tiers once no more store cascades can be submitted.

586 +

587 + Finalization means forwarding on_request_finished() to secondary tiers.

588 + It is delayed until pending GPU->primary stores finish, since their

589 + complete_store() callbacks may still submit primary->secondary stores.

590 + """

591 + state = self._req_state[req_id]

592 + if not state.is_finished:

593 + return

594 + if state.pending_primary_stores != 0:

595 + return

596 +

597 for tier in self.secondary_tiers:

566 - tier.on_request_finished(req_context)

567 - self._request_level_tiers.pop(req_context.req_id, None)

598 + tier.on_request_finished(state.req_context)

599 + del self._req_state[req_id]

600

601 @override

602 def on_schedule_end(self) -> None:

@@ -604,7 +636,7 @@ class TieringOffloadingManager(OffloadingManager):

636

637 @override

638 def reset_cache(self) -> None:

607 - """Drop all tracked state in the orchestrator and primary tier.

639 + """Reset transfer bookkeeping and primary-tier cache.

640

641 Called during sleep, weight update, or resume. Each secondary tier

642 drains its in-flight transfers via drain_jobs() so no tier I/O is

@@ -613,7 +645,9 @@ class TieringOffloadingManager(OffloadingManager):

645 from reusing primary slots while a transfer is mid-copy.

646

647 Secondary tiers are intentionally not reset: persistent stores

616 - (FS, network) keep their data across resets.

648 + (FS, network) keep their data across resets. Active request state is

649 + retained so those requests can continue after the reset; finished

650 + requests are finalized and removed.

651 """

652 for tier in self.secondary_tiers:

653 tier.drain_jobs()

@@ -626,9 +660,19 @@ class TieringOffloadingManager(OffloadingManager):

660 # called so no tier I/O is touching that memory.

661 self._pending_load_submissions.clear()

662

663 + finished_req_ids = []

664 + for req_id, state in self._req_state.items():

665 + state.pending_primary_stores = 0

666 + if not state.is_finished:

667 + continue

668 + for tier in self.secondary_tiers:

669 + tier.on_request_finished(state.req_context)

670 + finished_req_ids.append(req_id)

671 +

672 self.primary_tier.reset_cache()

673

631 - self._request_level_tiers.clear()

674 + for req_id in finished_req_ids:

675 + del self._req_state[req_id]

676 self._processed_jobs_this_step = False

677

678 @override