@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 2a2b5ca791

Fork 0

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, 06:00 UTC

[KV Offload] Add `on_schedule_end()` hook to separate step lifecycle from event draining (#44206)

Signed-off-by: Ronen Schaffer <ronen.schaffer@ibm.com>

Ronen Schaffer committed Jun 2, 2026 at 13:42 UTC 2a2b5ca7919014cb851bac690fc99d356520254b

5 files changed +54 -39

tests/v1/kv_offload/test_tiering_offloading.py

+16 -12

@@ -123,6 +123,11 @@ class TestTieringOffloadingManager:

123 secondary_tiers=[self.secondary_tier1, self.secondary_tier2],

124 )

125

126 + def _simulate_on_schedule_end(self):

127 + """Simulate end of scheduler step: lifecycle flush + drain events."""

128 + self.manager.on_schedule_end()

129 + list(self.manager.take_events())

130 +

131 def test_basic_store_to_primary(self, manager_setup):

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

133 blocks = to_keys(range(3))

@@ -185,10 +190,10 @@ class TestTieringOffloadingManager:

190 assert block.ref_cnt == 2

191

192 # End of step 1: _maybe_process_finished_jobs() was already called by

188 - # prepare_store() above (setting the per-step flag), so take_events()

193 + # prepare_store() above (setting the per-step flag), so on_schedule_end()

194 # does NOT poll get_finished_jobs() again — cascade completions remain

195 # unprocessed until the next step.

191 - list(self.manager.take_events())

196 + self._simulate_on_schedule_end()

197

198 # ref_cnt still held: cascade jobs finished (sync tier) but haven't

199 # been polled yet because the per-step guard skipped the second call.

@@ -202,7 +207,7 @@ class TestTieringOffloadingManager:

207

208 # End of step 2: flag was reset, so _maybe_process_finished_jobs()

209 # runs and processes the cascade completions (complete_read → ref_cnt--)

205 - list(self.manager.take_events())

210 + self._simulate_on_schedule_end()

211

212 # After cascade completes, ref_cnt should be 0

213 for block_hash in blocks:

@@ -238,10 +243,10 @@ class TestTieringOffloadingManager:

243 assert result is None # Retry later (promotion initiated)

244

245 # End of step 1: flushes deferred submit_load() calls

241 - list(self.manager.take_events())

246 + self._simulate_on_schedule_end()

247

248 # End of step 2: processes the completed promotion jobs

244 - list(self.manager.take_events())

249 + self._simulate_on_schedule_end()

250

251 # Now blocks should be in primary tier

252 assert count_hits(self.primary_tier, blocks) == 3

@@ -271,7 +276,7 @@ class TestTieringOffloadingManager:

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

277

278 # End of step: release ref_cnt from cascade

274 - list(self.manager.take_events())

279 + self._simulate_on_schedule_end()

280

281 # Now try to store 2 more blocks (should trigger eviction)

282 more_blocks = to_keys(range(5, 7))

@@ -289,7 +294,7 @@ class TestTieringOffloadingManager:

294 # Store blocks

295 self.manager.prepare_store(blocks, _CTX)

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

292 - list(self.manager.take_events())

297 + self._simulate_on_schedule_end()

298

299 self.secondary_tier1.touch = MagicMock(wraps=self.secondary_tier1.touch)

300 self.secondary_tier2.touch = MagicMock(wraps=self.secondary_tier2.touch)

@@ -328,7 +333,7 @@ class TestTieringOffloadingManager:

333 self.secondary_tier2.submit_store.assert_not_called()

334

335 def test_lookup_batches_submit_load_per_request(self, manager_setup):

331 - """lookup() defers submit_load until take_events(), one call per request.

336 + """lookup() defers submit_load until on_schedule_end(), one per request.

337

338 Blocks from different requests each get their own submit_load call, each

339 carrying the correct req_context.

@@ -354,7 +359,7 @@ class TestTieringOffloadingManager:

359 self.secondary_tier1.submit_load.assert_not_called()

360

361 # simulate end of step

357 - list(self.manager.take_events())

362 + self._simulate_on_schedule_end()

363

364 assert self.secondary_tier1.submit_load.call_count == 2

365 calls = self.secondary_tier1.submit_load.call_args_list

@@ -389,7 +394,7 @@ class TestTieringOffloadingManager:

394 assert result_a is None

395 assert result_b is None

396

392 - list(self.manager.take_events())

397 + self._simulate_on_schedule_end()

398

399 # Only one submit_load call despite two lookups

400 self.secondary_tier1.submit_load.assert_called_once()

@@ -449,8 +454,7 @@ class TestTieringOffloadingManager:

454 assert result is not None

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

456 # Drain cascade completions

452 - list(self.manager.take_events())

453 - list(self.manager.take_events())

457 + self._simulate_on_schedule_end()

458

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

460 self.secondary_tier1.on_new_request = (

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

@@ -897,6 +897,7 @@ class OffloadingConnectorScheduler:

897 self, scheduler_output: SchedulerOutput

898 ) -> KVConnectorMetadata:

899 self._update_req_states(scheduler_output)

900 + self.manager.on_schedule_end()

901

902 # Flush jobs for preempted requests.

903 for req_id in scheduler_output.preempted_req_ids or ():

vllm/v1/kv_offload/base.py

