diff --git a/tests/core/test_prefix_cache.py b/tests/core/test_prefix_cache.py index 171f36f71..9476a5140 100644 --- a/tests/core/test_prefix_cache.py +++ b/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 diff --git a/vllm_omni/core/prefix_cache/manager.py b/vllm_omni/core/prefix_cache/manager.py index 5294a25d4..a9aa381d7 100644 --- a/vllm_omni/core/prefix_cache/manager.py +++ b/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