[Diffusion] Prepare HunyuanImage3 for Scheduler-managed paged KV cache (#5541)

Signed-off-by: zwhzzz0821 <zwhzzz0821@users.noreply.github.com>
Co-authored-by: zwhzzz0821 <zwhzzz0821@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/diffusion/test_diffusion_scheduler.py

@@ -12,6 +12,8 @@ from pytest_mock import MockerFixture
 
 from vllm_omni.diffusion.data import DiffusionOutput, DiffusionRequestAbortedError
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVContext, DiffusionKVRequest
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched import (
     BaseScheduler,
@@ -259,6 +261,90 @@ class TestRequestScheduler:
         assert self.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_COMPLETED
         assert self.scheduler.has_requests() is False
 
+    def test_diffusion_kv_request_moves_to_scheduler_state(self) -> None:
+        self.scheduler.initialize(SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER))
+        request = _make_request("diffusion-kv")
+        prepared_layout = object()
+        request.prepared_layout = prepared_layout
+        text_context = DiffusionKVContext(context_id="text", cache_role="cross.text", num_tokens=8)
+        kv_request = DiffusionKVRequest(
+            "diffusion-kv/diffusion-kv/0",
+            sequence_id=0,
+            prefix_len=4,
+            target_len=8,
+            seq_len=16,
+            kv_contexts=(text_context,),
+        )
+        request.diffusion_kv_requests = (kv_request,)
+
+        self.scheduler.add_request(request)
+        state = self.scheduler.get_request_state(request.request_id)
+        scheduler_output = self.scheduler.schedule()
+
+        assert state is not None
+        assert state.diffusion_kv_requests == (kv_request,)
+        assert state.diffusion_kv_requests[0].kv_contexts == (text_context,)
+        assert request.diffusion_kv_requests is None
+        assert scheduler_output.scheduled_new_reqs[0].req.prepared_layout is prepared_layout
+
+    def test_paged_scheduler_rejects_missing_kv_request(self) -> None:
+        self.scheduler.initialize(SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER))
+        request = _make_request("diffusion-kv")
+
+        with pytest.raises(ValueError, match="did not produce DiffusionKVRequest"):
+            self.scheduler.add_request(request)
+
+    @pytest.mark.parametrize(
+        ("owner_name", "field_name"),
+        [
+            ("request", "past_key_values"),
+            ("sampling_params", "past_key_values"),
+            ("sampling_params", "cfg_text_past_key_values"),
+            ("sampling_params", "cfg_img_past_key_values"),
+            ("sampling_params", "cfg_branch_past_key_values"),
+        ],
+    )
+    def test_diffusion_kv_request_rejects_legacy_dense_kv(self, owner_name: str, field_name: str) -> None:
+        self.scheduler.initialize(SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER))
+        request = _make_request("diffusion-kv")
+        owner = request if owner_name == "request" else request.sampling_params
+        setattr(owner, field_name, object())
+        request.diffusion_kv_requests = (
+            DiffusionKVRequest(
+                "diffusion-kv/diffusion-kv/0",
+                sequence_id=0,
+                prefix_len=4,
+                target_len=8,
+                seq_len=16,
+            ),
+        )
+
+        with pytest.raises(ValueError, match=rf"{owner_name}\.{field_name}"):
+            self.scheduler.add_request(request)
+
+    def test_dense_scheduler_rejects_scheduler_only_kv_request(self) -> None:
+        request = _make_request("dense-kv-request")
+        request.diffusion_kv_requests = (
+            DiffusionKVRequest(
+                "dense-kv-request/diffusion-kv/0",
+                sequence_id=0,
+                prefix_len=4,
+                target_len=8,
+                seq_len=16,
+            ),
+        )
+
+        with pytest.raises(ValueError, match="dense_legacy request unexpectedly contains"):
+            self.scheduler.add_request(request)
+
+    def test_dense_scheduler_accepts_legacy_dense_kv(self) -> None:
+        request = _make_request("dense-kv")
+        request.sampling_params.past_key_values = object()
+
+        request_id = self.scheduler.add_request(request)
+
+        assert request_id == request.request_id
+
     def test_error_output_marks_finished_error(self) -> None:
         req_id = self.scheduler.add_request(_make_request("err"))
 
@@ -562,12 +648,20 @@ class TestDiffusionEngine:
         engine.abort_queue = queue.Queue()
 
         request = _make_request("engine")
+        prepared_layout = object()
+
+        def preprocess(req):
+            req.prepared_layout = prepared_layout
+            return req
+
+        engine.pre_process_func = mocker.Mock(side_effect=preprocess)
         runner_output = _make_request_output("engine")
         engine.execute_fn = mocker.Mock(return_value=runner_output)
 
         output = engine.add_req_and_wait_for_response(request)
 
         assert output is runner_output.result
+        engine.pre_process_func.assert_called_once_with(request)
         engine.execute_fn.assert_called_once()
 
     def test_supports_scheduler_interface_injection(self, mocker: MockerFixture) -> None:
@@ -631,17 +725,29 @@ class TestDiffusionEngine:
     async def test_step_streaming_raises_aborted_error(self, mocker: MockerFixture) -> None:
         engine = DiffusionEngine.__new__(DiffusionEngine)
         engine._check_and_start_background_loop = mocker.AsyncMock()
-        engine.pre_process_func = None
+        request = _make_request("req-abort")
+        prepared_layout = object()
+
+        def preprocess(req):
+            req.prepared_layout = prepared_layout
+            return req
+
+        engine.pre_process_func = mocker.Mock(side_effect=preprocess)
 
         async def _stream(_request):
             yield DiffusionOutput(aborted=True, abort_message="Request req-abort aborted.")
 
-        engine.async_add_req_and_stream_response = mocker.Mock(return_value=_stream(None))
+        engine._add_prepared_request = mocker.Mock(return_value=request.request_id)
+        engine.get_output_stream = mocker.Mock(return_value=_stream(None))
 
         with pytest.raises(DiffusionRequestAbortedError, match="Request req-abort aborted"):
-            async for _ in engine.step_streaming(_make_request("req-abort")):
+            async for _ in engine.step_streaming(request):
                 pass
 
+        engine.pre_process_func.assert_called_once_with(request)
+        engine._add_prepared_request.assert_called_once_with(request)
+        engine.get_output_stream.assert_called_once_with(request.request_id)
+
     def test_abort_queue_marks_request_finished_aborted(self) -> None:
         engine = DiffusionEngine.__new__(DiffusionEngine)
         engine._rpc_lock = threading.RLock()
@@ -774,6 +880,18 @@ class TestDiffusionEngine:
         with pytest.raises(RuntimeError, match="Dummy run failed: boom"):
             engine._dummy_run()
 
+    def test_dummy_run_delegates_preprocessing_to_sync_admission(self, mocker: MockerFixture) -> None:
+        engine = DiffusionEngine.__new__(DiffusionEngine)
+        engine.od_config = SimpleNamespace(model_class_name="mock_model", diffusion_load_format="default")
+        engine.pre_process_func = mocker.Mock()
+        engine.add_req_and_wait_for_response = mocker.Mock(return_value=DiffusionOutput(output=None))
+
+        engine._dummy_run()
+
+        engine.pre_process_func.assert_not_called()
+        admitted_request = engine.add_req_and_wait_for_response.call_args.args[0]
+        assert admitted_request.request_id == "dummy_req_id"
+
 
 class TestStepScheduler:
     def setup_method(self) -> None: