[Diffusion][HunyuanImage3] Add request paged prefix caching (#7719)

Signed-off-by: zwhzzz0821 <2831474076@qq.com>

tests/diffusion/test_diffusion_scheduler.py

@@ -26,7 +26,7 @@ from vllm_omni.diffusion.sched import (
     Scheduler,
     StepScheduler,
 )
-from vllm_omni.diffusion.sched.interface import CachedRequestData, NewRequestData
+from vllm_omni.diffusion.sched.interface import CachedRequestData, NewRequestData, SchedulerRequestState
 from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
@@ -96,6 +96,7 @@ def _initialize_paged_scheduler(
     *,
     num_blocks: int = 64,
     max_num_seqs: int = 1,
+    enable_prefix_caching: bool = False,
 ) -> None:
     native_kv_managers.register_all_kvcache_specs(None)
     spec = FullAttentionSpec(
@@ -121,22 +122,42 @@ def _initialize_paged_scheduler(
         kv_vllm_config=SimpleNamespace(
             model_config=SimpleNamespace(max_model_len=64),
             max_in_flight_tokens=64,
+            cache_config=SimpleNamespace(
+                enable_prefix_caching=enable_prefix_caching,
+                prefix_caching_hash_algo="sha256",
+            ),
         ),
     )
 
 
-def _attach_diffusion_kv(request: OmniDiffusionRequest, *, seq_len: int = 8) -> None:
+def _attach_diffusion_kv(
+    request: OmniDiffusionRequest,
+    *,
+    seq_len: int = 8,
+    prefix_len: int = 4,
+    cache_token_ids=(),
+) -> None:
     request.diffusion_kv_requests = (
         DiffusionKVRequest(
             f"{request.request_id}/diffusion-kv/0",
             sequence_id=0,
-            prefix_len=4,
+            prefix_len=prefix_len,
             target_len=4,
             seq_len=seq_len,
+            cache_token_ids=cache_token_ids,
         ),
     )
 
 
+def _make_aborted_request_output(req_id: str) -> RunnerOutput:
+    return RunnerOutput(
+        request_id=req_id,
+        step_index=None,
+        finished=True,
+        result=DiffusionOutput(output=None, aborted=True),
+    )
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("scheduler_cls", [RequestScheduler, StepScheduler])
 @pytest.mark.parametrize("failure", ["timeout", "registration"])
@@ -168,7 +189,7 @@ async def test_single_native_kv_failure_reaches_output_stream(mocker, scheduler_
     engine.abort_queue = queue.Queue()
     engine._cv = threading.Condition()
     engine.main_loop = asyncio.get_running_loop()
-    stream = asyncio.Queue()
+    stream: asyncio.Queue[DiffusionOutput] = asyncio.Queue()
     engine._out_streams = {"failed": stream}
     engine.executor = mocker.Mock()
     engine.executor.prepare_kv_for_forward.return_value = KVConnectorOutput()
@@ -192,7 +213,7 @@ class _StubScheduler:
         self._output = output
         self.initialized_with = None
         self._request_id = request.request_id
-        self._state = None
+        self._state: SchedulerRequestState | None = None
         self._scheduled = False
         self.max_num_running_reqs = 1
 
@@ -201,7 +222,7 @@ class _StubScheduler:
 
     def add_request(self, request: OmniDiffusionRequest) -> str:
         assert request is self._request
-        self._state = SimpleNamespace(request_id=self._request_id, req=request)
+        self._state = SchedulerRequestState(request_id=self._request_id, req=request)
         return self._request_id
 
     def schedule(self):
@@ -223,6 +244,7 @@ class _StubScheduler:
     def update_from_output(self, sched_output, output) -> set[str]:
         del sched_output
         assert output is self._output
+        assert self._state is not None
         self._state.status = DiffusionRequestStatus.FINISHED_COMPLETED
         return {self._request_id}
 
@@ -553,6 +575,197 @@ class TestRequestScheduler:
         assert metadata.request_id == request.request_id
         assert len(metadata.sequences[0].block_ids[0]) == 4
 
+    def test_diffusion_kv_publishes_only_after_successful_completion(self) -> None:
+        _initialize_paged_scheduler(self.scheduler, enable_prefix_caching=True)
+        first = _make_request("publish-success")
+        _attach_diffusion_kv(first, cache_token_ids=range(4))
+        self.scheduler.add_request(first)
+        first_schedule = self.scheduler.schedule()
+        assert self.scheduler.update_from_output(first_schedule, _make_request_output(first.request_id)) == {
+            first.request_id
+        }
+
+        warm = _make_request("publish-warm")
+        _attach_diffusion_kv(warm, cache_token_ids=range(4))
+        self.scheduler.add_request(warm)
+        warm_schedule = self.scheduler.schedule()
+        metadata = warm_schedule.scheduled_new_reqs[0].diffusion_kv_metadata
+        assert metadata is not None
+        assert metadata.sequences[0].cached_prefix_len == 4
+
+    @pytest.mark.parametrize("boundaries", [(8, 4), (8, 0), (0, 8), (4, 8), (8, 8), (4, 4), (0, 0)])
+    @pytest.mark.parametrize("num_branches", [1, 2])
+    def test_diffusion_kv_batches_only_matching_prefix_boundaries(self, boundaries, num_branches) -> None:
+        _initialize_paged_scheduler(self.scheduler, max_num_seqs=2, enable_prefix_caching=True)
+        manager = self.scheduler._diffusion_kv_manager
+        assert manager is not None
+        pool = manager.native_manager.block_pool
+        empty_free_blocks = pool.get_num_free_blocks()
+
+        warmup = _make_request("warmup")
+        _attach_diffusion_kv(warmup, seq_len=12, prefix_len=8, cache_token_ids=range(8))
+        self.scheduler.add_request(warmup)
+        self.scheduler.update_from_output(self.scheduler.schedule(), _make_request_output("warmup"))
+
+        for request_id, boundary in zip(("first", "second"), boundaries):
+            request = _make_request(request_id)
+            # Each request's CFG branches share a boundary, but requests may not.
+            tokens = tuple(range(boundary)) + tuple(range(100, 108 - boundary))
+            request.diffusion_kv_requests = tuple(
+                DiffusionKVRequest(
+                    f"{request_id}/diffusion-kv/{branch}",
+                    sequence_id=branch,
+                    prefix_len=8,
+                    target_len=4,
+                    seq_len=12,
+                    cache_token_ids=tokens,
+                )
+                for branch in range(num_branches)
+            )
+            self.scheduler.add_request(request)
+
+        scheduled = self.scheduler.schedule()
+        matching = boundaries[0] == boundaries[1]
+        expected_ids = ["first", "second"] if matching else ["first"]
+        assert _new_ids(scheduled) == expected_ids
+        assert {
+            sequence.cached_prefix_len
+            for req in scheduled.scheduled_new_reqs
+            for sequence in req.diffusion_kv_metadata.sequences
+        } == {boundaries[0]}
+
+        if not matching:
+            second_state = self.scheduler.get_request_state("second")
+            assert second_state.status is DiffusionRequestStatus.WAITING
+            assert not manager.has_request("second")
+            # Retrying admission cannot leak reservations or publish its suffix.
+            free_blocks = pool.get_num_free_blocks()
+            for _ in range(3):
+                retry = self.scheduler.schedule()
+                assert _new_ids(retry) == [] and _cached_ids(retry) == ["first"]
+                assert pool.get_num_free_blocks() == free_blocks
+                assert not manager.has_request("second")
+                for row in second_state.diffusion_kv_requests:
+                    assert row.num_computed_tokens == 0
+                    assert manager.native_manager.get_computed_blocks(row)[1] == boundaries[1]
+
+        self.scheduler.update_from_output(
+            scheduled,
+            BatchRunnerOutput.from_list([_make_request_output(req_id) for req_id in expected_ids]),
+        )
+        if not matching:
+            next_wave = self.scheduler.schedule()
+            assert _new_ids(next_wave) == ["second"]
+            assert next_wave.scheduled_new_reqs[0].diffusion_kv_metadata.sequences[0].cached_prefix_len == boundaries[1]
+            self.scheduler.update_from_output(next_wave, _make_request_output("second"))
+        assert not self.scheduler.has_requests()
+        assert pool.get_num_free_blocks() == empty_free_blocks
+
+    def test_deferred_allocation_is_not_registered_for_transfer(self, mocker) -> None:
+        _initialize_paged_scheduler(self.scheduler, max_num_seqs=2)
+        manager = self.scheduler._diffusion_kv_manager
+        connector = mocker.Mock()
+        connector.get_num_new_matched_tokens.return_value = (4, True)
+        self.scheduler._kv_connector = connector
+        # Model a compatibility decision that needs the completed lookup.
+        can_schedule = self.scheduler._can_schedule_waiting
+        mocker.patch.object(
+            self.scheduler,
+            "_can_schedule_waiting",
+            side_effect=lambda state: can_schedule(state)
+            and not (state.request_id == "deferred" and manager.has_request("deferred")),
+        )
+        for request_id in ("admitted", "deferred"):
+            request = _make_request(request_id)
+            _attach_diffusion_kv(request)
+            request.diffusion_kv_requests[0].prompt_token_ids = [1, 2, 3, 4]
+            request.kv_transfer_params = {"num_transfer_tokens": 4}
+            self.scheduler.add_request(request)
+
+        assert _new_ids(self.scheduler.schedule()) == ["admitted"]
+        assert connector.update_state_after_alloc.call_count == 1
+        assert not manager.has_request("deferred")
+        assert "deferred" not in self.scheduler._kv_request_generations
+        assert "deferred" not in self.scheduler._kv_loading_request_ids
+        assert "deferred/diffusion-kv/0" not in self.scheduler._kv_transfer_request_ids
+        assert self.scheduler.get_request_state("deferred").status is DiffusionRequestStatus.WAITING
+
+    def test_diffusion_kv_deferred_request_can_be_cancelled(self) -> None:
+        _initialize_paged_scheduler(self.scheduler, max_num_seqs=2, enable_prefix_caching=True)
+        manager = self.scheduler._diffusion_kv_manager
+        pool = manager.native_manager.block_pool
+        empty_free_blocks = pool.get_num_free_blocks()
+        for request_id, tokens in (("warmup", range(4)), ("hit", range(4)), ("miss", range(10, 14))):
+            request = _make_request(request_id)
+            _attach_diffusion_kv(request, cache_token_ids=tokens)
+            self.scheduler.add_request(request)
+            if request_id == "warmup":
+                self.scheduler.update_from_output(self.scheduler.schedule(), _make_request_output("warmup"))
+
+        scheduled = self.scheduler.schedule()
+        assert _new_ids(scheduled) == ["hit"]
+        self.scheduler.finish_requests("miss", DiffusionRequestStatus.FINISHED_ABORTED)
+        self.scheduler.update_from_output(scheduled, _make_request_output("hit", error="worker failed"))
+        assert not manager.has_request("miss")
+        assert not self.scheduler.has_requests()
+        assert pool.get_num_free_blocks() == empty_free_blocks
+
+    def test_diffusion_kv_preempted_request_keeps_reservation_when_deferred(self) -> None:
+        _initialize_paged_scheduler(self.scheduler, max_num_seqs=2, enable_prefix_caching=True)
+        manager = self.scheduler._diffusion_kv_manager
+        warmup = _make_request("warmup")
+        _attach_diffusion_kv(warmup, cache_token_ids=range(4))
+        self.scheduler.add_request(warmup)
+        self.scheduler.update_from_output(self.scheduler.schedule(), _make_request_output("warmup"))
+
+        hit = _make_request("hit")
+        _attach_diffusion_kv(hit, cache_token_ids=range(4))
+        self.scheduler.add_request(hit)
+        self.scheduler.schedule()
+        miss = _make_request("miss")
+        _attach_diffusion_kv(miss, cache_token_ids=range(10, 14))
+        self.scheduler.add_request(miss)
+        # Install retained state as if this request were resuming an earlier wave.
+        miss_state = self.scheduler.get_request_state("miss")
+        miss_state.status = DiffusionRequestStatus.PREEMPTED
+        miss_metadata = manager.reserve_request("miss", miss_state.diffusion_kv_requests)
+        assert _cached_ids(self.scheduler.schedule()) == ["hit"]
+        assert self.scheduler.get_request_state("miss").status is DiffusionRequestStatus.PREEMPTED
+        assert manager.get_metadata("miss") is miss_metadata
+        self.scheduler.finish_requests("hit", DiffusionRequestStatus.FINISHED_ABORTED)
+        assert _cached_ids(self.scheduler.schedule()) == ["miss"]
+        assert manager.get_metadata("miss") is miss_metadata
+        self.scheduler.close()
+
+    @pytest.mark.parametrize(
+        "terminal_output",
+        [
+            pytest.param(
+                _make_request_output("publish-error", error="worker failed"),
+                id="error",
+            ),
+            pytest.param(_make_aborted_request_output("publish-abort"), id="abort"),
+        ],
+    )
+    def test_diffusion_kv_does_not_publish_failed_or_aborted_request(
+        self,
+        terminal_output: RunnerOutput,
+    ) -> None:
+        _initialize_paged_scheduler(self.scheduler, enable_prefix_caching=True)
+        first = _make_request(terminal_output.request_id)
+        _attach_diffusion_kv(first, cache_token_ids=range(4))
+        self.scheduler.add_request(first)
+        first_schedule = self.scheduler.schedule()
+        assert self.scheduler.update_from_output(first_schedule, terminal_output) == {first.request_id}
+
+        warm = _make_request(f"{first.request_id}-warm")
+        _attach_diffusion_kv(warm, cache_token_ids=range(4))
+        self.scheduler.add_request(warm)
+        warm_schedule = self.scheduler.schedule()
+        metadata = warm_schedule.scheduled_new_reqs[0].diffusion_kv_metadata
+        assert metadata is not None
+        assert metadata.sequences[0].cached_prefix_len == 0
+
     def test_diffusion_kv_capacity_backpressures_fifo_until_blocks_are_freed(self) -> None:
         _initialize_paged_scheduler(self.scheduler, num_blocks=3, max_num_seqs=2)
         first_request = _make_request("first")