[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>

.buildkite/scripts/hardware_ci/run-amd-test.sh

@@ -64,7 +64,8 @@ container_name="rocm_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | hea
 docker pull "${image_name}"
 
 remove_docker_container() {
-   docker rm -f "${container_name}" || docker image rm -f "${image_name}" || true
+   docker ps -aq --filter "name=${container_name}" | xargs -r docker rm -f || true
+   docker image rm -f "${image_name}" || true
 }
 trap remove_docker_container EXIT
 
@@ -107,7 +108,6 @@ if [[ $commands == *"--shard-id="* ]]; then
         --network=host \
         --shm-size=16gb \
         --group-add "$render_gid" \
-        --rm \
         -e MIOPEN_DEBUG_CONV_DIRECT=0 \
         -e MIOPEN_DEBUG_CONV_GEMM=0 \
         -e VLLM_ROCM_USE_AITER=0 \
@@ -140,7 +140,6 @@ else
           --network=host \
           --shm-size=16gb \
           --group-add "$render_gid" \
-          --rm \
           -e MIOPEN_DEBUG_CONV_DIRECT=0 \
           -e MIOPEN_DEBUG_CONV_GEMM=0 \
           -e VLLM_ROCM_USE_AITER=0 \

tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_step_execution.py

@@ -0,0 +1,366 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 as hy3_module
+from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec
+from vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 import (
+    _STEP_AR_KV,
+    _STEP_CFG_FACTOR,
+    _STEP_GENERATOR,
+    _STEP_GUIDANCE_SCALE,
+    _STEP_INPUT_IDS,
+    _STEP_MODEL_KWARGS,
+    _STEP_PROMPT_KV,
+    HunyuanImage3Pipeline,
+)
+from vllm_omni.diffusion.worker.input_batch import InputBatch
+from vllm_omni.diffusion.worker.utils import DiffusionRequestState
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _pipeline():
+    pipeline = object.__new__(HunyuanImage3Pipeline)
+    pipeline._tkwrapper = SimpleNamespace(pad_token_id=0)
+    pipeline.od_config = SimpleNamespace(
+        diffusion_attention_config=AttentionConfig(default=AttentionSpec(backend="TORCH_SDPA")),
+        parallel_config=SimpleNamespace(sequence_parallel_size=1, cfg_parallel_size=1),
+        cache_backend=None,
+        diffusion_kv_cache_skip_step_indices=None,
+    )
+    pipeline._pipeline = SimpleNamespace()
+    return pipeline
+
+
+def _state(request_id: str, step_index: int) -> DiffusionRequestState:
+    state = DiffusionRequestState(
+        request_id=request_id,
+        sampling=SimpleNamespace(),
+        prompts=["prompt"],
+    )
+    state.step_index = step_index
+    state.timesteps = torch.tensor([1.0, 0.5, 0.25, 0.0])
+    state.latents = torch.zeros(1, 4, 8, 8)
+    state.extra = {
+        _STEP_CFG_FACTOR: 1,
+        _STEP_AR_KV: None,
+        _STEP_INPUT_IDS: None,
+        _STEP_GUIDANCE_SCALE: 1.0,
+        _STEP_MODEL_KWARGS: {
+            "num_image_tokens": 17,
+            "ar_kv_reuse_len": 0,
+        },
+    }
+    return state
+
+
+def _sampling_params(**extra_args):
+    return SimpleNamespace(
+        timesteps=None,
+        sigmas=None,
+        num_outputs_per_prompt=None,
+        extra_args=extra_args,
+        height=512,
+        width=512,
+        num_inference_steps=4,
+        guidance_scale=1.0,
+        guidance_scale_provided=True,
+        guidance_rescale=0.0,
+        generator=None,
+    )
+
+
+def test_hunyuan_step_group_key_ignores_step_index_for_later_steps():
+    pipeline = _pipeline()
+    states = [_state("req-0", 1), _state("req-1", 3)]
+
+    groups = pipeline._split_step_groups(states)
+
+    assert len(groups) == 1
+    assert [state.request_id for state in groups[0]] == ["req-0", "req-1"]
+
+
+@pytest.mark.parametrize(
+    ("sampling", "prompt_item", "expected_model_bot_task", "expected_system_bot_task"),
+    [
+        pytest.param(
+            _sampling_params(bot_task="think_recaption", use_system_prompt="dynamic"),
+            {"prompt": "prompt", "bot_task": "vanilla"},
+            "think",
+            "think",
+            id="extra-args-precedence",
+        ),
+        pytest.param(
+            _sampling_params(use_system_prompt="dynamic"),
+            {"prompt": "prompt", "bot_task": "vanilla"},
+            "image",
+            "image",
+            id="prompt-dict-fallback",
+        ),
+        pytest.param(
+            _sampling_params(use_system_prompt="dynamic"),
+            {"prompt": "prompt"},
+            "auto",
+            "image",
+            id="default-auto-system-prompt",
+        ),
+    ],
+)
+def test_prepare_encode_preserves_normal_hunyuan_bot_task_semantics(
+    monkeypatch,
+    sampling,
+    prompt_item,
+    expected_model_bot_task,
+    expected_system_bot_task,
+):
+    pipeline = _pipeline()
+    captured: dict[str, object] = {}
+
+    def fake_get_system_prompt(sys_type, bot_task, system_prompt=None):
+        del sys_type, system_prompt
+        captured["system_prompt_bot_task"] = bot_task
+        return "system prompt"
+
+    def fake_prepare_model_inputs(**kwargs):
+        captured.update(kwargs)
+        raise RuntimeError("stop after prepare_model_inputs")
+
+    monkeypatch.setattr(hy3_module, "get_system_prompt", fake_get_system_prompt)
+    pipeline.prepare_model_inputs = fake_prepare_model_inputs
+    state = DiffusionRequestState(
+        request_id="req-bot-task",
+        sampling=sampling,
+        prompts=[prompt_item],
+    )
+
+    with pytest.raises(RuntimeError, match="stop after prepare_model_inputs"):
+        pipeline.prepare_encode(state)
+
+    assert captured["bot_task"] == expected_model_bot_task
+    assert captured["system_prompt_bot_task"] == expected_system_bot_task
+
+
+def test_forward_uses_same_hunyuan_bot_task_semantics(monkeypatch):
+    pipeline = _pipeline()
+    captured: dict[str, object] = {}
+
+    def fake_get_system_prompt(sys_type, bot_task, system_prompt=None):
+        del sys_type, system_prompt
+        captured["system_prompt_bot_task"] = bot_task
+        return "system prompt"
+
+    def fake_prepare_model_inputs(**kwargs):
+        captured.update(kwargs)
+        raise RuntimeError("stop after prepare_model_inputs")
+
+    monkeypatch.setattr(hy3_module, "get_system_prompt", fake_get_system_prompt)
+    pipeline.prepare_model_inputs = fake_prepare_model_inputs
+    req = SimpleNamespace(
+        request_id="req-forward-bot-task",
+        sampling_params=_sampling_params(bot_task="think_recaption", use_system_prompt="dynamic"),
+        prompts=[{"prompt": "prompt", "bot_task": "vanilla"}],
+    )
+
+    with pytest.raises(RuntimeError, match="stop after prepare_model_inputs"):
+        pipeline.forward(req)
+
+    assert captured["bot_task"] == "think"
+    assert captured["system_prompt_bot_task"] == "think"
+
+
+def test_grouped_denoise_rejects_non_sdpa_attention_backend():
+    pipeline = _pipeline()
+    pipeline.od_config.diffusion_attention_config = AttentionConfig(default=AttentionSpec(backend="FLASH_ATTN"))
+
+    with pytest.raises(ValueError, match="only supports TORCH_SDPA"):
+        pipeline._ensure_grouped_attention_backend_supported(2)
+
+
+def test_single_denoise_allows_non_sdpa_attention_backend():
+    pipeline = _pipeline()
+    pipeline.od_config.diffusion_attention_config = AttentionConfig(default=AttentionSpec(backend="FLASH_ATTN"))
+
+    pipeline._ensure_grouped_attention_backend_supported(1)
+
+
+def test_grouped_denoise_allows_sdpa_attention_backend():
+    pipeline = _pipeline()
+
+    pipeline._ensure_grouped_attention_backend_supported(2)
+
+
+def test_step_scheduler_preserves_latent_dtype_for_mixed_progress_batches():
+    pipeline = _pipeline()
+    pipeline._pipeline = SimpleNamespace(prepare_extra_func_kwargs=lambda step, kwargs: {})
+
+    class FakeScheduler:
+        def step(self, noise_pred, timestep, latents, **kwargs):
+            del timestep, kwargs
+            return (latents.float() + noise_pred.float(),)
+
+    state = _state("req", 0)
+    state.timesteps = torch.tensor([1.0])
+    state.scheduler = FakeScheduler()
+    state.latents = torch.zeros(1, 4, 8, 8, dtype=torch.bfloat16)
+    state.extra[_STEP_GENERATOR] = None
+
+    pipeline.step_scheduler(state, torch.ones_like(state.latents, dtype=torch.float32))
+
+    assert state.latents.dtype == torch.bfloat16
+    assert state.step_index == 1
+
+
+def test_later_step_merge_shifts_spans_without_polluting_request_state():
+    pipeline = _pipeline()
+    states = [_state("short", 2), _state("long", 4)]
+    states[0].extra[_STEP_MODEL_KWARGS].update(
+        {
+            "attention_mask": torch.ones(1, 1, 3, 5, dtype=torch.bool),
+            "full_attn_spans": [[(2, 5)]],
+        }
+    )
+    states[1].extra[_STEP_MODEL_KWARGS].update(
+        {
+            "attention_mask": torch.ones(1, 1, 3, 7, dtype=torch.bool),
+            "full_attn_spans": [[(4, 7)]],
+        }
+    )
+    states[0].extra[_STEP_PROMPT_KV] = [{"lens": torch.tensor([2])}]
+    states[1].extra[_STEP_PROMPT_KV] = [{"lens": torch.tensor([4])}]
+
+    row_state_indexes = [0, 1]
+    row_branches = [0, 0]
+    _, merged = pipeline._merge_step_model_inputs(
+        states,
+        row_state_indexes,
+        row_branches,
+        first_step=False,
+    )
+
+    assert merged["attention_mask"].shape == (2, 1, 3, 7)
+    assert merged["full_attn_spans"] == [[(4, 7)], [(4, 7)]]
+
+    pipeline._split_merged_kwargs_to_states(states, merged, row_state_indexes, row_branches)
+
+    assert states[0].extra[_STEP_MODEL_KWARGS]["attention_mask"].shape == (1, 1, 3, 5)
+    assert states[1].extra[_STEP_MODEL_KWARGS]["attention_mask"].shape == (1, 1, 3, 7)
+    assert states[0].extra[_STEP_MODEL_KWARGS]["full_attn_spans"] == [[(2, 5)]]
+    assert states[1].extra[_STEP_MODEL_KWARGS]["full_attn_spans"] == [[(4, 7)]]
+
+
+def test_later_step_merge_allows_request_local_step_counts_and_guidance_values():
+    pipeline = _pipeline()
+    states = [_state("req-0", 1), _state("req-1", 3)]
+    for idx, state in enumerate(states):
+        state.extra[_STEP_MODEL_KWARGS].update(
+            {
+                "attention_mask": torch.ones(1, 1, 2, 4, dtype=torch.bool),
+                "full_attn_spans": [[(2, 4)]],
+                "guidance_scale": 3.0 + idx,
+                "num_inference_steps": 20 + idx,
+            }
+        )
+        state.extra[_STEP_PROMPT_KV] = [{"lens": torch.tensor([2])}]
+
+    _, merged = pipeline._merge_step_model_inputs(
+        states,
+        row_state_indexes=[0, 1],
+        row_branches=[0, 0],
+        first_step=False,
+    )
+
+    assert "guidance_scale" not in merged
+    assert "num_inference_steps" not in merged
+
+
+@pytest.mark.parametrize(
+    ("request_id", "mutate_state", "error_match"),
+    [
+        pytest.param(
+            "broken-req",
+            lambda state: state.extra.pop(_STEP_MODEL_KWARGS),
+            "broken-req",
+            id="missing-model-kwargs",
+        ),
+        pytest.param(
+            "bad-cfg",
+            lambda state: state.extra.__setitem__(_STEP_CFG_FACTOR, 3),
+            "bad-cfg",
+            id="unsupported-cfg-factor",
+        ),
+    ],
+)
+def test_denoise_step_reports_invalid_group_state_with_request_id(request_id, mutate_state, error_match):
+    pipeline = _pipeline()
+    state = _state(request_id, 0)
+    mutate_state(state)
+
+    with pytest.raises(ValueError, match=error_match):
+        pipeline.denoise_step(InputBatch.make_batch([state]))
+
+
+def test_denoise_step_uses_input_batch_group_order_and_splits_back(monkeypatch):
+    pipeline = _pipeline()
+    monkeypatch.setattr(HunyuanImage3Pipeline, "device", property(lambda self: torch.device("cpu")))
+    states = [_state("req-0", 1), _state("req-1", 3)]
+    for idx, state in enumerate(states):
+        prefix_len = 2 + idx * 2
+        state.latents = torch.full((1, 1), float(idx))
+        state.extra[_STEP_CFG_FACTOR] = 2
+        state.extra[_STEP_GUIDANCE_SCALE] = 1.0
+        state.extra[_STEP_INPUT_IDS] = None
+        state.extra[_STEP_MODEL_KWARGS].update(
+            {
+                "attention_mask": torch.ones(2, 1, 2, prefix_len + 2, dtype=torch.bool),
+                "full_attn_spans": [[(prefix_len, prefix_len + 2)], [(prefix_len, prefix_len + 2)]],
+            }
+        )
+        state.extra[_STEP_PROMPT_KV] = [
+            {
+                "key": torch.zeros(2, prefix_len, 1, 1),
+                "value": torch.zeros(2, prefix_len, 1, 1),
+                "lens": torch.tensor([prefix_len, prefix_len]),
+            }
+        ]
+
+    captured = {}
+
+    def fake_restore_prompt_kv_cache(states_arg, row_state_indexes, row_branches):
+        del states_arg
+        captured["row_state_indexes"] = list(row_state_indexes)
+        captured["row_branches"] = list(row_branches)
+
+    def fake_prepare_inputs_for_generation(input_ids, images, timestep, **model_kwargs):
+        captured["input_ids"] = input_ids
+        captured["images"] = images.clone()
+        captured["timestep"] = timestep.clone()
+        captured["merged_attention_mask_shape"] = tuple(model_kwargs["attention_mask"].shape)
+        captured["merged_full_attn_spans"] = model_kwargs["full_attn_spans"]
+        return {"model_kwargs": model_kwargs}
+
+    pipeline._restore_prompt_kv_cache = fake_restore_prompt_kv_cache
+    pipeline.prepare_inputs_for_generation = fake_prepare_inputs_for_generation
+    pipeline.forward_call = lambda **kwargs: {"diffusion_prediction": torch.tensor([[10.0], [20.0], [1.0], [2.0]])}
+    pipeline._update_model_kwargs_for_generation = lambda model_output, model_kwargs: model_kwargs
+    pipeline._pipeline = SimpleNamespace(cfg_operator=lambda cond, uncond, scale, step: cond + uncond)
+
+    batch = InputBatch.make_batch(states)
+    out = pipeline.denoise_step(batch)
+
+    assert captured["row_state_indexes"] == [0, 1, 0, 1]
+    assert captured["row_branches"] == [0, 0, 1, 1]
+    assert captured["input_ids"] is None
+    assert tuple(captured["images"].shape) == (4, 1)
+    assert captured["timestep"].tolist() == [0.5, 0.0, 0.5, 0.0]
+    assert captured["merged_attention_mask_shape"] == (4, 1, 2, 6)
+    assert captured["merged_full_attn_spans"] == [[(4, 6)], [(4, 6)], [(4, 6)], [(4, 6)]]
+    torch.testing.assert_close(out, torch.tensor([[11.0], [22.0]]))
+    assert states[0].extra[_STEP_MODEL_KWARGS]["attention_mask"].shape == (2, 1, 2, 4)
+    assert states[1].extra[_STEP_MODEL_KWARGS]["attention_mask"].shape == (2, 1, 2, 6)
+    assert states[0].extra[_STEP_MODEL_KWARGS]["full_attn_spans"] == [[(2, 4)], [(2, 4)]]
+    assert states[1].extra[_STEP_MODEL_KWARGS]["full_attn_spans"] == [[(4, 6)], [(4, 6)]]

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()
 

vllm_omni/diffusion/models/interface.py

@@ -16,6 +16,7 @@ if TYPE_CHECKING:
     import torch
 
     from vllm_omni.diffusion.data import DiffusionOutput
+    from vllm_omni.diffusion.worker.input_batch import InputBatch
     from vllm_omni.diffusion.worker.utils import DiffusionRequestState
 
 
@@ -56,8 +57,8 @@ class SupportsStepExecution(Protocol):
     def prepare_encode(self, state: DiffusionRequestState, **kwargs: Any) -> DiffusionRequestState:
         """Prepare request-level inputs and return initialized state."""
 
-    def denoise_step(self, state: DiffusionRequestState, **kwargs: Any) -> torch.Tensor | None:
-        """Run one denoise step."""
+    def denoise_step(self, input_batch: InputBatch, **kwargs: Any) -> torch.Tensor | None:
+        """Run one denoise forward on the runner-assembled batch."""
 
     def step_scheduler(self, state: DiffusionRequestState, noise_pred: torch.Tensor, **kwargs: Any) -> None:
         """Run one scheduler step."""

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -374,6 +374,13 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                     sampling=copy.deepcopy(req.sampling_params),
                     prompts=req.prompts,
                 )