@@ -257,6 +257,14 @@ class OffloadingManager(ABC):

257 """

258 return ()

259

260 + def on_schedule_end(self) -> None:

261 + """Called once at the end of each scheduler step.

262 +

263 + Managers may override this to flush deferred work accumulated

264 + during the step (e.g., batched promotions).

265 + """

266 + return

267 +

268 def reset_cache(self) -> None:

269 """Evict all tracked blocks and reset internal state."""

270 return

vllm/v1/kv_offload/tiering/base.py

@@ -185,6 +185,14 @@ class SecondaryTierManager(ABC):

185 """

186 return

187

188 + def on_schedule_end(self) -> None:

189 + """Called once at the end of each scheduler step.

190 +

191 + Secondary tiers may override this for per-step cleanup or

192 + deferred work submission.

193 + """

194 + return

195 +

196 def shutdown(self) -> None:

197 """Release resources held by this tier (threads, connections, etc.)."""

198 return

vllm/v1/kv_offload/tiering/manager.py

+21 -27

@@ -152,7 +152,7 @@ class TieringOffloadingManager(OffloadingManager):

152 self._transfer_jobs: dict[JobId, JobMetadata] = {}

153

154 # Pending promotion requests accumulated during lookup() calls; flushed

155 - # as one batched submit_load() per (tier, request) in take_events().

155 + # as one batched submit_load() per (tier, request) in on_schedule_end().

156 # Outer key: tier. Inner key: req_context.req_id — the same ReqContext

157 # object is reused for all block lookups of a given request per engine step.

158 self._pending_load_submissions: dict[

@@ -160,7 +160,7 @@ class TieringOffloadingManager(OffloadingManager):

160 ] = {}

161

162 # Gate for once-per-step execution of _maybe_process_finished_jobs().

163 - # Reset at the end of each step in take_events().

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

164 self._processed_jobs_this_step: bool = False

165

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

@@ -182,7 +182,7 @@ class TieringOffloadingManager(OffloadingManager):

182

183 Guarded by _processed_jobs_this_step: the first call in an engine step

184 does the actual polling; subsequent calls are no-ops. The flag is reset

185 - in take_events() at the end of each step.

185 + in on_schedule_end() at the end of each step.

186 """

187 if self._processed_jobs_this_step:

188 return

@@ -304,8 +304,8 @@ class TieringOffloadingManager(OffloadingManager):

304

305 store_spec = primary_write_result.store_spec

306 assert isinstance(store_spec, CPULoadStoreSpec)

307 - # Defer submit_load to take_events(). Group by (tier, request) so each

308 - # request's blocks are submitted as one batched job per tier.

307 + # Defer submit_load to on_schedule_end(). Group by (tier, request) so

308 + # each request's blocks are submitted as one batched job per tier.

309 tier_pending = self._pending_load_submissions.setdefault(tier, {})

310 ctx_id = req_context.req_id

311 if ctx_id not in tier_pending:

@@ -320,8 +320,8 @@ class TieringOffloadingManager(OffloadingManager):

320 def _flush_pending_promotions(self) -> None:

321 """Submit one batched submit_load() per (tier, request).

322

323 - Called from take_events() at the end of each engine step, flushing

324 - all promotion requests deferred during lookup().

323 + Called from on_schedule_end() at the end of each scheduler step,

324 + flushing all promotion requests deferred during lookup().

325 """

326 if not self._pending_load_submissions:

327 return

@@ -564,32 +564,26 @@ class TieringOffloadingManager(OffloadingManager):

564 self._request_level_tiers.pop(req_context.req_id, None)

565

566 @override

567 - def take_events(self) -> Iterable[OffloadingEvent]:

568 - """

569 - End-of-step hook: flush deferred work, yield events, reset per-step state.

567 + def on_schedule_end(self) -> None:

568 + """End-of-schedule hook: process finished jobs, flush deferred

569 + promotions, and reset the per-step gate.

570

571 - Called once per engine step from Scheduler.update_from_output() →

572 - connector.take_events(). Ensures _maybe_process_finished_jobs() has run

573 - at least once this step, flushes pending promotions, yields collected

574 - events, and resets the per-step flag.

575 -

576 - Yields:

577 - New OffloadingEvents collected since the last call.

571 + Called once per scheduler step from

572 + OffloadingConnectorScheduler.build_connector_meta().

573 """

579 - # TODO: Move _flush_pending_promotions() to a dedicated end_of_batch()

580 - # hook once one exists. For now, take_events() serves as the flush

581 - # point under the assumption that it is called at the end of each

582 - # engine step (Scheduler.update_from_output() → connector.take_events()).

583 - # When the dedicated hook is added, update tests that rely on

584 - # take_events() to signal end of step.

585 -

574 self._maybe_process_finished_jobs()

587 -

575 + self._processed_jobs_this_step = False

576 self._flush_pending_promotions()

577 + for tier in self.secondary_tiers:

578 + tier.on_schedule_end()

579

590 - # Reset the per-step gate so next step's first call does real work.

591 - self._processed_jobs_this_step = False

580 + @override

581 + def take_events(self) -> Iterable[OffloadingEvent]:

582 + """Yield offloading events collected since the last call.

583

584 + Yields:

585 + New OffloadingEvents collected since the last call.

586 + """

587 if self.events is not None:

588 yield from self.events

589 self.events.clear()