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

[Bugfix][KV Offload] Defer request-level cascade of in-flight primary keys (#53329)

Signed-off-by: almogtavor <almogtavor@gmail.com>

Almog Tavor committed Aug 26, 2026 at 01:20 UTC bc39ded3c4e7b0453af76c617883e5cd88092222

2 files changed +181 -8

tests/v1/kv_offload/tiering/test_tiering_offloading.py

+134

@@ -247,6 +247,35 @@ class TestExampleSecondaryTierManager:

247 assert tier.lookup(blocks[2], _CTX) is LookupResult.MISS

248

249

250 +# What a request-level cascade does with a key already present in the primary

251 +# tier, keyed by what primary lookup says about it. SUPPLY reaches the peer now,

252 +# DEFER is retried next step, DROP is never offered again, permanently so since

253 +# prepare_store counts the key as stored and the scheduler advances past its

254 +# chunk. A new LookupResult has to declare which it is, and the test named

255 +# beside it is where that behavior is actually exercised.

256 +_CASCADE_SUPPLY_BEHAVIOR = {

257 + # test_cascade_defers_keys_whose_primary_write_is_in_flight, after the

258 + # writer's complete_store lands.

259 + LookupResult.HIT: "supply",

260 + # The same test, before it lands.

261 + LookupResult.HIT_PENDING: "defer",

262 + # test_cascade_drops_deferred_keys_whose_primary_write_failed, where the

263 + # failed write frees the block.

264 + LookupResult.MISS: "drop",

265 + # test_cascade_rejects_retry_from_primary. The primary tier resolves every

266 + # key it holds, so RETRY cannot reach the cascade, and parking on it would

267 + # have no guarantee of draining, which is what makes parking HIT_PENDING

268 + # safe. It is the one case no real primary tier can produce.

269 + LookupResult.RETRY: "unreachable",

270 +}

271 +

272 +

273 +def test_every_lookup_result_has_a_declared_cascade_behavior():

274 + """A new LookupResult must state what the cascade does with it, rather

275 + than falling into whichever branch happens to catch it."""

276 + assert set(_CASCADE_SUPPLY_BEHAVIOR) == set(LookupResult)

277 +

278 +

279 class TestTieringOffloadingManager:

280 """Tests for TieringOffloadingManager."""

281

@@ -1079,6 +1108,111 @@ class TestTieringOffloadingManager:

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

1109 self.secondary_tier2.submit_store.assert_not_called()

1110

1111 + def _make_request_level_request(self, req_id: str) -> ReqContext:

1112 + """Start a request for which tier1 asks for request-level offloading,

1113 + with tier1's submit_store wrapped so the cascade can be observed."""

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

1115 + RequestOffloadingContext(policy=OffloadPolicy.REQUEST_LEVEL)

1116 + )

1117 + ctx = ReqContext(req_id=req_id)

1118 + self.manager.on_new_request(ctx)

1119 + self.secondary_tier1.submit_store = MagicMock(

1120 + wraps=self.secondary_tier1.submit_store

1121 + )

1122 + return ctx

1123 +

1124 + def _cascaded_keys_for(self, req_id: str) -> list[set[OffloadKey]]:

1125 + """Key sets tier1 was asked to store on behalf of `req_id`."""

1126 + return [

1127 + set(call.args[0].keys)

1128 + for call in self.secondary_tier1.submit_store.call_args_list

1129 + if call.args[0].req_context.req_id == req_id

1130 + ]

1131 +

1132 + def test_cascade_rejects_retry_from_primary(self, manager_setup):

1133 + """RETRY would be parked with no guarantee of ever draining, so the

1134 + cascade refuses it. No real primary tier returns it, so this is the one

1135 + disposition that has to be forced."""

1136 + keys = to_keys(range(3))

1137 + self._start_request()

1138 + assert self.manager.prepare_store(keys, _CTX) is not None

1139 + self.manager.complete_store(keys, _CTX, success=True)

1140 + self._simulate_on_schedule_end()

1141 +