+                state_req = copy.copy(req)
+                state_req.sampling_params = new_state.sampling
+                self.kv_transfer_manager.receive_multi_kv_cache_distributed(
+                    state_req,
+                    cfg_kv_collect_func=getattr(self.od_config, "cfg_kv_collect_func", None),
+                    target_device=getattr(self.pipeline, "device", None),
+                )
                 self.state_cache[request_id] = new_state
                 resolved.append(new_state)
 
@@ -452,8 +459,8 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         if not self.supports_step_mode():
             raise ValueError("Current pipeline does not support step execution.")
         # Stepwise mode only supports the basic state-driven denoise path for now.
-        # Request-mode extras such as cache backends, KV transfer, editing inputs,
-        # and similar features are not supported here yet.
+        # Request-mode extras such as cache backends, editing inputs, and
+        # similar features are not supported here yet.
         if self.od_config.cache_backend not in (None, "none"):
             raise ValueError("Step mode does not support cache_backend yet.")
 

vllm_omni/diffusion/worker/input_batch.py

@@ -12,7 +12,7 @@ step, and step outputs are scattered back into request states by
 from __future__ import annotations
 
 from collections.abc import Sequence
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 
 import numpy as np
 import torch
@@ -496,7 +496,7 @@ def _prepare_prompt_embeds(
     *,
     embeds_out: torch.Tensor | None = None,
     mask_out: torch.Tensor | None = None,
-) -> tuple[torch.Tensor, torch.Tensor | None]:
+) -> tuple[torch.Tensor | None, torch.Tensor | None]:
     prompt_embeds, prompt_embeds_mask = _prepare_padded_prompt_fields(
         states,
         embeds_attr="prompt_embeds",
@@ -505,8 +505,6 @@ def _prepare_prompt_embeds(
         embeds_out=embeds_out,
         mask_out=mask_out,
     )
-    if prompt_embeds is None:
-        raise ValueError("All requests must have `prompt_embeds` initialized.")
     return prompt_embeds, prompt_embeds_mask
 
 
@@ -581,6 +579,10 @@ class InputBatch:
     source of truth. ``InputBatch`` only assembles a contiguous view for the
     current step and refreshes dynamic fields in-place when composition is
     unchanged.
+
+    ``states`` is a narrow escape hatch for Hunyuan-style state-driven
+    pipelines that need request-private KV/cache metadata during denoise.
+    Other pipelines should continue to use the standard batch fields.
     """
 
     request_ids: list[str]
@@ -591,7 +593,7 @@ class InputBatch:
 
     latents: torch.Tensor
     timesteps: torch.Tensor
-    prompt_embeds: torch.Tensor
+    prompt_embeds: torch.Tensor | None
     prompt_embeds_mask: torch.Tensor | None
     negative_prompt_embeds: torch.Tensor | None
     negative_prompt_embeds_mask: torch.Tensor | None
@@ -604,6 +606,7 @@ class InputBatch:
     img_shapes: list | None = None
     txt_seq_lens: list[int] | None = None
     negative_txt_seq_lens: list[int] | None = None
+    states: Sequence[DiffusionRequestState] = field(default_factory=tuple)
 
     def __post_init__(self) -> None:
         if len(self.request_ids) != int(self.idx_mapping.numel()):
@@ -648,6 +651,9 @@ class InputBatch:
         self,
         selected_states: Sequence[DiffusionRequestState],
     ) -> None:
+        # Same-composition cache hits reuse static prompt fields; request ids
+        # must keep the same encoded prompt metadata for the batch lifetime.
+        self.states = tuple(selected_states)
         self._refresh_dynamic_fields(selected_states)
 
     def _rebuild(
@@ -662,6 +668,7 @@ class InputBatch:
         self.num_reqs_after_padding = len(request_ids)
         self.idx_mapping = idx_mapping
         self.idx_mapping_np = idx_mapping_np
+        self.states = tuple(selected_states)
         self.latents = _prepare_latents(selected_states, out=self.latents)
         self.timesteps = _prepare_timesteps(selected_states, out=self.timesteps)
         self._refresh_static_fields(selected_states)
@@ -718,6 +725,7 @@ class InputBatch:
                 selected_states,
                 "negative_txt_seq_lens",
             ),
+            states=tuple(selected_states),
         )