[Hotfix][Core] PrefixCache: re-plan same-step prefix hits whose rows are written after the early prefetch (#7877)

Signed-off-by: ZhengWG <zwg0606@gmail.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>

tests/core/test_prefix_cache.py

@@ -1363,6 +1363,43 @@ def test_same_step_hit_prefetch_starts_at_save(caplog):
     assert torch.equal(outs.hidden_states["b"][:8], expected_rows(view.slots_for("b", 0, 8)))
 
 
+@pytest.mark.parametrize("reuse_blocks", [False, True], ids=["fresh", "reused"])
+@pytest.mark.parametrize("deferred_mm", [False, True], ids=["immediate", "deferred"])
+def test_same_step_hit_refreshes_prefetch_after_write(reuse_blocks, deferred_mm):
+    policy = ModelCachePolicy(deferred_keys=frozenset({"mm"}) if deferred_mm else frozenset())
+    mgr, view = make_manager(policy=policy)
+    try:
+        blocks = [0, 1] if reuse_blocks else [8, 9]
+        sid = run_step(mgr, view, {"old": (blocks, 0, 8)}, mm={"mm": torch.full((8, 2), 100.0)})
+        mgr.materialize(sid, ["old"])
+        mgr.new_step_starts(FakeSchedOut(finished=["old"]))
+
+        mgr.new_step_starts(
+            FakeSchedOut(
+                new_reqs=[FakeNewReq("a", 0, [[0, 1]]), FakeNewReq("b", 8, [[0, 1, 2]])],
+                num_scheduled={"a": 8, "b": 4},
+            )
+        )
+        old_prefetch = dict(mgr._hit_prefetch["b"])
+        # Complete the early reads before A publishes the prefix B actually hits.
+        for future in old_prefetch.values():
+            future.result(timeout=5)
+
+        view.order = ["a", "b"]
+        view.req_blocks.update(a=[0, 1], b=[0, 1, 2])
+        view.computed.update(a=0, b=8)
+        hidden = torch.cat([torch.full((8, HIDDEN), 20.0), torch.full((4, HIDDEN), 30.0)])
+        mm = torch.cat([torch.full((8, 2), 200.0), torch.full((4, 2), 300.0)])
+        sid = mgr.save_outputs(hidden, {"mm": mm}, num_tokens_unpadded=12, num_tokens_padded=12)
+        outs = mgr.materialize(sid, ["a", "b"])
+
+        assert torch.equal(outs.hidden_states["b"][:8], hidden[:8])
+        assert torch.equal(outs.hidden_states["b"][8:], hidden[8:])
+        assert torch.equal(outs.mm_outputs["mm"]["b"], mm)
+    finally:
+        mgr.shutdown()
+
+
 def test_delayed_read_of_reassigned_hit_raises():
     """A version check with no COW copy still raises for live and finished.
     The production path registers the ref first so remap preserves the rows

vllm_omni/core/prefix_cache/manager.py

@@ -789,6 +789,19 @@ class OmniPrefixCacheManager:
         self._hit_spans.clear()
         self._hit_prefetch.clear()
 
+    def _invalidate_overwritten_prefetches(self) -> None:
+        """Caller holds ``_state_lock``. Refresh this step's hits whose
+        intended producer registered a write after the early prefetch."""
+        for future, src in self._prefetch_queue:
+            futs = self._hit_prefetch.get(src.req_id)
+            if futs is None or futs.get(src.key) is not future:
+                continue
+            versions = self._slot_status.get_slot_status(src.key).slot_version[src.slots]
+            if src.reserved_version is None or not torch.equal(versions, src.reserved_version):
+                # Let the old fetch finish and unregister its pending read.
+                # COW may have preserved the prior tenant, which this hit must not use.
+                del futs[src.key]
+
     def _prefetch_hit_spans(self) -> None:
         """Caller holds ``_state_lock``. Plan each hit span not yet planned
         and gather it on the prefetch thread.
@@ -917,6 +930,7 @@ class OmniPrefixCacheManager:
         if self._hit_spans:
             # Same-step hits: their rows are registered now (IN_TRANSIT on
             # this step's tasks); start the gather before the next step.
+            self._invalidate_overwritten_prefetches()
             self._prefetch_hit_spans()
         step_id = self._next_step_id
         self._next_step_id += 1