@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 48fcfc926c

[KV Offload] Add `ParentManager` ABC for secondary tier callbacks (#47274)

Signed-off-by: Ronen Schaffer <ronen.schaffer@ibm.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>

Ronen Schaffer committed Jul 7, 2026 at 13:51 UTC 48fcfc926c0a6e62a8f46557a59163702b7403b9

2 files changed +169 -44

vllm/v1/kv_offload/tiering/base.py

+48 -2

@@ -49,6 +49,43 @@ class JobResult:

49 success: bool

52 +class ParentManager(ABC):

53 + """Interface for secondary tiers to call back into the tiering manager.

54 +

55 + Passed to secondary tiers via serve_external_requests() each step.

56 + The _SecondaryTierFacingParent wrapper implements this, automatically

57 + excluding the calling tier from fan-out operations.

58 +

59 + Required call sequence for each remote request:

60 + 1. on_new_request(req_context) — set up per-request state

61 + 2. lookup(key, req_context) — check block availability

62 + (repeat per block)

63 + 3. create_store_job(keys, req_context) — pin blocks and get a

64 + job handle

65 + 4. on_request_finished(req_context) — clean up per-request state

66 +

67 + Steps 2-3 may be interleaved. Step 4 must be called even if no

68 + blocks were found, to avoid leaking async lookup state (e.g. in

69 + the fs tier's AsyncLookupManager).

70 + """

71 +

72 + @abstractmethod

73 + def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: ...

74 +

75 + @abstractmethod

76 + def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: ...

77 +

78 + @abstractmethod

79 + def create_store_job(

80 + self,

81 + keys: Collection[OffloadKey],

82 + req_context: ReqContext,

83 + ) -> JobMetadata: ...

84 +

85 + @abstractmethod

86 + def on_request_finished(self, req_context: ReqContext) -> None: ...

87 +

88 +

89 class SecondaryTierManager(ABC):

90 """

91 Abstract interface for managing a single non-primary offloading tier.

@@ -210,11 +247,20 @@ class SecondaryTierManager(ABC):

247 """

248 return

249

250 + def serve_external_requests(self, parent: ParentManager) -> None:

251 + """Process remotely-originated requests using the parent manager.

252 +

253 + Called once per scheduler step, BEFORE _flush_pending_promotions().

254 + The parent handle is valid only for the duration of this call.

255 + Tiers that don't serve external requests leave this as a no-op.

256 + """

257 + return

258 +

259 def on_schedule_end(self, context: ScheduleEndContext) -> None:

260 """Called once at the end of each scheduler step.

261

216 - Secondary tiers may override this for per-step cleanup or

217 - deferred work submission.

262 + Args:

263 + context: Per-step context from the scheduler.

264 """

265 return

266

vllm/v1/kv_offload/tiering/manager.py

+121 -42

@@ -48,6 +48,7 @@ from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion

48 from vllm.v1.kv_offload.tiering.base import (

49 JobId,

50 JobMetadata,

51 + ParentManager,

52 SecondaryTierManager,

53 )

@@ -120,6 +121,35 @@ class CPUPrimaryTierOffloadingManager(CPUOffloadingManager):

121 self._mmap_region.cleanup()

122

123

124 +class _SecondaryTierFacingParent(ParentManager):

125 + """Wrapper that implements ParentManager by delegating to the

126 + TieringOffloadingManager with exclude_tier set to the origin tier."""

127 +

128 + __slots__ = ("_m", "_origin")

129 +

130 + def __init__(

131 + self,

132 + manager: "TieringOffloadingManager",

133 + tier: SecondaryTierManager,

134 + ):

135 + self._m = manager

136 + self._origin = tier

137 +

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

139 + return self._m.on_new_request(req_context, exclude_tier=self._origin)

140 +

141 + def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult:

142 + return self._m.lookup(key, req_context, exclude_tier=self._origin)

143 +

144 + def create_store_job(

145 + self, keys: Collection[OffloadKey], req_context: ReqContext

146 + ) -> JobMetadata:

147 + return self._m.create_store_job(keys, req_context)

148 +

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

150 + return self._m.on_request_finished(req_context, exclude_tier=self._origin)

151 +

152 +

153 class TieringOffloadingManager(OffloadingManager):

154 """

155 Orchestrates multi-tier KV cache offloading.

@@ -180,6 +210,12 @@ class TieringOffloadingManager(OffloadingManager):

210 # complete_store(), since complete_store() can still submit cascades.

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

212

213 + # Cached ParentManager wrappers for each secondary tier.

214 + self._tier_parents: dict[SecondaryTierManager, _SecondaryTierFacingParent] = {

215 + tier: _SecondaryTierFacingParent(self, tier)

216 + for tier in self.secondary_tiers

217 + }

218 +

219 def _next_job_id(self) -> JobId:

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

221 job_id = self._job_id_counter

@@ -235,7 +271,13 @@ class TieringOffloadingManager(OffloadingManager):

271 )

272

273 @override

238 - def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult:

274 + def lookup(

275 + self,

276 + key: OffloadKey,

277 + req_context: ReqContext,

278 + *,

279 + exclude_tier: SecondaryTierManager | None = None,

280 + ) -> LookupResult:

281 """

282 Check whether a single block is offloaded and ready.

283

@@ -257,6 +299,10 @@ class TieringOffloadingManager(OffloadingManager):

299 MISS — block not found in any tier, or primary is full

300 and cannot accept a promotion.

301 """

302 + # Poll first so a promotion that finished since the last call is

303 + # already reflected as HIT (not stale HIT_PENDING/MISS) below, and

304 + # so blocks freed by cascade or promotion completions are evictable

305 + # in time for a promotion this lookup may initiate.

306 self._maybe_process_finished_jobs()

307

308 primary_hit = self.primary_tier.lookup(key, req_context)

@@ -267,6 +313,8 @@ class TieringOffloadingManager(OffloadingManager):

313

314 any_retry = False

315 for tier in self.secondary_tiers:

316 + if tier is exclude_tier:

317 + continue

318 result = tier.lookup(key, req_context)

319 if result is LookupResult.HIT:

320 if not self._initiate_promotion(tier, key, req_context):

@@ -359,8 +407,8 @@ class TieringOffloadingManager(OffloadingManager):

407 """

408 Prepare blocks to be loaded from primary tier to GPU.

409

362 - CRITICAL: This method calls _maybe_process_finished_jobs() FIRST to ensure

363 - that any completed promotions have been finalized and blocks are ready.

410 + Callers only pass keys already confirmed HIT by lookup() earlier this

411 + step.

412

413 This increments ref_cnt on the blocks in the primary tier, protecting

414 them from eviction during the transfer.

@@ -372,9 +420,6 @@ class TieringOffloadingManager(OffloadingManager):

420 Returns:

421 LoadStoreSpec for reading from primary tier.

422 """

375 - # Process completed promotions to ensure blocks are ready

376 - self._maybe_process_finished_jobs()

377 -

423 return self.primary_tier.prepare_load(keys, req_context)

424

425 @override

@@ -427,8 +472,15 @@ class TieringOffloadingManager(OffloadingManager):

472 evicted, or None if store cannot proceed.

473 """

474 # Step 1: Poll for completed async jobs FIRST

430 - # This decrements ref_cnt on primary blocks that have been

431 - # successfully transferred to secondary tiers.

475 + # _process_finished_jobs() handles two kinds of completions here:

476 + # - Cascade completions (store to a secondary tier, either a local

477 + # cascade or a store job created for a remote requester via

478 + # create_store_job()): decrements ref_cnt on the primary blocks

479 + # that were read, making them evictable again once ref_cnt hits 0.

480 + # - Promotion completions (secondary->primary loads): sets a

481 + # not-yet-ready block's ref_cnt from -1 to 0 via complete_write(),

482 + # making it evictable for the first time.

483 + # Both must be accounted for before the eviction decision below.

484 self._maybe_process_finished_jobs()

485

486 # Step 2: Store to primary tier (new blocks only).

@@ -478,20 +530,7 @@ class TieringOffloadingManager(OffloadingManager):

530 return

531

532 for tier in request_level_tiers:

481 - primary_blocks_spec = self.primary_tier.prepare_read(

482 - ready_keys, req_context

483 - )

484 -

485 - job_id = self._next_job_id()

486 - assert isinstance(primary_blocks_spec, CPULoadStoreSpec)

487 - job_metadata = JobMetadata(

488 - job_id=job_id,

489 - keys=ready_keys,

490 - block_ids=primary_blocks_spec.block_ids,

491 - is_promotion=False,

492 - req_context=req_context,

493 - )

494 - self._transfer_jobs[job_id] = job_metadata

533 + job_metadata = self.create_store_job(ready_keys, req_context)

534 tier.submit_store(job_metadata)

535

536 @override

@@ -529,22 +568,7 @@ class TieringOffloadingManager(OffloadingManager):

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

569 # secondary tier.

570 for tier in self.secondary_tiers:

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

533 -

534 - # Submit async store job: primary→secondary

535 - job_id = self._next_job_id()

536 -

537 - # Track this store job

538 - assert isinstance(primary_blocks_spec, CPULoadStoreSpec)

539 - job_metadata = JobMetadata(

540 - job_id=job_id,

541 - keys=keys,

542 - block_ids=primary_blocks_spec.block_ids,

543 - is_promotion=False,

544 - req_context=req_context,

545 - )

546 - self._transfer_jobs[job_id] = job_metadata

547 -

571 + job_metadata = self.create_store_job(keys, req_context)

572 tier.submit_store(job_metadata)

573

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

@@ -555,8 +579,40 @@ class TieringOffloadingManager(OffloadingManager):

579 state.pending_primary_stores -= 1

580 self._maybe_finalize_request(req_id)

581

582 + def create_store_job(

583 + self,

584 + keys: Collection[OffloadKey],

585 + req_context: ReqContext,

586 + ) -> JobMetadata:

587 + """Pin blocks in the primary tier and create a tracked store job.

588 +

589 + Calls prepare_read() to increment ref_cnt (protecting blocks

590 + from eviction during the async transfer), allocates a job ID,

591 + and registers the job in _transfer_jobs.

592 +

593 + The caller is responsible for the actual data transfer and

594 + reporting completion via get_finished_jobs().

595 + """

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

597 + assert isinstance(primary_blocks_spec, CPULoadStoreSpec)

598 + job_id = self._next_job_id()

599 + job_metadata = JobMetadata(

600 + job_id=job_id,

601 + keys=keys,

602 + block_ids=primary_blocks_spec.block_ids,

603 + is_promotion=False,

604 + req_context=req_context,

605 + )

606 + self._transfer_jobs[job_id] = job_metadata

607 + return job_metadata

608 +

609 @override

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

610 + def on_new_request(

611 + self,

612 + req_context: ReqContext,

613 + *,

614 + exclude_tier: SecondaryTierManager | None = None,

615 + ) -> RequestOffloadingContext:

616 """

617 Query each secondary tier for its offload policy preference.

618

@@ -565,6 +621,8 @@ class TieringOffloadingManager(OffloadingManager):

621 """

622 state = RequestState(req_context=req_context)

623 for tier in self.secondary_tiers:

624 + if tier is exclude_tier:

625 + continue

626 tier_ctx = tier.on_new_request(req_context)

627 if tier_ctx.policy == OffloadPolicy.REQUEST_LEVEL:

628 if state.request_level_tiers is None:

@@ -580,13 +638,22 @@ class TieringOffloadingManager(OffloadingManager):

638 return RequestOffloadingContext(policy=policy)

639

640 @override

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

641 + def on_request_finished(

642 + self,

643 + req_context: ReqContext,

644 + *,

645 + exclude_tier: SecondaryTierManager | None = None,

646 + ) -> None:

647 self.primary_tier.on_request_finished(req_context)

648 state = self._req_state[req_context.req_id]

649 state.is_finished = True

587 - self._maybe_finalize_request(req_context.req_id)

650 + self._maybe_finalize_request(req_context.req_id, exclude_tier)

651

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

652 + def _maybe_finalize_request(

653 + self,

654 + req_id: str,

655 + exclude_tier: SecondaryTierManager | None = None,

656 + ) -> None:

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

658

659 Finalization means forwarding on_request_finished() to secondary tiers.

@@ -600,6 +667,8 @@ class TieringOffloadingManager(OffloadingManager):

667 return

668

669 for tier in self.secondary_tiers:

670 + if tier is exclude_tier:

671 + continue

672 tier.on_request_finished(state.req_context)

673 del self._req_state[req_id]

674

@@ -611,8 +680,18 @@ class TieringOffloadingManager(OffloadingManager):

680 Called once per scheduler step from

681 OffloadingConnectorScheduler.build_connector_meta().

682 """

683 + # Catch-all poll: guarantees jobs are processed even on steps where

684 + # lookup()/prepare_store() were never called (e.g. no requests

685 + # scheduled but a tier still has_pending_work()).

686 self._maybe_process_finished_jobs()

687 +

688 + for tier in self.secondary_tiers:

689 + tier.serve_external_requests(self._tier_parents[tier])

690 +

691 + # Reset the per-step gate AFTER serve_external_requests so that

692 + # lookup() calls within it skip redundant _process_finished_jobs().

693 self._processed_jobs_this_step = False

694 +

695 self._flush_pending_promotions()

696 for tier in self.secondary_tiers:

697 tier.on_schedule_end(context)