From bc39ded3c4e7b0453af76c617883e5cd88092222 Mon Sep 17 00:00:00 2001 From: Almog Tavor <70065337+almogtavor@users.noreply.github.com> Date: Wed, 26 Aug 2026 01:20:03 +0200 Subject: [PATCH] [Bugfix][KV Offload] Defer request-level cascade of in-flight primary keys (#53329) Signed-off-by: almogtavor --- .../tiering/test_tiering_offloading.py | 134 ++++++++++++++++++ vllm/v1/kv_offload/tiering/manager.py | 55 +++++-- 2 files changed, 181 insertions(+), 8 deletions(-) diff --git a/tests/v1/kv_offload/tiering/test_tiering_offloading.py b/tests/v1/kv_offload/tiering/test_tiering_offloading.py index 9af714ac97..a0e26293f0 100644 --- a/tests/v1/kv_offload/tiering/test_tiering_offloading.py +++ b/tests/v1/kv_offload/tiering/test_tiering_offloading.py @@ -247,6 +247,35 @@ class TestExampleSecondaryTierManager: assert tier.lookup(blocks[2], _CTX) is LookupResult.MISS +# What a request-level cascade does with a key already present in the primary +# tier, keyed by what primary lookup says about it. SUPPLY reaches the peer now, +# DEFER is retried next step, DROP is never offered again, permanently so since +# prepare_store counts the key as stored and the scheduler advances past its +# chunk. A new LookupResult has to declare which it is, and the test named +# beside it is where that behavior is actually exercised. +_CASCADE_SUPPLY_BEHAVIOR = { + # test_cascade_defers_keys_whose_primary_write_is_in_flight, after the + # writer's complete_store lands. + LookupResult.HIT: "supply", + # The same test, before it lands. + LookupResult.HIT_PENDING: "defer", + # test_cascade_drops_deferred_keys_whose_primary_write_failed, where the + # failed write frees the block. + LookupResult.MISS: "drop", + # test_cascade_rejects_retry_from_primary. The primary tier resolves every + # key it holds, so RETRY cannot reach the cascade, and parking on it would + # have no guarantee of draining, which is what makes parking HIT_PENDING + # safe. It is the one case no real primary tier can produce. + LookupResult.RETRY: "unreachable", +} + + +def test_every_lookup_result_has_a_declared_cascade_behavior(): + """A new LookupResult must state what the cascade does with it, rather + than falling into whichever branch happens to catch it.""" + assert set(_CASCADE_SUPPLY_BEHAVIOR) == set(LookupResult) + + class TestTieringOffloadingManager: """Tests for TieringOffloadingManager.""" @@ -1079,6 +1108,111 @@ class TestTieringOffloadingManager: # tier2 (block-level) does not get existing blocks here. self.secondary_tier2.submit_store.assert_not_called() + def _make_request_level_request(self, req_id: str) -> ReqContext: + """Start a request for which tier1 asks for request-level offloading, + with tier1's submit_store wrapped so the cascade can be observed.""" + self.secondary_tier1.on_new_request = lambda req_context: ( + RequestOffloadingContext(policy=OffloadPolicy.REQUEST_LEVEL) + ) + ctx = ReqContext(req_id=req_id) + self.manager.on_new_request(ctx) + self.secondary_tier1.submit_store = MagicMock( + wraps=self.secondary_tier1.submit_store + ) + return ctx + + def _cascaded_keys_for(self, req_id: str) -> list[set[OffloadKey]]: + """Key sets tier1 was asked to store on behalf of `req_id`.""" + return [ + set(call.args[0].keys) + for call in self.secondary_tier1.submit_store.call_args_list + if call.args[0].req_context.req_id == req_id + ] + + def test_cascade_rejects_retry_from_primary(self, manager_setup): + """RETRY would be parked with no guarantee of ever draining, so the + cascade refuses it. No real primary tier returns it, so this is the one + disposition that has to be forced.""" + keys = to_keys(range(3)) + self._start_request() + assert self.manager.prepare_store(keys, _CTX) is not None + self.manager.complete_store(keys, _CTX, success=True) + self._simulate_on_schedule_end() + + ctx = self._make_request_level_request("req_cascade") + self.manager.primary_tier.lookup = lambda key, req_context: (LookupResult.RETRY) + + with pytest.raises(AssertionError): + self.manager._cascade_existing_blocks_to_request_level_tiers(keys, ctx, {0}) + + def _start_in_flight_primary_write(self, keys: list[OffloadKey]) -> ReqContext: + """Leave a GPU->primary write for `keys` open, so they look present to + prepare_store but are not yet readable.""" + writer_ctx = ReqContext(req_id="req_writer") + self._start_request(writer_ctx) + assert self.manager.prepare_store(keys, writer_ctx) is not None + return writer_ctx + + def test_cascade_defers_keys_whose_primary_write_is_in_flight(self, manager_setup): + """A key another request is still writing must reach the peer once the + write lands, not be dropped: prepare_store already counts it as stored + and the scheduler advances past its chunk, so nothing re-offers it.""" + keys = to_keys(range(3)) + writer_ctx = self._start_in_flight_primary_write(keys) + + ctx = self._make_request_level_request("req_cascade") + result = self.manager.prepare_store(keys, ctx) + assert result is not None + assert not result.keys_to_store + assert not self._cascaded_keys_for(ctx.req_id) + + self.manager.complete_store(keys, writer_ctx, success=True) + self._simulate_on_schedule_end() + + assert self._cascaded_keys_for(ctx.req_id) == [set(keys)] + + def test_deferred_cascade_holds_request_from_finalization(self, manager_setup): + """A request that finishes with keys still deferred must stay alive, or + its tiers are torn down before the peer is ever served.""" + keys = to_keys(range(3)) + writer_ctx = self._start_in_flight_primary_write(keys) + + ctx = self._make_request_level_request("req_cascade") + result = self.manager.prepare_store(keys, ctx) + assert result is not None + assert not result.keys_to_store + self.manager.on_request_finished(ctx) + + assert ctx.req_id in self.manager._req_state + assert self.manager.has_pending_work() + + self.manager.complete_store(keys, writer_ctx, success=True) + self._simulate_on_schedule_end() + + assert self._cascaded_keys_for(ctx.req_id) == [set(keys)] + assert ctx.req_id not in self.manager._req_state + + def test_cascade_drops_deferred_keys_whose_primary_write_failed( + self, manager_setup + ): + """A failed write frees the block, so the deferred key resolves to MISS + and the request finalizes instead of parking forever.""" + keys = to_keys(range(3)) + writer_ctx = self._start_in_flight_primary_write(keys) + + ctx = self._make_request_level_request("req_cascade") + result = self.manager.prepare_store(keys, ctx) + assert result is not None + assert not result.keys_to_store + self.manager.on_request_finished(ctx) + + self.manager.complete_store(keys, writer_ctx, success=False) + self._simulate_on_schedule_end() + + assert not self._cascaded_keys_for(ctx.req_id) + assert ctx.req_id not in self.manager._req_state + assert not self.manager.has_pending_work() + def test_reset_cache_clears_orchestrator_state(self, manager_setup): """reset_cache wipes every kind of orchestrator state and resets primary tier; pending submissions are dropped without being sent diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index eb5d5784d4..b2112ae165 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -74,6 +74,7 @@ class RequestState: pending_primary_stores: int = 0 is_finished: bool = False request_level_tiers: set[int] | None = None + pending_cascade_keys: list[OffloadKey] = field(default_factory=list) class JobMetadata(NamedTuple): @@ -610,13 +611,29 @@ class TieringOffloadingManager(OffloadingManager): """ For tiers that requested request-level policy, submit_store() for blocks that are already present in the primary tier. + + A key whose primary write is still in flight (HIT_PENDING) cannot be + dropped: prepare_store already excluded it as present, and the + scheduler advances past its chunk, so no path offers it again. Park it + instead. MISS keys are dropped, since nothing is there to read. + + The primary tier resolves every key it holds, so RETRY cannot reach + here. Parking on it would have no guarantee of ever draining, which is + what makes parking HIT_PENDING safe, so it is rejected rather than + guessed at. """ - # Filter out keys that are not ready in primary (e.g. in-flight) - ready_keys = tuple( - k - for k in keys - if self.primary_tier.lookup(k, req_context) is LookupResult.HIT - ) + state = self._req_state[req_context.req_id] + ready_keys = [] + for key in keys: + result = self.primary_tier.lookup(key, req_context) + if result is LookupResult.HIT: + ready_keys.append(key) + elif result is LookupResult.HIT_PENDING: + state.pending_cascade_keys.append(key) + else: + assert result is LookupResult.MISS, ( + f"primary tier returned {result} for a cascade key" + ) if not ready_keys: return @@ -625,6 +642,22 @@ class TieringOffloadingManager(OffloadingManager): tier = self.secondary_tiers[tier_idx] tier.submit_store(job_metadata) + def _flush_pending_cascades(self) -> None: + """Retry request-level cascades parked on an in-flight primary write. + + A parked key always resolves, to HIT or to MISS, so the set drains and + a request cannot be held from finalization forever. + """ + for req_id, state in list(self._req_state.items()): + if not state.pending_cascade_keys: + continue + assert state.request_level_tiers + keys, state.pending_cascade_keys = state.pending_cascade_keys, [] + self._cascade_existing_blocks_to_request_level_tiers( + keys, state.req_context, state.request_level_tiers + ) + self._maybe_finalize_request(req_id) + @override def complete_store( self, @@ -759,6 +792,8 @@ class TieringOffloadingManager(OffloadingManager): return if state.pending_primary_stores != 0: return + if state.pending_cascade_keys: + return for tier_idx, tier in enumerate(self.secondary_tiers): if tier_idx == exclude_tier_idx: @@ -788,6 +823,7 @@ class TieringOffloadingManager(OffloadingManager): self._processed_jobs_this_step = False self._flush_pending_promotions() + self._flush_pending_cascades() for tier in self.secondary_tiers: tier.on_schedule_end(context) @@ -802,8 +838,10 @@ class TieringOffloadingManager(OffloadingManager): # In-flight primary<->secondary transfers (pending promotions are # translated to transfer jobs in on_schedule_end), plus any work the # secondary tiers themselves still have outstanding. - return bool(self._jobs) or any( - tier.has_pending_work() for tier in self.secondary_tiers + return ( + bool(self._jobs) + or any(state.pending_cascade_keys for state in self._req_state.values()) + or any(tier.has_pending_work() for tier in self.secondary_tiers) ) @override @@ -848,6 +886,7 @@ class TieringOffloadingManager(OffloadingManager): finished_req_ids = [] for req_id, state in self._req_state.items(): state.pending_primary_stores = 0 + state.pending_cascade_keys.clear() if not state.is_finished: continue for tier in self.secondary_tiers: