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

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."""