[Feature] Add HunyuanImage3 DiT grouped step batching (#4041)

Signed-off-by: TaffyOfficial <2324465096@qq.com>
Signed-off-by: zuiho <2324465096@qq.com>
Co-authored-by: TaffyOfficial <2324465096@qq.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/diffusion/test_diffusion_step_pipeline.py

@@ -39,7 +39,8 @@ from vllm_omni.diffusion.sched.interface import (
 )
 from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
 from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
-from vllm_omni.diffusion.worker.utils import RunnerOutput
+from vllm_omni.diffusion.worker.input_batch import InputBatch
+from vllm_omni.diffusion.worker.utils import DiffusionRequestState, RunnerOutput
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
@@ -254,7 +255,9 @@ def _make_runner():
     runner.cache_backend = None
     runner.offload_backend = None
     runner.state_cache = {}
-    runner.kv_transfer_manager = SimpleNamespace()
+    runner.kv_transfer_manager = SimpleNamespace(
+        receive_multi_kv_cache_distributed=lambda req, cfg_kv_collect_func=None, target_device=None: None
+    )
     return runner
 
 
@@ -270,7 +273,9 @@ def _make_distributed_runner(mode: str, device: torch.device):
     runner.cache_backend = None
     runner.offload_backend = None
     runner.state_cache = {}
-    runner.kv_transfer_manager = SimpleNamespace()
+    runner.kv_transfer_manager = SimpleNamespace(
+        receive_multi_kv_cache_distributed=lambda req, cfg_kv_collect_func=None, target_device=None: None
+    )
     return runner
 
 
@@ -298,6 +303,17 @@ def _make_batch_scheduler_output(reqs, *, step_id=0, finished_req_ids=None):
     )
 
 
+def _make_input_batch_state(request_id: str, latent_value: float) -> DiffusionRequestState:
+    state = DiffusionRequestState(
+        request_id=request_id,
+        sampling=SimpleNamespace(),
+        prompts=None,
+    )
+    state.latents = torch.tensor([[latent_value]])
+    state.timesteps = torch.tensor([1.0])
+    return state
+
+
 def _make_cached_scheduler_output(request_id="req-1", step_id=1, finished_req_ids=None):
     return DiffusionSchedulerOutput(
         step_id=step_id,
@@ -374,6 +390,38 @@ def _distributed_step_worker(local_rank: int, world_size: int, mode: str, master
 # ---------------------------------------------------------------------------
 
 
+@pytest.mark.cpu
+def test_input_batch_cached_repack_refreshes_state_references_without_prompt_embeds():
+    first_state = _make_input_batch_state("req-1", 1.0)
+    batch = InputBatch.make_batch([first_state])
+    assert batch.prompt_embeds is None
+
+    replacement_state = _make_input_batch_state("req-1", 2.0)
+    repacked = InputBatch.make_batch([replacement_state], cached_batch=batch)
+
+    assert repacked is batch
+    assert repacked.states[0] is replacement_state
+    assert repacked.prompt_embeds is None
+    torch.testing.assert_close(repacked.latents, replacement_state.latents)
+
+
+@pytest.mark.cpu
+def test_input_batch_cached_repack_keeps_static_prompt_fields_for_same_composition():
+    first_state = _make_input_batch_state("req-1", 1.0)
+    first_state.prompt_embeds = torch.ones(1, 2, 3)
+    first_state.prompt_embeds_mask = torch.ones(1, 2, dtype=torch.bool)
+    batch = InputBatch.make_batch([first_state])
+
+    replacement_state = _make_input_batch_state("req-1", 2.0)
+    replacement_state.prompt_embeds = torch.full((1, 2, 3), 2.0)
+    replacement_state.prompt_embeds_mask = torch.ones(1, 2, dtype=torch.bool)
+    repacked = InputBatch.make_batch([replacement_state], cached_batch=batch)
+
+    assert repacked.states[0] is replacement_state
+    torch.testing.assert_close(repacked.latents, replacement_state.latents)
+    torch.testing.assert_close(repacked.prompt_embeds, torch.ones(1, 2, 3))
+
+
 @pytest.mark.cpu
 class TestRunner:
     """DiffusionModelRunner.execute_stepwise"""
@@ -427,6 +475,36 @@ class TestRunner:
         result = DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
         assert len(result) == 2
 
+    def test_receives_kv_payload_before_prepare_encode(self, monkeypatch):
+        runner = _make_runner()
+        captured: dict[str, object] = {}
+        kv_payload = object()
+
+        class _CapturingStepPipeline(_StepPipeline):
+            def prepare_encode(self, state, **kwargs):
+                captured["past_key_values"] = getattr(state.sampling, "past_key_values", None)
+                return super().prepare_encode(state, **kwargs)
+
+        class _KVTransferManager:
+            def receive_multi_kv_cache_distributed(self, req, cfg_kv_collect_func=None, target_device=None):
+                captured["cfg_kv_collect_func"] = cfg_kv_collect_func
+                captured["target_device"] = target_device
+                req.sampling_params.past_key_values = kv_payload
+
+        runner.pipeline = _CapturingStepPipeline()
+        runner.pipeline.device = torch.device("cpu")
+        runner.od_config.cfg_kv_collect_func = "collect-cfg"
+        runner.kv_transfer_manager = _KVTransferManager()
+        monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+        req = _make_step_request()
+
+        DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0))
+
+        assert captured["past_key_values"] is kv_payload
+        assert captured["cfg_kv_collect_func"] == "collect-cfg"
+        assert captured["target_device"] == torch.device("cpu")
+        assert getattr(req.sampling_params, "past_key_values", None) is None
+
     def test_rejects_missing_cached_state(self):
         runner = _make_runner()