[Feature] Port Bagel RDMA flow to latest main (#2000)

Signed-off-by: Jinheng Li <ahengljh@gmail.com>
Signed-off-by: natureofnature <wzliu@connect.hku.hk>
Co-authored-by: ahengljh <ahengljh@users.noreply.github.com>
Co-authored-by: natureofnature <wzliu@connect.hku.hk>

docs/design/feature/omni_connectors/mooncake_transfer_engine_connector.md

@@ -33,8 +33,8 @@ runtime:
         zmq_port: 50051               # ZMQ base port (see "Port Offset Scheme" below)
         protocol: "rdma"              # "rdma" or "tcp"
         device_name: ""               # RDMA device (e.g., "mlx5_0"), empty for auto-detect
-        memory_pool_size: 2147483648  # 2GB memory pool
-        memory_pool_device: "cpu"     # "cpu" for pinned memory, "cuda" for GPUDirect RDMA
+        memory_pool_size: 4294967296  # 4 GB (CPU); use 2147483648 (2 GB) for GPU
+        memory_pool_device: "cpu"     # "cpu" for pinned memory (recommended), "cuda" for GPUDirect RDMA
 ```
 
 Wire stages to the connector:
@@ -64,8 +64,8 @@ stage_args:
 
 | Parameter | Default | Description |
 |---|---|---|
-| `memory_pool_size` | 1 GB | Total size of the RDMA-registered memory pool in bytes. |
-| `memory_pool_device` | `"cpu"` | `"cpu"`: pinned host memory (recommended). `"cuda"`: GPU VRAM for GPUDirect RDMA (requires NIC-GPU direct PCIe connectivity). |
+| `memory_pool_size` | 4 GB (CPU) / 2 GB (GPU) | Total size of the RDMA-registered memory pool in bytes. Recommended 4 GB for CPU pinned memory; 2 GB for GPU VRAM to conserve device memory. |
+| `memory_pool_device` | `"cpu"` | `"cpu"`: pinned host memory (recommended, works on all topologies). `"cuda"`: GPU VRAM for GPUDirect RDMA (requires NIC-GPU direct PCIe connectivity, PIX topology). |
 
 ### Networking
 
@@ -107,10 +107,10 @@ receiver_connect  = remote_side_channel_port + tp_rank
 
 ## Memory Pool Modes
 
-| Mode | Config | Data Flow | Best For |
-|---|---|---|---|
-| CPU Pinned | `memory_pool_device: "cpu"` | GPU → CPU pool → RDMA → CPU pool → GPU | Most hardware topologies (recommended) |
-| GPUDirect | `memory_pool_device: "cuda"` | GPU → GPU pool → RDMA (NIC reads GPU BAR1) → GPU pool | NIC-GPU direct PCIe (PIX topology) |
+| Mode | Config | Recommended Pool Size | Data Flow | Best For |
+|---|---|---|---|---|
+| CPU Pinned | `memory_pool_device: "cpu"` | 4 GB | GPU → CPU pool → RDMA → CPU pool → GPU | Most hardware topologies (recommended) |
+| GPUDirect | `memory_pool_device: "cuda"` | 2 GB | GPU → GPU pool → RDMA (NIC reads GPU BAR1) → GPU pool | NIC-GPU direct PCIe (PIX topology) |
 
 > **Note**: GPUDirect RDMA requires the NIC and GPU to share a direct PCIe
 > switch (PIX topology). On systems where they are connected via PXB or NODE,

tests/diffusion/test_stage_diffusion_proc.py

@@ -0,0 +1,65 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import asyncio
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import asdict
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def test_process_batch_request_preserves_parent_request_id_and_kv_sender_info():
+    async def run_test():
+        captured = {}
+
+        def step(request):
+            captured["request"] = request
+            return [
+                SimpleNamespace(
+                    images=["img-1"],
+                    _multimodal_output={},
+                    metrics={},
+                    stage_durations={},
+                    peak_memory_mb=0.0,
+                    latents=None,
+                    final_output_type="image",
+                ),
+                SimpleNamespace(
+                    images=["img-2"],
+                    _multimodal_output={},
+                    metrics={},
+                    stage_durations={},
+                    peak_memory_mb=0.0,
+                    latents=None,
+                    final_output_type="image",
+                ),
+            ]
+
+        proc = object.__new__(StageDiffusionProc)
+        proc._engine = SimpleNamespace(step=step)
+        proc._executor = ThreadPoolExecutor(max_workers=1)
+
+        try:
+            result = await proc._process_batch_request(
+                request_id="req-parent",
+                prompts=["hello", "world"],
+                sampling_params_dict=asdict(OmniDiffusionSamplingParams()),
+                kv_sender_info={0: {"host": "10.0.0.2", "zmq_port": 50151}},
+            )
+        finally:
+            proc._executor.shutdown(wait=True)
+
+        request = captured["request"]
+        assert request.request_id == "req-parent"
+        assert request.request_ids == ["req-parent-0", "req-parent-1"]
+        assert request.kv_sender_info == {0: {"host": "10.0.0.2", "zmq_port": 50151}}
+        assert result.request_id == "req-parent"
+        assert result.images == ["img-1", "img-2"]
+
+    asyncio.run(run_test())

tests/distributed/omni_connectors/test_basic_connectors.py

@@ -120,3 +120,61 @@ def test_get_invalid_metadata(shm_connector):
 
     result = shm_connector.get("stage_0", "stage_1", "req_3", {"unknown": "format"})
     assert result is None
+
+
+def test_mooncake_connector_defaults_missing_host_to_detected_ip(monkeypatch: pytest.MonkeyPatch):
+    import vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector as mooncake_module
+
+    class _FakePool:
+        is_cuda = False
+
+        def pin_memory(self):
+            return self
+
+        def data_ptr(self):
+            return 1234
+
+    class _FakeTransferEngine:
+        def initialize(self, host, mode, protocol, device_name):
+            self.host = host
+            self.mode = mode
+            self.protocol = protocol
+            self.device_name = device_name
+            return 0
+
+        def get_rpc_port(self):
+            return 23456
+
+        def register_memory(self, base_ptr, pool_size):
+            del base_ptr, pool_size
+            return 0
+
+        def unregister_memory(self, base_ptr):
+            del base_ptr
+            return 0
+
+    monkeypatch.setattr(mooncake_module, "TransferEngine", _FakeTransferEngine)
+    monkeypatch.setattr(mooncake_module.torch, "empty", lambda *args, **kwargs: _FakePool())
+    monkeypatch.setattr(
+        mooncake_module.MooncakeTransferEngineConnector,
+        "_get_local_ip",
+        lambda self: "10.20.30.40",
+    )
+    monkeypatch.setattr(
+        mooncake_module.MooncakeTransferEngineConnector,
+        "_zmq_listener_loop",
+        lambda self: self._listener_ready.set(),
+    )
+
+    connector = mooncake_module.MooncakeTransferEngineConnector(
+        {
+            "zmq_port": 50051,
+            "memory_pool_size": 4096,
+        }
+    )
+    try:
+        assert connector.host == "10.20.30.40"
+        assert connector.engine.host == "10.20.30.40"
+        assert connector.get_connection_info()["host"] == "10.20.30.40"
+    finally:
+        connector.close()

tests/distributed/omni_connectors/test_kv_flow.py

@@ -1,8 +1,14 @@
+import json
+import struct
+
+import numpy as np
 import pytest
 import torch
 
+import vllm_omni.distributed.omni_connectors.kv_transfer_manager as kv_transfer_manager_module
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.distributed.omni_connectors.kv_transfer_manager import (
+    KVCacheTransferData,
     OmniKVCacheConfig,
     OmniKVTransferManager,
 )
@@ -60,6 +66,35 @@ def common_constants():
     }
 
 
+def _decode_stored_payload(data):
+    if isinstance(data, torch.Tensor) and data.dtype == torch.uint8 and data.dim() == 1:
+        return KVCacheTransferData.from_bytes(data.cpu().numpy().tobytes())
+
+    if isinstance(data, (bytes, bytearray, memoryview)):
+        return KVCacheTransferData.from_bytes(data)
+
+    return data
+
+
+def _make_serialized_payload() -> tuple[bytes, torch.Tensor]:
+    key_tensor = torch.arange(12, dtype=torch.float32).reshape(3, 4)
+    payload = KVCacheTransferData(
+        request_id="req-payload",
+        layer_blocks={"key_cache": [key_tensor], "value_cache": [None]},
+        block_ids=[1],
+        metadata={"seq_len": 3},
+    ).to_bytes()
+    return payload, key_tensor
+
+
+def _rewrite_serialized_header(payload: bytes, mutate_header) -> bytes:
+    header_len = struct.unpack(">I", payload[:4])[0]
+    header = json.loads(payload[4 : 4 + header_len])
+    mutate_header(header)
+    new_header = json.dumps(header, separators=(",", ":")).encode("utf-8")
+    return struct.pack(">I", len(new_header)) + new_header + payload[4 + header_len :]
+
+
 def test_manager_extraction(kv_config, mock_connector, common_constants):
     """Test extraction and sending logic in OmniKVTransferManager."""
     num_layers = common_constants["num_layers"]
@@ -95,7 +130,7 @@ def test_manager_extraction(kv_config, mock_connector, common_constants):
     expected_key = f"stage1->stage2:{full_request_id}"
     assert expected_key in mock_connector.store
 
-    data = mock_connector.store[expected_key]
+    data = _decode_stored_payload(mock_connector.store[expected_key])
     assert data["request_id"] == req_id
     assert "layer_blocks" in data
     assert len(data["layer_blocks"]["key_cache"]) == num_layers
@@ -106,6 +141,116 @@ def test_manager_extraction(kv_config, mock_connector, common_constants):
     assert data["layer_blocks"]["key_cache"][0].shape == expected_shape
 
 
+def test_from_bytes_rejects_out_of_bounds_header_len():
+    payload, _ = _make_serialized_payload()
+    bad_payload = struct.pack(">I", len(payload)) + payload[4:]
+
+    with pytest.raises(ValueError, match="header_len"):
+        KVCacheTransferData.from_bytes(bad_payload)
+
+    with pytest.raises(ValueError, match="header_len"):
+        KVCacheTransferData.from_bytes_gpu(torch.tensor(list(bad_payload), dtype=torch.uint8))
+
+
+def test_from_bytes_rejects_out_of_bounds_tensor_span():
+    payload, _ = _make_serialized_payload()
+    bad_payload = _rewrite_serialized_header(payload, lambda header: header["td"][0].update({"o": 4096}))
+
+    with pytest.raises(ValueError, match="tensor span"):
+        KVCacheTransferData.from_bytes(bad_payload)
+
+    with pytest.raises(ValueError, match="tensor span"):
+        KVCacheTransferData.from_bytes_gpu(torch.tensor(list(bad_payload), dtype=torch.uint8))
+
+
+def test_from_bytes_rejects_unsupported_dtype():
+    payload, _ = _make_serialized_payload()
+    bad_payload = _rewrite_serialized_header(payload, lambda header: header["td"][0].update({"d": "cuda"}))
+
+    with pytest.raises(ValueError, match="Unsupported dtype"):
+        KVCacheTransferData.from_bytes(bad_payload)
+
+    with pytest.raises(ValueError, match="Unsupported dtype"):
+        KVCacheTransferData.from_bytes_gpu(torch.tensor(list(bad_payload), dtype=torch.uint8))
+
+
+def test_from_bytes_uses_explicit_layer_index_descriptor():
+    payload, key_tensor = _make_serialized_payload()
+    payload_with_explicit_index = _rewrite_serialized_header(
+        payload,
+        lambda header: header["td"][0].update({"n": "key_cache_extra_suffix", "i": 0}),
+    )
+
+    data = KVCacheTransferData.from_bytes(payload_with_explicit_index)
+
+    assert torch.equal(data["layer_blocks"]["key_cache"][0], key_tensor)
+
+
+def test_update_sender_info_uses_configured_source_stage():
+    config = OmniKVCacheConfig(
+        connector_config={"type": "mock"},
+        stage_id=2,
+        engine_input_source=[1],
+        need_recv_cache=True,
+    )
+    manager = OmniKVTransferManager(config)
+
+    manager.update_sender_info(
+        {
+            0: {"host": "10.0.0.1", "zmq_port": 50151},
+            1: {"host": "10.0.0.2", "zmq_port": 50152},
+        }
+    )
+
+    assert manager.config.connector_config["sender_host"] == "10.0.0.2"
+    assert manager.config.connector_config["sender_zmq_port"] == 50152
+
+
+def test_clone_received_payload_tensors_breaks_buffer_alias():
+    payload, key_tensor = _make_serialized_payload()
+    raw = np.frombuffer(bytearray(payload), dtype=np.uint8)
+    data = KVCacheTransferData.from_bytes(memoryview(raw))
+
+    OmniKVTransferManager._clone_received_payload_tensors(data)
+    raw[:] = 0
+
+    assert torch.equal(data["layer_blocks"]["key_cache"][0], key_tensor)
+
+
+def test_receive_kv_cache_uses_exponential_backoff(monkeypatch):
+    config = OmniKVCacheConfig(
+        connector_config={"type": "mock"},
+        from_stage="sender",
+        stage_id="receiver",
+        need_recv_cache=True,
+        recv_timeout=0.3,
+    )
+    manager = OmniKVTransferManager(config)
+
+    class _NeverReadyConnector:
+        def get(self, **kwargs):
+            del kwargs
+            return None
+
+    manager._connector = _NeverReadyConnector()
+
+    now = {"value": 0.0}
+    sleep_intervals = []
+
+    monkeypatch.setattr(kv_transfer_manager_module.time, "time", lambda: now["value"])
+
+    def _fake_sleep(interval: float) -> None:
+        sleep_intervals.append(interval)
+        now["value"] += interval
+
+    monkeypatch.setattr(kv_transfer_manager_module.time, "sleep", _fake_sleep)
+
+    data, size = manager.receive_kv_cache_for_request("req-backoff")
+
+    assert (data, size) == (None, 0)
+    assert sleep_intervals == pytest.approx([0.01, 0.02, 0.04, 0.08, 0.16])
+
+
 def test_manager_extraction_tuple_layout(kv_config, mock_connector, common_constants):
     """Test extraction with tuple layout."""
     num_layers = common_constants["num_layers"]
@@ -135,7 +280,7 @@ def test_manager_extraction_tuple_layout(kv_config, mock_connector, common_const
     expected_key = f"stage1->stage2:{full_request_id}"
     assert expected_key in mock_connector.store
 
-    data = mock_connector.store[expected_key]
+    data = _decode_stored_payload(mock_connector.store[expected_key])
     expected_shape = (seq_len, num_heads, head_dim)
     for idx in range(len(kv_caches)):
         assert data["layer_blocks"]["key_cache"][idx].shape == expected_shape
@@ -165,7 +310,7 @@ def test_manager_extraction_mismatched_kv_block_counts(kv_config, mock_connector
     expected_key = f"stage1->stage2:{full_request_id}"
     assert expected_key in mock_connector.store
 
-    data = mock_connector.store[expected_key]
+    data = _decode_stored_payload(mock_connector.store[expected_key])
     expected_shape = (2 * block_size, num_heads, head_dim)
     assert data["layer_blocks"]["key_cache"][0].shape == expected_shape
     assert data["layer_blocks"]["value_cache"][0].shape == expected_shape
@@ -254,6 +399,82 @@ def test_manager_reception(kv_config, mock_connector, common_constants):
     assert req.kv_metadata["seq_len"] == seq_len
 
 
+def test_manager_reception_prefers_parent_request_id_for_batched_request(kv_config, mock_connector, common_constants):
+    """Batched diffusion requests must fetch KV using the parent/global request ID."""
+    num_layers = common_constants["num_layers"]
+    num_heads = common_constants["num_heads"]
+    head_dim = common_constants["head_dim"]
+    seq_len = common_constants["seq_len"]
+    parent_req_id = common_constants["req_id"]
+
+    expected_shape = (seq_len, num_heads, head_dim)
+    key_cache = [torch.randn(expected_shape) for _ in range(num_layers)]
+    value_cache = [torch.randn(expected_shape) for _ in range(num_layers)]
+
+    data_to_receive = {
+        "request_id": parent_req_id,
+        "layer_blocks": {"key_cache": key_cache, "value_cache": value_cache},
+        "metadata": {"seq_len": seq_len},
+        "block_ids": [],
+    }
+
+    manager = OmniKVTransferManager(kv_config)
+    manager._connector = mock_connector
+
+    full_request_id = f"omni_stage1_to_stage2_kv_cache_{parent_req_id}"
+    store_key = f"stage1->stage2:{full_request_id}"
+    mock_connector.store[store_key] = data_to_receive
+
+    req = OmniDiffusionRequest(
+        prompts=["prompt-a", "prompt-b"],
+        sampling_params=OmniDiffusionSamplingParams(),
+        request_ids=[f"{parent_req_id}-0", f"{parent_req_id}-1"],
+        request_id=parent_req_id,
+    )
+
+    success = manager.receive_kv_cache(req, target_device=torch.device("cpu"))
+
+    assert success
+    assert req.kv_metadata["seq_len"] == seq_len
+    assert torch.allclose(req.past_key_values.key_cache[0], key_cache[0])
+
+
+def test_receive_multi_kv_cache_uses_parent_request_id_for_cfg_collection(kv_config):
+    manager = OmniKVTransferManager(kv_config)
+
+    seen = {}
+
+    def collect_cfg(request_id, cfg_request_ids, kv_transfer_manager, target_device):
+        seen["request_id"] = request_id
+        seen["cfg_request_ids"] = cfg_request_ids
+        seen["kv_transfer_manager"] = kv_transfer_manager
+        seen["target_device"] = target_device
+        return {"cfg_text_kv_metadata": {"ok": True}}
+
+    req = OmniDiffusionRequest(
+        prompts=["prompt-a", "prompt-b"],
+        sampling_params=OmniDiffusionSamplingParams(),
+        request_ids=["req-parent-0", "req-parent-1"],
+        request_id="req-parent",
+    )
+    req.sampling_params.cfg_kv_request_ids = {"cfg_text": "req-parent__cfg_text"}
+
+    manager.receive_kv_cache = lambda request, target_device=None: request is req
+
+    success = manager.receive_multi_kv_cache(
+        req,
+        cfg_kv_collect_func=collect_cfg,
+        target_device=torch.device("cpu"),
+    )
+
+    assert success
+    assert seen["request_id"] == "req-parent"
+    assert seen["cfg_request_ids"] == {"cfg_text": "req-parent__cfg_text"}
+    assert seen["kv_transfer_manager"] is manager
+    assert seen["target_device"] == torch.device("cpu")
+    assert req.sampling_params.cfg_text_kv_metadata == {"ok": True}
+
+
 def test_integration_flow(common_constants):
     """Simulate extraction -> connector -> reception."""
     num_layers = common_constants["num_layers"]

tests/engine/test_orchestrator_kv_sender_info.py

@@ -0,0 +1,207 @@
+import asyncio
+from types import SimpleNamespace
+
+import pytest
+from vllm import SamplingParams
+
+from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState
+from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _DummySenderStage:
+    def __init__(self, sender_info):
+        self._sender_info = sender_info
+        self.engine_outputs = None
+
+    def set_engine_outputs(self, outputs):
+        self.engine_outputs = outputs
+
+    def get_kv_sender_info(self):
+        return self._sender_info
+
+
+class _DummyDiffusionStage:
+    stage_type = "diffusion"
+    custom_process_input_func = None
+
+    def __init__(self, engine_input_source=None):
+        self.engine_input_source = engine_input_source or [0]
+        self.calls = []
+
+    async def add_request_async(self, request_id, prompt, sampling_params, kv_sender_info=None):
+        self.calls.append(
+            {
+                "request_id": request_id,
+                "prompt": prompt,
+                "sampling_params": sampling_params,
+                "kv_sender_info": kv_sender_info,
+            }
+        )
+
+
+def test_stage_engine_core_client_builds_kv_sender_info_from_tcp_address():
+    client = object.__new__(StageEngineCoreClient)
+    client.stage_id = 0
+    client.client_addresses = {"input_address": "tcp://10.20.30.40:1234"}
+    client._omni_kv_config = None
+    client._kv_sender_info = None
+    client._kv_sender_initialized = False
+    client._kv_sender_host = client._resolve_contact_host()
+    client._initialize_kv_sender_endpoint()
+
+    assert client.get_kv_sender_info() == {
+        "host": "10.20.30.40",
+        "zmq_port": 50151,
+    }
+
+
+def test_stage_engine_core_client_falls_back_to_detected_ip_for_loopback(monkeypatch):
+    client = object.__new__(StageEngineCoreClient)
+    client.stage_id = 1
+    client.client_addresses = {"input_address": "tcp://127.0.0.1:1234"}
+    client._omni_kv_config = None
+    client._kv_sender_info = None
+    client._kv_sender_initialized = False
+    monkeypatch.setattr(client, "_detect_local_ip", lambda: "192.168.0.12")
+    client._kv_sender_host = client._resolve_contact_host()
+    client._initialize_kv_sender_endpoint()
+
+    assert client.get_kv_sender_info() == {
+        "host": "192.168.0.12",
+        "zmq_port": 50152,
+    }
+
+
+def test_stage_engine_core_client_uses_connector_config_for_sender_port():
+    client = object.__new__(StageEngineCoreClient)
+    client.stage_id = 3
+    client.client_addresses = {"input_address": "tcp://10.20.30.40:1234"}
+    client._kv_sender_info = None
+    client._kv_sender_initialized = False
+    client._omni_kv_config = {
+        "omni_from_stage": "3",
+        "connector_config": {
+            "type": "MooncakeTransferEngineConnector",
+            "role": "sender",
+            "host": "10.20.30.99",
+            "zmq_port": 51000,
+        },
+    }
+    client._kv_sender_host = client._resolve_contact_host()
+    client._initialize_kv_sender_endpoint()
+
+    assert client.get_kv_sender_info() == {
+        "host": "10.20.30.99",
+        "zmq_port": 51103,
+    }
+
+
+def test_stage_engine_core_client_preserves_explicit_loopback_sender_host():
+    client = object.__new__(StageEngineCoreClient)
+    client.stage_id = 2
+    client.client_addresses = {"input_address": "tcp://10.20.30.40:1234"}
+    client._kv_sender_info = None
+    client._kv_sender_initialized = False
+    client._omni_kv_config = {
+        "omni_from_stage": "2",
+        "connector_config": {
+            "type": "MooncakeTransferEngineConnector",
+            "role": "sender",
+            "host": "127.0.0.1",
+            "zmq_port": 51000,
+        },
+    }
+    client._kv_sender_host = client._resolve_contact_host()
+    client._initialize_kv_sender_endpoint()
+
+    assert client.get_kv_sender_info() == {
+        "host": "127.0.0.1",
+        "zmq_port": 51102,
+    }
+
+
+def test_forward_to_diffusion_attaches_kv_sender_info():
+    orchestrator = object.__new__(Orchestrator)
+    sender_stage = _DummySenderStage({"host": "10.0.0.2", "zmq_port": 50151})
+    diffusion_stage = _DummyDiffusionStage(engine_input_source=[0])
+
+    orchestrator.num_stages = 2
+    orchestrator.stage_clients = [sender_stage, diffusion_stage]
+    orchestrator._companion_map = {}
+    orchestrator.stage_vllm_configs = [None, None]
+    orchestrator.output_processors = [None, None]
+
+    params = OmniDiffusionSamplingParams()
+    req_state = OrchestratorRequestState(
+        request_id="req-1",
+        prompt={"prompt": "hello"},
+        sampling_params_list=[SamplingParams(max_tokens=4), params],
+        final_stage_id=1,
+    )
+
+    output = SimpleNamespace(request_id="req-1", finished=True)
+    asyncio.run(Orchestrator._forward_to_next_stage(orchestrator, "req-1", 0, output, req_state))
+
+    assert sender_stage.engine_outputs == [output]
+    assert diffusion_stage.calls[0]["request_id"] == "req-1"
+    assert diffusion_stage.calls[0]["kv_sender_info"] == {
+        0: {"host": "10.0.0.2", "zmq_port": 50151},
+    }
+    assert req_state.stage_submit_ts[1] > 0
+
+
+def test_forward_to_diffusion_uses_engine_input_source_for_kv_sender_info():
+    orchestrator = object.__new__(Orchestrator)
+    source_stage = _DummySenderStage({"host": "10.0.0.2", "zmq_port": 50151})
+    previous_stage = _DummySenderStage({"host": "10.0.0.9", "zmq_port": 59999})
+    diffusion_stage = _DummyDiffusionStage(engine_input_source=[0])
+
+    orchestrator.num_stages = 3
+    orchestrator.stage_clients = [source_stage, previous_stage, diffusion_stage]
+    orchestrator._companion_map = {}
+    orchestrator.stage_vllm_configs = [None, None, None]
+    orchestrator.output_processors = [None, None, None]
+
+    params = OmniDiffusionSamplingParams()
+    req_state = OrchestratorRequestState(
+        request_id="req-3",
+        prompt={"prompt": "hello"},
+        sampling_params_list=[SamplingParams(max_tokens=4), SamplingParams(max_tokens=4), params],
+        final_stage_id=2,
+    )
+
+    output = SimpleNamespace(request_id="req-3", finished=True)
+    asyncio.run(Orchestrator._forward_to_next_stage(orchestrator, "req-3", 1, output, req_state))
+
+    assert previous_stage.engine_outputs == [output]
+    assert diffusion_stage.calls[0]["kv_sender_info"] == {
+        0: {"host": "10.0.0.2", "zmq_port": 50151},
+    }
+
+
+def test_prewarm_diffusion_attaches_kv_sender_info():
+    orchestrator = object.__new__(Orchestrator)
+    sender_stage = _DummySenderStage({"host": "10.0.0.3", "zmq_port": 50151})
+    diffusion_stage = _DummyDiffusionStage(engine_input_source=[0])
+
+    orchestrator.stage_clients = [sender_stage, diffusion_stage]
+    orchestrator.num_stages = 2
+
+    req_state = OrchestratorRequestState(
+        request_id="req-2",
+        prompt={"prompt": "hello"},
+        sampling_params_list=[SamplingParams(max_tokens=4), OmniDiffusionSamplingParams()],
+        final_stage_id=1,
+    )
+
+    stage0_request = SimpleNamespace(prompt_token_ids=[1, 2, 3])
+    asyncio.run(Orchestrator._prewarm_async_chunk_stages(orchestrator, "req-2", stage0_request, req_state))
+
+    assert diffusion_stage.calls[0]["request_id"] == "req-2"
+    assert diffusion_stage.calls[0]["kv_sender_info"] == {
+        0: {"host": "10.0.0.3", "zmq_port": 50151},
+    }
+    assert req_state.stage_submit_ts[1] > 0

vllm_omni/diffusion/request.py

@@ -26,6 +26,8 @@ class OmniDiffusionRequest:
     sampling_params: OmniDiffusionSamplingParams
 
     request_ids: list[str] = field(default_factory=list)
+    request_id: str | None = None
+    kv_sender_info: dict | None = None
 
     def __post_init__(self):
         """Initialize dependent fields after dataclass initialization."""

vllm_omni/diffusion/stage_diffusion_client.py

@@ -179,6 +179,7 @@ class StageDiffusionClient:
         request_id: str,
         prompt: OmniPromptType,
         sampling_params: OmniDiffusionSamplingParams,
+        kv_sender_info: dict[int, dict[str, Any]] | None = None,
     ) -> None:
         self._request_socket.send(
             self._encoder.encode(
@@ -187,6 +188,7 @@ class StageDiffusionClient:
                     "request_id": request_id,
                     "prompt": prompt,
                     "sampling_params": self._sampling_params_to_dict(sampling_params),
+                    "kv_sender_info": kv_sender_info,
                 }
             )
         )
@@ -198,6 +200,7 @@ class StageDiffusionClient:
         request_id: str,
         prompts: list[OmniPromptType],
         sampling_params: OmniDiffusionSamplingParams,
+        kv_sender_info: dict[int, dict[str, Any]] | None = None,
     ) -> None:
         """Submit a list of prompts as a single batched engine call.
 
@@ -206,7 +209,12 @@ class StageDiffusionClient:
         *request_id*.
         """
         task = asyncio.create_task(
-            self._run_batch(request_id, prompts, sampling_params),
+            self._run_batch(
+                request_id,
+                prompts,
+                sampling_params,
+                kv_sender_info,
+            ),
             name=f"diffusion-batch-{request_id}",
         )
         self._tasks[request_id] = task
@@ -216,6 +224,7 @@ class StageDiffusionClient:
         request_id: str,
         prompts: list[OmniPromptType],
         sampling_params: OmniDiffusionSamplingParams,
+        kv_sender_info: dict[int, dict[str, Any]] | None = None,
     ) -> None:
         try:
             self._request_socket.send(
@@ -225,6 +234,7 @@ class StageDiffusionClient:
                         "request_id": request_id,
                         "prompts": prompts,
                         "sampling_params": self._sampling_params_to_dict(sampling_params),
+                        "kv_sender_info": kv_sender_info,
                     }
                 )
             )

vllm_omni/diffusion/stage_diffusion_proc.py

@@ -130,6 +130,7 @@ class StageDiffusionProc:
         request_id: str,
         prompt: Any,
         sampling_params_dict: dict,
+        kv_sender_info: dict[str, Any] | None = None,
     ) -> OmniRequestOutput:
         """Build a diffusion request and run DiffusionEngine.step()."""
         sampling_params = self._reconstruct_sampling_params(sampling_params_dict)
@@ -138,6 +139,8 @@ class StageDiffusionProc:
             prompts=[prompt],
             sampling_params=sampling_params,
             request_ids=[request_id],
+            request_id=request_id,
+            kv_sender_info=kv_sender_info,
         )
 
         loop = asyncio.get_running_loop()
@@ -152,6 +155,7 @@ class StageDiffusionProc:
         request_id: str,
         prompts: list[Any],
         sampling_params_dict: dict,
+        kv_sender_info: dict[str, Any] | None = None,
     ) -> OmniRequestOutput:
         """Build a batched diffusion request and run DiffusionEngine.step().
 
@@ -165,7 +169,9 @@ class StageDiffusionProc:
         request = OmniDiffusionRequest(
             prompts=prompts,
             sampling_params=sampling_params,
-            request_ids=[request_id] * len(prompts),
+            request_ids=[f"{request_id}-{i}" for i in range(len(prompts))],
+            request_id=request_id,
+            kv_sender_info=kv_sender_info,
         )
 
         loop = asyncio.get_running_loop()
@@ -346,10 +352,20 @@ class StageDiffusionProc:
 
         tasks: dict[str, asyncio.Task] = {}
 
-        async def _dispatch_request(request_id: str, prompt: Any, sampling_params_dict: dict) -> None:
+        async def _dispatch_request(
+            request_id: str,
+            prompt: Any,
+            sampling_params_dict: dict,
+            kv_sender_info: dict[str, Any] | None = None,
+        ) -> None:
             """Process a single diffusion request and send the response."""
             try:
-                result = await self._process_request(request_id, prompt, sampling_params_dict)
+                result = await self._process_request(
+                    request_id,
+                    prompt,
+                    sampling_params_dict,
+                    kv_sender_info=kv_sender_info,
+                )
                 await response_socket.send(encoder.encode({"type": "result", "output": result}))
             except DiffusionRequestAbortedError as e:
                 logger.info(
@@ -384,6 +400,7 @@ class StageDiffusionProc:
                             request_id,
                             msg["prompt"],
                             msg["sampling_params"],
+                            msg.get("kv_sender_info"),
                         )
                     )
                     tasks[request_id] = task
@@ -391,9 +408,19 @@ class StageDiffusionProc:
                 elif msg_type == "add_batch_request":
                     request_id = msg["request_id"]
 
-                    async def _dispatch_batch(rid: str, prompts: list, sp_dict: dict) -> None:
+                    async def _dispatch_batch(
+                        rid: str,
+                        prompts: list,
+                        sp_dict: dict,
+                        kv_sender_info: dict[str, Any] | None = None,
+                    ) -> None:
                         try:
-                            result = await self._process_batch_request(rid, prompts, sp_dict)
+                            result = await self._process_batch_request(
+                                rid,
+                                prompts,
+                                sp_dict,
+                                kv_sender_info=kv_sender_info,
+                            )
                             await response_socket.send(encoder.encode({"type": "result", "output": result}))
                         except DiffusionRequestAbortedError as e:
                             logger.info(
@@ -420,6 +447,7 @@ class StageDiffusionProc:
                             request_id,
                             msg["prompts"],
                             msg["sampling_params"],
+                            msg.get("kv_sender_info"),
                         )
                     )
                     tasks[request_id] = task

vllm_omni/distributed/omni_connectors/connectors/mooncake_transfer_engine_connector.py

@@ -277,13 +277,15 @@ class MooncakeTransferEngineConnector(OmniConnectorBase):
         }
 
         self.config = config
-        host_config = config.get("host", "127.0.0.1")
-        # Support "auto" to auto-detect local IP address
-        if host_config.lower() == "auto":
+        host_config = config.get("host")
+        host_value = "auto" if host_config is None else str(host_config)
+        # Default sender/receiver bootstrap to a routable local IP so the
+        # advertised endpoint matches the interface Mooncake binds.
+        if host_value.lower() == "auto" or host_value in {"", "*", "0.0.0.0", "::"}:
             self.host = self._get_local_ip()
             logger.info(f"Auto-detected local IP for RDMA: {self.host}")
         else:
-            self.host = host_config
+            self.host = host_value
         self.zmq_port = config.get("zmq_port", 50051)
         self.protocol = config.get("protocol", "rdma")
 

vllm_omni/distributed/omni_connectors/kv_transfer_manager.py

@@ -2,6 +2,8 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Unified OmniConnector and KV cache transfer management."""
 
+import json
+import struct
 import time
 from collections.abc import Callable
 from dataclasses import asdict, dataclass
@@ -12,12 +14,36 @@ 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
 
 logger = init_logger(__name__)
 
 LayerKV = torch.Tensor | tuple[torch.Tensor, torch.Tensor]
 
+_SAFE_TORCH_DTYPES = {
+    name: dtype
+    for name in (
+        "bool",
+        "uint8",
+        "int8",
+        "int16",
+        "int32",
+        "int64",
+        "float16",
+        "float32",
+        "float64",
+        "bfloat16",
+        "complex64",
+        "complex128",
+        "float8_e4m3fn",
+        "float8_e4m3fnuz",
+        "float8_e5m2",
+        "float8_e5m2fnuz",
+    )
+    if isinstance((dtype := getattr(torch, name, None)), torch.dtype)
+}
+
 
 @dataclass
 class OmniKVCacheConfig:
@@ -46,6 +72,242 @@ class KVCacheTransferData:
         """Convert to dictionary for serialization."""
         return asdict(self)
 
+    def to_bytes(self) -> bytes:
+        """Convert to compact binary format for fast transfer."""
+        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] = []
+        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):
+                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:
+                    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,
+                        "s": list(t.shape),
+                        "o": data_offset,
+                        "b": nbytes,
+                    }
+                )
+                gpu_tensors.append(t.view(torch.uint8).flatten())
+                data_offset += nbytes
+
+        if device is None:
+            raise RuntimeError("No CUDA tensors found, use to_bytes() instead")
+
+        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")
+
+        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)
+
+        pos = len(header_prefix)
+        for t_flat in gpu_tensors:
+            n = t_flat.numel()
+            output[pos : pos + n].copy_(t_flat)
+            pos += n
+
+        return output
+
+    @staticmethod
+    def _load_header_from_memoryview(raw_mv: memoryview) -> tuple[dict[str, Any], memoryview]:
+        if len(raw_mv) < 4:
+            raise ValueError("Corrupted KV payload: missing 4-byte header length")
+
+        header_len = struct.unpack(">I", raw_mv[:4])[0]
+        if header_len > len(raw_mv) - 4:
+            raise ValueError(f"Corrupted KV payload: header_len={header_len} exceeds buffer size={len(raw_mv)}")
+
+        return json.loads(bytes(raw_mv[4 : 4 + header_len])), raw_mv[4 + header_len :]
+
+    @staticmethod
+    def _load_header_from_tensor(gpu_tensor: torch.Tensor) -> tuple[dict[str, Any], int]:
+        if gpu_tensor.dtype != torch.uint8 or gpu_tensor.dim() != 1:
+            raise ValueError("Packed GPU KV payload must be a 1-D uint8 tensor")
+
+        total_bytes = int(gpu_tensor.numel())
+        if total_bytes < 4:
+            raise ValueError("Corrupted KV payload: missing 4-byte header length")
+
+        header_len = struct.unpack(">I", gpu_tensor[:4].cpu().numpy().tobytes())[0]
+        if header_len > total_bytes - 4:
+            raise ValueError(f"Corrupted KV payload: header_len={header_len} exceeds buffer size={total_bytes}")
+
+        header_bytes = gpu_tensor[4 : 4 + header_len].cpu().numpy().tobytes()
+        return json.loads(header_bytes), 4 + header_len
+
+    @staticmethod
+    def _validate_tensor_span(name: str, info: dict[str, Any], tensor_data_bytes: int) -> tuple[int, int]:
+        offset = info["o"]
+        nbytes = info["b"]
+        if offset < 0 or nbytes < 0 or offset + nbytes > tensor_data_bytes:
+            raise ValueError(
+                f"Corrupted KV payload tensor span for {name}: "
+                f"offset={offset}, bytes={nbytes}, tensor_data_bytes={tensor_data_bytes}"
+            )
+        return offset, nbytes
+
+    @staticmethod
+    def _resolve_torch_dtype(dtype_name: Any) -> torch.dtype:
+        torch_dtype = _SAFE_TORCH_DTYPES.get(str(dtype_name))
+        if torch_dtype is None:
+            raise ValueError(f"Unsupported dtype in KV payload: {dtype_name}")
+        return torch_dtype
+
+    @staticmethod
+    def _resolve_layer_idx(info: dict[str, Any], num_layers: int) -> int:
+        layer_idx = info.get("i")
+        if layer_idx is None:
+            name = info.get("n")
+            if isinstance(name, str) and name.startswith("key_cache_"):
+                layer_idx = int(name.removeprefix("key_cache_"))
+            elif isinstance(name, str) and name.startswith("value_cache_"):
+                layer_idx = int(name.removeprefix("value_cache_"))
+            else:
+                raise ValueError(f"Invalid KV tensor name in payload: {name}")
+
+        if not isinstance(layer_idx, int):
+            raise ValueError(f"Invalid layer index in KV payload: {layer_idx}")
+        if layer_idx < 0 or layer_idx >= num_layers:
+            raise ValueError(f"Invalid layer index in KV payload: {layer_idx} (num_layers={num_layers})")
+        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)
+
+        num_layers = header["nl"]
+        key_cache: list[torch.Tensor | None] = [None] * num_layers
+        value_cache: list[torch.Tensor | None] = [None] * num_layers
+
+        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"])
+            )
+            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"],
+        }
+
+    @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)
+
+        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
+
+        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"],
+        }
+
 
 class OmniKVTransferManager:
     """Unified management for OmniConnector and KV cache transfer.
@@ -79,6 +341,13 @@ class OmniKVTransferManager:
             else (None, None)
         )
 
+        if config.need_send_cache and config.connector_config:
+            try:
+                _ = self.connector
+                logger.info("Sender connector eagerly initialized")
+            except Exception as e:
+                logger.warning("Failed to eagerly initialize sender connector: %s", e)
+
     @classmethod
     def _create(cls, cfg: dict | None) -> "OmniKVTransferManager":
         """Create manager from raw config dict."""
@@ -140,8 +409,39 @@ class OmniKVTransferManager:
             cfg = self.config.connector_config
             if cfg and (c_type := cfg.get("type")):
                 try:
-                    logger.info(f"Initializing OmniConnector with config: {cfg}")
                     c_extra = {k: v for k, v in cfg.items() if k != "type"}
+                    if c_type == "MooncakeTransferEngineConnector":
+                        base_port = c_extra.get("zmq_port", 50051)
+                        c_extra["from_stage"] = (
+                            str(self.config.from_stage) if self.config.from_stage is not None else "0"
+                        )
+                        c_extra["to_stage"] = str(self.config.to_stage) if self.config.to_stage is not None else "1"
+
+                        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
+                        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)
+
+                    logger.info(
+                        "Initializing OmniConnector (purpose=kv_transfer) with config: %s, role: %s",
+                        cfg,
+                        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}")
@@ -157,6 +457,85 @@ class OmniKVTransferManager:
         """Get connector (compatibility wrapper for existing code)."""
         return self.connector
 
+    def _resolve_sender_info(
+        self, sender_info: dict[str, Any], sender_stage_id: str | int | None = None
+    ) -> dict[str, Any] | None:
+        if not sender_info:
+            return None
+
+        if "host" in sender_info:
+            return sender_info
+
+        if not isinstance(sender_info, dict):
+            return None
+
+        preferred_keys: list[str | int] = []
+        if sender_stage_id is None:
+            recv_from, _ = self.recv_stages
+            sender_stage_id = recv_from
+
+        if sender_stage_id is not None:
+            preferred_keys.append(sender_stage_id)
+            preferred_keys.append(str(sender_stage_id))
+            try:
+                preferred_keys.append(int(sender_stage_id))
+            except (TypeError, ValueError):
+                pass
+
+        for key in dict.fromkeys(preferred_keys):
+            info = sender_info.get(key)
+            if isinstance(info, dict) and "host" in info:
+                return info
+
+        candidates = [info for info in sender_info.values() if isinstance(info, dict) and "host" in info]
+        if len(candidates) == 1:
+            return candidates[0]
+
+        if candidates:
+            logger.warning(
+                "Ambiguous sender_info for sender_stage_id=%s: "
+                "expected caller to resolve a single sender entry, got %s",
+                sender_stage_id,
+                sender_info,
+            )
+        return None
+
+    @staticmethod
+    def _clone_received_payload_tensors(data: dict[str, Any]) -> dict[str, Any]:
+        if not isinstance(data, dict) or "layer_blocks" not in data:
+            return data
+
+        layer_blocks = data["layer_blocks"]
+        for cache_name in ("key_cache", "value_cache"):
+            cache_list = layer_blocks.get(cache_name, [])
+            for idx, tensor in enumerate(cache_list):
+                if isinstance(tensor, torch.Tensor):
+                    cache_list[idx] = tensor.clone()
+        return data
+
+    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."""
+        if not self.config.need_recv_cache:
+            return
+
+        actual_info = self._resolve_sender_info(sender_info, sender_stage_id=sender_stage_id)
+        if not actual_info or "host" not in actual_info:
+            logger.warning("Invalid sender_info format: %s", sender_info)
+            return
+
+        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")
+
+        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"))
+            except Exception:
+                if hasattr(self._connector, "sender_host"):
+                    self._connector.sender_host = actual_info.get("host")
+                if hasattr(self._connector, "sender_zmq_port"):
+                    self._connector.sender_zmq_port = actual_info.get("zmq_port")
+
     def handle_finished_requests_kv_transfer(
         self,
         finished_reqs: dict[str, dict[str, Any]],
@@ -203,7 +582,8 @@ class OmniKVTransferManager:
 
                 custom_metadata = data.get("custom_metadata")
 
-                # Extract KV cache from GPU blocks -> CPU tensors
+                # Extract KV cache from GPU blocks and keep it on-device when
+                # possible so raw-data connectors can use the fast path.
                 kv_data = self._extract_kv_cache(
                     req_id, block_ids, seq_len, kv_caches, block_size, cache_dtype, custom_metadata
                 )
@@ -280,9 +660,8 @@ class OmniKVTransferManager:
                 flat_k = flat_k[:seq_len]
                 flat_v = flat_v[:seq_len]
 
-            # Move to CPU
-            key_cache[layer_idx] = flat_k.detach().cpu().contiguous()
-            value_cache[layer_idx] = flat_v.detach().cpu().contiguous()
+            key_cache[layer_idx] = flat_k.detach().contiguous()
+            value_cache[layer_idx] = flat_v.detach().contiguous()
 
         if not any(k is not None for k in key_cache):
             return None
@@ -311,14 +690,40 @@ class OmniKVTransferManager:
         if not from_stage or not to_stage:
             raise ValueError("Transfer stages (omni_from_stage, omni_to_stage) not configured")
 
-        # Prepare data and transfer with retry
-        data_dict = kv_data.to_dict()
-        data_dict["request_id"] = transfer_req_id
+        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)
 
-        success, size, _ = self._transfer_with_retry(from_stage, to_stage, f"kv_cache_{transfer_req_id}", data_dict)
+        try:
+            if supports_raw:
+                transfer_data = kv_data.to_gpu_tensor()
+            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
+
+        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)
+        elapsed = time.perf_counter() - transfer_start
 
         if success:
-            logger.info(f"KV transfer OK: {transfer_req_id}, {size} bytes")
+            mbps = (size / 1024 / 1024) / elapsed if elapsed > 0 else 0
+            logger.info(
+                "KV transfer OK: %s, %s bytes, %.3fs, %.1f MB/s",
+                transfer_req_id,
+                size,
+                elapsed,
+                mbps,
+            )
         else:
             logger.error(f"KV transfer FAILED: {transfer_req_id}")
 
@@ -327,7 +732,7 @@ class OmniKVTransferManager:
         from_stage: str,
         to_stage: str,
         request_id: str,
-        data: dict[str, Any],
+        data: "dict[str, Any] | bytes | torch.Tensor",
         max_retries: int = 3,
     ) -> tuple[bool, int, dict[str, Any] | None]:
         """Transfer data with retry and exponential backoff.
@@ -393,6 +798,8 @@ class OmniKVTransferManager:
 
         timeout = self.config.recv_timeout
         start_time = time.time()
+        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}...")
 
@@ -400,33 +807,74 @@ class OmniKVTransferManager:
             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:
-                    data, size = result
-                    logger.info(f"Successfully received KV cache for {request_id}, {size} bytes")
+                    raw_data, size = result
+                    elapsed = time.time() - start_time
+                    link_ms = (time.perf_counter() - link_start) * 1000
+                    managed_buffer = None
 
-                    # Move tensors to target device if specified
-                    if target_device is not None and isinstance(data, dict) and "layer_blocks" in data:
-                        layer_blocks = data["layer_blocks"]
-                        for cache_list in [
-                            layer_blocks.get("key_cache", []),
-                            layer_blocks.get("value_cache", []),
-                        ]:
-                            for i, tensor in enumerate(cache_list):
-                                if isinstance(tensor, torch.Tensor) and tensor.device != target_device:
-                                    cache_list[i] = tensor.to(target_device).contiguous()
+                    if hasattr(raw_data, "tensor") and hasattr(raw_data, "release"):
+                        managed_buffer = raw_data
+                        try:
+                            buf_tensor = raw_data.tensor
+                            if buf_tensor.is_cuda:
+                                data = KVCacheTransferData.from_bytes_gpu(buf_tensor)
+                                raw_data.release()
+                                managed_buffer = None
+                            else:
+                                data = KVCacheTransferData.from_bytes(memoryview(buf_tensor.numpy()))
+                                data = self._clone_received_payload_tensors(data)
+                                raw_data.release()
+                                managed_buffer = None
+                        except Exception as e:
+                            logger.error("Failed to deserialize KV cache from ManagedBuffer: %s", e)
+                            if managed_buffer is not None:
+                                raw_data.release()
+                            return None, 0
+                    elif isinstance(raw_data, (bytes, bytearray)):
+                        data = KVCacheTransferData.from_bytes(raw_data)
+                    elif isinstance(raw_data, torch.Tensor) and raw_data.dtype == torch.uint8 and raw_data.dim() == 1:
+                        data = KVCacheTransferData.from_bytes(raw_data.cpu().numpy().tobytes())
+                    else:
+                        data = raw_data
 
+                    try:
+                        if isinstance(data, dict) and "layer_blocks" in data:
+                            layer_blocks = data["layer_blocks"]
+                            for cache_list in [
+                                layer_blocks.get("key_cache", []),
+                                layer_blocks.get("value_cache", []),
+                            ]:
+                                for i, tensor in enumerate(cache_list):
+                                    if not isinstance(tensor, torch.Tensor):
+                                        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()
+
+                    logger.info(
+                        "Successfully received KV cache for %s, %s bytes, wait=%.3fs, link=%.1fms",
+                        request_id,
+                        size,
+                        elapsed,
+                        link_ms,
+                    )
                     return data, size
 
                 if time.time() - start_time > timeout:
                     logger.error(f"Timeout waiting for KV cache for request {request_id} after {timeout}s")
                     return None, 0
 
-                time.sleep(0.5)
+                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}")
@@ -459,6 +907,16 @@ class OmniKVTransferManager:
             if hasattr(req, "sampling_params") and req.sampling_params is not None:
                 req.sampling_params.kv_metadata = data["metadata"]
 
+    @staticmethod
+    def _resolve_request_id(req: Any) -> str | None:
+        """Resolve the logical request ID used for KV transfer lookups."""
+        request_id = getattr(req, "request_id", None)
+        if request_id:
+            return request_id
+        if hasattr(req, "request_ids") and req.request_ids:
+            return req.request_ids[0]
+        return None
+
     # Legacy compatibility method
     def receive_kv_cache(self, req: Any, target_device: torch.device | None = None) -> bool:
         """Receive KV cache and populate request object (legacy interface).
@@ -470,11 +928,11 @@ class OmniKVTransferManager:
         Returns:
             True if successful, False otherwise
         """
-        request_id = getattr(req, "request_id", None)
-        if not request_id and hasattr(req, "request_ids") and req.request_ids:
-            # Adaptation for new OmniDiffusionRequest which has list of prompts/ids
-            request_id = req.request_ids[0]
+        kv_sender_info = getattr(req, "kv_sender_info", None)
+        if kv_sender_info:
+            self.update_sender_info(kv_sender_info, sender_stage_id=self.recv_stages[0])
 
+        request_id = self._resolve_request_id(req)
         if not request_id:
             logger.warning("Request has no ID, cannot receive KV cache")
             return False
@@ -513,9 +971,7 @@ class OmniKVTransferManager:
 
         cfg_ids = getattr(getattr(req, "sampling_params", None), "cfg_kv_request_ids", None)
         if cfg_ids and cfg_kv_collect_func:
-            request_id = getattr(req, "request_id", None) or (
-                req.request_ids[0] if hasattr(req, "request_ids") and req.request_ids else None
-            )
+            request_id = self._resolve_request_id(req)
             try:
                 cfg_kvs = cfg_kv_collect_func(
                     request_id,

vllm_omni/distributed/omni_connectors/utils/initialization.py

@@ -19,9 +19,17 @@ else:
 
 logger = get_connector_logger(__name__)
 
+# Reserve a separate port range for KV-transfer sockets so they do not
+# collide with request-forwarding endpoints that share the same base port.
+KV_TRANSFER_PORT_OFFSET = 100
+
 
 def initialize_connectors_from_config(
-    config_path: str | Path | None = None, default_shm_threshold: int = 65536
+    config_path: str | Path | None = None,
+    default_shm_threshold: int = 65536,
+    purpose: str = "request_forwarding",
+    caller_stage_id: int | str | None = None,
+    is_sender: bool | None = None,
 ) -> tuple[OmniTransferConfig | None, dict[tuple[str, str], OmniConnectorBase]]:
     """
     Initialize connectors from configuration file.
@@ -36,12 +44,20 @@ def initialize_connectors_from_config(
         return None, {}
 
     # create connectors from config
-    connectors = create_connectors_from_config(transfer_config.connectors)
+    connectors = create_connectors_from_config(
+        transfer_config.connectors,
+        purpose=purpose,
+        caller_stage_id=caller_stage_id,
+        is_sender=is_sender,
+    )
     return transfer_config, connectors
 
 
 def create_connectors_from_config(
     connectors_config: dict[tuple[str, str], ConnectorSpec],
+    purpose: str = "request_forwarding",
+    caller_stage_id: int | str | None = None,
+    is_sender: bool | None = None,
 ) -> dict[tuple[str, str], OmniConnectorBase]:
     """
     Create connectors from config.
@@ -52,12 +68,59 @@ def create_connectors_from_config(
     Returns:
         A dictionary of connectors.
     """
+    purpose_port_offsets = {
+        "request_forwarding": 0,
+        "kv_transfer": KV_TRANSFER_PORT_OFFSET,
+    }
+    port_offset = purpose_port_offsets.get(purpose, 0)
+    orchestrator_port_offset = 200
+
     connectors = {}
     for edge_key, connector_spec in connectors_config.items():
+        from_stage, to_stage = edge_key
         try:
-            connector = OmniConnectorFactory.create_connector(connector_spec)
+            if connector_spec.name == "MooncakeTransferEngineConnector":
+                extra = dict(connector_spec.extra) if connector_spec.extra else {}
+                base_port = extra.get("zmq_port", 50051)
+                try:
+                    stage_offset = int(from_stage)
+                except (TypeError, ValueError):
+                    stage_offset = 0
+
+                if str(caller_stage_id) == "orchestrator":
+                    adjusted_port = base_port + orchestrator_port_offset + stage_offset
+                else:
+                    adjusted_port = base_port + port_offset + stage_offset
+                extra["zmq_port"] = adjusted_port
+
+                if is_sender is not None:
+                    extra["role"] = "sender" if is_sender else "receiver"
+                    if not is_sender:
+                        extra.setdefault("sender_host", extra.get("host", "127.0.0.1"))
+                        extra.setdefault("sender_zmq_port", adjusted_port)
+                elif caller_stage_id is not None:
+                    caller_str = str(caller_stage_id)
+                    if caller_str == from_stage:
+                        extra["role"] = "sender"
+                    elif caller_str == to_stage:
+                        extra["role"] = "receiver"
+                        extra.setdefault("sender_host", extra.get("host", "127.0.0.1"))
+                        extra.setdefault("sender_zmq_port", adjusted_port)
+                    else:
+                        extra["role"] = "sender"
+                else:
+                    extra["role"] = extra.get("role", "auto")
+
+                connector = OmniConnectorFactory.create_connector(ConnectorSpec(name=connector_spec.name, extra=extra))
+            else:
+                connector = OmniConnectorFactory.create_connector(connector_spec)
             connectors[edge_key] = connector
-            logger.info(f"Created connector for {edge_key[0]} -> {edge_key[1]}: {type(connector).__name__}")
+            logger.info(
+                "Created connector for %s -> %s: %s",
+                from_stage,
+                to_stage,
+                type(connector).__name__,
+            )
         except Exception as e:
             raise RuntimeError(f"Failed to initialize connector for edge {edge_key}: {e}") from e
 
@@ -289,7 +352,11 @@ def initialize_orchestrator_connectors(
     else:
         default_shm_threshold = max(0, shm_threshold_bytes)
     transfer_config, connectors = initialize_connectors_from_config(
-        config_path, default_shm_threshold=default_shm_threshold
+        config_path,
+        default_shm_threshold=default_shm_threshold,
+        purpose="request_forwarding",
+        caller_stage_id="orchestrator",
+        is_sender=True,
     )
     return transfer_config, connectors
 
@@ -316,6 +383,7 @@ def get_stage_connector_config(
 def build_stage_connectors(
     stage_id: int,
     connectors_config: dict[str, Any],
+    purpose: str = "request_forwarding",
 ) -> dict[tuple[str, str], Any] | None:
     """Instantiate OmniConnectors for a stage based on config."""
     if not connectors_config:
@@ -352,7 +420,12 @@ def build_stage_connectors(
 
     try:
         # Use unified connector creation logic
-        connectors = create_connectors_from_config(stage_connector_specs)
+        connectors = create_connectors_from_config(
+            stage_connector_specs,
+            purpose=purpose,
+            caller_stage_id=stage_id,
+            is_sender=False,
+        )
     except Exception as exc:  # pragma: no cover - defensive logging
         # Fail fast so the stage does not start with missing connectors.
         logger.exception("[Stage-%s] Failed to initialize connectors: %s", stage_id, exc)

vllm_omni/engine/async_omni_engine.py

@@ -395,15 +395,17 @@ class AsyncOmniEngine:
                         proc=proc,
                         addresses=addresses,
                     )
+                    logger.info("[AsyncOmniEngine] Stage %s engine launch started", metadata.stage_id)
+                    # Keep the stage-specific device visibility until vLLM
+                    # finishes starting all child processes.
+                    complete_stage_handshake(proc, handshake_address, addresses, vllm_config)
+                    logger.info("[AsyncOmniEngine] Stage %s engine startup completed", metadata.stage_id)
                 finally:
                     if previous_visible_devices is None:
                         current_omni_platform.unset_device_control_env_var()
                     else:
                         current_omni_platform.set_device_control_env_var(previous_visible_devices)
 
-            logger.info("[AsyncOmniEngine] Stage %s engine launch started", metadata.stage_id)
-            complete_stage_handshake(proc, handshake_address, addresses, vllm_config)
-            logger.info("[AsyncOmniEngine] Stage %s engine startup completed", metadata.stage_id)
             assert started_stage is not None
             return started_stage
         except Exception:

vllm_omni/engine/orchestrator.py

@@ -477,6 +477,30 @@ class Orchestrator:
             ),
         )
 
+    def _build_kv_sender_info(self, sender_stage_ids: list[int]) -> dict[int, dict[str, Any]] | None:
+        """Build per-request sender info for diffusion KV-transfer receivers."""
+        sender_infos: dict[int, dict[str, Any]] = {}
+        for sender_stage_id in dict.fromkeys(sender_stage_ids):
+            if sender_stage_id < 0 or sender_stage_id >= self.num_stages:
+                continue
+
+            sender_stage = self.stage_clients[sender_stage_id]
+            get_sender_info = getattr(sender_stage, "get_kv_sender_info", None)
+            if not callable(get_sender_info):
+                continue
+
+            sender_info = get_sender_info()
+            if not sender_info:
+                logger.warning(
+                    "[Orchestrator] Stage-%s has no KV sender info available",
+                    sender_stage_id,
+                )
+                continue
+
+            sender_infos[sender_stage_id] = sender_info
+
+        return sender_infos or None
+
     async def _forward_to_next_stage(
         self,
         req_id: str,
@@ -522,14 +546,22 @@ class Orchestrator:
                         req_id,
                     )
 
+            source_stage_ids = list(getattr(next_client, "engine_input_source", None) or [stage_id])
+            kv_sender_info = self._build_kv_sender_info(sender_stage_ids=source_stage_ids)
             if isinstance(diffusion_prompt, list):
                 await next_client.add_batch_request_async(
                     req_id,
                     diffusion_prompt,
                     params,
+                    kv_sender_info=kv_sender_info,
                 )
             else:
-                await next_client.add_request_async(req_id, diffusion_prompt, params)
+                await next_client.add_request_async(
+                    req_id,
+                    diffusion_prompt,
+                    params,
+                    kv_sender_info=kv_sender_info,
+                )
             req_state.stage_submit_ts[next_stage_id] = _time.time()
             return
 
@@ -731,7 +763,14 @@ class Orchestrator:
             params = req_state.sampling_params_list[next_stage_id]
 
             if next_client.stage_type == "diffusion":
-                await next_client.add_request_async(request_id, req_state.prompt, params)
+                source_stage_ids = list(getattr(next_client, "engine_input_source", None) or [next_stage_id - 1])
+                kv_sender_info = self._build_kv_sender_info(sender_stage_ids=source_stage_ids)
+                await next_client.add_request_async(
+                    request_id,
+                    req_state.prompt,
+                    params,
+                    kv_sender_info=kv_sender_info,
+                )
                 req_state.stage_submit_ts[next_stage_id] = _time.time()
                 continue
 

vllm_omni/engine/stage_engine_core_client.py

@@ -6,12 +6,15 @@ Directly inherits from vLLM's AsyncMPClient to reuse EngineCore architecture.
 
 from __future__ import annotations
 
+import socket
 from typing import TYPE_CHECKING, Any
+from urllib.parse import urlparse
 
 from vllm.logger import init_logger
 from vllm.v1.engine import EngineCoreRequest
 from vllm.v1.engine.core_client import AsyncMPClient
 
+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:
@@ -76,6 +79,11 @@ class StageEngineCoreClient(AsyncMPClient):
 
         self.engine_outputs: Any = None
         self._proc = proc
+        self.client_addresses = dict(client_addresses or {})
+        self._omni_kv_config = getattr(getattr(vllm_config, "model_config", None), "omni_kv_config", None)
+        self._kv_sender_host = self._resolve_contact_host()
+        self._kv_sender_info: dict[str, Any] | None = None
+        self._kv_sender_initialized = False
 
         logger.info(
             "[StageEngineCoreClient] Stage-%s initializing EngineCore",
@@ -104,6 +112,7 @@ class StageEngineCoreClient(AsyncMPClient):
                     shutdown_error,
                 )
             raise
+        self._initialize_kv_sender_endpoint()
         logger.info(
             "[StageEngineCoreClient] Stage-%s EngineCore running",
             self.stage_id,
@@ -118,6 +127,118 @@ class StageEngineCoreClient(AsyncMPClient):
 
     # ==================== Stage Methods ====================
 
+    @staticmethod
+    def _detect_local_ip() -> str | None:
+        """Best-effort local IP detection for cross-node connector bootstrap."""
+        try:
+            with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock:
+                sock.connect(("8.8.8.8", 80))
+                return sock.getsockname()[0]
+        except Exception:
+            try:
+                return socket.gethostbyname(socket.gethostname())
+            except Exception:
+                return None
+
+    def _resolve_contact_host(self) -> str | None:
+        """Resolve a routable host for this stage from its client addresses."""
+        for key in ("input_address", "output_address", "stats_update_address"):
+            address = self.client_addresses.get(key)
+            if not address:
+                continue
+            host = urlparse(address).hostname
+            if host in {None, "", "*", "0.0.0.0", "::"}:
+                continue
+            if host in {"localhost", "127.0.0.1"}:
+                detected = self._detect_local_ip()
+                if detected:
+                    return detected
+                continue
+            return host
+        return self._detect_local_ip()
+
+    def _get_kv_connector_config(self) -> dict[str, Any] | None:
+        omni_kv_config = getattr(self, "_omni_kv_config", None)
+        if not isinstance(omni_kv_config, dict):
+            return None
+        connector_config = omni_kv_config.get("connector_config")
+        if not isinstance(connector_config, dict):
+            return None
+        return connector_config
+
+    def _resolve_sender_host_from_config(self, connector_config: dict[str, Any]) -> str | None:
+        host = connector_config.get("sender_host") or connector_config.get("host")
+        if host in {None, "", "auto", "*", "0.0.0.0", "::"}:
+            return self._resolve_contact_host()
+        return str(host)
+
+    def _initialize_kv_sender_endpoint(self) -> None:
+        if self._kv_sender_initialized:
+            return
+        self._kv_sender_initialized = True
+        connector_config = self._get_kv_connector_config()
+        if connector_config is None or connector_config.get("role") != "sender":
+            return
+
+        sender_host = self._resolve_sender_host_from_config(connector_config)
+        if sender_host is not None:
+            self._kv_sender_host = sender_host
+
+        sender_port = connector_config.get("sender_zmq_port")
+        if sender_port is None:
+            base_port = connector_config.get("zmq_port")
+            if base_port is None:
+                return
+
+            omni_kv_config = getattr(self, "_omni_kv_config", None)
+            from_stage = self.stage_id
+            if isinstance(omni_kv_config, dict):
+                from_stage = omni_kv_config.get("omni_from_stage", from_stage)
+
+            try:
+                sender_port = int(base_port) + KV_TRANSFER_PORT_OFFSET + int(from_stage)
+            except (TypeError, ValueError):
+                logger.warning(
+                    "[StageEngineCoreClient] Stage-%s could not resolve sender_zmq_port "
+                    "from base_port=%s and from_stage=%s",
+                    self.stage_id,
+                    base_port,
+                    from_stage,
+                )
+                return
+
+        if self._kv_sender_host is None:
+            return
+
+        self._kv_sender_info = {
+            "host": str(self._kv_sender_host),
+            "zmq_port": int(sender_port),
+        }
+
+    def get_kv_sender_info(
+        self,
+        *,
+        base_port: int = 50051,
+        kv_transfer_port_offset: int = KV_TRANSFER_PORT_OFFSET,
+    ) -> dict[str, Any] | None:
+        """Build sender bootstrap info for diffusion KV transfer receivers.
+
+        ``base_port`` and ``kv_transfer_port_offset`` are only used by the
+        legacy fallback path when no connector-level sender endpoint is
+        configured in ``omni_kv_config``.
+        """
+        if self._kv_sender_info is not None:
+            return dict(self._kv_sender_info)
+
+        if self._kv_sender_host is None:
+            self._kv_sender_host = self._resolve_contact_host()
+        if self._kv_sender_host is None:
+            return None
+        return {
+            "host": self._kv_sender_host,
+            "zmq_port": base_port + kv_transfer_port_offset + int(self.stage_id),
+        }
+
     def set_engine_outputs(self, engine_outputs: EngineCoreOutput) -> None:
         """Set engine outputs (called by orchestrator)."""
         self.engine_outputs = engine_outputs

vllm_omni/model_executor/stage_configs/bagel.yaml

@@ -1,5 +1,9 @@
 # Stage 0: Thinker (multimodal understanding + text generation)
 
+# By default this config uses the shared-memory connector for stage-0 -> stage-1 forwarding.
+# To switch to RDMA, add output_connectors/input_connectors that point to
+# rdma_connector and keep the rest of the pipeline unchanged.
+
 stage_args:
   - stage_id: 0
     stage_type: llm
@@ -36,6 +40,9 @@ stage_args:
       seed: 52
       detokenize: True
       repetition_penalty: 1.05
+    # Optional RDMA override:
+    # output_connectors:
+    #   to_stage_1: rdma_connector
 
   - stage_id: 1
     stage_type: diffusion
@@ -62,6 +69,9 @@ stage_args:
     is_comprehension: false
     default_sampling_params:
       seed: 52
+    # Optional RDMA override:
+    # input_connectors:
+    #   from_stage_0: rdma_connector
 
 # Runtime edges
 runtime:
@@ -78,6 +88,23 @@ runtime:
       extra:
         shm_threshold_bytes: 65536 # 64KB threshold
 
+    # Optional RDMA connector template for Bagel. To enable it, point
+    # stage-0 output_connectors/to_stage_1 and stage-1 input_connectors/from_stage_0
+    # to rdma_connector instead of relying on the default shared-memory path.
+    rdma_connector:
+      name: MooncakeTransferEngineConnector
+      extra:
+        host: "auto"
+        zmq_port: 50051
+        protocol: "rdma"
+        device_name: ""
+        # Memory pool for RDMA-registered buffers.
+        # Supports both CPU pinned memory ("cpu") and GPU VRAM ("cuda").
+        # CPU mode works on all topologies; GPU mode (GPUDirect RDMA) requires
+        # NIC-GPU direct PCIe connectivity (PIX topology).
+        # Recommended: 4 GB for CPU, 2 GB for GPU (to conserve VRAM).
+        memory_pool_size: 4294967296   # 4 GB
+        memory_pool_device: "cpu"
 
   edges:
     - from: 0