1142 + ctx = self._make_request_level_request("req_cascade")

1143 + self.manager.primary_tier.lookup = lambda key, req_context: (LookupResult.RETRY)

1144 +

1145 + with pytest.raises(AssertionError):

1146 + self.manager._cascade_existing_blocks_to_request_level_tiers(keys, ctx, {0})

1147 +

1148 + def _start_in_flight_primary_write(self, keys: list[OffloadKey]) -> ReqContext:

1149 + """Leave a GPU->primary write for `keys` open, so they look present to

1150 + prepare_store but are not yet readable."""

1151 + writer_ctx = ReqContext(req_id="req_writer")

1152 + self._start_request(writer_ctx)

1153 + assert self.manager.prepare_store(keys, writer_ctx) is not None

1154 + return writer_ctx

1155 +

1156 + def test_cascade_defers_keys_whose_primary_write_is_in_flight(self, manager_setup):

1157 + """A key another request is still writing must reach the peer once the

1158 + write lands, not be dropped: prepare_store already counts it as stored

1159 + and the scheduler advances past its chunk, so nothing re-offers it."""

1160 + keys = to_keys(range(3))

1161 + writer_ctx = self._start_in_flight_primary_write(keys)

1162 +

1163 + ctx = self._make_request_level_request("req_cascade")

1164 + result = self.manager.prepare_store(keys, ctx)

1165 + assert result is not None

1166 + assert not result.keys_to_store

1167 + assert not self._cascaded_keys_for(ctx.req_id)

1168 +

1169 + self.manager.complete_store(keys, writer_ctx, success=True)

1170 + self._simulate_on_schedule_end()

1171 +

1172 + assert self._cascaded_keys_for(ctx.req_id) == [set(keys)]

1173 +

1174 + def test_deferred_cascade_holds_request_from_finalization(self, manager_setup):

1175 + """A request that finishes with keys still deferred must stay alive, or

1176 + its tiers are torn down before the peer is ever served."""

1177 + keys = to_keys(range(3))

1178 + writer_ctx = self._start_in_flight_primary_write(keys)

1179 +

1180 + ctx = self._make_request_level_request("req_cascade")

1181 + result = self.manager.prepare_store(keys, ctx)

1182 + assert result is not None

1183 + assert not result.keys_to_store

1184 + self.manager.on_request_finished(ctx)

1185 +

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

1187 + assert self.manager.has_pending_work()

1188 +

1189 + self.manager.complete_store(keys, writer_ctx, success=True)

1190 + self._simulate_on_schedule_end()

1191 +

1192 + assert self._cascaded_keys_for(ctx.req_id) == [set(keys)]

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

1194 +

1195 + def test_cascade_drops_deferred_keys_whose_primary_write_failed(

1196 + self, manager_setup

1197 + ):

1198 + """A failed write frees the block, so the deferred key resolves to MISS

1199 + and the request finalizes instead of parking forever."""

1200 + keys = to_keys(range(3))

1201 + writer_ctx = self._start_in_flight_primary_write(keys)

1202 +

1203 + ctx = self._make_request_level_request("req_cascade")

1204 + result = self.manager.prepare_store(keys, ctx)

1205 + assert result is not None

1206 + assert not result.keys_to_store

1207 + self.manager.on_request_finished(ctx)

1208 +

1209 + self.manager.complete_store(keys, writer_ctx, success=False)

1210 + self._simulate_on_schedule_end()

1211 +

1212 + assert not self._cascaded_keys_for(ctx.req_id)

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

1214 + assert not self.manager.has_pending_work()

1215 +

1216 def test_reset_cache_clears_orchestrator_state(self, manager_setup):

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

1218 primary tier; pending submissions are dropped without being sent

vllm/v1/kv_offload/tiering/manager.py

+47 -8

@@ -74,6 +74,7 @@ class RequestState:

74 pending_primary_stores: int = 0

75 is_finished: bool = False

76 request_level_tiers: set[int] | None = None

77 + pending_cascade_keys: list[OffloadKey] = field(default_factory=list)

80 class JobMetadata(NamedTuple):

@@ -610,13 +611,29 @@ class TieringOffloadingManager(OffloadingManager):

611 """

612 For tiers that requested request-level policy, submit_store() for

613 blocks that are already present in the primary tier.

614 +

615 + A key whose primary write is still in flight (HIT_PENDING) cannot be

616 + dropped: prepare_store already excluded it as present, and the

617 + scheduler advances past its chunk, so no path offers it again. Park it

618 + instead. MISS keys are dropped, since nothing is there to read.

619 +

620 + The primary tier resolves every key it holds, so RETRY cannot reach

621 + here. Parking on it would have no guarantee of ever draining, which is

622 + what makes parking HIT_PENDING safe, so it is rejected rather than

623 + guessed at.

624 """

614 - # Filter out keys that are not ready in primary (e.g. in-flight)

615 - ready_keys = tuple(

616 - k

617 - for k in keys

618 - if self.primary_tier.lookup(k, req_context) is LookupResult.HIT

619 - )

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

626 + ready_keys = []

627 + for key in keys:

628 + result = self.primary_tier.lookup(key, req_context)

629 + if result is LookupResult.HIT:

630 + ready_keys.append(key)

631 + elif result is LookupResult.HIT_PENDING:

632 + state.pending_cascade_keys.append(key)

633 + else:

634 + assert result is LookupResult.MISS, (

635 + f"primary tier returned {result} for a cascade key"

636 + )

637 if not ready_keys:

638 return

639

@@ -625,6 +642,22 @@ class TieringOffloadingManager(OffloadingManager):

642 tier = self.secondary_tiers[tier_idx]

643 tier.submit_store(job_metadata)

644

645 + def _flush_pending_cascades(self) -> None:

646 + """Retry request-level cascades parked on an in-flight primary write.

647 +

648 + A parked key always resolves, to HIT or to MISS, so the set drains and

649 + a request cannot be held from finalization forever.

650 + """

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

652 + if not state.pending_cascade_keys:

653 + continue

654 + assert state.request_level_tiers

655 + keys, state.pending_cascade_keys = state.pending_cascade_keys, []

656 + self._cascade_existing_blocks_to_request_level_tiers(

657 + keys, state.req_context, state.request_level_tiers

658 + )

659 + self._maybe_finalize_request(req_id)

660 +

661 @override

662 def complete_store(

663 self,

@@ -759,6 +792,8 @@ class TieringOffloadingManager(OffloadingManager):

792 return

793 if state.pending_primary_stores != 0:

794 return

795 + if state.pending_cascade_keys:

796 + return

797

798 for tier_idx, tier in enumerate(self.secondary_tiers):

799 if tier_idx == exclude_tier_idx:

@@ -788,6 +823,7 @@ class TieringOffloadingManager(OffloadingManager):

823 self._processed_jobs_this_step = False

824

825 self._flush_pending_promotions()

826 + self._flush_pending_cascades()

827 for tier in self.secondary_tiers:

828 tier.on_schedule_end(context)

829

@@ -802,8 +838,10 @@ class TieringOffloadingManager(OffloadingManager):

838 # In-flight primary<->secondary transfers (pending promotions are

839 # translated to transfer jobs in on_schedule_end), plus any work the

840 # secondary tiers themselves still have outstanding.

805 - return bool(self._jobs) or any(

806 - tier.has_pending_work() for tier in self.secondary_tiers

841 + return (

842 + bool(self._jobs)

843 + or any(state.pending_cascade_keys for state in self._req_state.values())

844 + or any(tier.has_pending_work() for tier in self.secondary_tiers)

845 )

846

847 @override

@@ -848,6 +886,7 @@ class TieringOffloadingManager(OffloadingManager):

886 finished_req_ids = []

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

888 state.pending_primary_stores = 0

889 + state.pending_cascade_keys.clear()

890 if not state.is_finished:

891 continue

892 for tier in self.secondary_tiers: