From f84f0e72fd10f6e61a6db4284cc6a8f5dec0a629 Mon Sep 17 00:00:00 2001 From: Honghan Zhu Date: Wed, 15 Jul 2026 13:08:15 +0800 Subject: [PATCH] [Bug Fix] Hunyuan image cfg bugfix (#4752) Signed-off-by: zhuhh97 Signed-off-by: dph97 Co-authored-by: zhuhh97 Co-authored-by: Claude Opus 4.7 Co-authored-by: Gao Han Co-authored-by: Hongsheng Liu --- docker/Dockerfile.cuda | 2 +- docker/Dockerfile.rocm | 2 +- .../diffusion/test_diffusion_model_runner.py | 19 +++++ .../omni_connectors/test_tp_rank_aware.py | 70 +++++++++++++++++++ .../distributed/group_coordinator.py | 28 +++++--- .../hunyuan_image3_transformer.py | 10 +++ .../omni_connectors/kv_transfer_manager.py | 14 +++- 7 files changed, 132 insertions(+), 13 deletions(-) diff --git a/docker/Dockerfile.cuda b/docker/Dockerfile.cuda index 55d5859a8..4e9075f2a 100644 --- a/docker/Dockerfile.cuda +++ b/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 diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 7c45cb770..49c6f02aa 100644 --- a/docker/Dockerfile.rocm +++ b/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. diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py index 55584609d..5d4ef5d72 100644 --- a/tests/diffusion/test_diffusion_model_runner.py +++ b/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): diff --git a/tests/distributed/omni_connectors/test_tp_rank_aware.py b/tests/distributed/omni_connectors/test_tp_rank_aware.py index 8f8783db2..261de0b32 100644 --- a/tests/distributed/omni_connectors/test_tp_rank_aware.py +++ b/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.""" diff --git a/vllm_omni/diffusion/distributed/group_coordinator.py b/vllm_omni/diffusion/distributed/group_coordinator.py index 5c5b092ef..0a67fd941 100644 --- a/vllm_omni/diffusion/distributed/group_coordinator.py +++ b/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 = [] diff --git a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py index 0416fdcf1..f38d1ca7c 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py +++ b/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 diff --git a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py index 304489bbb..51e604336 100644 --- a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py +++ b/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]