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

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

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

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: