[Feature] Bagel: Support tp+cfg parallel using mooncake transfer engine connector (#2705)

Signed-off-by: natureofnature <wzliu@connect.hku.hk>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/distributed/omni_connectors/test_tp_rank_aware.py

@@ -0,0 +1,716 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for rank-aware KV transfer (TP > 1) and heterogeneous TP support.
+
+Covers:
+- _build_rank_aware_send_keys / _build_rank_aware_recv_keys
+- _get_kv_source_ranks / _get_kv_target_ranks / get_kv_connector_key
+- update_sender_info storing base host/port
+- receive path constructing per-rank metadata for connector.get()
+- Mooncake connector _query_metadata_at and partial-metadata get() path
+"""
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+
+from vllm_omni.distributed.omni_connectors.kv_transfer_manager import (
+    KVCacheTransferData,
+    OmniKVCacheConfig,
+    OmniKVTransferManager,
+)
+from vllm_omni.distributed.omni_connectors.utils.initialization import (
+    KV_RANK_PORT_STRIDE,
+)
+from vllm_omni.distributed.omni_connectors.utils.kv_utils import (
+    KVTPTopology,
+    build_rank_aware_recv_keys,
+    build_rank_aware_send_keys,
+    get_kv_connector_key,
+    get_kv_source_ranks,
+    get_kv_target_ranks,
+    merge_received_rank_shards,
+    slice_received_rank_shard,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_manager(
+    from_tp: int = 1,
+    to_tp: int = 1,
+    local_rank: int = 0,
+    from_stage: str = "stage0",
+    to_stage: str = "stage1",
+    stage_id: str = "stage1",
+    need_recv: bool = True,
+    need_send: bool = False,
+    recv_timeout: float = 0.3,
+) -> OmniKVTransferManager:
+    """Build a manager with TP params injected, bypassing torch.distributed."""
+    config = OmniKVCacheConfig(
+        connector_config={"type": "mock"},
+        from_stage=from_stage,
+        to_stage=to_stage,
+        stage_id=stage_id,
+        need_recv_cache=need_recv,
+        need_send_cache=need_send,
+        recv_timeout=recv_timeout,
+        from_tp=from_tp,
+        to_tp=to_tp,
+    )
+    with (
+        patch("vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_local_tp_rank", return_value=local_rank),
+        patch(
+            "vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_tp_world_size",
+            return_value=max(from_tp, to_tp),
+        ),
+    ):
+        mgr = OmniKVTransferManager(config)
+    return mgr
+
+
+def _make_payload(head_values: list[float], request_id: str = "req-1") -> dict:
+    head_tensor = torch.tensor(head_values, dtype=torch.float32).view(1, len(head_values), 1).repeat(2, 1, 1)
+    return {
+        "request_id": request_id,
+        "layer_blocks": {
+            "key_cache": [head_tensor.clone()],
+            "value_cache": [(head_tensor + 100).clone()],
+        },
+        "block_ids": [0],
+        "metadata": {"seq_len": 2},
+    }
+
+
+def _make_transfer_data(head_values: list[float], request_id: str = "req-1") -> KVCacheTransferData:
+    payload = _make_payload(head_values, request_id=request_id)
+    return KVCacheTransferData(
+        request_id=request_id,
+        layer_blocks=payload["layer_blocks"],
+        block_ids=payload["block_ids"],
+        metadata=payload["metadata"],
+    )
+
+
+# ── Key format helper ────────────────────────────────────────────────
+
+
+class TestConnectorKeyFormat:
+    def test_key_format_matches_pr2677(self):
+        key = get_kv_connector_key("req-1", "stage0", 0, 1, 2)
+        assert key == "req-1_stage0_0_1_2"
+
+    def test_key_fields_are_positional(self):
+        key = get_kv_connector_key("r", "s", 5, 3, 7)
+        parts = key.split("_")
+        assert parts == ["r", "s", "5", "3", "7"]
+
+
+# ── Source / target rank mapping ─────────────────────────────────────
+
+
+class TestRankMapping:
+    """Verify get_kv_target_ranks and get_kv_source_ranks for various TP configs."""
+
+    def test_homogeneous_tp2_rank0(self):
+        topo = KVTPTopology(source_tp_size=2, target_tp_size=2, local_rank=0)
+        assert get_kv_target_ranks(topo) == [0]
+        assert get_kv_source_ranks(topo) == [0]
+
+    def test_homogeneous_tp2_rank1(self):
+        topo = KVTPTopology(source_tp_size=2, target_tp_size=2, local_rank=1)
+        assert get_kv_target_ranks(topo) == [1]
+        assert get_kv_source_ranks(topo) == [1]
+
+    def test_homogeneous_tp4_rank3(self):
+        topo = KVTPTopology(source_tp_size=4, target_tp_size=4, local_rank=3)
+        assert get_kv_target_ranks(topo) == [3]
+        assert get_kv_source_ranks(topo) == [3]
+
+    def test_sender_gt_receiver_tp4_to_tp2_rank0(self):
+        """Receiver rank 0 should receive from sender rank 0 and 1."""
+        topo = KVTPTopology(source_tp_size=4, target_tp_size=2, local_rank=0)
+        assert get_kv_source_ranks(topo) == [0, 1]
+
+    def test_sender_gt_receiver_tp4_to_tp2_rank1(self):
+        """Receiver rank 1 should receive from sender rank 2 and 3."""
+        topo = KVTPTopology(source_tp_size=4, target_tp_size=2, local_rank=1)
+        assert get_kv_source_ranks(topo) == [2, 3]
+
+    def test_sender_lt_receiver_tp2_to_tp4_rank0(self):
+        """Sender rank 0 should send to receiver ranks 0 and 1."""
+        topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=0)
+        assert get_kv_target_ranks(topo) == [0, 1]
+
+    def test_sender_lt_receiver_tp2_to_tp4_rank1(self):
+        topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=1)
+        assert get_kv_target_ranks(topo) == [2, 3]
+
+    def test_receiver_lt_sender_source_ranks(self):
+        """Receiver rank 0 with tp2_to_tp4 should source from rank 0 only."""
+        topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=0)
+        assert get_kv_source_ranks(topo) == [0]
+
+    def test_invalid_topology_raises(self):
+        topo = KVTPTopology(source_tp_size=3, target_tp_size=2, local_rank=0)
+        with pytest.raises(ValueError, match="divisible"):
+            get_kv_source_ranks(topo)
+
+
+# ── _build_rank_aware_recv_keys ──────────────────────────────────────
+
+
+class TestBuildRankAwareRecvKeys:
+    """Verify build_rank_aware_recv_keys returns (key, from_rank) tuples."""
+
+    def test_tp1_returns_legacy_key_with_none_rank(self):
+        topo = KVTPTopology(source_tp_size=1, target_tp_size=1, local_rank=0)
+        pairs = build_rank_aware_recv_keys("req-1", "stage0", "stage1", topo)
+        assert len(pairs) == 1
+        key, rank = pairs[0]
+        assert key == "omni_stage0_to_stage1_kv_cache_req-1"
+        assert rank is None
+
+    def test_homogeneous_tp2_rank0(self):
+        topo = KVTPTopology(source_tp_size=2, target_tp_size=2, local_rank=0)
+        pairs = build_rank_aware_recv_keys("req-1", "stage0", "stage1", topo)
+        assert len(pairs) == 1
+        key, rank = pairs[0]
+        assert key == "req-1_stage0_0_0_0"
+        assert rank == 0
+
+    def test_homogeneous_tp2_rank1(self):
+        topo = KVTPTopology(source_tp_size=2, target_tp_size=2, local_rank=1)
+        pairs = build_rank_aware_recv_keys("req-1", "stage0", "stage1", topo)
+        assert len(pairs) == 1
+        key, rank = pairs[0]
+        assert key == "req-1_stage0_0_1_1"
+        assert rank == 1
+
+    def test_heterogeneous_tp4_to_tp2_rank0_gets_two_keys(self):
+        """Receiver rank 0 with source_tp=4, target_tp=2 should get 2 keys."""
+        topo = KVTPTopology(source_tp_size=4, target_tp_size=2, local_rank=0)
+        pairs = build_rank_aware_recv_keys("req-1", "stage0", "stage1", topo)
+        assert len(pairs) == 2
+
+        keys = [k for k, _ in pairs]
+        ranks = [r for _, r in pairs]
+        assert keys == ["req-1_stage0_0_0_0", "req-1_stage0_0_1_0"]
+        assert ranks == [0, 1]
+
+    def test_heterogeneous_tp4_to_tp2_rank1_gets_two_keys(self):
+        topo = KVTPTopology(source_tp_size=4, target_tp_size=2, local_rank=1)
+        pairs = build_rank_aware_recv_keys("req-1", "stage0", "stage1", topo)
+        assert len(pairs) == 2
+
+        ranks = [r for _, r in pairs]
+        assert ranks == [2, 3]
+
+    def test_heterogeneous_tp2_to_tp4_rank2_gets_one_key(self):
+        """Receiver rank 2 with source_tp=2, target_tp=4 should get 1 key from sender rank 1."""
+        topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=2)
+        pairs = build_rank_aware_recv_keys("req-1", "stage0", "stage1", topo)
+        assert len(pairs) == 1
+        key, rank = pairs[0]
+        assert rank == 1
+        assert key == "req-1_stage0_0_1_2"
+
+
+# ── _build_rank_aware_send_keys ──────────────────────────────────────
+
+
+class TestBuildRankAwareSendKeys:
+    def test_tp1_returns_legacy_key(self):
+        topo = KVTPTopology(source_tp_size=1, target_tp_size=1, local_rank=0)
+        keys = build_rank_aware_send_keys("req-1", "stage0", "stage1", topo)
+        assert keys == ["omni_stage0_to_stage1_kv_cache_req-1"]
+
+    def test_homogeneous_tp2_rank0(self):
+        topo = KVTPTopology(source_tp_size=2, target_tp_size=2, local_rank=0)
+        keys = build_rank_aware_send_keys("req-1", "stage0", "stage1", topo)
+        assert keys == ["req-1_stage0_0_0_0"]
+
+    def test_sender_lt_receiver_tp2_to_tp4_rank0_sends_two_keys(self):
+        topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=0)
+        keys = build_rank_aware_send_keys("req-1", "stage0", "stage1", topo)
+        assert len(keys) == 2
+        assert keys == ["req-1_stage0_0_0_0", "req-1_stage0_0_0_1"]
+
+
+# ── update_sender_info stores base host/port ─────────────────────────
+
+
+class TestUpdateSenderInfoBase:
+    def test_stores_base_host_and_port(self):
+        mgr = _make_manager(from_tp=2, to_tp=2, local_rank=0)
+        mgr.update_sender_info({"host": "10.0.0.1", "zmq_port": 50151})
+
+        assert mgr._sender_base_host == "10.0.0.1"
+        assert mgr._sender_base_zmq_port == 50151
+
+    def test_rank1_adjusts_default_port_but_preserves_base(self):
+        mgr = _make_manager(from_tp=2, to_tp=2, local_rank=1)
+        mgr.update_sender_info({"host": "10.0.0.1", "zmq_port": 50151})
+
+        assert mgr._sender_base_host == "10.0.0.1"
+        assert mgr._sender_base_zmq_port == 50151
+        expected_adjusted = 50151 + 1 * KV_RANK_PORT_STRIDE
+        assert mgr.config.connector_config["sender_zmq_port"] == expected_adjusted
+
+    def test_nested_sender_info_resolves_correctly(self):
+        """Nested sender_info keyed by integer stage id should resolve
+        using recv_stages (engine_input_source → recv_from)."""
+        config = OmniKVCacheConfig(
+            connector_config={"type": "mock"},
+            stage_id=2,
+            engine_input_source=[1],
+            need_recv_cache=True,
+            from_tp=2,
+            to_tp=2,
+        )
+        with (
+            patch("vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_local_tp_rank", return_value=0),
+            patch("vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_tp_world_size", return_value=2),
+        ):
+            mgr = OmniKVTransferManager(config)
+        mgr.update_sender_info(
+            {
+                0: {"host": "10.0.0.1", "zmq_port": 50151},
+                1: {"host": "10.0.0.2", "zmq_port": 50152},
+            }
+        )
+        assert mgr._sender_base_host == "10.0.0.2"
+        assert mgr._sender_base_zmq_port == 50152
+
+
+# ── receive path constructs per-rank metadata ────────────────────────
+
+
+class TestReceiveConstructsMetadata:
+    """Verify that receive_kv_cache_for_request passes metadata with
+    correct (host, port) to connector.get() for heterogeneous TP."""
+
+    def test_tp1_no_metadata_passed(self):
+        """TP=1: connector.get() should be called WITHOUT metadata."""
+        mgr = _make_manager(from_tp=1, to_tp=1, local_rank=0, recv_timeout=0.05)
+        mgr.update_sender_info({"host": "10.0.0.1", "zmq_port": 50151})
+
+        calls = []
+
+        class _Connector:
+            def get(self, from_stage, to_stage, get_key, metadata=None):
+                calls.append({"key": get_key, "metadata": metadata})
+                return None
+
+        mgr._connector = _Connector()
+        mgr.receive_kv_cache_for_request("req-1")
+
+        assert len(calls) > 0
+        assert calls[0]["metadata"] is None
+
+    def test_homogeneous_tp2_rank0_passes_metadata(self):
+        """TP=2 rank 0: metadata should point to sender rank 0's port."""
+        mgr = _make_manager(from_tp=2, to_tp=2, local_rank=0, recv_timeout=0.05)
+        mgr.update_sender_info({"host": "10.0.0.1", "zmq_port": 50151})
+
+        calls = []
+
+        class _Connector:
+            def get(self, from_stage, to_stage, get_key, metadata=None):
+                calls.append({"key": get_key, "metadata": metadata})
+                return None
+
+        mgr._connector = _Connector()
+        mgr.receive_kv_cache_for_request("req-1")
+
+        assert len(calls) > 0
+        meta = calls[0]["metadata"]
+        assert meta is not None
+        assert meta["source_host"] == "10.0.0.1"
+        assert meta["source_port"] == 50151 + 0 * KV_RANK_PORT_STRIDE
+
+    def test_homogeneous_tp2_rank1_passes_metadata_with_offset(self):
+        mgr = _make_manager(from_tp=2, to_tp=2, local_rank=1, recv_timeout=0.05)
+        mgr.update_sender_info({"host": "10.0.0.1", "zmq_port": 50151})
+
+        calls = []
+
+        class _Connector:
+            def get(self, from_stage, to_stage, get_key, metadata=None):
+                calls.append({"key": get_key, "metadata": metadata})
+                return None
+
+        mgr._connector = _Connector()
+        mgr.receive_kv_cache_for_request("req-1")
+
+        meta = calls[0]["metadata"]
+        assert meta["source_port"] == 50151 + 1 * KV_RANK_PORT_STRIDE
+
+    def test_heterogeneous_tp4_to_tp2_rank0_multiple_metadata(self):
+        """Receiver rank 0 with source_tp=4, target_tp=2 should call get() with
+        two different metadata entries for sender ranks 0 and 1."""
+        mgr = _make_manager(from_tp=4, to_tp=2, local_rank=0, recv_timeout=0.05)
+        mgr.update_sender_info({"host": "10.0.0.1", "zmq_port": 50151})
+
+        calls = []
+
+        class _Connector:
+            def get(self, from_stage, to_stage, get_key, metadata=None):
+                calls.append({"key": get_key, "metadata": metadata})
+                return None
+
+        mgr._connector = _Connector()
+        mgr.receive_kv_cache_for_request("req-1")
+
+        seen_ports = set()
+        for c in calls:
+            if c["metadata"]:
+                seen_ports.add(c["metadata"]["source_port"])
+        expected_ports = {
+            50151 + 0 * KV_RANK_PORT_STRIDE,
+            50151 + 1 * KV_RANK_PORT_STRIDE,
+        }
+        assert expected_ports.issubset(seen_ports)
+
+
+# ── Mooncake connector _query_metadata_at ────────────────────────────
+
+
+class TestMooncakeQueryMetadataAt:
+    """Test the connector's _query_metadata_at method and partial-metadata
+    path in get() without requiring real RDMA/Mooncake."""
+
+    def test_query_metadata_at_returns_full_metadata(self):
+        """Mock the ZMQ interaction to verify _query_metadata_at returns
+        complete metadata including data_size."""
+
+        try:
+            from vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector import (
+                MooncakeTransferEngineConnector,
+                QueryResponse,
+            )
+        except ImportError:
+            pytest.skip("Mooncake not available")
+
+        import msgspec
+
+        connector = MagicMock(spec=MooncakeTransferEngineConnector)
+        connector._get_req_socket = MagicMock()
+
+        mock_socket = MagicMock()
+        resp = QueryResponse(request_id="test_key@s0_s1", data_size=4096, is_fast_path=True)
+        mock_socket.recv.return_value = msgspec.msgpack.encode(resp)
+        connector._get_req_socket.return_value = mock_socket
+
+        result = MooncakeTransferEngineConnector._query_metadata_at(
+            connector,
+            "test_key@s0_s1",
+            "10.0.0.1",
+            50151,
+        )
+
+        assert result is not None
+        assert result["source_host"] == "10.0.0.1"
+        assert result["source_port"] == 50151
+        assert result["data_size"] == 4096
+        assert result["is_fast_path"] is True
+
+    def test_query_metadata_at_returns_none_on_not_found(self):
+        try:
+            from vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector import (
+                INFO_NOT_FOUND,
+                MooncakeTransferEngineConnector,
+            )
+        except ImportError:
+            pytest.skip("Mooncake not available")
+
+        connector = MagicMock(spec=MooncakeTransferEngineConnector)
+        mock_socket = MagicMock()
+        mock_socket.recv.return_value = INFO_NOT_FOUND
+        connector._get_req_socket.return_value = mock_socket
+
+        result = MooncakeTransferEngineConnector._query_metadata_at(
+            connector,
+            "test_key@s0_s1",
+            "10.0.0.1",
+            50151,
+        )
+        assert result is None
+
+
+# ── Merge / slice hooks ──────────────────────────────────────────────
+
+
+class TestMergeSliceHooks:
+    def test_single_shard_passes_through(self):
+        payload = {"layer_blocks": {"key_cache": [1]}}
+        assert merge_received_rank_shards([payload]) == payload
+
+    def test_default_merger_concats_head_dim(self):
+        p0 = _make_payload([0.0])
+        p1 = _make_payload([1.0])
+        result = merge_received_rank_shards([p0, p1])
+        key_cache = result["layer_blocks"]["key_cache"][0]
+        value_cache = result["layer_blocks"]["value_cache"][0]
+        assert key_cache.shape == (2, 2, 1)
+        assert value_cache.shape == (2, 2, 1)
+        assert torch.equal(key_cache[:, :, 0], torch.tensor([[0.0, 1.0], [0.0, 1.0]]))
+        assert torch.equal(value_cache[:, :, 0], torch.tensor([[100.0, 101.0], [100.0, 101.0]]))
+
+    def test_custom_merger_hook_called(self):
+        merged = {"merged": True}
+        assert merge_received_rank_shards([{}, {}], merger=lambda payloads: merged) == merged
+
+    def test_slicer_hook_called(self):
+        topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=0)
+        sliced = {"sliced": True}
+        assert slice_received_rank_shard({"full": True}, topo, slicer=lambda payload: sliced) == sliced
+
+    def test_default_slicer_extracts_rank_local_heads(self):
+        topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=1)
+        payload = _make_payload([0.0, 1.0])
+        result = slice_received_rank_shard(payload, topo)
+        key_cache = result["layer_blocks"]["key_cache"][0]
+        value_cache = result["layer_blocks"]["value_cache"][0]
+        assert key_cache.shape == (2, 1, 1)
+        assert value_cache.shape == (2, 1, 1)
+        assert torch.equal(key_cache[:, :, 0], torch.tensor([[1.0], [1.0]]))
+        assert torch.equal(value_cache[:, :, 0], torch.tensor([[101.0], [101.0]]))
+
+    def test_presliced_payload_is_not_sliced_twice(self):
+        topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=1)
+        payload = _make_payload([1.0])
+        payload["metadata"]["tp_head_slice"] = {"applied": True, "target_rank": 1}
+        result = slice_received_rank_shard(payload, topo)
+        assert result is payload
+
+    def test_round_trip_merge_from_tp4_to_tp2(self):
+        topo = KVTPTopology(source_tp_size=4, target_tp_size=2, local_rank=1)
+        source_ranks = get_kv_source_ranks(topo)
+        payloads = [_make_payload([float(rank)]) for rank in source_ranks]
+        result = merge_received_rank_shards(payloads)
+        key_cache = result["layer_blocks"]["key_cache"][0]
+        assert torch.equal(key_cache[:, :, 0], torch.tensor([[2.0, 3.0], [2.0, 3.0]]))
+
+    def test_round_trip_slice_from_tp2_to_tp4(self):
+        topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=3)
+        payload = _make_payload([2.0, 3.0])
+        result = slice_received_rank_shard(payload, topo)
+        key_cache = result["layer_blocks"]["key_cache"][0]
+        assert torch.equal(key_cache[:, :, 0], torch.tensor([[3.0], [3.0]]))
+
+
+class TestSenderSideSlicing:
+    def test_transfer_slices_before_sending_to_multiple_targets(self):
+        mgr = _make_manager(
+            from_tp=2,
+            to_tp=4,
+            local_rank=0,
+            need_send=True,
+            need_recv=False,
+        )
+        sent_payloads = []
+
+        class _Connector:
+            supports_raw_data = False
+
+            def put(self, from_stage, to_stage, put_key, data):
+                sent_payloads.append((put_key, KVCacheTransferData.from_bytes(data)))
+                return True, len(data), {}
+
+        mgr._connector = _Connector()
+        mgr._transfer_kv_cache(_make_transfer_data([0.0, 1.0]), "req-1")
+
+        assert [key for key, _ in sent_payloads] == ["req-1_stage0_0_0_0", "req-1_stage0_0_0_1"]
+        assert sent_payloads[0][1]["layer_blocks"]["key_cache"][0].shape == (2, 1, 1)
+        assert sent_payloads[1][1]["layer_blocks"]["key_cache"][0].shape == (2, 1, 1)
+        assert torch.equal(
+            sent_payloads[0][1]["layer_blocks"]["key_cache"][0][:, :, 0],
+            torch.tensor([[0.0], [0.0]]),
+        )
+        assert torch.equal(
+            sent_payloads[1][1]["layer_blocks"]["key_cache"][0][:, :, 0],
+            torch.tensor([[1.0], [1.0]]),
+        )
+        assert sent_payloads[0][1]["metadata"]["tp_head_slice"]["target_rank"] == 0
+        assert sent_payloads[1][1]["metadata"]["tp_head_slice"]["target_rank"] == 1
+
+
+class _MockBroadcastGroup:
+    def __init__(self, world_size: int, rank_in_group: int, broadcast_value=None, recv_value=None):
+        self.world_size = world_size
+        self.rank_in_group = rank_in_group
+        self.broadcast_value = broadcast_value
+        self.recv_value = recv_value
+        self.broadcast_calls = []
+        self.send_calls = []
+        self.recv_calls = []
+        self.shm_broadcaster = None
+
+    def broadcast_object(self, obj=None, src: int = 0):
+        self.broadcast_calls.append((obj, src))
+        return self.broadcast_value if self.broadcast_value is not None else obj
+
+    def send_object(self, obj, dst: int):
+        self.send_calls.append((dst, obj))
+
+    def recv_object(self, src: int):
+        self.recv_calls.append(src)
+        return self.recv_value
+
+
+class TestDistributedReceive:
+    def test_tp_cfg_leader_receives_then_sends_branch_local_payloads(self):
+        mgr = _make_manager(from_tp=2, to_tp=4, local_rank=0)
+        req = SimpleNamespace(request_id="req-1", sampling_params=SimpleNamespace())
+        world_group = _MockBroadcastGroup(world_size=4, rank_in_group=2)
+        cfg_group = _MockBroadcastGroup(world_size=3, rank_in_group=0)
+
+        def _receive(req_obj, cfg_func, target_device):
+            req_obj.past_key_values = SimpleNamespace(key_cache=[torch.tensor([1.0])])
+            req_obj.kv_metadata = {"source": "leader"}
+            req_obj.sampling_params.past_key_values = req_obj.past_key_values
+            req_obj.sampling_params.kv_metadata = req_obj.kv_metadata
+            req_obj.sampling_params.cfg_text_past_key_values = SimpleNamespace(key_cache=[torch.tensor([2.0])])
+            req_obj.sampling_params.cfg_text_kv_metadata = {"source": "cfg_text"}
+            req_obj.sampling_params.cfg_img_past_key_values = SimpleNamespace(key_cache=[torch.tensor([3.0])])
+            req_obj.sampling_params.cfg_img_kv_metadata = {"source": "cfg_img"}
+            return True
+
+        mgr.receive_multi_kv_cache = MagicMock(side_effect=_receive)
+        with (
+            patch("vllm_omni.diffusion.distributed.parallel_state.get_world_group", return_value=world_group),
+            patch(
+                "vllm_omni.diffusion.distributed.parallel_state.get_classifier_free_guidance_world_size",
+                return_value=3,
+            ),
+            patch(
+                "vllm_omni.diffusion.distributed.parallel_state.get_classifier_free_guidance_rank",
+                return_value=0,
+            ),
+            patch("vllm_omni.diffusion.distributed.parallel_state.get_cfg_group", return_value=cfg_group),
+        ):
+            assert mgr.receive_multi_kv_cache_distributed(req) is True
+
+        mgr.receive_multi_kv_cache.assert_called_once()
+        assert mgr.receive_multi_kv_cache.call_args.args[2] == torch.device("cpu")
+        assert req.kv_metadata == {"source": "leader"}
+        assert cfg_group.broadcast_calls == []
+        assert [dst for dst, _ in cfg_group.send_calls] == [1, 2]
+        rank1_payload = cfg_group.send_calls[0][1]
+        rank2_payload = cfg_group.send_calls[1][1]
+        assert torch.equal(rank1_payload["past_key_values"].key_cache[0], torch.tensor([1.0]))
+        assert torch.equal(rank2_payload["past_key_values"].key_cache[0], torch.tensor([1.0]))
+        assert rank1_payload["sp.cfg_active_branch"] == "cfg_text"
+        assert rank2_payload["sp.cfg_active_branch"] == "cfg_img"
+        assert rank1_payload["sp.cfg_branch_roles"] == ["cfg_text", "cfg_img"]
+        assert rank2_payload["sp.cfg_branch_roles"] == ["cfg_text", "cfg_img"]
+        assert "sp.cfg_branch_past_key_values" in rank1_payload
+        assert "sp.cfg_branch_past_key_values" in rank2_payload
+        assert list(rank1_payload["sp.cfg_branch_past_key_values"].keys()) == ["cfg_text"]
+        assert list(rank2_payload["sp.cfg_branch_past_key_values"].keys()) == ["cfg_img"]
+        assert "sp.cfg_text_past_key_values" in rank1_payload
+        assert "sp.cfg_img_past_key_values" not in rank1_payload
+        assert "sp.cfg_img_past_key_values" in rank2_payload
+        assert "sp.cfg_text_past_key_values" not in rank2_payload
+
+    def test_tp_cfg_follower_receives_local_payload_without_receiving(self):
+        mgr = _make_manager(from_tp=2, to_tp=4, local_rank=1)
+        req = SimpleNamespace(request_id="req-1", sampling_params=SimpleNamespace())
+        world_group = _MockBroadcastGroup(world_size=4, rank_in_group=3)
+        cfg_payload = {
+            "past_key_values": SimpleNamespace(key_cache=[torch.tensor([1.0])]),
+            "kv_metadata": {"source": "main"},
+            "sp.past_key_values": SimpleNamespace(key_cache=[torch.tensor([1.0])]),
+            "sp.kv_metadata": {"source": "main"},
+            "sp.cfg_active_branch": "cfg_text",
+            "sp.cfg_branch_roles": ["cfg_text", "cfg_img"],
+            "sp.cfg_branch_past_key_values": {
+                "cfg_text": SimpleNamespace(key_cache=[torch.tensor([2.0])]),
+            },
+            "sp.cfg_branch_kv_metadata": {"cfg_text": {"source": "cfg-text"}},
+            "sp.cfg_text_past_key_values": SimpleNamespace(key_cache=[torch.tensor([2.0])]),
+        }
+        cfg_group = _MockBroadcastGroup(world_size=2, rank_in_group=1, recv_value=cfg_payload)
+
+        mgr.receive_multi_kv_cache = MagicMock(return_value=True)
+        with (
+            patch("vllm_omni.diffusion.distributed.parallel_state.get_world_group", return_value=world_group),
+            patch(
+                "vllm_omni.diffusion.distributed.parallel_state.get_classifier_free_guidance_world_size",
+                return_value=2,
+            ),
+            patch(
+                "vllm_omni.diffusion.distributed.parallel_state.get_classifier_free_guidance_rank",
+                return_value=1,
+            ),
+            patch("vllm_omni.diffusion.distributed.parallel_state.get_cfg_group", return_value=cfg_group),
+        ):
+            assert mgr.receive_multi_kv_cache_distributed(req) is True
+
+        mgr.receive_multi_kv_cache.assert_not_called()
+        assert req.kv_metadata == {"source": "main"}
+        assert torch.equal(req.past_key_values.key_cache[0], torch.tensor([1.0]))
+        assert torch.equal(req.sampling_params.past_key_values.key_cache[0], torch.tensor([1.0]))
+        assert req.sampling_params.cfg_active_branch == "cfg_text"
+        assert req.sampling_params.cfg_branch_roles == ["cfg_text", "cfg_img"]
+        assert torch.equal(
+            req.sampling_params.cfg_branch_past_key_values["cfg_text"].key_cache[0],
+            torch.tensor([2.0]),
+        )
+        assert req.sampling_params.cfg_branch_kv_metadata == {"cfg_text": {"source": "cfg-text"}}
+        assert torch.equal(req.sampling_params.cfg_text_past_key_values.key_cache[0], torch.tensor([2.0]))
+        assert cfg_group.broadcast_calls == []
+        assert cfg_group.recv_calls == [0]
+
+    def test_tp_without_cfg_keeps_independent_receive_path(self):
+        mgr = _make_manager(from_tp=2, to_tp=2, local_rank=1)
+        req = SimpleNamespace(request_id="req-1", sampling_params=SimpleNamespace())
+        world_group = _MockBroadcastGroup(world_size=2, rank_in_group=1)
+        mgr.receive_multi_kv_cache = MagicMock(return_value=True)
+
+        with patch("vllm_omni.diffusion.distributed.parallel_state.get_world_group", return_value=world_group):
+            assert mgr.receive_multi_kv_cache_distributed(req, target_device=torch.device("cpu")) is True
+
+        mgr.receive_multi_kv_cache.assert_called_once_with(req, None, torch.device("cpu"))
+
+
+# ── TP auto-detect ───────────────────────────────────────────────────
+
+
+class TestAutoDetectTP:
+    def test_auto_detect_when_config_defaults(self):
+        """When config from_tp/to_tp == 1 (default), manager should auto-detect."""
+        config = OmniKVCacheConfig(
+            connector_config={"type": "mock"},
+            from_stage="s0",
+            stage_id="s1",
+            need_recv_cache=True,
+        )
+        with (
+            patch("vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_local_tp_rank", return_value=0),
+            patch("vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_tp_world_size", return_value=4),
+        ):
+            mgr = OmniKVTransferManager(config)
+        assert mgr._tp_topo.source_tp_size == 4
+        assert mgr._tp_topo.target_tp_size == 4
+
+    def test_explicit_tp_overrides_auto_detect(self):
+        config = OmniKVCacheConfig(
+            connector_config={"type": "mock"},
+            from_stage="s0",
+            stage_id="s1",
+            need_recv_cache=True,
+            from_tp=2,
+            to_tp=4,
+        )
+        with (
+            patch("vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_local_tp_rank", return_value=0),
+            patch("vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_tp_world_size", return_value=8),
+        ):
+            mgr = OmniKVTransferManager(config)
+        assert mgr._tp_topo.source_tp_size == 2
+        assert mgr._tp_topo.target_tp_size == 4

tests/engine/test_async_omni_engine_stage_init.py

@@ -183,6 +183,7 @@ def test_launch_llm_stage_passes_stage_init_timeout_to_complete_stage_handshake(
     engine.model = "dummy-model"
     engine.single_stage_mode = False
     engine._omni_master_server = None
+    engine.stage_configs = []
 
     metadata = types.SimpleNamespace(stage_id=0, runtime_cfg={"devices": "0"})
     fake_vllm_config = types.SimpleNamespace()
@@ -238,6 +239,7 @@ def test_launch_llm_stage_releases_launch_lock_before_complete_stage_handshake(m
     engine.model = "dummy-model"
     engine.single_stage_mode = False
     engine._omni_master_server = None
+    engine.stage_configs = []
 
     fake_vllm_config = types.SimpleNamespace()
     fake_addresses = types.SimpleNamespace()
@@ -378,3 +380,70 @@ def test_attach_llm_stage_uses_omni_input_preprocessor(monkeypatch):
     assert input_processor is not None
     assert isinstance(input_processor.input_preprocessor, DummyOmniInputPreprocessor)
     assert input_processor.input_preprocessor.renderer is input_processor.renderer
+
+
+def test_inject_kv_stage_info_infers_sender_tp_topology():
+    from vllm_omni.engine.stage_init_utils import inject_kv_stage_info
+
+    stage0 = types.SimpleNamespace(
+        stage_id=0,
+        engine_args={
+            "tensor_parallel_size": 4,
+            "omni_kv_config": {
+                "need_send_cache": True,
+                "omni_from_stage": "0",
+                "omni_to_stage": "1",
+            },
+        },
+        engine_input_source=[],
+    )
+    stage1 = types.SimpleNamespace(
+        stage_id=1,
+        engine_args={
+            "parallel_config": {
+                "tensor_parallel_size": 2,
+                "cfg_parallel_size": 1,
+            },
+            "omni_kv_config": {"need_recv_cache": True},
+        },
+        engine_input_source=[0],
+    )
+
+    inject_kv_stage_info(stage0, 0, [stage0, stage1])
+
+    assert stage0.engine_args["omni_kv_config"]["stage_id"] == 0
+    assert stage0.engine_args["omni_kv_config"]["rank_mapping"] == {"from_tp": 4, "to_tp": 2}
+
+
+def test_inject_kv_stage_info_infers_receiver_tp_topology():
+    from vllm_omni.engine.stage_init_utils import inject_kv_stage_info
+
+    stage0 = types.SimpleNamespace(
+        stage_id=0,
+        engine_args={
+            "tensor_parallel_size": 4,
+            "omni_kv_config": {"need_send_cache": True},
+        },
+        engine_input_source=[],
+    )
+    stage1 = types.SimpleNamespace(
+        stage_id=1,
+        engine_args={
+            "parallel_config": {
+                "tensor_parallel_size": 2,
+                "cfg_parallel_size": 1,
+            },
+            "omni_kv_config": {
+                "need_recv_cache": True,
+                "omni_from_stage": "0",
+                "omni_to_stage": "1",
+            },
+        },
+        engine_input_source=[0],
+    )
+
+    inject_kv_stage_info(stage1, 1, [stage0, stage1])
+
+    assert stage1.engine_args["omni_kv_config"]["stage_id"] == 1
+    assert stage1.engine_args["omni_kv_config"]["engine_input_source"] == [0]
+    assert stage1.engine_args["omni_kv_config"]["rank_mapping"] == {"from_tp": 4, "to_tp": 2}

tests/engine/test_single_stage_mode.py

@@ -1555,6 +1555,7 @@ class TestLaunchLlmStageSingleStageMode:
         engine.single_stage_mode = True
         engine._single_stage_id_filter = 0
         engine._llm_stage_launch_lock = threading.Lock()
+        engine.stage_configs = []
         mock_oms = mocker.Mock(spec=OmniMasterServer)
         mock_oms.address = "127.0.0.1"
         mock_oms.port = 25000
@@ -1629,6 +1630,7 @@ class TestLaunchLlmStageSingleStageMode:
         engine.single_stage_mode = False
         engine._omni_master_server = None
         engine._llm_stage_launch_lock = threading.Lock()
+        engine.stage_configs = []
 
         fake_vllm_config = mocker.Mock()
         fake_executor_cls = mocker.Mock()

vllm_omni/diffusion/distributed/group_coordinator.py

@@ -104,6 +104,7 @@ class GroupCoordinator:
         self.local_rank = local_rank
         self.device_group = None
         self.cpu_group = None
+        self.shm_broadcaster = None
 
         for ranks in group_ranks:
             device_group = torch.distributed.new_group(ranks, backend=torch_distributed_backend)
@@ -316,7 +317,7 @@ class GroupCoordinator:
 
         assert dst < self.world_size, f"Invalid dst rank ({dst})"
 
-        assert dst != self.rank, "Invalid destination rank. Destination rank is the same as the current rank."
+        assert dst != self.rank_in_group, "Invalid destination rank. Destination rank is the same as the current rank."
 
         # Serialize object to tensor and get the size as well
         object_tensor = torch.frombuffer(pickle.dumps(obj), dtype=torch.uint8)
@@ -338,7 +339,7 @@ class GroupCoordinator:
 
         assert src < self.world_size, f"Invalid src rank ({src})"
 
-        assert src != self.rank, "Invalid source rank. Source rank is the same as the current rank."
+        assert src != self.rank_in_group, "Invalid source rank. Source rank is the same as the current rank."
 
         size_tensor = torch.empty(1, dtype=torch.long, device="cpu")
 

vllm_omni/diffusion/models/bagel/pipeline_bagel.py

@@ -365,28 +365,52 @@ class BagelPipeline(nn.Module, DiffusionPipelineProfilerMixin):
             if req.sampling_params.kv_metadata and "image_shape" in req.sampling_params.kv_metadata:
                 image_shape = tuple(req.sampling_params.kv_metadata["image_shape"])
 
-            cfg_text_kv = getattr(req.sampling_params, "cfg_text_past_key_values", None)
+            branch_kvs = getattr(req.sampling_params, "cfg_branch_past_key_values", None) or {}
+            branch_metadata = getattr(req.sampling_params, "cfg_branch_kv_metadata", None) or {}
+            active_branch = getattr(req.sampling_params, "cfg_active_branch", None)
+            branch_roles = getattr(req.sampling_params, "cfg_branch_roles", None) or list(branch_kvs.keys())
+
+            cfg_text_kv = getattr(req.sampling_params, "cfg_text_past_key_values", None) or branch_kvs.get("cfg_text")
+            cfg_text_metadata = getattr(req.sampling_params, "cfg_text_kv_metadata", None) or branch_metadata.get(
+                "cfg_text"
+            )
+            cfg_img_kv = getattr(req.sampling_params, "cfg_img_past_key_values", None) or branch_kvs.get("cfg_img")
+            cfg_img_metadata = getattr(req.sampling_params, "cfg_img_kv_metadata", None) or branch_metadata.get(
+                "cfg_img"
+            )
+
+            cfg_parallel_contract = (
+                active_branch is not None or bool(branch_roles) or cfg_text_kv is not None or cfg_img_kv is not None
+            )
+            if cfg_parallel_contract:
+                logger.info(
+                    "CFG enabled with injected branch KV context roles=%s active=%s",
+                    branch_roles,
+                    active_branch,
+                )
+
             if cfg_text_kv is not None:
-                logger.info("CFG enabled with multi-KV: using injected cfg_text KV Cache")
                 cfg_text_seq_len = cfg_text_kv.key_cache[0].shape[0]
                 cfg_text_context["past_key_values"] = cfg_text_kv
                 cfg_text_context["kv_lens"] = [cfg_text_seq_len]
-                cfg_text_metadata = getattr(req.sampling_params, "cfg_text_kv_metadata", None)
                 if cfg_text_metadata and "ropes" in cfg_text_metadata:
                     cfg_text_context["ropes"] = cfg_text_metadata["ropes"]
                 else:
                     cfg_text_context["ropes"] = [cfg_text_seq_len]
 
-                cfg_img_kv = getattr(req.sampling_params, "cfg_img_past_key_values", None) or injected_kv
+            if cfg_img_kv is None and cfg_text_kv is not None:
+                cfg_img_kv = injected_kv
+
+            if cfg_img_kv is not None:
                 cfg_img_seq_len = cfg_img_kv.key_cache[0].shape[0]
                 cfg_img_context["past_key_values"] = cfg_img_kv
                 cfg_img_context["kv_lens"] = [cfg_img_seq_len]
-                cfg_img_metadata = getattr(req.sampling_params, "cfg_img_kv_metadata", None)
                 if cfg_img_metadata and "ropes" in cfg_img_metadata:
                     cfg_img_context["ropes"] = cfg_img_metadata["ropes"]
                 else:
                     cfg_img_context["ropes"] = [cfg_img_seq_len]
-            else:
+
+            if not cfg_parallel_contract:
                 logger.warning("CFG is disabled: only single KV cache available")
                 gen_params = BagelGenParams(
                     num_timesteps=gen_params.num_timesteps,

vllm_omni/distributed/omni_connectors/kv_transfer_manager.py

@@ -14,8 +14,20 @@ from vllm.logger import init_logger
 
 from .factory import OmniConnectorFactory
 from .utils.config import ConnectorSpec
-from .utils.initialization import KV_TRANSFER_PORT_OFFSET
-from .utils.kv_utils import normalize_layer_kv
+from .utils.initialization import KV_RANK_PORT_STRIDE
+from .utils.kv_utils import (
+    KVTPTopology,
+    build_rank_aware_recv_keys,
+    build_rank_aware_send_keys,
+    get_kv_target_ranks,
+    get_local_tp_rank,
+    get_tp_world_size,
+    kv_zmq_port,
+    merge_received_rank_shards,
+    normalize_layer_kv,
+    slice_layer_blocks,
+    slice_received_rank_shard,
+)
 
 logger = init_logger(__name__)
 
@@ -57,6 +69,8 @@ class OmniKVCacheConfig:
     need_recv_cache: bool = False
     need_send_cache: bool = False
     recv_timeout: float = 30.0
+    from_tp: int = 1
+    to_tp: int = 1
 
 
 @dataclass
@@ -72,82 +86,44 @@ class KVCacheTransferData:
         """Convert to dictionary for serialization."""
         return asdict(self)
 
-    def to_bytes(self) -> bytes:
-        """Convert to compact binary format for fast transfer."""
+    def _build_tensors_desc(self, *, cpu: bool) -> tuple[list[dict[str, Any]], list, int, torch.device | None]:
+        """Iterate layer blocks and build tensor descriptors + data chunks.
+
+        Returns ``(tensors_desc, chunks, total_bytes, device)``.
+        *chunks* contains ``bytes`` when *cpu* is True, flat uint8 GPU tensors otherwise.
+        """
         tensors_desc: list[dict[str, Any]] = []
-        tensor_bufs: list[bytes] = []
-        data_offset = 0
-
-        for cache_name in ("key_cache", "value_cache"):
-            cache_list = self.layer_blocks.get(cache_name, [])
-            for layer_idx, tensor in enumerate(cache_list):
-                if tensor is None:
-                    tensors_desc.append({"n": f"{cache_name}_{layer_idx}", "x": True})
-                    continue
-
-                t = tensor.detach().cpu().contiguous()
-                dtype_str = str(t.dtype).removeprefix("torch.")
-                raw = t.view(torch.uint8).numpy().tobytes()
-                tensors_desc.append(
-                    {
-                        "n": f"{cache_name}_{layer_idx}",
-                        "i": layer_idx,
-                        "d": dtype_str,
-                        "s": list(t.shape),
-                        "o": data_offset,
-                        "b": len(raw),
-                    }
-                )
-                tensor_bufs.append(raw)
-                data_offset += len(raw)
-
-        header = json.dumps(
-            {
-                "rid": self.request_id,
-                "bids": self.block_ids,
-                "meta": self.metadata,
-                "td": tensors_desc,
-                "nl": len(self.layer_blocks.get("key_cache", [])),
-            },
-            separators=(",", ":"),
-        ).encode("utf-8")
-        return b"".join([struct.pack(">I", len(header)), header] + tensor_bufs)
-
-    def to_gpu_tensor(self) -> torch.Tensor:
-        """Convert to a packed GPU tensor for raw-data connectors."""
-        tensors_desc: list[dict[str, Any]] = []
-        gpu_tensors: list[torch.Tensor] = []
+        chunks: list = []
         data_offset = 0
         device = None
 
         for cache_name in ("key_cache", "value_cache"):
-            cache_list = self.layer_blocks.get(cache_name, [])
-            for layer_idx, tensor in enumerate(cache_list):
+            for layer_idx, tensor in enumerate(self.layer_blocks.get(cache_name, [])):
                 if tensor is None:
                     tensors_desc.append({"n": f"{cache_name}_{layer_idx}", "x": True})
                     continue
-
                 t = tensor.detach().contiguous()
-                if device is None and t.is_cuda:
+                if cpu:
+                    t = t.cpu()
+                elif device is None and t.is_cuda:
                     device = t.device
-                dtype_str = str(t.dtype).removeprefix("torch.")
                 nbytes = t.numel() * t.element_size()
                 tensors_desc.append(
                     {
                         "n": f"{cache_name}_{layer_idx}",
                         "i": layer_idx,
-                        "d": dtype_str,
+                        "d": str(t.dtype).removeprefix("torch."),
                         "s": list(t.shape),
                         "o": data_offset,
                         "b": nbytes,
                     }
                 )
-                gpu_tensors.append(t.view(torch.uint8).flatten())
+                chunks.append(t.view(torch.uint8).numpy().tobytes() if cpu else t.view(torch.uint8).flatten())
                 data_offset += nbytes
 
-        if device is None:
-            raise RuntimeError("No CUDA tensors found, use to_bytes() instead")
+        return tensors_desc, chunks, data_offset, device
 
+    def _build_header_bytes(self, tensors_desc: list[dict[str, Any]]) -> bytes:
         header = json.dumps(
             {
                 "rid": self.request_id,
@@ -158,19 +134,26 @@ class KVCacheTransferData:
             },
             separators=(",", ":"),
         ).encode("utf-8")
+        return struct.pack(">I", len(header)) + header
 
-        header_prefix = struct.pack(">I", len(header)) + header
-        total_size = len(header_prefix) + data_offset
-        output = torch.empty(total_size, dtype=torch.uint8, device=device)
-        header_tensor = torch.frombuffer(bytearray(header_prefix), dtype=torch.uint8)
-        output[: len(header_prefix)].copy_(header_tensor)
+    def to_bytes(self) -> bytes:
+        """Convert to compact binary format for fast transfer."""
+        tensors_desc, chunks, _, _ = self._build_tensors_desc(cpu=True)
+        return b"".join([self._build_header_bytes(tensors_desc)] + chunks)
 
+    def to_gpu_tensor(self) -> torch.Tensor:
+        """Convert to a packed GPU tensor for raw-data connectors."""
+        tensors_desc, chunks, data_offset, device = self._build_tensors_desc(cpu=False)
+        if device is None:
+            raise RuntimeError("No CUDA tensors found, use to_bytes() instead")
+        header_prefix = self._build_header_bytes(tensors_desc)
+        output = torch.empty(len(header_prefix) + data_offset, dtype=torch.uint8, device=device)
+        output[: len(header_prefix)].copy_(torch.frombuffer(bytearray(header_prefix), dtype=torch.uint8))
         pos = len(header_prefix)
-        for t_flat in gpu_tensors:
+        for t_flat in chunks:
             n = t_flat.numel()
             output[pos : pos + n].copy_(t_flat)
             pos += n
-
         return output
 
     @staticmethod
@@ -237,11 +220,8 @@ class KVCacheTransferData:
         return layer_idx
 
     @staticmethod
-    def from_bytes(raw: "bytes | bytearray | memoryview") -> dict[str, Any]:
-        """Reconstruct KV cache data from the packed bytes format."""
-        raw_mv = memoryview(raw) if not isinstance(raw, memoryview) else raw
-        header, tensor_data_mv = KVCacheTransferData._load_header_from_memoryview(raw_mv)
-
+    def _populate_caches(header: dict[str, Any], get_tensor: callable) -> dict[str, Any]:
+        """Shared deserialization loop for both CPU and GPU paths."""
         num_layers = header["nl"]
         key_cache: list[torch.Tensor | None] = [None] * num_layers
         value_cache: list[torch.Tensor | None] = [None] * num_layers
@@ -249,20 +229,9 @@ class KVCacheTransferData:
         for info in header["td"]:
             if info.get("x"):
                 continue
-
             name: str = info["n"]
             torch_dtype = KVCacheTransferData._resolve_torch_dtype(info["d"])
-            offset, nbytes = KVCacheTransferData._validate_tensor_span(name, info, len(tensor_data_mv))
-            t = (
-                torch.frombuffer(
-                    tensor_data_mv,
-                    dtype=torch.uint8,
-                    offset=offset,
-                    count=nbytes,
-                )
-                .view(torch_dtype)
-                .reshape(info["s"])
-            )
+            t = get_tensor(info).view(torch_dtype).reshape(info["s"])
             layer_idx = KVCacheTransferData._resolve_layer_idx(info, num_layers)
             if name.startswith("key_cache_"):
                 key_cache[layer_idx] = t
@@ -276,37 +245,30 @@ class KVCacheTransferData:
             "metadata": header["meta"],
         }
 
+    @staticmethod
+    def from_bytes(raw: "bytes | bytearray | memoryview") -> dict[str, Any]:
+        """Reconstruct KV cache data from the packed bytes format."""
+        raw_mv = memoryview(raw) if not isinstance(raw, memoryview) else raw
+        header, tensor_data_mv = KVCacheTransferData._load_header_from_memoryview(raw_mv)
+        data_len = len(tensor_data_mv)
+
+        def _get(info: dict) -> torch.Tensor:
+            offset, nbytes = KVCacheTransferData._validate_tensor_span(info["n"], info, data_len)
+            return torch.frombuffer(tensor_data_mv, dtype=torch.uint8, offset=offset, count=nbytes)
+
+        return KVCacheTransferData._populate_caches(header, _get)
+
     @staticmethod
     def from_bytes_gpu(gpu_tensor: torch.Tensor) -> dict[str, Any]:
         """Reconstruct KV cache data from a packed GPU tensor."""
         header, data_start = KVCacheTransferData._load_header_from_tensor(gpu_tensor)
+        data_len = int(gpu_tensor.numel()) - data_start
 
-        num_layers = header["nl"]
-        key_cache: list[torch.Tensor | None] = [None] * num_layers
-        value_cache: list[torch.Tensor | None] = [None] * num_layers
-        tensor_data_bytes = int(gpu_tensor.numel()) - data_start
+        def _get(info: dict) -> torch.Tensor:
+            offset, nbytes = KVCacheTransferData._validate_tensor_span(info["n"], info, data_len)
+            return gpu_tensor[data_start + offset : data_start + offset + nbytes].clone()
 
-        for info in header["td"]:
-            if info.get("x"):
-                continue
-
-            name: str = info["n"]
-            torch_dtype = KVCacheTransferData._resolve_torch_dtype(info["d"])
-            offset, nbytes = KVCacheTransferData._validate_tensor_span(name, info, tensor_data_bytes)
-            t = gpu_tensor[data_start + offset : data_start + offset + nbytes].clone()
-            t = t.view(torch_dtype).reshape(info["s"])
-            layer_idx = KVCacheTransferData._resolve_layer_idx(info, num_layers)
-            if name.startswith("key_cache_"):
-                key_cache[layer_idx] = t
-            elif name.startswith("value_cache_"):
-                value_cache[layer_idx] = t
-
-        return {
-            "request_id": header["rid"],
-            "layer_blocks": {"key_cache": key_cache, "value_cache": value_cache},
-            "block_ids": header["bids"],
-            "metadata": header["meta"],
-        }
+        return KVCacheTransferData._populate_caches(header, _get)
 
 
 class OmniKVTransferManager:
@@ -341,6 +303,30 @@ class OmniKVTransferManager:
             else (None, None)
         )
 
+        local_rank = get_local_tp_rank()
+
+        if config.from_tp <= 1 and config.to_tp <= 1:
+            detected_tp = get_tp_world_size()
+            from_tp = detected_tp
+            to_tp = detected_tp
+        else:
+            from_tp = config.from_tp
+            to_tp = config.to_tp
+
+        self._tp_topo = KVTPTopology(source_tp_size=from_tp, target_tp_size=to_tp, local_rank=local_rank)
+
+        # Injectable hooks (compatible with PR #2677 OmniConnectorModelRunnerMixin).
+        self.kv_send_key_builder: Callable | None = None
+        self.kv_recv_key_builder: Callable | None = None
+        self.kv_payload_merger: Callable | None = None
+        self.kv_payload_slicer: Callable | None = None
+
+        # Base sender endpoint (rank-0 host/port) stored during
+        # update_sender_info().  Used by the receive path to construct
+        # per-rank metadata for heterogeneous TP without querying a registry.
+        self._sender_base_host: str | None = None
+        self._sender_base_zmq_port: int | None = None
+
         if config.need_send_cache and config.connector_config:
             try:
                 _ = self.connector
@@ -348,11 +334,20 @@ class OmniKVTransferManager:
             except Exception as e:
                 logger.warning("Failed to eagerly initialize sender connector: %s", e)
 
+    # ------------------------------------------------------------------ #
+    #  Factory helpers
+    # ------------------------------------------------------------------ #
+
     @classmethod
     def _create(cls, cfg: dict | None) -> "OmniKVTransferManager":
         """Create manager from raw config dict."""
         if not cfg or not isinstance(cfg, dict):
             return cls(OmniKVCacheConfig())
+
+        rank_mapping = cfg.get("rank_mapping", {})
+        if not isinstance(rank_mapping, dict):
+            rank_mapping = {}
+
         return cls(
             OmniKVCacheConfig(
                 connector_config=cfg.get("connector_config"),
@@ -363,18 +358,17 @@ class OmniKVTransferManager:
                 need_recv_cache=cfg.get("need_recv_cache", False),
                 need_send_cache=cfg.get("need_send_cache", False),
                 recv_timeout=cfg.get("recv_timeout", 30.0),
+                from_tp=int(rank_mapping.get("from_tp", 1)),
+                to_tp=int(rank_mapping.get("to_tp", 1)),
             )
         )
 
     @classmethod
-    def from_model_config(cls, config: Any) -> "OmniKVTransferManager":
-        """Create from model config (for AR model runner)."""
+    def from_od_config(cls, config: Any) -> "OmniKVTransferManager":
+        """Create from model or OmniDiffusion config."""
         return cls._create(getattr(config, "omni_kv_config", None))
 
-    @classmethod
-    def from_od_config(cls, config: Any) -> "OmniKVTransferManager":
-        """Create from OmniDiffusion config (for diffusion runner)."""
-        return cls._create(getattr(config, "omni_kv_config", None))
+    from_model_config = from_od_config
 
     @classmethod
     def from_vllm_config(cls, vllm_config: Any, model_config: Any) -> "OmniKVTransferManager":
@@ -417,45 +411,33 @@ class OmniKVTransferManager:
                         )
                         c_extra["to_stage"] = str(self.config.to_stage) if self.config.to_stage is not None else "1"
 
+                        try:
+                            stage_int = int(self.config.from_stage) if self.config.from_stage is not None else 0
+                        except (TypeError, ValueError):
+                            stage_int = 0
+                        zmq_port = kv_zmq_port(base_port, stage_int, self._tp_topo.local_rank)
+
                         if self.config.need_send_cache:
                             c_extra["role"] = "sender"
-                            from_stage = self.config.from_stage
-                            if from_stage is not None:
-                                try:
-                                    c_extra["zmq_port"] = base_port + KV_TRANSFER_PORT_OFFSET + int(from_stage)
-                                except (TypeError, ValueError):
-                                    c_extra["zmq_port"] = base_port + KV_TRANSFER_PORT_OFFSET
+                            c_extra["zmq_port"] = zmq_port
                         elif self.config.need_recv_cache:
                             c_extra["role"] = "receiver"
-                            from_stage = self.config.from_stage
-                            sender_port = base_port + KV_TRANSFER_PORT_OFFSET
-                            if from_stage is not None:
-                                try:
-                                    sender_port = base_port + KV_TRANSFER_PORT_OFFSET + int(from_stage)
-                                except (TypeError, ValueError):
-                                    pass
                             c_extra.setdefault("sender_host", c_extra.get("host", "127.0.0.1"))
-                            c_extra.setdefault("sender_zmq_port", sender_port)
+                            c_extra.setdefault("sender_zmq_port", zmq_port)
 
                     logger.info(
-                        "Initializing OmniConnector (purpose=kv_transfer) with config: %s, role: %s",
-                        cfg,
+                        "Initializing OmniConnector type=%s role=%s",
+                        c_type,
                         c_extra.get("role", "N/A"),
                     )
                     self._connector = OmniConnectorFactory.create_connector(ConnectorSpec(name=c_type, extra=c_extra))
