[Bug Fix] Hunyuan image cfg bugfix (#4752)

Signed-off-by: zhuhh97 <zhuhonghan@huawei.com>
Signed-off-by: dph97 <nenbaying@163.com>
Co-authored-by: zhuhh97 <z00806815@gmail.com>
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

docker/Dockerfile.cuda

@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai:v0.24.0
+ARG BASE_IMAGE=vllm/vllm-openai:v0.25.0
 FROM ${BASE_IMAGE}
 
 ARG COMMON_WORKDIR=/app

docker/Dockerfile.rocm

@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.24.0
+ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.25.0
 FROM ${BASE_IMAGE} AS base
 
 # Declare a variable to know if we want to use the nightly build or the stable build.

tests/diffusion/test_diffusion_model_runner.py

@@ -509,6 +509,25 @@ def test_split_diffusion_output_by_request_slices_tuple_outputs():
     assert torch.equal(outputs[1].output[1], torch.tensor([20]))
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_execute_model_runs_forward_after_kv_receive(monkeypatch):
+    """execute_model runs the pipeline forward after the KV receive step."""
+    runner = _make_runner(cache_backend=None, cache_backend_name=None)
+    runner.kv_transfer_manager.receive_multi_kv_cache_distributed = lambda *a, **k: True
+    req = _make_request(skip_cache_refresh=True)
+
+    monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "reset_peak_memory_stats", lambda: None)
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0)
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0)
+
+    output = DiffusionModelRunner.execute_model(runner, req)
+
+    assert output.output == "ok"
+    assert runner.pipeline.forward_calls == 1
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):

tests/distributed/omni_connectors/test_tp_rank_aware.py

@@ -937,6 +937,76 @@ class TestDistributedReceive:
 # ── TP auto-detect ───────────────────────────────────────────────────
 
 
+class TestBuildCfgRankLocalPayloads:
+    """Guard _build_cfg_rank_local_payloads padding logic for the
+    cfg_size==2 / no branch_roles scenario (hunyuan-image shared-prefix reuse)."""
+
+    def test_cfg_size2_no_branch_roles_copies_main_payload_to_follower(self):
+        """When AR doesn't split KV by branch, cfg follower (rank 1) should
+        receive a copy of the positive KV, not None."""
+        req = SimpleNamespace(
+            request_id="req-1",
+            past_key_values=SimpleNamespace(key_cache=[torch.tensor([1.0])]),
+            kv_metadata={"seq_len": 10},
+            sampling_params=SimpleNamespace(),
+        )
+        payloads = OmniKVTransferManager._build_cfg_rank_local_payloads(req, cfg_size=2)
+
+        assert len(payloads) == 2
+        # Rank 0: main payload (always present)
+        assert payloads[0] is not None
+        assert "past_key_values" in payloads[0]
+        # Rank 1: must also be non-None — same positive KV for shared-prefix reuse
+        assert payloads[1] is not None
+        assert payloads[1]["past_key_values"] is req.past_key_values
+        assert payloads[1]["kv_metadata"] == {"seq_len": 10}
+
+    def test_cfg_size3_no_branch_roles_keeps_none_padding(self):
+        """When cfg_size > 2 and no branch_roles, padding stays None (safety guard)."""
+        req = SimpleNamespace(
+            request_id="req-1",
+            past_key_values=SimpleNamespace(key_cache=[torch.tensor([1.0])]),
+            kv_metadata={"seq_len": 10},
+            sampling_params=SimpleNamespace(),
+        )
+        payloads = OmniKVTransferManager._build_cfg_rank_local_payloads(req, cfg_size=3)
+
+        assert len(payloads) == 3
+        assert payloads[0] is not None
+        assert payloads[1] is None
+        assert payloads[2] is None
+
+    def test_empty_main_payload_keeps_none_padding(self):
+        """When main_payload is empty, padding stays None even with cfg_size==2."""
+        req = SimpleNamespace(
+            request_id="req-1",
+            sampling_params=SimpleNamespace(),
+        )
+        payloads = OmniKVTransferManager._build_cfg_rank_local_payloads(req, cfg_size=2)
+
+        assert len(payloads) == 2
+        # No past_key_values → main_payload is empty dict → falsy → padding stays None
+        assert payloads[1] is None
+
+    def test_with_branch_roles_uses_original_behavior(self):
+        """When branch_roles exist, follower payload comes from branch KV, not padding."""
+        req = SimpleNamespace(
+            request_id="req-1",
+            past_key_values=SimpleNamespace(key_cache=[torch.tensor([1.0])]),
+            kv_metadata={"seq_len": 10},
+            sampling_params=SimpleNamespace(
+                cfg_text_past_key_values=SimpleNamespace(key_cache=[torch.tensor([2.0])]),
+                cfg_text_kv_metadata={"source": "cfg_text"},
+            ),
+        )
+        payloads = OmniKVTransferManager._build_cfg_rank_local_payloads(req, cfg_size=2)
+
+        assert len(payloads) == 2
+        assert payloads[0] is not None
+        assert payloads[1] is not None
+        assert payloads[1]["sp.cfg_active_branch"] == "cfg_text"
+
+
 class TestAutoDetectTP:
     def test_auto_detect_when_config_defaults(self):
         """When config from_tp/to_tp == 1 (default), manager should auto-detect."""

vllm_omni/diffusion/distributed/group_coordinator.py

@@ -589,11 +589,7 @@ class GroupCoordinator:
         if dst is None:
             dst = self.group_next_rank
 
-        torch.distributed.send(
-            tensor,
-            self.ranks[dst],
-            group=(self.device_groups[self.rank_in_group % 2] if self.world_size == 2 else self.device_group),
-        )
+        torch.distributed.send(tensor, self.ranks[dst], group=self.device_group)
 
     def recv(self, size: torch.Size, dtype: torch.dtype, src: int | None = None) -> torch.Tensor:
         """Receives a tensor from the src rank."""
@@ -602,11 +598,7 @@ class GroupCoordinator:
             src = self.group_prev_rank
 
         tensor = torch.empty(size, dtype=dtype, device=self.device)
-        torch.distributed.recv(
-            tensor,
-            self.ranks[src],
-            (self.device_groups[(self.rank_in_group + 1) % 2] if self.world_size == 2 else self.device_group),
-        )
+        torch.distributed.recv(tensor, self.ranks[src], self.device_group)
         return tensor
 
     def destroy(self):
@@ -710,6 +702,22 @@ class PipelineGroupCoordinator(GroupCoordinator):
                 self.skip_device_group = skip_device_group
         assert self.skip_device_group is not None
 
+    def send(self, tensor: torch.Tensor, dst: int | None = None) -> None:
+        if dst is None:
+            dst = self.group_next_rank
+
+        group = self.device_groups[self.rank_in_group % 2] if self.world_size == 2 else self.device_group
+        torch.distributed.send(tensor, self.ranks[dst], group=group)
+
+    def recv(self, size: torch.Size, dtype: torch.dtype, src: int | None = None) -> torch.Tensor:
+        if src is None:
+            src = self.group_prev_rank
+
+        tensor = torch.empty(size, dtype=dtype, device=self.device)
+        group = self.device_groups[(self.rank_in_group + 1) % 2] if self.world_size == 2 else self.device_group
+        torch.distributed.recv(tensor, self.ranks[src], group)
+        return tensor
+
     def reset_buffer(self):
         self.recv_tasks_queue = []
         self.receiving_tasks = []

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py

@@ -1099,6 +1099,12 @@ class ImageKVCacheManager:
 
         neg_kv = (key.reshape(-1, num_kv_heads, head_dim), value.reshape(-1, num_kv_heads, head_dim))
         self._injected_ar_kv = [self._injected_ar_kv[0], neg_kv]
+        logger.debug(
+            "[AR KV Reuse] neg branch: reused shared_prefix_len=%d, neg_kv_total_len=%d (q_len_actual=%d)",
+            shared_prefix_len,
+            key.shape[1],
+            q_len_actual,
+        )
         return key, value
 
     def __call__(
@@ -2919,6 +2925,10 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
     ):
         ar_kv_data = model_kwargs.pop("ar_kv_data", None)
         if ar_kv_data is None:
+            logger.debug(
+                "[AR KV Reuse] cfg_rank=%s: no AR KV received, fallback to full recompute (reuse_len=0)",
+                cfg_rank,
+            )
             return input_ids, 0
 
         # 1. positive prefix len

vllm_omni/distributed/omni_connectors/kv_transfer_manager.py

@@ -856,8 +856,20 @@ class OmniKVTransferManager:
 
             payloads.append(local_payload)
 
+        padding = None
+        if not branch_roles and cfg_size == 2 and main_payload:
+            # AR didn't split KV by branch; cfg follower needs the same
+            # positive KV for shared-prefix reuse. Only safe for cfg_size == 2:
+            # for cfg_size > 2 every rank must get its own KV, which we can't
+            # assign here, so padding stays None and the runner raises.
+            # Shallow dict copy: shares PKV tensors with main_payload. Safe
+            # only because the send path (cfg_scatter -> pack -> to_gpu_tensor)
+            # reads tensors to serialize, never mutates in place. Any async/
+            # KV-pool path that mutates req.past_key_values mid-scatter must
+            # clone tensors here instead.
+            padding = dict(main_payload)
         while len(payloads) < cfg_size:
-            payloads.append(None)
+            payloads.append(padding)
 
         return payloads[:cfg_size]