-                except Exception as e:
-                    logger.error(f"Failed to initialize OmniConnector: {e}")
-                    import traceback
-
-                    traceback.print_exc()
-                    # Cache failure sentinel to avoid repeated initialization attempts in hot paths.
+                except Exception:
+                    logger.exception("Failed to initialize OmniConnector")
                     self._connector = False
 
         return self._connector if self._connector else None
 
-    def get_connector(self):
-        """Get connector (compatibility wrapper for existing code)."""
-        return self.connector
+    get_connector = property(lambda self: self.connector)
 
     def _resolve_sender_info(
         self, sender_info: dict[str, Any], sender_stage_id: str | int | None = None
@@ -513,8 +495,187 @@ class OmniKVTransferManager:
                     cache_list[idx] = tensor.clone()
         return data
 
+    def _slice_transfer_data_for_target(self, kv_data: KVCacheTransferData, target_rank: int) -> KVCacheTransferData:
+        """Pre-slice sender payload for one target rank when sender TP < receiver TP."""
+        topo = self._tp_topo
+        ratio = topo.target_tp_size // topo.source_tp_size
+        offset_in_sender = target_rank % ratio
+        metadata = dict(kv_data.metadata) if isinstance(kv_data.metadata, dict) else {}
+        metadata["tp_head_slice"] = {
+            "applied": True,
+            "side": "sender",
+            "target_rank": target_rank,
+            "source_rank": topo.local_rank,
+            "from_tp": topo.source_tp_size,
+            "to_tp": topo.target_tp_size,
+            "offset_in_shard": offset_in_sender,
+            "num_slices": ratio,
+        }
+        return KVCacheTransferData(
+            request_id=kv_data.request_id,
+            layer_blocks=slice_layer_blocks(kv_data.layer_blocks, offset_in_sender, ratio),
+            block_ids=list(kv_data.block_ids),
+            metadata=metadata,
+        )
+
+    def _serialize_transfer_payload(self, kv_data: KVCacheTransferData) -> torch.Tensor | bytes | dict[str, Any]:
+        """Serialize KV transfer data using the connector's fastest supported path."""
+        if getattr(self.connector, "supports_raw_data", False):
+            try:
+                return kv_data.to_gpu_tensor()
+            except Exception:
+                pass
+        try:
+            return kv_data.to_bytes()
+        except Exception:
+            return kv_data.to_dict()
+
+    @staticmethod
+    def _collect_request_kv_payload(req: Any) -> dict[str, object]:
+        """Collect request-side KV objects for object broadcast."""
+        kv_payload: dict[str, object] = {}
+        for attr in ("past_key_values", "kv_metadata"):
+            val = getattr(req, attr, None)
+            if val is not None:
+                kv_payload[attr] = val
+
+        if hasattr(req, "sampling_params") and req.sampling_params is not None:
+            for key in list(vars(req.sampling_params).keys()):
+                if key in ("past_key_values", "kv_metadata") or (
+                    key.startswith("cfg_")
+                    and (
+                        key.endswith("_past_key_values")
+                        or key.endswith("_kv_metadata")
+                        or key
+                        in (
+                            "cfg_kv_request_ids",
+                            "cfg_active_branch",
+                            "cfg_branch_roles",
+                            "cfg_branch_past_key_values",
+                            "cfg_branch_kv_metadata",
+                        )
+                    )
+                ):
+                    val = getattr(req.sampling_params, key, None)
+                    if val is not None:
+                        kv_payload[f"sp.{key}"] = val
+
+        return kv_payload
+
+    @staticmethod
+    def _apply_request_kv_payload(
+        req: Any,
+        kv_payload: dict[str, object],
+        target_device: torch.device | None = None,
+    ) -> None:
+        """Apply a broadcast KV payload back onto a request object."""
+        for attr in ("past_key_values", "kv_metadata"):
+            val = kv_payload.get(attr)
+            if val is not None:
+                if target_device is not None:
+                    val = _move_to_device(val, target_device)
+                setattr(req, attr, val)
+
+        if hasattr(req, "sampling_params") and req.sampling_params is not None:
+            for key, val in kv_payload.items():
+                if key.startswith("sp."):
+                    if target_device is not None:
+                        val = _move_to_device(val, target_device)
+                    setattr(req.sampling_params, key[3:], val)
+
+    @staticmethod
+    def _discover_cfg_branch_roles(req: Any) -> list[str]:
+        """Discover CFG branch roles in a stable order."""
+        sampling_params = getattr(req, "sampling_params", None)
+        if sampling_params is None:
+            return []
+
+        roles: list[str] = []
+        branch_map = getattr(sampling_params, "cfg_branch_past_key_values", None) or {}
+        for preferred_role in ("cfg_text", "cfg_img"):
+            if (
+                preferred_role in branch_map
+                or getattr(sampling_params, f"{preferred_role}_past_key_values", None) is not None
+            ):
+                roles.append(preferred_role)
+
+        for role in branch_map.keys():
+            if role not in roles and branch_map.get(role) is not None:
+                roles.append(role)
+
+        for key in vars(sampling_params).keys():
+            if not (key.startswith("cfg_") and key.endswith("_past_key_values")):
+                continue
+            role = key.removesuffix("_past_key_values")
+            if role in ("cfg_branch",) or role in roles:
+                continue
+            if getattr(sampling_params, key, None) is not None:
+                roles.append(role)
+
+        return roles
+
+    @classmethod
+    def _build_cfg_rank_local_payloads(cls, req: Any, cfg_size: int) -> list[dict[str, object] | None]:
+        """Build per-cfg-rank payloads so each rank receives only its branch KV."""
+        full_payload = cls._collect_request_kv_payload(req)
+        payloads: list[dict[str, object] | None] = []
+
+        main_payload = {
+            key: value
+            for key, value in full_payload.items()
+            if key in ("past_key_values", "kv_metadata", "sp.past_key_values", "sp.kv_metadata")
+        }
+        branch_roles = cls._discover_cfg_branch_roles(req)
+        if branch_roles:
+            main_payload["sp.cfg_branch_roles"] = list(branch_roles)
+            main_payload["sp.cfg_active_branch"] = None
+        payloads.append(main_payload or None)
+
+        sampling_params = getattr(req, "sampling_params", None)
+        branch_map = getattr(sampling_params, "cfg_branch_past_key_values", None) or {}
+        branch_metadata_map = getattr(sampling_params, "cfg_branch_kv_metadata", None) or {}
+
+        for role in branch_roles:
+            if sampling_params is None:
+                payloads.append(None)
+                continue
+
+            branch_kv = branch_map.get(role)
+            if branch_kv is None:
+                branch_kv = getattr(sampling_params, f"{role}_past_key_values", None)
+            branch_metadata = branch_metadata_map.get(role)
+            if branch_metadata is None:
+                branch_metadata = getattr(sampling_params, f"{role}_kv_metadata", None)
+            if branch_kv is None:
+                payloads.append(None)
+                continue
+
+            local_payload = dict(main_payload)
+            local_payload["sp.cfg_active_branch"] = role
+            local_payload["sp.cfg_branch_roles"] = list(branch_roles)
+            local_payload["sp.cfg_branch_past_key_values"] = {role: branch_kv}
+            local_payload[f"sp.{role}_past_key_values"] = branch_kv
+            if branch_metadata is not None:
+                local_payload["sp.cfg_branch_kv_metadata"] = {role: branch_metadata}
+                local_payload[f"sp.{role}_kv_metadata"] = branch_metadata
+
+            payloads.append(local_payload)
+
+        while len(payloads) < cfg_size:
+            payloads.append(None)
+
+        return payloads[:cfg_size]
+
     def update_sender_info(self, sender_info: dict[str, Any], sender_stage_id: str | int | None = None) -> None:
-        """Update receiver-side sender info before loading remote KV cache."""
+        """Update receiver-side sender info before loading remote KV cache.
+
+        The orchestrator always reports rank-0's ZMQ port.  When TP > 1 the
+        receiver must offset the port so that each TP rank connects to the
+        corresponding sender rank's port.
+
+        The base host/port are also stored so that the receive path can
+        construct per-rank metadata for heterogeneous TP scenarios.
+        """
         if not self.config.need_recv_cache:
             return
 
@@ -523,18 +684,39 @@ class OmniKVTransferManager:
             logger.warning("Invalid sender_info format: %s", sender_info)
             return
 
+        sender_host = actual_info.get("host")
+        base_zmq_port = actual_info.get("zmq_port")
+
+        # Store base sender info for per-rank metadata construction.
+        self._sender_base_host = sender_host
+        if base_zmq_port is not None:
+            self._sender_base_zmq_port = int(base_zmq_port)
+
+        # --- Default sender: offset to match this receiver's corresponding sender rank ---
+        zmq_port = base_zmq_port
+        if zmq_port is not None and self._tp_topo.local_rank > 0:
+            zmq_port = int(zmq_port) + self._tp_topo.local_rank * KV_RANK_PORT_STRIDE
+
         if self.config.connector_config:
-            self.config.connector_config["sender_host"] = actual_info.get("host")
-            self.config.connector_config["sender_zmq_port"] = actual_info.get("zmq_port")
+            self.config.connector_config["sender_host"] = sender_host
+            self.config.connector_config["sender_zmq_port"] = zmq_port
 
         if self._connector and hasattr(self._connector, "update_sender_info"):
             try:
-                self._connector.update_sender_info(actual_info.get("host"), actual_info.get("zmq_port"))
+                self._connector.update_sender_info(sender_host, zmq_port)
             except Exception:
                 if hasattr(self._connector, "sender_host"):
-                    self._connector.sender_host = actual_info.get("host")
+                    self._connector.sender_host = sender_host
                 if hasattr(self._connector, "sender_zmq_port"):
-                    self._connector.sender_zmq_port = actual_info.get("zmq_port")
+                    self._connector.sender_zmq_port = zmq_port
+
+        logger.info(
+            "Sender info updated: host=%s, base_port=%s, adjusted_port=%s (local_rank=%s)",
+            sender_host,
+            base_zmq_port,
+            zmq_port,
+            self._tp_topo.local_rank,
+        )
 
     def handle_finished_requests_kv_transfer(
         self,
@@ -692,35 +874,54 @@ class OmniKVTransferManager:
 
         kv_data.request_id = transfer_req_id
         serialization_start = time.perf_counter()
-        transfer_data: torch.Tensor | bytes | dict[str, Any]
-        supports_raw = getattr(self.connector, "supports_raw_data", False)
+        topo = self._tp_topo
+        send_keys = build_rank_aware_send_keys(
+            transfer_req_id, from_stage, to_stage, topo, hook=self.kv_send_key_builder
+        )
+        sender_slice_active = (
+            topo.source_tp_size < topo.target_tp_size and len(send_keys) > 1 and not callable(self.kv_send_key_builder)
+        )
+        per_key_payloads: list[tuple[str, torch.Tensor | bytes | dict[str, Any]]] = []
 
-        try:
-            if supports_raw:
-                transfer_data = kv_data.to_gpu_tensor()
+        if sender_slice_active:
+            target_ranks = get_kv_target_ranks(topo)
+            if len(target_ranks) != len(send_keys):
+                logger.warning(
+                    "Skip sender-side KV slicing because target rank count does not match send key count: "
+                    "target_ranks=%s send_keys=%s",
+                    len(target_ranks),
+                    len(send_keys),
+                )
+                sender_slice_active = False
             else:
-                raise RuntimeError("Connector does not support raw tensor")
-        except Exception:
-            try:
-                transfer_data = kv_data.to_bytes()
-            except Exception:
-                data_dict = kv_data.to_dict()
-                data_dict["request_id"] = transfer_req_id
-                transfer_data = data_dict
+                for put_key, target_rank in zip(send_keys, target_ranks, strict=False):
+                    sliced_kv_data = self._slice_transfer_data_for_target(kv_data, target_rank)
+                    per_key_payloads.append((put_key, self._serialize_transfer_payload(sliced_kv_data)))
+
+        if not per_key_payloads:
+            transfer_data = self._serialize_transfer_payload(kv_data)
+            per_key_payloads = [(put_key, transfer_data) for put_key in send_keys]
 
         serialization_ms = (time.perf_counter() - serialization_start) * 1000
         logger.info("KV cache serialized for %s in %.1f ms", transfer_req_id, serialization_ms)
 
         transfer_start = time.perf_counter()
-        success, size, _ = self._transfer_with_retry(from_stage, to_stage, f"kv_cache_{transfer_req_id}", transfer_data)
+        total_size = 0
+        all_succeeded = True
+        for put_key, transfer_data in per_key_payloads:
+            success, size, _ = self._transfer_with_retry(from_stage, to_stage, put_key, transfer_data)
+            total_size += size
+            all_succeeded = all_succeeded and success
+
         elapsed = time.perf_counter() - transfer_start
 
-        if success:
-            mbps = (size / 1024 / 1024) / elapsed if elapsed > 0 else 0
+        if all_succeeded:
+            mbps = (total_size / 1024 / 1024) / elapsed if elapsed > 0 else 0
             logger.info(
-                "KV transfer OK: %s, %s bytes, %.3fs, %.1f MB/s",
+                "KV transfer OK: %s, %s bytes across %s key(s), %.3fs, %.1f MB/s",
                 transfer_req_id,
-                size,
+                total_size,
+                len(send_keys),
                 elapsed,
                 mbps,
             )
@@ -731,7 +932,7 @@ class OmniKVTransferManager:
         self,
         from_stage: str,
         to_stage: str,
-        request_id: str,
+        put_key: str,
         data: "dict[str, Any] | bytes | torch.Tensor",
         max_retries: int = 3,
     ) -> tuple[bool, int, dict[str, Any] | None]:
@@ -740,7 +941,7 @@ class OmniKVTransferManager:
         Args:
             from_stage: Source stage identifier
             to_stage: Target stage identifier
-            request_id: Request identifier for the key
+            put_key: Pre-built connector key (rank-aware when TP > 1)
             data: Data to transfer
             max_retries: Maximum number of retry attempts
 
@@ -749,14 +950,12 @@ class OmniKVTransferManager:
         """
         for attempt in range(max_retries):
             try:
-                # Build the full key for connector
-                full_request_id = f"omni_{from_stage}_to_{to_stage}_{request_id}"
                 success, size, metadata = self.connector.put(
-                    from_stage=from_stage, to_stage=to_stage, put_key=full_request_id, data=data
+                    from_stage=from_stage, to_stage=to_stage, put_key=put_key, data=data
                 )
                 if success:
                     return success, size, metadata
-                logger.warning(f"Transfer attempt {attempt + 1} failed for {request_id}")
+                logger.warning(f"Transfer attempt {attempt + 1} failed for {put_key}")
             except Exception as e:
                 logger.warning(f"Transfer attempt {attempt + 1} exception: {e}")
 
@@ -801,22 +1000,46 @@ class OmniKVTransferManager:
         poll_interval = 0.01
         max_poll_interval = 0.5
 
-        logger.info(f"Wait for KV cache for request {request_id} from stage {from_stage} to {to_stage}...")
+        topo = self._tp_topo
+        recv_key_pairs = build_rank_aware_recv_keys(
+            request_id, from_stage, to_stage, topo, hook=self.kv_recv_key_builder
+        )
+        pending_pairs = list(recv_key_pairs)
+        received_payloads: dict[str, tuple[dict[str, Any], int]] = {}
+
+        logger.info(
+            "Wait for KV cache for request %s from stage %s to %s via %s key(s)...",
+            request_id,
+            from_stage,
+            to_stage,
+            len(recv_key_pairs),
+        )
 
         try:
             while True:
-                # Build the full key for connector
-                full_request_id = f"omni_{from_stage}_to_{to_stage}_kv_cache_{request_id}"
                 link_start = time.perf_counter()
-                result = self.connector.get(
-                    from_stage=from_stage,
-                    to_stage=to_stage,
-                    get_key=full_request_id,
-                )
-                if result:
+                for get_key, from_rank in list(pending_pairs):
+                    # Construct per-rank metadata so the connector queries
+                    # the correct sender endpoint (heterogeneous TP path).
+                    # When from_rank is None (TP<=1), metadata stays None
+                    # and the connector falls back to its default sender.
+                    rank_metadata: dict[str, Any] | None = None
+                    if from_rank is not None and self._sender_base_host and self._sender_base_zmq_port is not None:
+                        rank_metadata = {
+                            "source_host": self._sender_base_host,
+                            "source_port": self._sender_base_zmq_port + from_rank * KV_RANK_PORT_STRIDE,
+                        }
+
+                    result = self.connector.get(
+                        from_stage=from_stage,
+                        to_stage=to_stage,
+                        get_key=get_key,
+                        metadata=rank_metadata,
+                    )
+                    if not result:
+                        continue
+
                     raw_data, size = result
-                    elapsed = time.time() - start_time
-                    link_ms = (time.perf_counter() - link_start) * 1000
                     managed_buffer = None
 
                     if hasattr(raw_data, "tensor") and hasattr(raw_data, "release"):
@@ -844,6 +1067,21 @@ class OmniKVTransferManager:
                     else:
                         data = raw_data
 
+                    received_payloads[get_key] = (data, size)
+                    pending_pairs.remove((get_key, from_rank))
+
+                if not pending_pairs and received_payloads:
+                    elapsed = time.time() - start_time
+                    link_ms = (time.perf_counter() - link_start) * 1000
+                    ordered_payloads = [received_payloads[key][0] for key, _ in recv_key_pairs]
+                    total_size = sum(received_payloads[key][1] for key, _ in recv_key_pairs)
+
+                    if len(ordered_payloads) == 1:
+                        data = ordered_payloads[0]
+                    else:
+                        data = merge_received_rank_shards(ordered_payloads, merger=self.kv_payload_merger)
+                    data = slice_received_rank_shard(data, topo, slicer=self.kv_payload_slicer)
+
                     try:
                         if isinstance(data, dict) and "layer_blocks" in data:
                             layer_blocks = data["layer_blocks"]
@@ -856,18 +1094,18 @@ class OmniKVTransferManager:
                                         continue
                                     if target_device is not None and tensor.device != target_device:
                                         cache_list[i] = tensor.to(target_device).contiguous()
-                    finally:
-                        if managed_buffer is not None:
-                            managed_buffer.release()
+                    except Exception:
+                        logger.exception("Failed to move KV cache tensors to target device")
 
                     logger.info(
-                        "Successfully received KV cache for %s, %s bytes, wait=%.3fs, link=%.1fms",
+                        "Successfully received KV cache for %s, %s bytes across %s key(s), wait=%.3fs, link=%.1fms",
                         request_id,
-                        size,
+                        total_size,
+                        len(recv_key_pairs),
                         elapsed,
                         link_ms,
                     )
-                    return data, size
+                    return data, total_size
 
                 if time.time() - start_time > timeout:
                     logger.error(f"Timeout waiting for KV cache for request {request_id} after {timeout}s")
@@ -876,11 +1114,8 @@ class OmniKVTransferManager:
                 time.sleep(poll_interval)
                 poll_interval = min(poll_interval * 2, max_poll_interval)
 
-        except Exception as e:
-            logger.error(f"Error receiving KV cache for {request_id}: {e}")
-            import traceback
-
-            traceback.print_exc()
+        except Exception:
+            logger.exception("Error receiving KV cache for %s", request_id)
             return None, 0
 
     def apply_kv_cache_to_request(self, req: Any, data: dict[str, Any]) -> None:
@@ -994,73 +1229,79 @@ class OmniKVTransferManager:
         cfg_kv_collect_func: Callable | None = None,
         target_device: torch.device | None = None,
     ) -> bool:
-        """Broadcast-aware wrapper around :meth:`receive_multi_kv_cache`.
+        """Distributed wrapper around :meth:`receive_multi_kv_cache`.
 
-        SharedMemory connector is single-reader: once rank 0 consumes the
-        segment it is deleted.  For multi-GPU stages (e.g. sequence-parallel)
-        only rank 0 receives; the result is then broadcast to every other
-        rank via the world process-group.
-
-        For single-worker stages this is equivalent to calling
-        :meth:`receive_multi_kv_cache` directly.
+        TP-aware path selection:
+        - world size 1: direct receive
+        - TP active, cfg size 1: each rank independently receives
+        - TP active, cfg size > 1: cfg-rank 0 receives, then broadcasts to
+          peers that share the same TP rank
+        - TP inactive: legacy rank-0 receive then world broadcast
         """
-        from vllm_omni.diffusion.distributed.parallel_state import get_world_group
+        from vllm_omni.diffusion.distributed.parallel_state import (
+            get_cfg_group,
+            get_classifier_free_guidance_rank,
+            get_classifier_free_guidance_world_size,
+            get_world_group,
+        )
 
         world = get_world_group()
 
         if world.world_size <= 1:
             return self.receive_multi_kv_cache(req, cfg_kv_collect_func, target_device)
 
-        # --- rank 0: receive to CPU (needed for pickle-based broadcast) ---
+        topo = self._tp_topo
+        tp_active = topo.source_tp_size > 1 or topo.target_tp_size > 1
+        cfg_size = 1
+        cfg_rank = 0
+        cfg_group = None
+        try:
+            cfg_size = get_classifier_free_guidance_world_size()
+            cfg_rank = get_classifier_free_guidance_rank()
+            cfg_group = get_cfg_group()
+        except Exception:
+            cfg_size = 1
+            cfg_rank = 0
+            cfg_group = None
+
+        if tp_active and cfg_size <= 1:
+            logger.info(
+                "Rank-aware KV receive: rank %s independently receiving (from_tp=%s, to_tp=%s)",
+                topo.local_rank,
+                topo.source_tp_size,
+                topo.target_tp_size,
+            )
+            return self.receive_multi_kv_cache(req, cfg_kv_collect_func, target_device)
+
+        if tp_active and cfg_size > 1 and cfg_group is not None:
+            kv_payload: dict[str, object] | None = None
+            if cfg_rank == 0:
+                received = self.receive_multi_kv_cache(req, cfg_kv_collect_func, torch.device("cpu"))
+                rank_payloads = self._build_cfg_rank_local_payloads(req, cfg_size) if received else [None] * cfg_size
+                kv_payload = rank_payloads[0]
+                for dst_rank in range(1, cfg_size):
+                    cfg_group.send_object(rank_payloads[dst_rank], dst_rank)
+            else:
+                kv_payload = cfg_group.recv_object(0)
+
+            if not kv_payload:
+                return False
+
+            self._apply_request_kv_payload(req, kv_payload, target_device)
+            return True
+
+        kv_payload: dict[str, object] | None = None
         if world.rank_in_group == 0:
-            self.receive_multi_kv_cache(req, cfg_kv_collect_func, torch.device("cpu"))
+            received = self.receive_multi_kv_cache(req, cfg_kv_collect_func, torch.device("cpu"))
+            if received:
+                kv_payload = self._collect_request_kv_payload(req)
 
-            kv_payload: dict[str, object] = {}
-            for attr in ("past_key_values", "kv_metadata"):
-                val = getattr(req, attr, None)
-                if val is not None:
-                    kv_payload[attr] = val
+        kv_payload = world.broadcast_object(kv_payload, src=0)
 
-            if hasattr(req, "sampling_params") and req.sampling_params is not None:
-                for key in list(vars(req.sampling_params).keys()):
-                    if (key.startswith("cfg_") and key.endswith("_past_key_values")) or key in (
-                        "past_key_values",
-                        "kv_metadata",
-                    ):
-                        val = getattr(req.sampling_params, key, None)
-                        if val is not None:
-                            kv_payload[f"sp.{key}"] = val
-
-            payload_list = [kv_payload]
-            # Use broadcast_object_list (pickle-based) instead of broadcast_tensor_dict
-            # because the KV cache is a heterogeneous nested structure (NaiveCache objects
-            # with metadata + tensors), not a flat tensor dict.  This runs once before
-            # the denoising loop so the serialization cost is negligible.
-            torch.distributed.broadcast_object_list(payload_list, src=world.ranks[0], group=world.cpu_group)
-            kv_payload = payload_list[0]
-        else:
-            payload_list: list[dict[str, object] | None] = [None]
-            torch.distributed.broadcast_object_list(payload_list, src=world.ranks[0], group=world.cpu_group)
-            kv_payload = payload_list[0]
-
-        # --- apply on ALL ranks (rank 0 also needs CPU→GPU move) ---
         if not kv_payload:
             return False
 
-        for attr in ("past_key_values", "kv_metadata"):
-            val = kv_payload.get(attr)
-            if val is not None:
-                if target_device is not None:
-                    val = _move_to_device(val, target_device)
-                setattr(req, attr, val)
-
-        if hasattr(req, "sampling_params") and req.sampling_params is not None:
-            for key, val in kv_payload.items():
-                if key.startswith("sp."):
-                    if target_device is not None:
-                        val = _move_to_device(val, target_device)
-                    setattr(req.sampling_params, key[3:], val)
-
+        self._apply_request_kv_payload(req, kv_payload, target_device)
         return True
 
 

vllm_omni/distributed/omni_connectors/utils/kv_utils.py

@@ -1,15 +1,380 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Utility helpers for KV cache manipulation."""
+"""Utility helpers for KV cache manipulation, TP routing, and merge/slice."""
+
+from __future__ import annotations
+
+import os
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
 
 import torch
+from vllm.distributed.parallel_state import (
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+)
 from vllm.logger import init_logger
 
+from .initialization import KV_RANK_PORT_STRIDE, KV_TRANSFER_PORT_OFFSET
+
 logger = init_logger(__name__)
 
 LayerKV = torch.Tensor | tuple[torch.Tensor, torch.Tensor]
 
 
+# ------------------------------------------------------------------ #
+#  TP Topology
+# ------------------------------------------------------------------ #
+
+
+@dataclass(frozen=True)
+class KVTPTopology:
+    """Immutable descriptor for a KV-transfer parallel mapping.
+
+    Captures sender/receiver parallel sizes and the local rank within
+    that parallel dimension.  Works for any divisible parallel dimension
+    (TP, SP, Ring Attention).
+    """
+
+    source_tp_size: int
+    target_tp_size: int
+    local_rank: int
+
+    def __post_init__(self) -> None:
+        if self.source_tp_size <= 0 or self.target_tp_size <= 0:
+            raise ValueError(
+                f"Parallel sizes must be positive: "
+                f"source_tp_size={self.source_tp_size}, target_tp_size={self.target_tp_size}"
+            )
+        if self.local_rank < 0:
+            raise ValueError(f"local_rank must be non-negative, got {self.local_rank}")
+
+    @property
+    def is_heterogeneous(self) -> bool:
+        return self.source_tp_size != self.target_tp_size
+
+    @property
+    def ratio(self) -> int:
+        """Larger parallel size divided by smaller. Always >= 1."""
+        return max(self.source_tp_size, self.target_tp_size) // min(self.source_tp_size, self.target_tp_size)
+
+
+# ------------------------------------------------------------------ #
+#  Runtime TP detection
+# ------------------------------------------------------------------ #
+
+
+def get_local_tp_rank() -> int:
+    """Return the TP-local rank of this worker process.
+
+    Uses ``get_tensor_model_parallel_rank()`` which returns the rank
+    within the TP group only, not the stage-global rank.
+    """
+    try:
+        return get_tensor_model_parallel_rank()
+    except Exception:
+        logger.debug("TP parallel state not initialized, falling back to LOCAL_RANK env", exc_info=True)
+    try:
+        return int(os.environ.get("LOCAL_RANK", "0"))
+    except (ValueError, TypeError):
+        return 0
+
+
+def get_tp_world_size() -> int:
+    """Return the TP world size (tensor-parallel dimension only).
+
+    Uses ``get_tensor_model_parallel_world_size()`` so that
+    cfg_parallel, SP, PP etc. are not included in the count.
+    """
+    try:
+        return get_tensor_model_parallel_world_size()
+    except Exception:
+        logger.debug("TP parallel state not initialized, defaulting world_size=1", exc_info=True)
+    return 1
+
+
+# ------------------------------------------------------------------ #
+#  ZMQ port computation
+# ------------------------------------------------------------------ #
+
+
+def kv_zmq_port(base_port: int, from_stage: int, local_rank: int = 0) -> int:
+    """Compute the ZMQ port for a KV-transfer connector.
+
+    Each TP rank gets its own port so that TP > 1 deployments do not
+    cause ``EADDRINUSE`` when multiple sender workers bind on the same
+    host.  The formula is backward-compatible: rank 0 produces the same
+    port as the previous ``base + OFFSET + stage`` formula.
+    """
+    return base_port + KV_TRANSFER_PORT_OFFSET + local_rank * KV_RANK_PORT_STRIDE + from_stage
+
+
+# ------------------------------------------------------------------ #
+#  TP topology validation and rank routing
+# ------------------------------------------------------------------ #
+
+
+def validate_kv_tp_topology(topo: KVTPTopology) -> None:
+    """Reject heterogeneous TP mappings that cannot be routed losslessly."""
+    larger = max(topo.source_tp_size, topo.target_tp_size)
+    smaller = min(topo.source_tp_size, topo.target_tp_size)
+    if larger % smaller != 0:
+        raise ValueError(
+            f"KV TP mapping must be divisible: "
+            f"source_tp_size={topo.source_tp_size}, "
+            f"target_tp_size={topo.target_tp_size}"
+        )
+
+
+def get_kv_target_ranks(topo: KVTPTopology) -> list[int]:
+    """Which remote ranks this local rank sends KV shards to (send side)."""
+    validate_kv_tp_topology(topo)
+    if topo.source_tp_size == topo.target_tp_size:
+        return [topo.local_rank]
+    if topo.source_tp_size > topo.target_tp_size:
+        return [topo.local_rank // (topo.source_tp_size // topo.target_tp_size)]
+    ratio = topo.target_tp_size // topo.source_tp_size
+    return [topo.local_rank * ratio + i for i in range(ratio)]
+
+
+def get_kv_source_ranks(topo: KVTPTopology) -> list[int]:
+    """Which remote ranks this local rank receives KV shards from (recv side)."""
+    validate_kv_tp_topology(topo)
+    if topo.source_tp_size == topo.target_tp_size:
+        return [topo.local_rank]
+    if topo.source_tp_size > topo.target_tp_size:
+        ratio = topo.source_tp_size // topo.target_tp_size
+        return [topo.local_rank * ratio + i for i in range(ratio)]
+    return [topo.local_rank // (topo.target_tp_size // topo.source_tp_size)]
+
+
+# ------------------------------------------------------------------ #
+#  Rank-aware connector key building
+# ------------------------------------------------------------------ #
+
+
+def get_kv_connector_key(
+    req_id: str,
+    from_stage: int | str,
+    chunk_id: int,
+    from_rank: int,
+    to_rank: int,
+) -> str:
+    """Build connector key that includes rank info for KV transfers.
+
+    Format matches PR #2677: ``{req_id}_{from_stage}_{chunk_id}_{from_rank}_{to_rank}``
+    """
+    return f"{req_id}_{from_stage}_{chunk_id}_{from_rank}_{to_rank}"
+
+
+def build_rank_aware_send_keys(
+    request_id: str,
+    from_stage: str,
+    to_stage: str,
+    topo: KVTPTopology,
+    hook: Callable | None = None,
+) -> list[str]:
+    """Build send-side connector keys, checking injectable hook first."""
+    if callable(hook):
+        keys = list(hook(request_id, from_stage, to_stage))
+        if keys:
+            return keys
+    if topo.source_tp_size <= 1 and topo.target_tp_size <= 1:
+        return [f"omni_{from_stage}_to_{to_stage}_kv_cache_{request_id}"]
+    target_ranks = get_kv_target_ranks(topo)
+    return [get_kv_connector_key(request_id, from_stage, 0, topo.local_rank, r) for r in target_ranks]
+
+
+def build_rank_aware_recv_keys(
+    request_id: str,
+    from_stage: str,
+    to_stage: str,
+    topo: KVTPTopology,
+    hook: Callable | None = None,
+) -> list[tuple[str, int | None]]:
+    """Build recv-side connector keys with sender rank info.
+
+    Returns a list of ``(key, from_rank)`` tuples.  ``from_rank`` is
+    ``None`` when TP <= 1 (single sender, no per-rank routing needed).
+    For TP > 1, ``from_rank`` identifies which sender rank owns the
+    key so that the connector can route metadata queries to the
+    correct endpoint.
+    """
+    if callable(hook):
+        raw = list(hook(request_id, from_stage, to_stage))
+        if raw:
+            if isinstance(raw[0], tuple):
+                return raw
+            # Hook returned plain strings (e.g. OmniConnectorModelRunnerMixin.
+            # get_rank_aware_kv_keys). Reconstruct from_rank from topology so
+            # Mooncake connector can route metadata queries to the correct
+            # sender endpoint in heterogeneous TP.
+            # TODO: have the mixin return (key, from_rank) tuples directly
+            # to avoid this indirect reconstruction.
+            source_ranks = get_kv_source_ranks(topo)
+            if len(raw) == len(source_ranks):
+                return list(zip(raw, source_ranks))
+            return [(k, None) for k in raw]
+    if topo.source_tp_size <= 1 and topo.target_tp_size <= 1:
+        return [(f"omni_{from_stage}_to_{to_stage}_kv_cache_{request_id}", None)]
+    source_ranks = get_kv_source_ranks(topo)
+    return [(get_kv_connector_key(request_id, from_stage, 0, r, topo.local_rank), r) for r in source_ranks]
+
+
+# ------------------------------------------------------------------ #
+#  KV tensor head slicing (heterogeneous TP)
+# ------------------------------------------------------------------ #
+
+
+def slice_kv_tensor_heads(
+    tensor: torch.Tensor | None,
+    offset_in_shard: int,
+    num_slices: int,
+) -> torch.Tensor | None:
+    """Slice one KV tensor along its head dimension (dim 1)."""
+    if tensor is None:
+        return None
+    if not isinstance(tensor, torch.Tensor):
+        return tensor
+    if tensor.dim() < 2:
+        raise ValueError(f"Expected KV tensor with a head dimension, got shape={tuple(tensor.shape)}")
+    if num_slices <= 0:
+        raise ValueError(f"num_slices must be > 0, got {num_slices}")
+    if not (0 <= offset_in_shard < num_slices):
+        raise ValueError(f"offset_in_shard must be in [0, {num_slices}), got {offset_in_shard}")
+
+    heads_in_shard = tensor.shape[1]
+    if heads_in_shard % num_slices != 0:
+        raise ValueError(
+            "KV head count must be divisible for heterogeneous TP slicing: "
+            f"heads_in_shard={heads_in_shard}, num_slices={num_slices}"
+        )
+
+    heads_per_slice = heads_in_shard // num_slices
+    start = offset_in_shard * heads_per_slice
+    end = start + heads_per_slice
+    return tensor[:, start:end, ...].contiguous()
+
+
+def slice_layer_blocks(
+    layer_blocks: dict[str, Any],
+    offset_in_shard: int,
+    num_slices: int,
+) -> dict[str, list[torch.Tensor | None]]:
+    """Slice all KV layers for one logical receiver rank."""
+    sliced_blocks: dict[str, list[torch.Tensor | None]] = {}
+    for cache_name in ("key_cache", "value_cache"):
+        cache_list = layer_blocks.get(cache_name, [])
+        sliced_blocks[cache_name] = [
+            slice_kv_tensor_heads(tensor, offset_in_shard, num_slices) for tensor in cache_list
+        ]
+    return sliced_blocks
+
+
+# ------------------------------------------------------------------ #
+#  Multi-rank merge and receiver-side slice
+# ------------------------------------------------------------------ #
+
+
+def merge_received_rank_shards(
+    payloads: list[dict[str, Any]],
+    merger: Callable | None = None,
+) -> dict[str, Any] | None:
+    """Merge multiple source-rank KV shards for one target rank.
+
+    When *merger* is provided (injectable hook), it is called directly.
+    Otherwise the default merges along the head dimension (dim 1).
+    """
+    if callable(merger):
+        return merger(payloads)
+    if not payloads:
+        return None
+    if len(payloads) == 1:
+        return payloads[0]
+
+    base_payload = payloads[0]
+    if not isinstance(base_payload, dict) or "layer_blocks" not in base_payload:
+        return base_payload
+
+    merged: dict[str, Any] = {
+        "request_id": base_payload.get("request_id"),
+        "block_ids": list(base_payload.get("block_ids", [])),
+        "metadata": dict(base_payload.get("metadata", {})),
+    }
+    merged_layer_blocks: dict[str, list[torch.Tensor | None]] = {}
+
+    for cache_name in ("key_cache", "value_cache"):
+        cache_lists = [payload.get("layer_blocks", {}).get(cache_name, []) for payload in payloads]
+        num_layers = max((len(cache_list) for cache_list in cache_lists), default=0)
+        merged_cache: list[torch.Tensor | None] = []
+
+        for layer_idx in range(num_layers):
+            layer_tensors = [
+                cache_list[layer_idx]
+                for cache_list in cache_lists
+                if layer_idx < len(cache_list) and cache_list[layer_idx] is not None
+            ]
+            if not layer_tensors:
+                merged_cache.append(None)
+            elif len(layer_tensors) == 1 or not isinstance(layer_tensors[0], torch.Tensor):
+                merged_cache.append(layer_tensors[0])
+            else:
+                merged_cache.append(torch.cat(layer_tensors, dim=1).contiguous())
+
+        merged_layer_blocks[cache_name] = merged_cache
+
+    merged["layer_blocks"] = merged_layer_blocks
+    return merged
+
+
+def slice_received_rank_shard(
+    payload: dict[str, Any] | None,
+    topo: KVTPTopology,
+    slicer: Callable | None = None,
+) -> dict[str, Any] | None:
+    """Optionally slice a received payload to extract this rank's portion.
+
+    Used when ``to_tp > from_tp``: the sender sent full heads and each
+    receiver rank slices out its own subset.
+    """
+    if callable(slicer):
+        return slicer(payload)
+    if not payload or topo.target_tp_size <= topo.source_tp_size or "layer_blocks" not in payload:
+        return payload
+
+    metadata = payload.get("metadata", {})
+    slice_metadata = metadata.get("tp_head_slice") if isinstance(metadata, dict) else None
+    if isinstance(slice_metadata, dict) and slice_metadata.get("applied"):
+        tagged_rank = slice_metadata.get("target_rank")
+        if tagged_rank is not None and tagged_rank != topo.local_rank:
+            logger.warning(
+                "Received pre-sliced KV payload for unexpected target rank: expected=%s got=%s",
+                topo.local_rank,
+                tagged_rank,
+            )
+        return payload
+
+    ratio = topo.target_tp_size // topo.source_tp_size
+    offset_in_sender = topo.local_rank % ratio
+    updated_metadata = dict(metadata) if isinstance(metadata, dict) else {}
+    updated_metadata["tp_head_slice"] = {
+        "applied": True,
+        "side": "receiver",
+        "target_rank": topo.local_rank,
+        "from_tp": topo.source_tp_size,
+        "to_tp": topo.target_tp_size,
+        "offset_in_shard": offset_in_sender,
+        "num_slices": ratio,
+    }
+    return {
+        "request_id": payload.get("request_id"),
+        "layer_blocks": slice_layer_blocks(payload["layer_blocks"], offset_in_sender, ratio),
+        "block_ids": list(payload.get("block_ids", [])),
+        "metadata": updated_metadata,
+    }
+
+
 def normalize_layer_kv(
     layer_kv: LayerKV,
     *,

vllm_omni/engine/async_omni_engine.py

@@ -61,6 +61,7 @@ from vllm_omni.engine.stage_engine_startup import (
 )
 from vllm_omni.engine.stage_init_utils import (
     StartedLlmStage,
+    _inject_inferred_kv_tp_topology,
     acquire_device_locks,
     build_diffusion_config,
     build_engine_args_dict,
@@ -78,7 +79,10 @@ from vllm_omni.engine.stage_init_utils import (
     setup_stage_devices,
     terminate_alive_proc,
 )
-from vllm_omni.entrypoints.utils import load_and_resolve_stage_configs
+from vllm_omni.entrypoints.utils import (
+    inject_omni_kv_config,
+    load_and_resolve_stage_configs,
+)
 from vllm_omni.inputs.preprocess import OmniInputPreprocessor
 from vllm_omni.platforms import current_omni_platform
 
@@ -378,6 +382,12 @@ class AsyncOmniEngine:
                             omni_kv["omni_to_stage"] = omni_to
                             omni_kv.setdefault("stage_id", metadata.stage_id)
                             engine_args_dict["omni_kv_config"] = omni_kv
+                        if self.stage_configs:
+                            _inject_inferred_kv_tp_topology(
+                                engine_args_dict.get("omni_kv_config"),
+                                metadata.stage_id,
+                                self.stage_configs,
+                            )
                         vllm_config, executor_class = build_vllm_config(
                             stage_cfg,
                             self.model,
@@ -747,10 +757,8 @@ class AsyncOmniEngine:
                                 setup_stage_devices(configured_stage_id, metadata.runtime_cfg)
                                 omni_conn_cfg, omni_from, omni_to = omni_kv_connector
                                 if omni_conn_cfg:
-                                    from vllm_omni.entrypoints.utils import inject_omni_kv_config
-
                                     inject_omni_kv_config(stage_cfg, omni_conn_cfg, omni_from, omni_to)
-                                inject_kv_stage_info(stage_cfg, configured_stage_id)
+                                inject_kv_stage_info(stage_cfg, configured_stage_id, self.stage_configs)
                                 if self.single_stage_mode:
                                     assert self._omni_master_server is not None
                                     stage_clients[stage_idx] = self._launch_diffusion_stage(

vllm_omni/engine/stage_engine_core_client.py

@@ -14,7 +14,9 @@ from vllm.logger import init_logger
 from vllm.v1.engine import EngineCoreRequest
 from vllm.v1.engine.core_client import AsyncMPClient, DPLBAsyncMPClient
 
-from vllm_omni.distributed.omni_connectors.utils.initialization import KV_TRANSFER_PORT_OFFSET
+from vllm_omni.distributed.omni_connectors.utils.initialization import (
+    KV_TRANSFER_PORT_OFFSET,
+)
 from vllm_omni.engine.stage_init_utils import StageMetadata
 
 if TYPE_CHECKING:
@@ -246,6 +248,8 @@ class StageEngineCoreClientBase:
                 from_stage = omni_kv_config.get("omni_from_stage", from_stage)
 
             try:
+                # Orchestrator always reports rank-0's port; receiver
+                # workers add their own local_rank * KV_RANK_PORT_STRIDE.
                 sender_port = int(base_port) + KV_TRANSFER_PORT_OFFSET + int(from_stage)
             except (TypeError, ValueError):
                 logger.warning(
@@ -284,6 +288,7 @@ class StageEngineCoreClientBase:
             self._kv_sender_host = self._resolve_contact_host()
         if self._kv_sender_host is None:
             return None
+        # rank-0 base port; receiver workers adjust per KV_RANK_PORT_STRIDE.
         return {
             "host": self._kv_sender_host,
             "zmq_port": base_port + kv_transfer_port_offset + int(self.stage_id),

vllm_omni/engine/stage_init_utils.py

@@ -13,7 +13,7 @@ import importlib
 import multiprocessing as mp
 import os
 import time
-from collections.abc import Callable
+from collections.abc import Callable, Sequence
 from dataclasses import dataclass
 from typing import Any, Literal
 
@@ -101,8 +101,110 @@ def resolve_worker_cls(engine_args: dict[str, Any]) -> None:
         raise ValueError(f"Unknown worker_type: {worker_type}")
 
 
-def inject_kv_stage_info(stage_cfg: Any, stage_id: int) -> None:
-    """Inject stage metadata into omni_kv_config when present."""
+def _get_attr_or_item(obj: Any, key: str, default: Any = None) -> Any:
+    """Read *key* from *obj* regardless of whether it's a dict or object."""
+    if hasattr(obj, "get"):
+        return obj.get(key, default)
+    return getattr(obj, key, default)
+
+
+def _tp_size_for_stage(stage_configs: Sequence[Any], stage_id: Any) -> int | None:
+    """Resolve tensor_parallel_size for *stage_id* from the loaded stage configs."""
+    id_strs = {str(stage_id)}
+    try:
+        id_strs.add(str(int(stage_id)))
+    except (TypeError, ValueError):
+        pass
+
+    for stage_cfg in stage_configs:
+        if str(getattr(stage_cfg, "stage_id", None)) not in id_strs:
+            continue
+        engine_args = getattr(stage_cfg, "engine_args", None)
+        if engine_args is None:
+            return 1
+        parallel_config = _get_attr_or_item(engine_args, "parallel_config")
+        if parallel_config is not None:
+            tp = _get_attr_or_item(parallel_config, "tensor_parallel_size", 1)
+        else:
+            tp = _get_attr_or_item(engine_args, "tensor_parallel_size", 1)
+        try:
+            return max(1, int(tp))
+        except (TypeError, ValueError):
+            return 1
+    return None
+
+
+def _inject_inferred_kv_tp_topology(
+    omni_kv: Any,
+    stage_id: int,
+    stage_configs: Sequence[Any],
+    engine_input_source: Sequence[int] | None = None,
+) -> None:
+    """Infer adjacent-stage TP topology and inject it into omni_kv_config.
+
+    This keeps heterogeneous TP working without requiring user-authored
+    rank_mapping blocks in config files.
+    """
+    if omni_kv is None:
+        return
+
+    if hasattr(omni_kv, "get"):
+        need_send = bool(omni_kv.get("need_send_cache", False))
+        need_recv = bool(omni_kv.get("need_recv_cache", False))
+        omni_from_stage = omni_kv.get("omni_from_stage")
+        omni_to_stage = omni_kv.get("omni_to_stage")
+        rank_mapping = omni_kv.get("rank_mapping")
+    else:
+        need_send = bool(getattr(omni_kv, "need_send_cache", False))
+        need_recv = bool(getattr(omni_kv, "need_recv_cache", False))
+        omni_from_stage = getattr(omni_kv, "omni_from_stage", None)
+        omni_to_stage = getattr(omni_kv, "omni_to_stage", None)
+        rank_mapping = getattr(omni_kv, "rank_mapping", None)
+
+    if not need_send and not need_recv:
+        return
+
+    current_tp = _tp_size_for_stage(stage_configs, stage_id)
+    if current_tp is None:
+        return
+
+    peer_stage_id = None
+    from_tp = None
+    to_tp = None
+    if str(omni_from_stage) == str(stage_id):
+        peer_stage_id = omni_to_stage
+        from_tp = current_tp
+        to_tp = _tp_size_for_stage(stage_configs, peer_stage_id)
+    elif str(omni_to_stage) == str(stage_id):
+        peer_stage_id = omni_from_stage
+        from_tp = _tp_size_for_stage(stage_configs, peer_stage_id)
+        to_tp = current_tp
+    elif need_recv and engine_input_source:
+        peer_stage_id = engine_input_source[0]
+        from_tp = _tp_size_for_stage(stage_configs, peer_stage_id)
+        to_tp = current_tp
+
+    if from_tp is None or to_tp is None:
+        return
+
+    if not isinstance(rank_mapping, dict):
+        rank_mapping = {}
+    rank_mapping.setdefault("from_tp", int(from_tp))
+    rank_mapping.setdefault("to_tp", int(to_tp))
+
+    if hasattr(omni_kv, "__setitem__"):
+        omni_kv["rank_mapping"] = rank_mapping
+    else:
+        setattr(omni_kv, "rank_mapping", rank_mapping)
+
+
+def inject_kv_stage_info(stage_cfg: Any, stage_id: int, stage_configs: Sequence[Any] | None = None) -> None:
+    """Inject stage_id, engine_input_source, and inferred TP topology into omni_kv_config.
+
+    When *stage_configs* is provided, also infers from_tp/to_tp for
+    heterogeneous TP topologies so the KV transfer manager can compute
+    rank mappings automatically.
+    """
     try:
         engine_args = stage_cfg.engine_args
         if hasattr(engine_args, "get"):
@@ -125,6 +227,14 @@ def inject_kv_stage_info(stage_cfg: Any, stage_id: int) -> None:
                 omni_kv.setdefault("engine_input_source", list(engine_input_source))
             elif hasattr(omni_kv, "__setitem__") and "engine_input_source" not in omni_kv:
                 omni_kv["engine_input_source"] = list(engine_input_source)
+
+        if stage_configs:
+            _inject_inferred_kv_tp_topology(
+                omni_kv,
+                stage_id=stage_id,
+                stage_configs=stage_configs,
+                engine_input_source=engine_input_source,
+            )
     except Exception as e:
         logger.debug("Failed to inject stage info into omni_kv_config: %s", e)
 

vllm_omni/entrypoints/openai/serving_chat.py

@@ -86,6 +86,7 @@ from vllm_omni.entrypoints.openai.image_api_utils import validate_layered_layers
 from vllm_omni.entrypoints.openai.protocol import OmniChatCompletionStreamResponse
 from vllm_omni.entrypoints.openai.protocol.audio import AudioResponse, CreateAudio
 from vllm_omni.entrypoints.openai.utils import (
+    get_stage_type,
     get_supported_speakers_from_hf_config,
     parse_lora_request,
     validate_requested_speaker,
@@ -294,6 +295,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         )
 
         num_inference_steps = None
+        cfg_text_scale = None
+        cfg_img_scale = None
         # Omni multistage image generation: Stage-0 (AR) should receive a clean
         # text prompt (and optional conditioning image/size) so the model's own
         # processor can construct the correct inputs.
@@ -342,6 +345,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                     except Exception:
                         pass
                 negative_prompt = extra_body.get("negative_prompt")
+                cfg_text_scale = extra_body.get("cfg_text_scale")
+                cfg_img_scale = extra_body.get("cfg_img_scale")
 
                 engine_prompt_image: dict[str, Any] | None = None
                 is_img2img = False
@@ -397,14 +402,18 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                     sampling_params_list = self._build_sampling_params_list_from_request(request)
 
                 # Apply user-specified overrides to diffusion stage(s) for image generation
-                if _image_gen_height is not None or _image_gen_width is not None or num_inference_steps is not None:
-                    for idx, sp in enumerate(sampling_params_list):
-                        if hasattr(sp, "height") and _image_gen_height is not None:
-                            sp.height = _image_gen_height
-                        if hasattr(sp, "width") and _image_gen_width is not None:
-                            sp.width = _image_gen_width
-                        if hasattr(sp, "num_inference_steps") and num_inference_steps is not None:
-                            sp.num_inference_steps = num_inference_steps
+                for idx, sp in enumerate(sampling_params_list):
+                    if hasattr(sp, "height") and _image_gen_height is not None:
+                        sp.height = _image_gen_height
+                    if hasattr(sp, "width") and _image_gen_width is not None:
+                        sp.width = _image_gen_width
+                    if hasattr(sp, "num_inference_steps") and num_inference_steps is not None:
+                        sp.num_inference_steps = num_inference_steps
+                    if hasattr(sp, "extra_args") and sp.extra_args is not None:
+                        if cfg_text_scale is not None:
+                            sp.extra_args["cfg_text_scale"] = cfg_text_scale
+                        if cfg_img_scale is not None:
+                            sp.extra_args["cfg_img_scale"] = cfg_img_scale
 
                 self._log_inputs(
                     request_id,
@@ -2108,6 +2117,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
             num_inference_steps = extra_body.get("num_inference_steps")
             guidance_scale = extra_body.get("guidance_scale")
             true_cfg_scale = extra_body.get("true_cfg_scale") or extra_body.get("cfg_scale")
+            cfg_text_scale = extra_body.get("cfg_text_scale")
+            cfg_img_scale = extra_body.get("cfg_img_scale")
             seed = extra_body.get("seed")
             negative_prompt = extra_body.get("negative_prompt")
             num_outputs_per_prompt = extra_body.get("num_outputs_per_prompt", 1)
@@ -2162,6 +2173,10 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 gen_params.guidance_scale = guidance_scale
             if true_cfg_scale is not None:
                 gen_params.true_cfg_scale = true_cfg_scale
+            if cfg_text_scale is not None:
+                gen_params.extra_args["cfg_text_scale"] = cfg_text_scale
+            if cfg_img_scale is not None:
+                gen_params.extra_args["cfg_img_scale"] = cfg_img_scale
             if num_frames is not None:
                 gen_params.num_frames = num_frames
             if guidance_scale_2 is not None:
@@ -2206,10 +2221,30 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
 
             # Generate image
             diffusion_engine = cast(AsyncOmni, self._diffusion_engine)
+            stage_configs = list(getattr(diffusion_engine, "stage_configs", []) or [])
+            default_params_list = list(getattr(diffusion_engine, "default_sampling_params_list", []) or [])
+
+            sampling_params_list: list[Any] = []
+            for idx, stage_cfg in enumerate(stage_configs):
+                if get_stage_type(stage_cfg) == "diffusion":
+                    sampling_params_list.append(gen_params)
+                    continue
+
+                default_stage_params = default_params_list[idx] if idx < len(default_params_list) else SamplingParams()
+                if hasattr(default_stage_params, "clone"):
+                    try:
+                        default_stage_params = default_stage_params.clone()
+                    except Exception:
+                        pass
+                sampling_params_list.append(default_stage_params)
+
+            if not sampling_params_list:
+                sampling_params_list = [gen_params]
+
             result = None
             async for output in diffusion_engine.generate(
                 prompt=gen_prompt,
-                sampling_params_list=[gen_params],  # Pass as single-stage params
+                sampling_params_list=sampling_params_list,
                 request_id=request_id,
             ):
                 result = output

vllm_omni/inputs/data.py

@@ -267,6 +267,10 @@ class OmniDiffusionSamplingParams:
     cfg_text_kv_metadata: dict[str, Any] | None = None
     cfg_img_kv_metadata: dict[str, Any] | None = None
     cfg_kv_request_ids: dict[str, str] | None = None
+    cfg_active_branch: str | None = None
+    cfg_branch_roles: list[str] | None = None
+    cfg_branch_past_key_values: dict[str, Any] | None = None
+    cfg_branch_kv_metadata: dict[str, dict[str, Any]] | None = None
 
     # Component modules
     modules: dict[str, Any] = field(default_factory=dict)