[Core] Add NIXL omni connector (#6093)

Signed-off-by: yuanwu <yuan.wu@intel.com>
Signed-off-by: Yuan Wu <yuan.wu@intel.com>

docs/.nav.yml

@@ -136,6 +136,7 @@ nav:
           - Mooncake Store Connector: design/feature/omni_connectors/mooncake_store_connector.md
           - Mooncake Transfer Engine Connector: design/feature/omni_connectors/mooncake_transfer_engine_connector.md
           - Mori Transfer Engine Connector: design/feature/omni_connectors/mori_transfer_engine_connector.md
+          - NIXL Connector: design/feature/omni_connectors/nixl_connector.md
           - Shared Memory Connector: design/feature/omni_connectors/shared_memory_connector.md
           - Yuanrong Store Connector: design/feature/omni_connectors/yuanrong_connector.md
           - Yuanrong Transfer Engine Connector: design/feature/omni_connectors/yuanrong_transfer_engine_connector.md

docs/configuration/environment_variables.md

@@ -90,6 +90,13 @@ depends on the installed kernels and model path.
 | `VLLM_VIDEO_ASYNC_CHUNK` | `on` or `off`; default `on` | Streaming video output; read on attribute access | Environment-only setting. Values are trimmed and case-normalized; an invalid value warns once and uses `on`. | Experimental |
 | `VLLM_VIDEO_AUDIO_DELTA_MODE` | `fast` or `slow`; default `fast` | Streaming video audio deltas; read on attribute access | Environment-only setting. Values are trimmed and case-normalized; an invalid value warns once and uses `fast`. | Experimental |
 
+### NIXL stage transfer
+
+| Name | Type and default | Applies to and read time | Precedence and invalid values | Lifecycle |
+| --- | --- | --- | --- | --- |
+| `VLLM_OMNI_NIXL_LEASE_S` | Float seconds; default `3600` | Unclaimed producer payload expiry; connector construction | Nonempty environment value overrides `lease_seconds`; empty falls back to config/default. Invalid floats raise `ValueError`. Claimed READ allocations never expire by time. | Experimental |
+| `VLLM_OMNI_NIXL_XFER_TIMEOUT_S` | Float seconds; default `300` | Receiver READ wait; connector construction | Nonempty environment value overrides `transfer_timeout_s`; empty falls back to config/default. Invalid floats raise `ValueError`. Timeout does not cancel DMA or release active allocations. | Experimental |
+
 ### Server storage
 
 Storage names use Pydantic's nested-settings delimiter (`__`). They are read

docs/design/feature/disaggregated_inference.md

@@ -9,6 +9,7 @@ Backend-specific setup lives in separate docs:
 - [MooncakeStoreConnector](omni_connectors/mooncake_store_connector.md)
 - [MooncakeTransferEngineConnector](omni_connectors/mooncake_transfer_engine_connector.md)
 - [MoriTransferEngineConnector](omni_connectors/mori_transfer_engine_connector.md)
+- [NixlConnector](omni_connectors/nixl_connector.md)
 - [YuanrongConnector](omni_connectors/yuanrong_connector.md)
 - [YuanrongTransferEngineConnector](omni_connectors/yuanrong_transfer_engine_connector.md)
 
@@ -25,6 +26,7 @@ Current connectors operate in D2H2D (device to host to device) mode.
 | Multi node (Mooncake Store) | MooncakeStoreConnector | TCP-based, requires Mooncake Master + metadata server. |
 | Multi node (Mooncake RDMA) | MooncakeTransferEngineConnector | RDMA/TCP direct transfer with managed memory pool. Fastest. |
 | Multi node (Mori RDMA) | MoriTransferEngineConnector | RDMA direct transfer via Mori IOEngine. |
+| Multi node (NIXL) | NixlConnector | NIXL `READ` over UCX; the only backend covering Intel XPU. |
 | Multi node (Yuanrong) | YuanrongConnector | Requires Yuanrong Datasystem + etcd. |
 | Ascend NPU P2P (Yuanrong TE) | YuanrongTransferEngineConnector | Uses Yuanrong TransferEngine directly. Configure NPU device IPv4 and `memory_pool_device: "npu"`. |
 

docs/design/feature/omni_connectors/nixl_connector.md

@@ -0,0 +1,135 @@
+# NixlConnector
+
+## When to Use
+
+Multi-node or intra-node stage transfer over NIXL, which brokers RDMA/shared-memory
+transports through a single API. Useful when the deployment already standardises on
+NIXL for KV transfer, or when the target accelerator has no Mooncake/Mori backend --
+NIXL reaches Intel XPU through UCX's Level Zero support.
+
+## Mechanism
+
+Uses vLLM's `NixlWrapper` (`vllm.distributed.nixl_utils`) to register the producer's
+tensors and let the consumer pull them with a NIXL `READ`.
+
+- Data Plane: NIXL agent-to-agent `READ`, GPU-to-GPU where the backend allows it.
+- Control Plane: a ZMQ ROUTER socket atomically claims the source allocation before
+  returning descriptors. Forwarded `put()` metadata includes the generation and
+  endpoint; the consumer claims that generation before submitting any READ.
+
+Transfer metadata uses schema version 1. Tensor descriptors are grouped by NIXL
+memory type (`DRAM`, `VRAM`, and so on), while each descriptor retains its global
+tensor index so mixed CPU/accelerator payloads can be reconstructed in order.
+
+Payloads are not restricted to tensors. A single tensor or a list of tensors is
+transferred as-is; a nested structure has its tensor leaves extracted and shipped
+alongside a msgpack-encoded skeleton; anything else is msgpack-encoded into one
+uint8 tensor. Mapping and sequence structure is preserved. `msgspec.Struct` values,
+including `OmniPayloadStruct`, are normalised to mappings so their tensor leaves use
+the direct NIXL path rather than CPU object serialisation.
+
+Each NIXL agent uses strict thread synchronization. A consumer performs one short
+metadata query per scheduler poll (10ms by default), which prevents an unavailable
+producer from blocking unrelated requests. The regular transfer-completion handshake
+keeps the configured transfer timeout.
+
+## Installation
+
+CUDA hosts can install the published wheel:
+
+```bash
+pip install nixl
+```
+
+For Intel XPU, use `docker/Dockerfile.xpu`. Its upstream vLLM XPU base image
+provides a matching torch-xpu, UCX with Level Zero support, and NIXL runtime.
+
+## Configuration
+
+```yaml
+connectors:
+  nixl_connector:
+    name: NixlConnector
+    extra:
+      host: "auto"
+      zmq_port: 50061
+      backends: ["UCX"]
+
+stages:
+  - stage_id: 0
+    output_connectors:
+      to_stage_1: nixl_connector
+
+  - stage_id: 1
+    input_connectors:
+      from_stage_0: nixl_connector
+```
+
+Parameters:
+
+- `host`: address the producer binds its handshake socket to (`"auto"` to detect).
+- `zmq_port`: handshake port. If a producer omits it, an ephemeral port is opened
+  for ownership claims and returned in `put()` metadata. The value in the deploy YAML is a
+  base port. Purpose, replica, rank, and producer-stage offsets are added centrally,
+  so colocated connectors and tensor-parallel workers receive distinct endpoints
+  without one config entry per rank.
+- `sender_host` / `sender_zmq_port`: consumer-side override naming the producer's
+  handshake endpoint. Only needed when the consumer cannot learn it from metadata.
+  In replicated deployments, the producer endpoint is carried with each request and
+  selected at `get()` time; it is not mutable connector-wide state.
+  Explicit sender ports are preserved, not recomputed using the receiver replica.
+  Intermediate NIXL stages retain their incoming receiver configuration and a
+  separate outgoing bind configuration, so one worker can receive and publish.
+- `backends`: NIXL backends to register memory with. Defaults to `["UCX"]`.
+- `receive_device`: forces where received tensors land. By default the consumer
+  keeps the producer's device *type* but uses its own current device of that type.
+- `memory_type`: overrides the NIXL memory type, otherwise `DRAM` for CPU tensors
+  and `VRAM` for accelerator tensors.
+- `lease_seconds`: how long a `put()` payload stays registered while waiting to be
+  read (default 3600). The consumer reports completion, so this only bounds payloads
+  nobody has claimed. An internal reaper releases only unclaimed registrations.
+  `VLLM_OMNI_NIXL_LEASE_S` overrides it.
+- `transfer_timeout_s`: how long a `get()` waits for its `READ` to complete
+  (default 300). NIXL 1.3 has no transfer cancellation API, so a timed-out transfer's
+  buffers and registrations remain owned by the connector until NIXL reports a
+  terminal state; `close()` waits for that state before releasing them.
+  `VLLM_OMNI_NIXL_XFER_TIMEOUT_S` overrides it.
+
+### Source ownership and failure limits
+
+Claims and expiry use the same lock. Completion identifies both a payload
+generation and a unique READ claim; duplicate or stale ACKs cannot release a
+different reader or replacement payload. Replacing a claimed key is rejected.
+Terminal errors also release claims; timeout, unknown status, and active sibling
+transfers retain the complete destination bundle and source claim until terminal.
+Zero-element tensors retain their shape/dtype/skeleton positions but never create
+DMA descriptors or registrations.
+
+A lost metadata response, lost completion ACK, or abandoned consumer can retain a
+claim indefinitely. NIXL 1.3 cannot prove remote cancellation, so neither TTL nor
+`cleanup()` frees those allocations. Producer `close()` rejects new work and
+retains its agent, listener and claimed allocations; a later `close()` finishes
+teardown after claims drain. Permanently abandoned claims remain until process
+exit. This favors memory safety over bounded shutdown/memory usage. Only trusted
+peers may access this unauthenticated control plane. Both endpoints must use the
+claim-aware protocol; rolling compatibility with older GET_META clients is not
+provided. Legacy externally-owned metadata without a generation is accepted only
+under the caller's lifetime guarantee, not managed as a leased allocation here.
+
+## Validation
+
+The native CUDA smoke test requires NIXL and at least two visible CUDA devices:
+
+```bash
+UCX_RCACHE_MAX_UNRELEASED=1024 pytest -o addopts='' -q -s \
+  tests/distributed/omni_connectors/test_nixl_connector_native.py
+```
+
+It starts independent producer and consumer processes, loads the real UCX backend,
+performs the ZMQ `GET_META` handshake and a native NIXL `READ` from GPU 0 to GPU 1,
+checks mixed CPU/GPU tensor leaves from an `OmniPayloadStruct`, sends `XFER_DONE`,
+and verifies that producer registrations are released. This validates same-host,
+cross-GPU operation. Inter-node UCX transport is supported by the design but was not
+validated by this test.
+
+For more details, refer to the [NIXL repository](https://github.com/ai-dynamo/nixl).

docs/design/index.md

@@ -36,6 +36,7 @@ implementation contract; it is not, by itself, a general support claim.
 - [Mooncake Store Connector](feature/omni_connectors/mooncake_store_connector.md)
 - [Mooncake Transfer Engine Connector](feature/omni_connectors/mooncake_transfer_engine_connector.md)
 - [Mori Transfer Engine Connector](feature/omni_connectors/mori_transfer_engine_connector.md)
+- [NIXL Connector](feature/omni_connectors/nixl_connector.md)
 - [Shared Memory Connector](feature/omni_connectors/shared_memory_connector.md)
 - [Yuanrong Store Connector](feature/omni_connectors/yuanrong_connector.md)
 - [Yuanrong Transfer Engine Connector](feature/omni_connectors/yuanrong_transfer_engine_connector.md)

tests/config/test_environment_variables.py

@@ -164,7 +164,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
     """Make an inventory expansion an explicit review decision."""
     category_counts = Counter(item.category for item in ENVIRONMENT_VARIABLE_INVENTORY.values())
     assert category_counts == {
-        EnvironmentVariableCategory.PUBLIC_OMNI: 25,
+        EnvironmentVariableCategory.PUBLIC_OMNI: 27,
         EnvironmentVariableCategory.INHERITED_VLLM: 20,
         EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27,
         EnvironmentVariableCategory.MODEL_SPECIFIC: 63,

tests/core/sched/test_omni_scheduling_coordinator.py

@@ -58,6 +58,7 @@ def _make_request(req_id: str, status: str = "waiting") -> SimpleNamespace:
         num_computed_tokens=0,
         _all_token_ids=[],
         _output_token_ids=[],
+        payload_sender_info=None,
     )
 
 
@@ -270,6 +271,19 @@ class TestWaitingForInputTransition(unittest.TestCase):
         self.assertEqual(len(coord.pending_input_registrations), 1)
         self.assertEqual(coord.pending_input_registrations[0].request_id, "r1")
 
+    def test_pending_input_registration_carries_payload_sender_info(self):
+        coord = OmniSchedulingCoordinator(stage_id=1)
+        req = _make_request("r1", status=RequestStatus.WAITING_FOR_INPUT)
+        req.payload_sender_info = {"host": "10.0.0.1", "zmq_port": 50051}
+        waiting = MockQueue([req])
+
+        coord.process_pending_full_payload_inputs(waiting, set())
+
+        self.assertEqual(
+            coord.pending_input_registrations[0].payload_sender_info,
+            {"host": "10.0.0.1", "zmq_port": 50051},
+        )
+
     def test_idle_cycles_retain_received_marker_before_request_appears(self):
         coord = OmniSchedulingCoordinator(stage_id=1)
         coord._full_payload_input_received.add("late")

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -4,7 +4,7 @@
 import threading
 from collections import deque
 from types import MethodType, SimpleNamespace
-from unittest.mock import Mock, patch
+from unittest.mock import Mock, call, patch
 
 import pytest
 import torch
@@ -543,7 +543,7 @@ def test_create_connector_config_parsing(monkeypatch, raw_cfg, expected_name, ex
     assert connector == "ok"
     assert isinstance(captured["spec"], ConnectorSpec)
     assert captured["spec"].name == expected_name
-    assert captured["spec"].extra == expected_extra
+    assert captured["spec"].extra == {**expected_extra, "stage_id": 0}
 
 
 def test_load_poll(build_adapter):
@@ -570,8 +570,51 @@ def test_load_poll(build_adapter):
     assert request.additional_information == payload
     assert adapter.get_req_chunk["req-1"] == 1
     assert "req-1" not in adapter._finished_load_reqs
-    assert "req-1" in adapter.upstream_exhausted_requests
-    assert "req-1" not in adapter._pending_load_reqs
+
+
+def test_load_poll_uses_request_scoped_payload_sender_endpoint(build_adapter):
+    adapter, connector = build_adapter(stage_id=2, model_mode="ar")
+    request = _req("req-1", RequestStatus.WAITING, external_req_id="external-1")
+    request.payload_sender_info = {"host": "10.0.0.1", "zmq_port": 50051}
+    connector.get.return_value = None
+
+    adapter._poll_single_request(_dequeue_load_entry(adapter, request))
+
+    connector.get.assert_called_once_with(
+        "1",
+        "2",
+        "req-1_1_0",
+        {"source_host": "10.0.0.1", "source_port": 50051},
+    )
+
+
+def test_load_poll_keeps_concurrent_payload_sender_endpoints_distinct(build_adapter):
+    adapter, connector = build_adapter(stage_id=2, model_mode="ar")
+    first = _req("req-1", RequestStatus.WAITING)
+    first.payload_sender_info = {"host": "10.0.0.1", "zmq_port": 50051}
+    second = _req("req-2", RequestStatus.WAITING)
+    second.payload_sender_info = {"host": "10.0.0.2", "zmq_port": 51051}
+    connector.get.return_value = None
+
+    second_entry = _dequeue_load_entry(adapter, second)
+    first_entry = _dequeue_load_entry(adapter, first)
+    adapter._poll_single_request(second_entry)
+    adapter._poll_single_request(first_entry)
+
+    assert connector.get.call_args_list == [
+        call(
+            "1",
+            "2",
+            "req-2_1_0",
+            {"source_host": "10.0.0.2", "source_port": 51051},
+        ),
+        call(
+            "1",
+            "2",
+            "req-1_1_0",
+            {"source_host": "10.0.0.1", "source_port": 50051},
+        ),
+    ]
 
 
 def test_load_async_does_not_requeue_registered_inflight_request(build_adapter):
@@ -2770,7 +2813,8 @@ def test_successful_put_keeps_terminal_chunk_sequence(build_adapter, finish_insi
                 scheduler.add_request(final_update)
         return True, 1, {}
 
-    def get_by_key(_from_stage, _to_stage, key):
+    def get_by_key(_from_stage, _to_stage, key, metadata=None):
+        assert metadata is None
         payload = payloads.pop(key, None)
         return (payload, 1) if payload is not None else None
 

tests/distributed/omni_connectors/test_nixl_connector.py

@@ -0,0 +1,891 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""
+Unit tests for the NixlConnector metadata handshake.
+
+NIXL itself is stubbed out, so these tests exercise the ZMQ control plane and
+the payload normalisation rather than any actual RDMA transfer.
+"""
+
+import ctypes
+import sys
+import time
+import types
+
+import msgspec
+import pytest
+import torch
+import zmq
+
+from vllm_omni.data_entry_keys import HiddenStatesStruct, MetaStruct, OmniPayloadStruct
+
+pytestmark = [pytest.mark.cpu, pytest.mark.parallel, pytest.mark.core_model]
+
+PORT = 47431
+
+
+def _wait_for_metadata(consumer, key, metadata=None, *, timeout=2.0):
+    """Retry outside the connector's deliberately bounded metadata query."""
+    deadline = time.monotonic() + timeout
+    while time.monotonic() < deadline:
+        resolved = consumer._resolve_metadata(key, metadata)
+        if resolved is not None:
+            return resolved
+        time.sleep(0.01)
+    pytest.fail(f"Timed out waiting for metadata for {key!r}")
+
+
+@pytest.mark.parametrize("misses", [0, 2])
+def test_wait_for_metadata_retries_until_success(monkeypatch, misses):
+    expected = {"claim_id": "claim"}
+    forwarded = {"source_host": "127.0.0.1", "source_port": PORT}
+    calls = []
+
+    def resolve(key, metadata):
+        calls.append((key, metadata))
+        return None if len(calls) <= misses else expected
+
+    ticks = iter(range(10))
+    monkeypatch.setattr(time, "monotonic", lambda: next(ticks))
+    monkeypatch.setattr(time, "sleep", lambda _: None)
+    consumer = types.SimpleNamespace(_resolve_metadata=resolve)
+    assert _wait_for_metadata(consumer, "retry", forwarded, timeout=5) is expected
+    assert calls == [("retry", forwarded)] * (misses + 1)
+
+
+def test_wait_for_metadata_stops_at_deadline(monkeypatch):
+    calls = []
+    ticks = iter([0.0, 0.0, 1.0, 2.0])
+    monkeypatch.setattr(time, "monotonic", lambda: next(ticks))
+    monkeypatch.setattr(time, "sleep", lambda _: None)
+    consumer = types.SimpleNamespace(_resolve_metadata=lambda *args: calls.append(args))
+    with pytest.raises(pytest.fail.Exception, match="Timed out waiting for metadata for 'missing'"):
+        _wait_for_metadata(consumer, "missing")
+    assert calls == [("missing", None)] * 2
+
+
+@pytest.mark.parametrize("async_chunk", [False, True])
+def test_three_stage_incoming_and_outgoing_endpoints(nixl_connector_cls, async_chunk):
+    from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec, OmniTransferConfig
+    from vllm_omni.distributed.omni_connectors.utils.initialization import resolve_connector_spec
+    from vllm_omni.engine.stage_init_utils import get_stage_connector_spec
+
+    config = OmniTransferConfig(
+        connectors={
+            ("0", "1"): ConnectorSpec(name="NixlConnector", extra={"host": "127.0.0.1", "zmq_port": PORT}),
+            ("1", "2"): ConnectorSpec(name="NixlConnector", extra={"host": "127.0.0.1", "zmq_port": PORT + 10}),
+        }
+    )
+    connectors = []
+    try:
+        for stage in range(3):
+            spec = get_stage_connector_spec(config, stage, async_chunk)
+            resolved = resolve_connector_spec(ConnectorSpec(**spec), stage_id=stage, role=spec["extra"]["role"])
+            connectors.append(nixl_connector_cls(resolved.extra))
+        first, middle, last = connectors
+        assert middle._sender_zmq_port == PORT
+        assert middle._serving_handshake
+        assert middle._zmq_port == PORT + 11
+        first.put("0", "1", "incoming", torch.ones(1))
+        incoming = _wait_for_metadata(middle, "incoming")
+        assert incoming["sender_zmq_port"] == PORT
+        assert incoming["generation"] == first._published["incoming"]["generation"]
+        middle.put("1", "2", "outgoing", torch.ones(1))
+        outgoing = _wait_for_metadata(last, "outgoing")
+        assert outgoing["sender_zmq_port"] == PORT + 11
+        assert outgoing["generation"] == middle._published["outgoing"]["generation"]
+    finally:
+        for connector in reversed(connectors):
+            for pending in connector._pending.values():
+                pending.claims.clear()  # No DMA in this metadata-only probe.
+            connector.close()
+
+
+@pytest.mark.parametrize("name", ["NixlConnector", "MooncakeTransferEngineConnector"])
+def test_explicit_sender_port_is_not_derived(name):
+    from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec
+    from vllm_omni.distributed.omni_connectors.utils.initialization import resolve_connector_spec
+
+    spec = ConnectorSpec(name=name, extra={"sender_zmq_port": 55000, "from_stage": 0})
+    resolved = resolve_connector_spec(spec, stage_id=1, role="receiver", replica_id=3, local_rank=2)
+    assert resolved.extra["sender_zmq_port"] == 55000
+    assert spec.extra == {"sender_zmq_port": 55000, "from_stage": 0}
+    spec.extra["zmq_port"] = "${UNUSED_NIXL_BASE_PORT}"
+    assert resolve_connector_spec(spec, stage_id=1, role="receiver").extra["sender_zmq_port"] == 55000
+
+
+@pytest.mark.parametrize("need_recv_cache", [False, True])
+@pytest.mark.parametrize("update_before_init", [False, True])
+def test_manager_receiver_does_not_bind_shared_producer_port(producer, need_recv_cache, update_before_init):
+    """Payload-only and KV receivers must dial, not bind, the incoming edge."""
+    from vllm_omni.distributed.omni_connectors.kv_transfer_manager import (
+        OmniKVCacheConfig,
+        OmniKVTransferManager,
+    )
+
+    manager = OmniKVTransferManager(
+        OmniKVCacheConfig(
+            connector_config={
+                "type": "NixlConnector",
+                "role": "receiver",
+                "host": "127.0.0.1",
+                "zmq_port": PORT,
+                "backends": ["UCX"],
+            },
+            from_stage="0",
+            to_stage="1",
+            stage_id=1,
+            need_recv_cache=need_recv_cache,
+        )
+    )
+    sender_info = {"host": producer.host, "zmq_port": producer._zmq_port}
+    if update_before_init:
+        manager.update_sender_info(sender_info)
+    receiver = manager.connector
+    try:
+        assert receiver is not None, "Receiver incorrectly tried to bind the producer's occupied port"
+        assert receiver._zmq_port is None
+        assert not receiver._serving_handshake
+        if not update_before_init:
+            assert receiver._sender_zmq_port is None
+            manager.update_sender_info(sender_info)
+        assert receiver._sender_zmq_port == producer._zmq_port
+        assert manager.config.connector_config["zmq_port"] == PORT  # Do not mutate the shared specification.
+        success, _, _ = producer.put("0", "1", "manager-receive", torch.ones(1))
+        assert success
+        metadata = _wait_for_metadata(receiver, "manager-receive")
+        assert metadata["generation"] == producer._published["manager-receive"]["generation"]
+        receiver._notify_transfer_done("manager-receive", metadata)
+        assert not producer._pending
+    finally:
+        if receiver is not None:
+            receiver.close()
+
+
+@pytest.mark.parametrize("need_recv_cache", [False, True])
+def test_manager_receiver_preserves_explicit_standalone_sender(producer, need_recv_cache):
+    from vllm_omni.distributed.omni_connectors.kv_transfer_manager import (
+        OmniKVCacheConfig,
+        OmniKVTransferManager,
+    )
+
+    manager = OmniKVTransferManager(
+        OmniKVCacheConfig(
+            connector_config={
+                "type": "NixlConnector",
+                "role": "receiver",
+                "zmq_port": PORT,
+                "sender_host": producer.host,
+                "sender_zmq_port": producer._zmq_port,
+            },
+            need_recv_cache=need_recv_cache,
+        )
+    )
+    receiver = manager.connector
+    try:
+        assert receiver is not None
+        assert receiver._zmq_port is None and not receiver._serving_handshake
+        assert (receiver._sender_host, receiver._sender_zmq_port) == (producer.host, producer._zmq_port)
+        success, _, _ = producer.put("0", "1", "standalone-receive", torch.ones(1))
+        assert success
+        metadata = _wait_for_metadata(receiver, "standalone-receive")
+        assert metadata["generation"] == producer._published["standalone-receive"]["generation"]
+        receiver._notify_transfer_done("standalone-receive", metadata)
+        assert not producer._pending
+    finally:
+        if receiver is not None:
+            receiver.close()
+
+
+def test_middle_stage_advertises_its_outgoing_replica_endpoint():
+    from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient
+
+    client = object.__new__(StageEngineCoreClient)
+    client.stage_id = 1
+    client.replica_id = 3
+    client.vllm_config = types.SimpleNamespace(
+        model_config=types.SimpleNamespace(
+            stage_connector_config={
+                "name": "NixlConnector",
+                "extra": {
+                    "role": "receiver",
+                    "host": "10.0.0.1",
+                    "zmq_port": 47000,
+                    "from_stage": 0,
+                    "outgoing": {"host": "10.0.0.2", "zmq_port": 48000, "from_stage": 1},
+                },
+            }
+        )
+    )
+    assert client._build_payload_sender_info() == {"host": "10.0.0.2", "zmq_port": 51073}
+
+
+@pytest.mark.parametrize("direct", [False, True])
+@pytest.mark.usefixtures("reliable_claim_queries")
+def test_claimed_source_survives_expiry_and_cleanup(producer, consumer, direct):
+    _, _, metadata = producer.put("0", "1", "claimed", torch.ones(1))
+    resolved = consumer._resolve_metadata("claimed", metadata if direct else None)
+    pending = producer._pending["claimed"]
+    pending.deadline = 0
+    producer._cleanup_expired_pending()
+    assert producer._pending.get("claimed") is pending
+    producer.cleanup("claimed")
+    assert producer._pending.get("claimed") is pending
+    assert producer._agent.registered
+    consumer._notify_transfer_done("claimed", resolved)
+    consumer._notify_transfer_done("claimed", resolved)
+    assert producer._pending == {}
+    assert producer._agent.registered == []
+
+
+@pytest.mark.usefixtures("reliable_claim_queries")
+def test_claims_are_generation_scoped_and_duplicate_ack_cannot_release_sibling(producer, consumer):
+    _, _, original = producer.put("0", "1", "owners", torch.ones(1))
+    first = consumer._resolve_metadata("owners", original)
+    second = consumer._resolve_metadata("owners", original)
+    assert first["claim_id"] != second["claim_id"]
+    assert not producer.put("0", "1", "owners", torch.zeros(1))[0]
+    consumer._notify_transfer_done("owners", first)
+    consumer._notify_transfer_done("owners", first)
+    assert len(producer._pending["owners"].claims) == 1
+    consumer._notify_transfer_done("owners", second)
+    assert not producer._pending
+    producer.put("0", "1", "owners", torch.zeros(1))
+    consumer._notify_transfer_done("owners", second)
+    assert "owners" in producer._pending
+    assert consumer._resolve_metadata("owners", original) is None
+
+
+@pytest.mark.parametrize("get_metadata", [None, {"schema_version": 1, "tensor_specs": [], "descriptor_groups": []}])
+@pytest.mark.usefixtures("reliable_claim_queries")
+def test_abandoned_claim_survives_close_and_late_completion(producer, consumer, get_metadata):
+    _, _, metadata = producer.put("0", "1", "abandoned", torch.ones(1))
+    claimed = consumer._resolve_metadata("abandoned", metadata)
+    producer._pending["abandoned"].deadline = 0
+    producer.close()
+    assert producer._closing and not producer._closed
+    assert producer._agent.registered
+    assert producer._listener_thread.is_alive()
+    with pytest.raises(RuntimeError, match="closed"):
+        producer.put("0", "1", "new", torch.ones(1))
+    metrics = dict(producer._metrics)
+    with pytest.raises(RuntimeError, match="Cannot get data: NixlConnector is closed"):
+        producer.get("0", "1", "new", get_metadata)
+    assert producer._metrics == metrics
+    assert producer._pending["abandoned"].claims == {claimed["claim_id"]}
+    assert consumer._resolve_metadata("abandoned", metadata) is None
+    consumer._notify_transfer_done("abandoned", claimed)
+    assert not producer._pending
+    assert not producer._agent.registered
+    assert producer._listener_thread.is_alive()
+    producer.close()
+    assert not producer._agent.registered
+    assert producer._closed
+
+
+@pytest.fixture
+def copying_native_agent(consumer, monkeypatch):
+    """Strict fake descriptors and actual CPU copies; not native NIXL evidence."""
+    calls = []
+
+    def descriptors(regions, memory_type):
+        assert regions and all(region[0] > 0 and region[1] > 0 for region in regions)
+        return regions
+
+    def prepare(operation, local, local_ids, remote, remote_ids):
+        assert operation == "READ"
+        assert local_ids == remote_ids == list(range(len(local)))
+        return list(zip(local, remote, strict=True))
+
+    def transfer(pairs):
+        calls.append(pairs)
+        for destination, source in pairs:
+            assert destination[1] == source[1] > 0
+            ctypes.memmove(destination[0], source[0], source[1])
+
+    agent = consumer._agent
+    monkeypatch.setattr(agent, "add_remote_agent", lambda metadata: "producer", raising=False)
+    monkeypatch.setattr(agent, "get_xfer_descs", descriptors, raising=False)
+    monkeypatch.setattr(agent, "prep_xfer_dlist", lambda agent, descs: descs, raising=False)
+    monkeypatch.setattr(agent, "make_prepped_xfer", prepare, raising=False)
+    monkeypatch.setattr(agent, "transfer", transfer, raising=False)
+    monkeypatch.setattr(agent, "check_xfer_state", lambda handle: "DONE", raising=False)
+    for method in ("release_xfer_handle", "release_dlist_handle", "remove_remote_agent"):
+        monkeypatch.setattr(agent, method, lambda handle: None, raising=False)
+    return calls
+
+
+@pytest.mark.parametrize("case", ["empty", "all_empty", "mixed", "scalar", "structured_empty", "structured_mixed"])
+def test_zero_byte_leaves_roundtrip_without_native_descriptors(producer, consumer, copying_native_agent, case):
+    empty = torch.empty((2, 0, 3), dtype=torch.float64)
+    other = torch.empty((0,), dtype=torch.int64)
+    scalar = torch.tensor(7)
+    vector = torch.arange(6, dtype=torch.float32)
+    payload = {
+        "empty": empty,
+        "all_empty": [empty, other],
+        "mixed": [empty, scalar, other, vector],
+        "scalar": scalar,
+        "structured_empty": {"leaves": (empty, [other]), "label": "empty"},
+        "structured_mixed": {"leaves": (empty, [scalar, other, vector]), "label": "mixed"},
+    }[case]
+    ok, size, metadata = producer.put("0", "1", "empty-slots", payload)
+    assert ok
+    indices = [index for group in metadata["descriptor_groups"] for index in group["tensor_indices"]]
+    assert sorted(indices) == [i for i, spec in enumerate(metadata["tensor_specs"]) if spec["size"]]
+    actual, received_size = consumer.get("0", "1", "empty-slots", metadata)
+    if isinstance(payload, dict):
+        assert actual["label"] == payload["label"]
+        torch.testing.assert_close(actual["leaves"], payload["leaves"], rtol=0, atol=0)
+    else:
+        torch.testing.assert_close(actual, payload, rtol=0, atol=0)
+    assert size == received_size
+    if case in ("empty", "all_empty"):
+        assert not copying_native_agent
+    assert not producer._pending and not producer._agent.registered
+    assert not consumer._agent.registered
+
+
+@pytest.mark.parametrize("direct", [False, True])
+@pytest.mark.parametrize("outcome", ["done", "error", "timeout", "unknown"])
+def test_read_ownership_through_terminal_and_deferred_paths(
+    producer, consumer, copying_native_agent, monkeypatch, direct, outcome
+):
+    import threading
+
+    _, _, metadata = producer.put("0", "1", "active-read", torch.ones(2))
+    entered, proceed = threading.Event(), threading.Event()
+    state = ["PROC"]
+    monkeypatch.setattr(consumer._agent, "check_xfer_state", lambda handle: state[0])
+
+    def wait(handle, key):
+        entered.set()
+        assert proceed.wait(5)
+        if outcome == "done":
+            state[0] = "DONE"
+        elif outcome == "error":
+            state[0] = "ERR"
+            raise RuntimeError("terminal error")
+        else:
+            if outcome == "unknown":
+                state[0] = "UNKNOWN"
+            raise TimeoutError("still owned")
+
+    monkeypatch.setattr(consumer, "_wait_for_transfer", wait)
+    worker = threading.Thread(target=consumer.get, args=("0", "1", "active-read", metadata if direct else None))
+    worker.start()
+    try:
+        assert entered.wait(5)
+        pending = producer._pending["active-read"]
+        pending.deadline = 0
+        producer._cleanup_expired_pending()
+        producer.cleanup("active-read")
+        assert producer._pending["active-read"] is pending
+        assert producer._agent.registered
+    finally:
+        proceed.set()
+        worker.join(5)
+    assert not worker.is_alive()
+    if outcome in ("timeout", "unknown"):
+        assert producer._agent.registered
+        assert consumer._deferred_transfers
+        consumer._reap_deferred_transfers()
+        assert producer._agent.registered
+        state[0] = "DONE"
+        deadline = time.monotonic() + 5
+        while producer._pending and time.monotonic() < deadline:
+            time.sleep(0.01)
+    assert not producer._pending
+    assert not producer._agent.registered
+
+
+class _FakeNixlAgent:
+    def __init__(self, name, config=None):
+        self.name = name
+        self.config = config
+        self.registered = []
+
+    def get_reg_descs(self, regions, memory_type):
+        return ("reg", tuple(regions), memory_type)
+
+    def register_memory(self, descs, backends=None):
+        self.registered.append(descs)
+
+    def deregister_memory(self, descs):
+        if descs in self.registered:
+            self.registered.remove(descs)
+
+    def get_agent_metadata(self):
+        return f"agent-metadata-{self.name}".encode()
+
+
+@pytest.fixture
+def nixl_connector_cls(monkeypatch):
+    """Import NixlConnector with vLLM's optional NIXL dependency stubbed out."""
+    stub = types.ModuleType("vllm.distributed.nixl_utils")
+    stub.NixlWrapper = _FakeNixlAgent  # type: ignore[attr-defined]
+    config_module = types.ModuleType("test_nixl_config_api")
+    config_sync_type = types.SimpleNamespace(NIXL_THREAD_SYNC_STRICT="strict")
+    config_module.nixl_thread_sync_t = config_sync_type  # type: ignore[attr-defined]
+
+    def agent_config(**kwargs):
+        return kwargs
+
+    agent_config.__module__ = config_module.__name__
+    config_module.nixl_agent_config = agent_config  # type: ignore[attr-defined]
+    stub.nixl_agent_config = agent_config  # type: ignore[attr-defined]
+    monkeypatch.setitem(sys.modules, "vllm.distributed.nixl_utils", stub)
+    monkeypatch.setitem(sys.modules, config_module.__name__, config_module)
+    sync_type = types.SimpleNamespace(NIXL_THREAD_SYNC_STRICT="wrong-public-enum")
+    nixl_module = types.ModuleType("nixl")
+    nixl_module.nixl_thread_sync_t = sync_type  # type: ignore[attr-defined]
+    monkeypatch.setitem(sys.modules, "nixl", nixl_module)
+
+    from vllm_omni.distributed.omni_connectors.connectors.nixl_connector import NixlConnector
+
+    return NixlConnector
+
+
+@pytest.fixture
+def producer(nixl_connector_cls):
+    connector = nixl_connector_cls({"role": "sender", "host": "127.0.0.1", "zmq_port": PORT})
+    yield connector
+    # These control-plane tests do not submit DMA. Explicitly complete claims
+    # made by metadata-only probes before tearing down the shared port.
+    for key, pending in list(connector._pending.items()):
+        for claim in list(pending.claims):
+            from vllm_omni.distributed.omni_connectors.connectors.nixl_connector import _XFER_DONE_MSG
+
+            connector._handle_handshake_message(
+                _XFER_DONE_MSG
+                + msgspec.msgpack.encode({"key": key, "generation": pending.generation, "claim_id": claim})
+            )
+    connector.close()
+
+
+@pytest.fixture
+def consumer(nixl_connector_cls):
+    connector = nixl_connector_cls({"role": "receiver", "sender_host": "127.0.0.1", "sender_zmq_port": PORT})
+    yield connector
+    connector.close()
+
+
+@pytest.fixture
+def reliable_claim_queries(producer, consumer, monkeypatch):
+    """Ownership probes need exact claims, not claims retained after lost replies.
+
+    Exercise the real resolver, wire encoding and producer handler synchronously;
+    discovery tests separately cover bounded queries over real ZMQ sockets.
+    """
+    import uuid
+
+    from vllm_omni.distributed.omni_connectors.connectors.nixl_connector import _GET_META_MSG, _META_NOT_FOUND
+
+    def query(key, host, port, *, generation=None):
+        assert (host, port) == (producer.host, producer._zmq_port)
+        request = {"key": key, "generation": generation, "claim_id": uuid.uuid4().hex}
+        reply = producer._handle_handshake_message(_GET_META_MSG + msgspec.msgpack.encode(request))
+        return None if reply == _META_NOT_FOUND else msgspec.msgpack.decode(reply)
+
+    monkeypatch.setattr(consumer, "_query_metadata_at", query)
+
+
+def test_put_publishes_its_handshake_endpoint(producer):
+    ok, size, metadata = producer.put("0", "1", "req-0", torch.arange(8, dtype=torch.float32))
+
+    assert ok is True
+    assert size == 32
+    assert metadata["sender_host"] == "127.0.0.1"
+    assert metadata["sender_zmq_port"] == PORT
+
+
+def test_agent_uses_strict_thread_synchronization(producer):
+    assert producer._agent.config["sync_mode"] == "strict"
+
+
+def test_handshake_serves_metadata_when_caller_has_none(producer, consumer):
+    _, _, published = producer.put("0", "1", "req-1", torch.arange(4, dtype=torch.float32))
+
+    resolved = _wait_for_metadata(consumer, "req-1")
+
+    # msgpack has no tuple type, so region descriptors arrive as lists; get()
+    # re-tuples them before handing them to NIXL.
+    assert resolved.pop("claim_id")
+    assert resolved == msgspec.msgpack.decode(msgspec.msgpack.encode(published))
+    assert [tuple(region) for group in resolved["descriptor_groups"] for region in group["regions"]] == [
+        region for group in published["descriptor_groups"] for region in group["regions"]
+    ]
+
+
+def test_legacy_metadata_without_generation_skips_the_handshake(consumer):
+    """Externally owned metadata without a generation needs no ownership claim."""
+    direct = {"schema_version": 1, "kind": "tensors", "tensor_specs": []}
+
+    assert consumer._resolve_metadata("req-2", direct) is direct
+
+
+def test_source_endpoint_metadata_overrides_configured_sender(consumer, monkeypatch):
+    queried = []
+    expected = {"schema_version": 1, "kind": "tensors", "tensor_specs": []}
+
+    def query(get_key, host, port):
+        queried.append((get_key, host, port))
+        return expected
+
+    monkeypatch.setattr(consumer, "_query_metadata_at", query)
+
+    resolved = consumer._resolve_metadata(
+        "req-tp4-rank3",
+        {"source_host": "10.0.0.3", "source_port": PORT + 3 * 16},
+    )
+
+    assert resolved is expected
+    assert queried == [("req-tp4-rank3", "10.0.0.3", PORT + 3 * 16)]
+
+
+def test_unknown_key_is_queried_once(producer, consumer, monkeypatch):
+    socket = consumer._get_req_socket(f"tcp://127.0.0.1:{PORT}")
+    send_count = 0
+    original_send = socket.send
+
+    def count_send(message):
+        nonlocal send_count
+        send_count += 1
+        return original_send(message)
+
+    monkeypatch.setattr(socket, "send", count_send)
+
+    assert consumer._resolve_metadata("never-published", None) is None
+    assert send_count == 1
+    assert socket.getsockopt(zmq.RCVTIMEO) == 10
+
+
+@pytest.mark.usefixtures("reliable_claim_queries")
+def test_transfer_done_releases_the_producer_buffer(producer, consumer):
+    _, _, metadata = producer.put("0", "1", "req-3", torch.arange(4, dtype=torch.float32))
+    assert producer._pending and producer._agent.registered
+
+    metadata = consumer._resolve_metadata("req-3", metadata)
+    consumer._notify_transfer_done("req-3", metadata)
+
+    assert producer._pending == {}
+    assert producer._published == {}
+    assert producer._agent.registered == []
+
+
+def test_direct_metadata_has_ephemeral_ownership_endpoint(nixl_connector_cls):
+    connector = nixl_connector_cls({})
+    try:
+        assert connector._zmq_ctx is not None
+        assert connector._listener_thread is not None
+        _, _, metadata = connector.put("0", "1", "req-4", torch.zeros(2))
+        assert metadata["sender_host"]
+        assert metadata["sender_zmq_port"] > 0
+    finally:
+        connector.close()
+
+
+def test_idle_producer_lease_expires_without_another_put(nixl_connector_cls):
+    connector = nixl_connector_cls({"role": "sender", "lease_seconds": 0.01})
+    try:
+        connector.put("0", "1", "req-expire", torch.zeros(2))
+
+        deadline = time.monotonic() + 1.0
+        while connector._pending and time.monotonic() < deadline:
+            time.sleep(0.01)
+
+        assert connector._pending == {}
+        assert connector._published == {}
+        assert connector._agent.registered == []
+        assert connector._registered_descs == []
+    finally:
+        connector.close()
+
+
+def test_close_stops_lease_reaper(nixl_connector_cls):
+    connector = nixl_connector_cls({"role": "sender"})
+    lease_thread = connector._lease_thread
+
+    connector.close()
+
+    assert lease_thread is not None
+    assert not lease_thread.is_alive()
+    assert connector._lease_thread is None
+
+
+def test_deferred_transfer_is_retained_while_active(nixl_connector_cls):
+    from vllm_omni.distributed.omni_connectors.connectors.nixl_connector import _DeferredTransfer
+
+    connector = nixl_connector_cls({"role": "receiver"})
+    released = []
+    connector._agent.check_xfer_state = lambda handle: "PROC"
+    connector._agent.release_xfer_handle = lambda handle: released.append(("handle", handle))
+    connector._agent.release_dlist_handle = lambda handle: released.append(("dlist", handle))
+    connector._agent.remove_remote_agent = lambda agent: released.append(("agent", agent))
+    transfer = _DeferredTransfer(
+        tensors=[torch.zeros(1)],
+        registrations=["registration"],
+        dlists=["local", "remote"],
+        handles=["transfer"],
+        remote_agent="producer",
+    )
+    connector._defer_transfer(transfer)
+    try:
+        connector._reap_deferred_transfers()
+
+        assert released == []
+        assert transfer in connector._deferred_transfers
+        assert transfer.tensors
+    finally:
+        connector._agent.check_xfer_state = lambda handle: "DONE"
+        connector.close()
+
+
+def test_deferred_transfer_releases_exactly_once_after_done(nixl_connector_cls):
+    from vllm_omni.distributed.omni_connectors.connectors.nixl_connector import _DeferredTransfer
+
+    connector = nixl_connector_cls({"role": "receiver"})
+    released = []
+    connector._agent.check_xfer_state = lambda handle: "DONE"
+    connector._agent.release_xfer_handle = lambda handle: released.append(("handle", handle))
+    connector._agent.release_dlist_handle = lambda handle: released.append(("dlist", handle))
+    connector._agent.remove_remote_agent = lambda agent: released.append(("agent", agent))
+    connector._agent.deregister_memory = lambda descs: released.append(("registration", descs))
+    transfer = _DeferredTransfer(
+        tensors=[torch.zeros(1)],
+        registrations=["registration"],
+        dlists=["local", "remote"],
+        handles=["transfer"],
+        remote_agent="producer",
+    )
+    connector._defer_transfer(transfer)
+    try:
+        connector._reap_deferred_transfers()
+        connector._reap_deferred_transfers()
+
+        assert released == [
+            ("handle", "transfer"),
+            ("dlist", "local"),
+            ("dlist", "remote"),
+            ("agent", "producer"),
+            ("registration", "registration"),
+        ]
+        assert connector._deferred_transfers == []
+        assert transfer.tensors == []
+    finally:
+        connector.close()
+
+
+def test_active_sibling_transfer_requires_deferred_ownership(nixl_connector_cls):
+    connector = nixl_connector_cls({"role": "receiver"})
+    connector._agent.check_xfer_state = lambda handle: {"failed": "ERR", "active": "PROC"}[handle]
+    try:
+        assert connector._transfers_may_be_active(["failed", "active"]) is True
+        assert connector._transfers_may_be_active(["failed"]) is False
+    finally:
+        connector.close()
+
+
+def test_get_defers_complete_ownership_when_sibling_transfer_is_active(nixl_connector_cls):
+    connector = nixl_connector_cls({"role": "receiver"})
+    connector._agent.add_remote_agent = lambda metadata: "producer"
+    connector._agent.get_xfer_descs = lambda regions, memory_type: (regions, memory_type)
+    connector._agent.prep_xfer_dlist = lambda agent, descs: (agent, tuple(descs[0]))
+    handles = iter(["failed", "active"])
+    connector._agent.make_prepped_xfer = lambda *_args: next(handles)
+    connector._agent.transfer = lambda handle: None
+    connector._agent.check_xfer_state = lambda handle: {"failed": "ERR", "active": "PROC"}[handle]
+    released = []
+    connector._agent.release_xfer_handle = lambda handle: released.append(("handle", handle))
+    connector._agent.release_dlist_handle = lambda handle: released.append(("dlist", handle))
+    connector._agent.remove_remote_agent = lambda agent: released.append(("agent", agent))
+    connector._agent.deregister_memory = lambda descs: released.append(("registration", descs))
+    metadata = {
+        "schema_version": 1,
+        "kind": "tensors",
+        "agent_metadata": b"producer-metadata",
+        "tensor_specs": [
+            {"shape": [1], "dtype": "torch.float32", "device": "cpu", "size": 4},
+            {"shape": [1], "dtype": "torch.float32", "device": "cpu", "size": 4},
+        ],
+        "descriptor_groups": [
+            {"memory_type": "DRAM", "tensor_indices": [0], "regions": [(1, 4, 0, "")]},
+            {"memory_type": "DRAM", "tensor_indices": [1], "regions": [(2, 4, 0, "")]},
+        ],
+        "size": 8,
+    }
+
+    try:
+        assert connector.get("0", "1", "req-active-sibling", metadata) is None
+        assert released == []
+        assert len(connector._deferred_transfers) == 1
+        deferred = connector._deferred_transfers[0]
+        assert deferred.handles == ["failed", "active"]
+        assert len(deferred.registrations) == 2
+        assert len(deferred.dlists) == 4
+        assert deferred.remote_agent == "producer"
+        assert len(deferred.tensors) == 2
+    finally:
+        connector._agent.check_xfer_state = lambda handle: "DONE"
+        connector.close()
+
+
+@pytest.mark.parametrize(
+    "payload,expected_kind",
+    [
+        (torch.zeros(4), "tensors"),
+        ([torch.zeros(2), torch.ones(3)], "tensors"),
+        ({"hidden": torch.zeros(4), "meta": {"token_role_ids": [1, 2]}}, "structured"),
+        ({"prompt": "a cat", "steps": 8}, "object"),
+    ],
+)
+def test_payload_kinds_round_trip_through_the_handshake(producer, consumer, payload, expected_kind):
+    _, _, published = producer.put("0", "1", "req-kind", payload)
+    assert published["kind"] == expected_kind
+
+    resolved = _wait_for_metadata(consumer, "req-kind")
+
+    assert resolved["kind"] == expected_kind
+    assert resolved["tensor_specs"] == published["tensor_specs"]
+
+
+def test_structured_payload_groups_descriptors_by_memory_type(producer, monkeypatch):
+    monkeypatch.setattr(
+        producer,
+        "_resolve_memory_type",
+        lambda tensor: "DRAM" if tensor.dtype == torch.uint8 else "VRAM",
+    )
+
+    _, _, metadata = producer.put(
+        "0",
+        "1",
+        "req-mixed",
+        {"meta": "value", "hidden": torch.ones(4, dtype=torch.float32)},
+    )
+
+    assert metadata["schema_version"] == 1
+    assert metadata["descriptor_groups"] == [
+        {
+            "memory_type": "DRAM",
+            "tensor_indices": [0],
+            "regions": metadata["descriptor_groups"][0]["regions"],
+        },
+        {
+            "memory_type": "VRAM",
+            "tensor_indices": [1],
+            "regions": metadata["descriptor_groups"][1]["regions"],
+        },
+    ]
+    assert len(producer._agent.registered) == 2
+
+
+def test_omni_payload_struct_uses_structured_tensor_path(producer):
+    payload = OmniPayloadStruct(
+        hidden_states=HiddenStatesStruct(output=torch.ones(4, dtype=torch.float32)),
+        meta=MetaStruct(left_context_size=2),
+    )
+
+    _, _, metadata = producer.put("0", "1", "req-struct", payload)
+    skeleton, tensors = producer._extract_tensor_leaves(payload)
+
+    assert metadata["kind"] == "structured"
+    assert len(tensors) == 1
+    assert torch.equal(tensors[0], payload.hidden_states.output)
+    assert skeleton["meta"]["left_context_size"] == 2
+
+
+def test_partial_group_registration_failure_rolls_back(producer, monkeypatch):
+    monkeypatch.setattr(
+        producer,
+        "_resolve_memory_type",
+        lambda tensor: "DRAM" if tensor.dtype == torch.uint8 else "VRAM",
+    )
+    original_register = producer._agent.register_memory
+
+    def fail_second_group(descs, backends=None):
+        if descs[2] == "VRAM":
+            raise RuntimeError("registration failed")
+        original_register(descs, backends=backends)
+
+    monkeypatch.setattr(producer._agent, "register_memory", fail_second_group)
+
+    success, _, metadata = producer.put(
+        "0",
+        "1",
+        "req-partial",
+        {"meta": "value", "hidden": torch.ones(4, dtype=torch.float32)},
+    )
+
+    assert success is False
+    assert metadata is None
+    assert producer._agent.registered == []
+    assert producer._registered_descs == []
+    assert "req-partial" not in producer._pending
+
+
+def test_reusing_put_key_releases_previous_registration(producer):
+    producer.put("0", "1", "req-reuse", torch.zeros(2))
+    previous_registration = producer._agent.registered[0]
+
+    producer.put("0", "1", "req-reuse", torch.ones(2))
+
+    assert previous_registration not in producer._agent.registered
+    assert len(producer._agent.registered) == 1
+    assert len(producer._registered_descs) == 1
+
+
+def test_expired_snapshot_cannot_claim_replacement(producer):
+    from vllm_omni.distributed.omni_connectors.connectors.nixl_connector import _PendingPayload
+
+    old = _PendingPayload([torch.zeros(1)], ["old"], 0.0)
+    replacement = _PendingPayload([torch.ones(1)], ["new"], time.monotonic() + 60)
+    producer._pending["req-race"] = replacement
+
+    assert producer._take_pending("req-race", expected=old) is None
+    assert producer._pending["req-race"] is replacement
+
+
+@pytest.mark.parametrize(
+    "indices",
+    [
+        [0, 0],
+        [0],
+        [0, 2],
+        [-1, 0],
+    ],
+)
+def test_descriptor_groups_require_exact_tensor_index_partition(nixl_connector_cls, indices):
+    metadata = {
+        "schema_version": 1,
+        "descriptor_groups": [
+            {
+                "memory_type": "DRAM",
+                "tensor_indices": indices,
+                "regions": [(1, 4, 0, "")] * len(indices),
+            }
+        ],
+    }
+
+    with pytest.raises(RuntimeError, match="exact partition"):
+        nixl_connector_cls._validated_descriptor_groups(
+            metadata, [{"shape": [1], "dtype": "torch.float32", "size": 4}] * 2
+        )
+
+
+def test_receive_device_ignores_the_producer_index(nixl_connector_cls):
+    """A producer on cuda:3 must not dictate the consumer's card."""
+    connector = nixl_connector_cls({})
+    try:
+        assert connector._resolve_receive_device("cpu") == torch.device("cpu")
+        assert connector._resolve_receive_device(None) == torch.device("cpu")
+    finally:
+        connector.close()
+
+
+def test_receive_device_config_wins(nixl_connector_cls):
+    connector = nixl_connector_cls({"receive_device": "cpu"})
+    try:
+        assert connector._resolve_receive_device("cuda:3") == torch.device("cpu")
+    finally:
+        connector.close()

tests/distributed/omni_connectors/test_nixl_connector_native.py

@@ -0,0 +1,204 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import multiprocessing as mp
+import socket
+import time
+from typing import Any
+
+import pytest
+import torch
+
+from vllm_omni.data_entry_keys import HiddenStatesStruct, MetaStruct, OmniPayloadStruct
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda, pytest.mark.parallel]
+
+
+def _native_nixl_available() -> bool:
+    try:
+        from vllm.distributed.nixl_utils import NixlWrapper
+    except ImportError:
+        return False
+    return NixlWrapper is not None and torch.cuda.is_available() and torch.accelerator.device_count() >= 2
+
+
+def _free_port() -> int:
+    with socket.socket() as sock:
+        sock.bind(("127.0.0.1", 0))
+        return int(sock.getsockname()[1])
+
+
+def _producer(port: int, ready: Any, consumed: Any, result: Any, claimed: Any, proceed: Any, wire: Any) -> None:
+    from vllm_omni.distributed.omni_connectors.connectors.nixl_connector import NixlConnector
+
+    current_omni_platform.set_device(0)
+    connector = NixlConnector(
+        {
+            "role": "sender",
+            "host": "127.0.0.1",
+            "zmq_port": port,
+            "agent_name": "native-smoke-producer",
+        }
+    )
+    try:
+        payload = OmniPayloadStruct(
+            hidden_states=HiddenStatesStruct(output=torch.arange(8, dtype=torch.float32, device="cuda:0")),
+            meta=MetaStruct(token_role_ids=torch.tensor([1, 2, 3], dtype=torch.int64, device="cpu")),
+            kv_metadata={"empty": torch.empty((2, 0, 3), device="cuda:0")},
+            request_id="native-smoke",
+        )
+        success, size, metadata = connector.put("0", "1", "native-smoke", payload)
+        result.put(("put", success, size, metadata["kind"] if metadata else None))
+        wire.put(metadata)
+        ready.set()
+        if not claimed.wait(timeout=30):
+            raise TimeoutError("consumer did not claim source")
+        pending = connector._pending["native-smoke"]
+        pending.deadline = 0
+        connector._cleanup_expired_pending()
+        connector.cleanup("native-smoke")
+        assert connector._pending["native-smoke"] is pending
+        assert pending.claims and connector._registered_descs
+        result.put(("source_retained", len(pending.claims)))
+        proceed.set()
+        if not consumed.wait(timeout=30):
+            raise TimeoutError("consumer did not complete the native NIXL transfer")
+        deadline = time.monotonic() + 10
+        while connector._pending and time.monotonic() < deadline:
+            time.sleep(0.01)
+        result.put(("cleanup", len(connector._pending), len(connector._registered_descs)))
+    except Exception as error:
+        result.put(("producer_error", repr(error)))
+        raise
+    finally:
+        connector.close()
+
+
+def _consumer(
+    port: int, ready: Any, consumed: Any, result: Any, claimed: Any, proceed: Any, wire: Any, direct: bool
+) -> None:
+    from vllm_omni.distributed.omni_connectors.kv_transfer_manager import (
+        OmniKVCacheConfig,
+        OmniKVTransferManager,
+    )
+
+    if not ready.wait(timeout=30):
+        raise TimeoutError("producer did not publish native NIXL metadata")
+    current_omni_platform.set_device(1)
+    manager = OmniKVTransferManager(
+        OmniKVCacheConfig(
+            connector_config={
+                "type": "NixlConnector",
+                "role": "receiver",
+                "host": "127.0.0.1",
+                "zmq_port": port,  # Shared deployment edge: already bound by producer.
+                "backends": ["UCX"],
+                "receive_device": "cuda",
+                "agent_name": "native-smoke-consumer",
+            },
+            from_stage="0",
+            to_stage="1",
+            stage_id=1,
+            need_recv_cache=direct,  # Exercise both payload-only and KV receiver setup.
+        )
+    )
+    manager.update_sender_info({"host": "127.0.0.1", "zmq_port": port})
+    connector = manager.connector
+    assert connector is not None, "Receiver must not bind the producer's occupied port"
+    assert connector._zmq_port is None and not connector._serving_handshake
+    assert connector._backends == ["UCX"]
+    try:
+        original_wait = connector._wait_for_transfer
+        states = []
+
+        def wait(handle, key):
+            states.append(connector._agent.check_xfer_state(handle))
+            claimed.set()
+            if not proceed.wait(timeout=30):
+                raise TimeoutError("producer did not verify active ownership")
+            original_wait(handle, key)
+
+        connector._wait_for_transfer = wait
+        metadata = wire.get(timeout=5)
+        received = None
+        deadline = time.monotonic() + 30
+        while received is None and time.monotonic() < deadline:
+            received = connector.get(
+                "0",
+                "1",
+                "native-smoke",
+                metadata if direct else {"source_host": "127.0.0.1", "source_port": port},
+            )
+        if received is None:
+            raise TimeoutError("native NIXL transfer did not complete")
+        payload, size = received
+        hidden = payload["hidden_states"]["output"]
+        token_roles = payload["meta"]["token_role_ids"]
+        empty = payload["kv_metadata"]["empty"]
+        assert empty.shape == (2, 0, 3) and str(empty.device) == "cuda:1"
+        result.put(("native_states", states))
+        result.put(
+            (
+                "get",
+                size,
+                hidden.cpu().tolist(),
+                str(hidden.dtype),
+                list(hidden.shape),
+                str(hidden.device),
+                token_roles.cpu().tolist(),
+                str(token_roles.dtype),
+                list(token_roles.shape),
+                str(token_roles.device),
+            )
+        )
+        consumed.set()
+    except Exception as error:
+        result.put(("consumer_error", repr(error)))
+        raise
+    finally:
+        connector.close()
+
+
+@pytest.mark.skipif(not _native_nixl_available(), reason="requires NIXL and at least two CUDA devices")
+@pytest.mark.parametrize("direct", [False, True])
+def test_native_two_process_structured_mixed_device_transfer(direct):
+    context = mp.get_context("spawn")
+    ready = context.Event()
+    consumed = context.Event()
+    claimed = context.Event()
+    proceed = context.Event()
+    wire = context.Queue()
+    result = context.Queue()
+    port = _free_port()
+    producer = context.Process(target=_producer, args=(port, ready, consumed, result, claimed, proceed, wire))
+    consumer = context.Process(target=_consumer, args=(port, ready, consumed, result, claimed, proceed, wire, direct))
+
+    producer.start()
+    consumer.start()
+    producer.join(timeout=60)
+    consumer.join(timeout=60)
+    if producer.is_alive():
+        producer.terminate()
+    if consumer.is_alive():
+        consumer.terminate()
+
+    records = [result.get(timeout=5) for _ in range(5)]
+    print("Native ownership records:", records)
+    assert producer.exitcode == 0, records
+    assert consumer.exitcode == 0, records
+    assert records[0][0] == "put"
+    assert records[0][1] is True
+    assert records[0][2] >= 56
+    assert records[0][3] == "structured"
+    get_record = next(record for record in records if record[0] == "get")
+    assert get_record[1] == records[0][2]
+    assert get_record[2:6] == (
+        list(range(8)),
+        "torch.float32",
+        [8],
+        "cuda:1",
+    )
+    assert get_record[6:] == ([1, 2, 3], "torch.int64", [3], "cuda:1")
+    assert next(record for record in records if record[0] == "cleanup") == ("cleanup", 0, 0)
+    assert next(record for record in records if record[0] == "source_retained") == ("source_retained", 1)

tests/engine/test_orchestrator_kv_sender_info.py

@@ -1,13 +1,22 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 import asyncio
+from collections.abc import Callable
 from types import SimpleNamespace
+from typing import Any
 
 import pytest
 from vllm import SamplingParams
+from vllm.v1.engine.core_client import AsyncMPClient, DPLBAsyncMPClient
 
 from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker
 from vllm_omni.engine.messages import OutputMessage
 from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState
-from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient
+from vllm_omni.engine.stage_engine_core_client import (
+    DPLBStageEngineCoreClient,
+    StageEngineCoreClient,
+)
 from vllm_omni.engine.stage_pool import StagePool
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
@@ -28,7 +37,7 @@ class _DummySenderStage:
 class _DummyDiffusionStage:
     stage_type = "diffusion"
     final_output = True
-    custom_process_input_func = None
+    custom_process_input_func: Callable[..., Any] | None = None
 
     def __init__(self, engine_input_source=None):
         self.engine_input_source = engine_input_source or [0]
@@ -54,6 +63,97 @@ def _build_sender_pool(stage_id: int, sender_info: dict[str, object]) -> StagePo
     )
 
 
+@pytest.mark.parametrize(
+    ("client_class", "base_client_class"),
+    [
+        (StageEngineCoreClient, AsyncMPClient),
+        (DPLBStageEngineCoreClient, DPLBAsyncMPClient),
+    ],
+)
+def test_stage_engine_core_client_builds_payload_sender_info_after_base_init(
+    monkeypatch, client_class, base_client_class
+):
+    vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(
+            omni_kv_config=None,
+            hf_config=None,
+            stage_connector_config=None,
+        )
+    )
+    metadata = SimpleNamespace(
+        stage_id=0,
+        replica_id=0,
+        stage_type="llm",
+        model_stage="main",
+        is_comprehension=False,
+        requires_multimodal_data=False,
+        engine_input_source=[],
+        final_output=False,
+        final_output_type=None,
+        default_sampling_params=None,
+        prompt_transform_func=None,
+        prompt_expand_func=None,
+        custom_process_input_func=None,
+    )
+
+    def fake_base_init(self, config, *_args, **_kwargs):
+        self.vllm_config = config
+        self.resources = SimpleNamespace(engine_dead=False)
+
+    monkeypatch.setattr(base_client_class, "__init__", fake_base_init)
+
+    client = client_class(vllm_config, object, metadata=metadata)
+
+    assert client.vllm_config is vllm_config
+    assert client.get_payload_sender_info() is None
+
+
+@pytest.mark.parametrize("outgoing", [False, True])
+@pytest.mark.parametrize("base_port", [None, 48000])
+def test_payload_sender_endpoint_matches_resolver_with_unequal_replicas(outgoing, base_port):
+    from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec
+    from vllm_omni.distributed.omni_connectors.utils.initialization import resolve_connector_spec
+
+    edge = {"host": "10.0.0.2", "from_stage": 1}
+    if base_port is not None:
+        edge["zmq_port"] = base_port
+    extra = (
+        {
+            "role": "receiver",
+            "host": "10.0.0.1",
+            "zmq_port": 47000,
+            "from_stage": 0,
+            "outgoing": edge,
+        }
+        if outgoing
+        else {"role": "sender", **edge}
+    )
+    client = object.__new__(StageEngineCoreClient)
+    client.stage_id = 1
+    client.replica_id = 3
+    client.vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(stage_connector_config={"name": "NixlConnector", "extra": extra})
+    )
+    producer = resolve_connector_spec(
+        ConnectorSpec(name="NixlConnector", extra=extra),
+        stage_id=1,
+        role=extra["role"],
+        replica_id=client.replica_id,
+    )
+    consumer = resolve_connector_spec(
+        ConnectorSpec(name="NixlConnector", extra=edge), stage_id=2, role="receiver", replica_id=7
+    )
+
+    sender_info = client._build_payload_sender_info()
+
+    assert sender_info == {
+        "host": "10.0.0.2",
+        "zmq_port": (50051 if base_port is None else base_port) + 3 * 1024 + 1,
+    }
+    assert sender_info["zmq_port"] == producer.extra["zmq_port"]
+    assert sender_info["zmq_port"] != consumer.extra["sender_zmq_port"]
+
+
 def test_stage_engine_core_client_builds_kv_sender_info_from_tcp_address():
     client = object.__new__(StageEngineCoreClient)
     client.stage_id = 0
@@ -259,3 +359,58 @@ def test_prewarm_diffusion_attaches_kv_sender_info():
         0: {"host": "10.0.0.3", "zmq_port": 50151},
     }
     assert req_state.stage_submit_ts[1] > 0
+
+
+def test_prewarm_submits_bound_payload_endpoint_for_concurrent_replicas():
+    orchestrator = object.__new__(Orchestrator)
+    endpoints = {"a": {"host": "10.0.0.2", "zmq_port": 52099}, "b": {"host": "10.0.0.3", "zmq_port": 54147}}
+    source = SimpleNamespace(
+        get_bound_client=lambda key: SimpleNamespace(get_payload_sender_info=lambda: endpoints[key]),
+        get_bound_replica_id=lambda key: {"a": 2, "b": 4}[key],
+    )
+    submitted = {}
+
+    async def submit(key, state, request, **kwargs):
+        await asyncio.sleep(0)
+        submitted[key] = request
+
+    target = SimpleNamespace(
+        stage_type="llm",
+        stage_client=SimpleNamespace(engine_input_source=[0]),
+        stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(hf_config=None, max_model_len=64)),
+        submit_initial=submit,
+        get_bound_replica_id=lambda key: {"a": 7, "b": 1}[key],
+    )
+    orchestrator.stage_pools = [source, target]
+    orchestrator._stage_receives_async_chunks = lambda stage: True
+    orchestrator._record_duplex_stage_submission = lambda *args: None
+    orchestrator._emit_tx_edge = lambda **kwargs: None
+
+    async def dispatch(callback, **kwargs):
+        await callback()
+        return True
+
+    orchestrator._dispatch_or_fail_request = dispatch
+
+    async def run():
+        await asyncio.gather(
+            *[
+                orchestrator._prewarm_async_chunk_stages(
+                    key,
+                    SimpleNamespace(prompt_token_ids=[1]),
+                    OrchestratorRequestState(
+                        request_id=key,
+                        prompt={},
+                        sampling_params_list=[SamplingParams(), SamplingParams()],
+                        final_stage_id=1,
+                    ),
+                )
+                for key in endpoints
+            ]
+        )
+
+    asyncio.run(run())
+    from vllm_omni.engine import OmniEngineCoreRequest
+
+    assert all(isinstance(request, OmniEngineCoreRequest) for request in submitted.values())
+    assert {key: request.payload_sender_info for key, request in submitted.items()} == endpoints

tests/engine/test_orchestrator_stage_input_bridge.py

@@ -104,9 +104,13 @@ class FakeInputProcessor:
 
 
 class FakePrewarmPool:
+    """Single-replica pool double that binds requests only on submission."""
+
     stage_type = "llm"
 
     def __init__(self, role: str) -> None:
+        self.stage_client = SimpleNamespace()
+        self._bound_request_ids: set[str] = set()
         self.stage_vllm_config = SimpleNamespace(
             model_config=SimpleNamespace(
                 max_model_len=64,
@@ -115,12 +119,16 @@ class FakePrewarmPool:
         )
         self.submitted: list[Any] = []
 
-    async def submit_initial(self, _request_id, _req_state, request, prompt_text=None):
+    async def submit_initial(self, request_id, _req_state, request, prompt_text=None):
         self.submitted.append(request)
+        self._bound_request_ids.add(request_id)
         return 0
 
-    def get_bound_replica_id(self, _request_id):
-        return 0
+    def get_bound_replica_id(self, request_id):
+        return 0 if request_id in self._bound_request_ids else None
+
+    def get_bound_client(self, request_id):
+        return self.stage_client if self.get_bound_replica_id(request_id) is not None else None
 
 
 def _duplex_stage_port_submission():
@@ -261,11 +269,14 @@ async def test_forward_text_prompt_uses_target_stage_input_processor() -> None:
 
 
 @pytest.mark.asyncio
-async def test_async_prewarm_skips_outgoing_only_stage() -> None:
+@pytest.mark.parametrize("payload_sender_info", [None, {"host": "10.0.0.2", "zmq_port": 52099}])
+async def test_async_prewarm_skips_outgoing_only_stage(payload_sender_info) -> None:
     orchestrator = object.__new__(Orchestrator)
     stage0 = FakePrewarmPool("sender")
     stage1 = FakePrewarmPool("sender")
     stage2 = FakePrewarmPool("receiver")
+    if payload_sender_info is not None:
+        stage1.stage_client.get_payload_sender_info = MagicMock(return_value=payload_sender_info)
     orchestrator.stage_pools = [stage0, stage1, stage2]
     orchestrator._emit_tx_edge = lambda **_kwargs: None
     orchestrator._record_duplex_stage_submission = MagicMock()
@@ -285,7 +296,14 @@ async def test_async_prewarm_skips_outgoing_only_stage() -> None:
 
     assert prewarmed is True
     assert stage1.submitted == []
+    assert stage1.get_bound_client("req-prewarm") is None
     assert len(stage2.submitted) == 1
+    assert stage2.get_bound_client("req-prewarm") is stage2.stage_client
+    assert stage2.submitted[0].payload_sender_info == payload_sender_info
+    if payload_sender_info is not None:
+        stage1.stage_client.get_payload_sender_info.assert_called_once_with()
+    assert stage2.submitted[0].external_req_id == "req-prewarm"
+    assert stage2.submitted[0].resumable is True
     assert 1 not in req_state.stage_submit_ts
     assert 2 in req_state.stage_submit_ts
     orchestrator._record_duplex_stage_submission.assert_called_once_with(

tests/worker/test_omni_connector_mixin.py

@@ -10,6 +10,7 @@ from __future__ import annotations
 
 import time
 import unittest
+from copy import deepcopy
 from types import SimpleNamespace
 from typing import Any
 from unittest.mock import MagicMock, patch
@@ -20,6 +21,10 @@ import torch
 from vllm_omni.distributed.omni_connectors.kv_transfer_manager import (
     OmniKVTransferManager,
 )
+from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec
+from vllm_omni.distributed.omni_connectors.utils.initialization import (
+    resolve_connector_spec,
+)
 from vllm_omni.outputs import OmniConnectorOutput
 from vllm_omni.worker.omni_connector_model_runner_mixin import (
     OmniConnectorModelRunnerMixin,
@@ -564,6 +569,55 @@ class TestChunkStreamCompletedGuard(unittest.TestCase):
 
         host.shutdown_omni_connectors()
 
+    def test_poll_uses_request_scoped_payload_sender_endpoint(self):
+        host = self._make_host(stage_id=1)
+        host._omni_connector.get = MagicMock(return_value=None)
+        req = _make_request("req-1", "ext-req-1")
+        req.payload_sender_info = {"host": "10.0.0.1", "zmq_port": 50051}
+
+        host.register_chunk_recv(req)
+        host._poll_single_request("req-1")
+
+        host._omni_connector.get.assert_called_once_with(
+            "0",
+            "1",
+            "ext-req-1_0_0",
+            {"source_host": "10.0.0.1", "source_port": 50051},
+        )
+        host.shutdown_omni_connectors()
+
+    def test_concurrent_requests_keep_distinct_payload_sender_endpoints(self):
+        host = self._make_host(stage_id=1)
+        host._omni_connector.get = MagicMock(return_value=None)
+        first = _make_request("req-1", "ext-req-1")
+        first.payload_sender_info = {"host": "10.0.0.1", "zmq_port": 50051}
+        second = _make_request("req-2", "ext-req-2")
+        second.payload_sender_info = {"host": "10.0.0.2", "zmq_port": 51051}
+
+        host.register_chunk_recv(first)
+        host.register_chunk_recv(second)
+        host._poll_single_request("req-2")
+        host._poll_single_request("req-1")
+
+        self.assertEqual(
+            host._omni_connector.get.call_args_list,
+            [
+                unittest.mock.call(
+                    "0",
+                    "1",
+                    "ext-req-2_0_0",
+                    {"source_host": "10.0.0.2", "source_port": 51051},
+                ),
+                unittest.mock.call(
+                    "0",
+                    "1",
+                    "ext-req-1_0_0",
+                    {"source_host": "10.0.0.1", "source_port": 50051},
+                ),
+            ],
+        )
+        host.shutdown_omni_connectors()
+
     def test_finish_sentinel_populates_completed_set(self):
         """Receiving is_finished=True adds to _chunk_stream_completed."""
         host = self._make_host(stage_id=1)
@@ -1381,6 +1435,144 @@ class TestConnectorConfigValidation(unittest.TestCase):
             host.init_omni_connectors(model_config=model_config)
 
 
+class TestRankAwareHandshakePort(unittest.TestCase):
+    """A sender must bind the port the orchestrator advertises for its rank."""
+
+    @staticmethod
+    def _resolve(name, extra, local_rank=0):
+        spec = resolve_connector_spec(
+            ConnectorSpec(name=name, extra=extra),
+            stage_id=int(extra.get("stage_id", 0)),
+            role=extra.get("role"),
+            local_rank=local_rank,
+        )
+        return spec.extra
+
+    def test_sender_port_matches_the_orchestrator_formula(self):
+        for local_rank in (0, 1, 3):
+            with self.subTest(local_rank=local_rank):
+                extra = {"zmq_port": 50071, "stage_id": 2, "role": "sender"}
+                resolved = self._resolve("NixlConnector", extra, local_rank=local_rank)
+                self.assertEqual(resolved["zmq_port"], 50071 + 2 + local_rank * 16)
+                self.assertEqual(extra["zmq_port"], 50071, "the source config must not be mutated")
+
+    def test_receiver_gets_sender_endpoint_without_a_bind_port(self):
+        receiver = self._resolve(
+            "NixlConnector",
+            {"zmq_port": 50071, "stage_id": 3, "from_stage": 2, "role": "receiver"},
+            local_rank=1,
+        )
+        self.assertNotIn("zmq_port", receiver)
+        self.assertEqual(receiver["sender_zmq_port"], 50071 + 2 + 16)
+
+    def test_handshake_less_connector_keeps_its_config(self):
+        shm = self._resolve("SharedMemoryConnector", {"zmq_port": 50071, "role": "sender"}, local_rank=1)
+        self.assertEqual(shm["zmq_port"], 50071)
+
+    def test_unresolvable_port_is_rejected(self):
+        extra = {"zmq_port": "${MISSING_PORT_VAR}", "stage_id": 0, "role": "sender"}
+        with self.assertRaises(ValueError):
+            self._resolve("NixlConnector", extra)
+
+
+@pytest.mark.parametrize(
+    ("extra", "expected"),
+    [
+        (
+            {"role": "sender", "from_stage": 2, "zmq_port": 50071},
+            {"zmq_port": 52153},
+        ),
+        (
+            {"role": "receiver", "from_stage": 1, "zmq_port": 50071, "host": "producer"},
+            {"sender_zmq_port": 52152, "sender_host": "producer"},
+        ),
+        (
+            {
+                "role": "receiver",
+                "from_stage": 1,
+                "zmq_port": 50071,
+                "sender_zmq_port": 49000,
+                "sender_host": "bound-producer",
+                "outgoing": {"from_stage": 2, "zmq_port": 51071, "host": "local-worker"},
+            },
+            {"sender_zmq_port": 49000, "sender_host": "bound-producer", "zmq_port": 53153, "host": "local-worker"},
+        ),
+    ],
+)
+def test_split_runtime_factory_resolves_nixl_endpoints(extra, expected):
+    """The runner must pass the resolved spec, not the raw config, to the factory."""
+    from vllm_omni.distributed.omni_connectors.model_runner import omni_connector_runtime as runtime
+
+    original = deepcopy(extra)
+    config = SimpleNamespace(stage_id=2, stage_connector_config={"name": "NixlConnector", "extra": extra})
+    with (
+        patch.object(runtime, "get_local_tp_rank", return_value=2, create=True),
+        patch.object(runtime, "get_omni_replica_id", return_value=2, create=True),
+        patch.object(runtime.OmniConnectorFactory, "create_connector") as factory,
+    ):
+        assert MixinHost._create_connector(config) is factory.return_value
+
+    spec = factory.call_args.args[0]
+    assert spec.name == "NixlConnector"
+    assert spec.extra["stage_id"] == 2
+    for key, value in expected.items():
+        assert spec.extra[key] == value
+    if extra["role"] == "receiver" and "outgoing" not in extra:
+        assert "zmq_port" not in spec.extra
+    assert extra == original
+
+
+@pytest.mark.parametrize("async_chunk", [False, True])
+def test_split_payload_endpoint_retry_cleanup_and_reuse(async_chunk):
+    host = MixinHost()
+    host.init_omni_connectors(_make_model_config(async_chunk=async_chunk))
+    host._stage_id = 1
+    host._omni_connector = MagicMock()
+    host._omni_connector.get.return_value = None
+    first = _make_request("req-1", "ext-req-1")
+    first.payload_sender_info = {"host": "producer-a", "zmq_port": "50051"}
+    second = _make_request("req-2", "ext-req-2")
+    second.payload_sender_info = {"host": "producer-b", "zmq_port": 51051}
+
+    try:
+        host.register_chunk_recv(first)
+        host.register_chunk_recv(second)
+        for req_id in ("req-1", "req-2", "req-1"):
+            assert not host._poll_single_request(req_id)
+        assert host._omni_connector.get.call_args_list == [
+            unittest.mock.call("0", "1", "ext-req-1_0_0", {"source_host": "producer-a", "source_port": 50051}),
+            unittest.mock.call("0", "1", "ext-req-2_0_0", {"source_host": "producer-b", "source_port": 51051}),
+            unittest.mock.call("0", "1", "ext-req-1_0_0", {"source_host": "producer-a", "source_port": 50051}),
+        ]
+        host.cleanup_finished_request("req-1")
+        assert "req-1" not in host._pending_load_reqs
+        assert host._pending_load_reqs["req-2"] is second
+        host.register_chunk_recv(_make_request("req-1", "reused"))
+        assert not host._poll_single_request("req-1")
+        host._omni_connector.get.assert_called_with("0", "1", "reused_0_0")
+    finally:
+        host.shutdown_omni_connectors()
+
+
+@pytest.mark.parametrize("method", ["_recv_full_payload_result", "_recv_async_chunk_result"])
+@pytest.mark.parametrize("rank", [0, 1])
+def test_split_payload_sender_metadata_respects_tp_leader(method, rank):
+    host = MixinHost()
+    connector = MagicMock()
+    metadata = {"source_host": "producer", "source_port": 50051}
+    with (
+        patch.object(host, "_get_local_tp_group", return_value=_FakeTPGroup(world_size=2, rank_in_group=rank)),
+        patch.object(host, "is_data_transfer_rank", return_value=rank == 0),
+    ):
+        result = getattr(host, method)(connector, "0", "1", "req_0_0", metadata)
+    if rank == 0:
+        connector.get.assert_called_once_with("0", "1", "req_0_0", metadata)
+        assert result is connector.get.return_value
+    else:
+        connector.get.assert_not_called()
+        assert result is None
+
+
 class _FailingConnector:
     """Connector whose put() fails a configurable number of times."""
 

vllm_omni/config/environment_variable_inventory.py

@@ -67,6 +67,8 @@ _PUBLIC_OMNI = (
     "VLLM_OMNI_EVENT_DRIVEN_ORCH",
     "VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS",
     "VLLM_OMNI_INPUT_WAIT_TIMEOUT_S",
+    "VLLM_OMNI_NIXL_LEASE_S",
+    "VLLM_OMNI_NIXL_XFER_TIMEOUT_S",
     "VLLM_OMNI_ORCH_MONITOR_PATH",
     "VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY",
     "VLLM_OMNI_SERVER_STORAGE__FILE_TTL",

vllm_omni/core/sched/omni_scheduling_coordinator.py

@@ -121,6 +121,7 @@ class OmniSchedulingCoordinator:
                     OmniChunkRecvHandle(
                         request_id=request.request_id,
                         external_req_id=getattr(request, "external_req_id", None),
+                        payload_sender_info=getattr(request, "payload_sender_info", None),
                     )
                 )
             elif request.status == RequestStatus.WAITING_FOR_INPUT:
@@ -134,6 +135,7 @@ class OmniSchedulingCoordinator:
                         OmniChunkRecvHandle(
                             request_id=request.request_id,
                             external_req_id=getattr(request, "external_req_id", None),
+                            payload_sender_info=getattr(request, "payload_sender_info", None),
                         )
                     )
         if to_remove:

vllm_omni/core/sched/output.py

@@ -105,6 +105,7 @@ class OmniChunkRecvHandle:
 
     request_id: str
     external_req_id: str | None = None
+    payload_sender_info: dict[str, object] | None = None
 
 
 @dataclass

vllm_omni/distributed/__init__.py

@@ -7,6 +7,7 @@ from .omni_connectors import (
     MooncakeStoreConnector,
     MooncakeTransferEngineConnector,
     MoriTransferEngineConnector,
+    NixlConnector,
     OmniConnectorBase,
     OmniConnectorFactory,
     OmniTransferConfig,
@@ -27,6 +28,7 @@ __all__ = [
     "MooncakeStoreConnector",
     "MooncakeTransferEngineConnector",
     "MoriTransferEngineConnector",
+    "NixlConnector",
     "SharedMemoryConnector",
     "YuanrongConnector",
     "YuanrongTransferEngineConnector",

vllm_omni/distributed/omni_connectors/__init__.py

@@ -22,6 +22,11 @@ try:
     from .connectors.mori_transfer_engine_connector import MoriTransferEngineConnector
 except ImportError:
     MoriTransferEngineConnector = None  # RDMA deps (msgspec/zmq/mori) not installed
+
+try:
+    from .connectors.nixl_connector import NixlConnector
+except ImportError:
+    NixlConnector = None  # NIXL deps not installed
 from .factory import OmniConnectorFactory
 from .utils.config import ConnectorSpec, OmniTransferConfig
 from .utils.initialization import (
@@ -50,6 +55,7 @@ __all__ = [
     "MooncakeStoreConnector",
     "MooncakeTransferEngineConnector",
     "MoriTransferEngineConnector",
+    "NixlConnector",
     "SharedMemoryConnector",
     "YuanrongConnector",
     "YuanrongTransferEngineConnector",

vllm_omni/distributed/omni_connectors/factory.py

@@ -125,6 +125,17 @@ def _create_mori_transfer_engine_connector(config: dict[str, Any]) -> OmniConnec
     return MoriTransferEngineConnector(config)
 
 
+def _create_nixl_connector(config: dict[str, Any]) -> OmniConnectorBase:
+    try:
+        from .connectors.nixl_connector import NixlConnector
+    except ImportError:
+        import sys
+
+        sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
+        from omni_connectors.connectors.nixl_connector import NixlConnector
+    return NixlConnector(config)
+
+
 # Register connectors
 OmniConnectorFactory.register_connector("MooncakeStoreConnector", _create_mooncake_store_connector)
 OmniConnectorFactory.register_connector("MooncakeTransferEngineConnector", _create_mooncake_transfer_engine_connector)
@@ -132,5 +143,6 @@ OmniConnectorFactory.register_connector("SharedMemoryConnector", _create_shm_con
 OmniConnectorFactory.register_connector("YuanrongConnector", _create_yuanrong_connector)
 OmniConnectorFactory.register_connector("YuanrongTransferEngineConnector", _create_yuanrong_transfer_engine_connector)
 OmniConnectorFactory.register_connector("MoriTransferEngineConnector", _create_mori_transfer_engine_connector)
+OmniConnectorFactory.register_connector("NixlConnector", _create_nixl_connector)
 # Backward-compatible aliases – will be removed in the future
 OmniConnectorFactory.register_connector("MooncakeConnector", _create_mooncake_store_connector)

vllm_omni/distributed/omni_connectors/kv_transfer_manager.py

@@ -564,6 +564,13 @@ class OmniKVTransferManager:
                             # the wrong process. Explicit sender_* YAML values
                             # are still preserved for standalone connector use.
 
+                    if c_type == "NixlConnector" and c_extra.get("role") == "receiver":
+                        # This manager receives on the shared incoming edge;
+                        # its YAML zmq_port belongs to the producer. NIXL treats
+                        # an explicit zmq_port as a local listener, even for a
+                        # receiver. Keep only the request-scoped sender_* endpoint.
+                        c_extra.pop("zmq_port", None)
+
                     logger.info(
                         "Initializing OmniConnector type=%s role=%s",
                         c_type,
@@ -922,7 +929,9 @@ class OmniKVTransferManager:
         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:
+        # A stage can sit on the receiving end of a payload-only edge, where the
+        # sender endpoint still has to be applied even though no KV is expected.
+        if not self.config.need_recv_cache and (self.config.connector_config or {}).get("role") != "receiver":
             return
 
         actual_info = self._resolve_sender_info(sender_info, sender_stage_id=sender_stage_id)

vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py

@@ -182,14 +182,19 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
         from_stage: str,
         to_stage: str,
         connector_get_key: str,
+        metadata: dict[str, Any] | None = None,
     ) -> Any:
         """Receive one ordinary non-KV stage payload on the local leader rank only."""
         tp_group = self._get_local_tp_group()
         if tp_group is None or getattr(tp_group, "world_size", 1) <= 1:
-            return connector.get(from_stage, to_stage, connector_get_key)
+            if metadata is None:
+                return connector.get(from_stage, to_stage, connector_get_key)
+            return connector.get(from_stage, to_stage, connector_get_key, metadata)
         if not self.is_data_transfer_rank():
             return None
-        return connector.get(from_stage, to_stage, connector_get_key)
+        if metadata is None:
+            return connector.get(from_stage, to_stage, connector_get_key)
+        return connector.get(from_stage, to_stage, connector_get_key, metadata)
 
     def _recv_full_payload_result(
         self,
@@ -197,6 +202,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
         from_stage: str,
         to_stage: str,
         connector_get_key: str,
+        metadata: dict[str, Any] | None = None,
     ) -> Any:
         """Receive one full-payload transfer on the local leader rank only."""
         return self._recv_ordinary_stage_result(
@@ -204,6 +210,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
             from_stage,
             to_stage,
             connector_get_key,
+            metadata,
         )
 
     def _recv_async_chunk_result(
@@ -212,6 +219,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
         from_stage: str,
         to_stage: str,
         connector_get_key: str,
+        metadata: dict[str, Any] | None = None,
     ) -> Any:
         """Receive one ordinary async chunk on the local leader rank only."""
         return self._recv_ordinary_stage_result(
@@ -219,6 +227,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
             from_stage,
             to_stage,
             connector_get_key,
+            metadata,
         )
 
     @staticmethod
@@ -1027,6 +1036,14 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
         chunk_id = self._get_req_chunk[req_id]
         external_req_id = self._request_ids_mapping.get(req_id, req_id)
         connector_get_key = f"{external_req_id}_{target_stage_id}_{chunk_id}"
+        request = self._pending_load_reqs.get(req_id)
+        sender_info = getattr(request, "payload_sender_info", None)
+        metadata = None
+        if isinstance(sender_info, dict):
+            host = sender_info.get("host")
+            port = sender_info.get("zmq_port")
+            if host and port:
+                metadata = {"source_host": str(host), "source_port": int(port)}
 
         if self._async_chunk:
             result = self._recv_async_chunk_result(
@@ -1034,6 +1051,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
                 str(target_stage_id),
                 str(self._stage_id),
                 connector_get_key,
+                metadata,
             )
         else:
             result = self._recv_full_payload_result(
@@ -1041,6 +1059,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
                 str(target_stage_id),
                 str(self._stage_id),
                 connector_get_key,
+                metadata,
             )
 
         if result is None:

vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py

@@ -19,6 +19,8 @@ from vllm_omni.distributed.omni_connectors.utils.config import (
     ConnectorSpec,
     get_stage_connector_role,
 )
+from vllm_omni.distributed.omni_connectors.utils.initialization import resolve_connector_spec
+from vllm_omni.distributed.omni_connectors.utils.kv_utils import get_local_tp_rank, get_omni_replica_id
 
 logger = init_logger("vllm_omni.worker.omni_connector_model_runner_mixin")
 
@@ -186,7 +188,11 @@ class _OmniConnectorRuntimeMixin:
                 getattr(topology, "target_tp_size", None),
                 getattr(topology, "local_rank", None),
             )
-            if all(isinstance(value, int) for value in effective_mapping):
+            if (
+                isinstance(effective_mapping[0], int)
+                and isinstance(effective_mapping[1], int)
+                and isinstance(effective_mapping[2], int)
+            ):
                 rank_cfg = {
                     "from_tp": effective_mapping[0],
                     "to_tp": effective_mapping[1],
@@ -521,7 +527,15 @@ class _OmniConnectorRuntimeMixin:
         elif not isinstance(extra, dict):
             raise RuntimeError(f"Invalid extra config for connector {name}: expected dict, got {type(extra).__name__}")
 
-        spec = ConnectorSpec(name=name, extra=extra)
+        stage_id = int(getattr(model_config, "stage_id", extra.get("stage_id", 0)))
+        role = extra.get("role")
+        spec = resolve_connector_spec(
+            ConnectorSpec(name=name, extra=extra),
+            stage_id=stage_id,
+            role=str(role) if role is not None else None,
+            local_rank=get_local_tp_rank(),
+            replica_id=get_omni_replica_id(),
+        )
         try:
             return OmniConnectorFactory.create_connector(spec)
         except Exception as exc:

vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py

@@ -20,6 +20,8 @@ from vllm_omni.data_entry_keys import MetaStruct, OmniPayloadStruct, unflatten_p
 from ..adapter import construct_next_stage_streaming_input_prompt
 from ..factory import OmniConnectorFactory
 from ..utils.config import ConnectorSpec, stage_receives_chunks
+from ..utils.initialization import resolve_connector_spec
+from ..utils.kv_utils import get_local_tp_rank, get_omni_replica_id
 from ..utils.logging import get_connector_logger
 from .base import OmniTransferAdapterBase
 
@@ -47,10 +49,20 @@ class _SenderGeneration:
 class _LoadEntry:
     """Identify one receiver registration across queue and I/O boundaries."""
 
-    __slots__ = ("request",)
+    __slots__ = ("request", "source_metadata")
 
     def __init__(self, request: Request) -> None:
         self.request = request
+        sender_info = getattr(request, "payload_sender_info", None)
+        self.source_metadata = None
+        if isinstance(sender_info, dict):
+            host = sender_info.get("host")
+            port = sender_info.get("zmq_port")
+            if host and port:
+                self.source_metadata = {
+                    "source_host": str(host),
+                    "source_port": int(port),
+                }
 
     @property
     def request_id(self) -> str:
@@ -280,11 +292,20 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
                 "extra": getattr(connector_config, "extra", {}),
             }
 
-        connector_specs = ConnectorSpec(
+        connector_spec = ConnectorSpec(
             name=connector_config.get("name", "SharedMemoryConnector"),
             extra=connector_config.get("extra", {}),
         )
-        return OmniConnectorFactory.create_connector(connector_specs)
+        stage_id = int(getattr(model_config, "stage_id", connector_spec.extra.get("stage_id", 0)))
+        role = connector_spec.extra.get("role")
+        resolved_spec = resolve_connector_spec(
+            connector_spec,
+            stage_id=stage_id,
+            role=str(role) if role is not None else None,
+            local_rank=get_local_tp_rank(),
+            replica_id=get_omni_replica_id(),
+        )
+        return OmniConnectorFactory.create_connector(resolved_spec)
 
     def load_async(self, request: Request):
         """Register a request for asynchronous chunk retrieval.
@@ -484,6 +505,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
                 str(target_stage_id),
                 str(stage_id),
                 connector_get_key,
+                entry.source_metadata,
             )
         except Exception as e:
             logger.error(f"SharedMemoryConnector get failed for req {connector_get_key}: {e}")

vllm_omni/distributed/omni_connectors/utils/config.py

@@ -13,6 +13,7 @@ TRANSFER_ENGINE_CONNECTOR_NAMES = frozenset(
         "MooncakeTransferEngineConnector",
         "MoriTransferEngineConnector",
         "YuanrongTransferEngineConnector",
+        "NixlConnector",
     }
 )
 

vllm_omni/distributed/omni_connectors/utils/initialization.py

@@ -4,9 +4,10 @@
 """Utilities for OmniConnector configuration and validation."""
 
 import json
+import os
 import warnings
 from pathlib import Path
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Literal
 
 from ..factory import OmniConnectorFactory
 from .config import TRANSFER_ENGINE_CONNECTOR_NAMES, ConnectorSpec, OmniTransferConfig
@@ -38,6 +39,91 @@ KV_RANK_PORT_STRIDE = 16
 # comfortably sized block per replica for TP-rank and stage offsets.
 KV_REPLICA_PORT_STRIDE = 1024
 
+REQUEST_FORWARDING_PORT_OFFSET = 0
+ORCHESTRATOR_PORT_OFFSET = 200
+
+ConnectorPurpose = Literal["request_forwarding", "kv_transfer", "orchestrator"]
+
+
+def compute_connector_zmq_port(
+    base_port: int,
+    *,
+    purpose: ConnectorPurpose,
+    from_stage: int,
+    local_rank: int = 0,
+    replica_id: int = 0,
+) -> int:
+    """Return a collision-free connector endpoint for an edge and worker."""
+    purpose_offsets = {
+        "request_forwarding": REQUEST_FORWARDING_PORT_OFFSET,
+        "kv_transfer": KV_TRANSFER_PORT_OFFSET,
+        "orchestrator": ORCHESTRATOR_PORT_OFFSET,
+    }
+    if purpose not in purpose_offsets:
+        raise ValueError(f"Unsupported connector purpose: {purpose}")
+    if not 0 <= from_stage < KV_RANK_PORT_STRIDE:
+        raise ValueError(f"from_stage must be in [0, {KV_RANK_PORT_STRIDE}), got {from_stage}")
+    if local_rank < 0:
+        raise ValueError(f"local_rank must be non-negative, got {local_rank}")
+    if replica_id < 0:
+        raise ValueError(f"replica_id must be non-negative, got {replica_id}")
+
+    port = (
+        base_port
+        + purpose_offsets[purpose]
+        + replica_id * KV_REPLICA_PORT_STRIDE
+        + local_rank * KV_RANK_PORT_STRIDE
+        + from_stage
+    )
+    if not 1 <= port <= 65535:
+        raise ValueError(f"Resolved connector port must be in [1, 65535], got {port}")
+    return port
+
+
+def resolve_connector_spec(
+    spec: ConnectorSpec,
+    *,
+    stage_id: int,
+    role: str | None,
+    purpose: ConnectorPurpose = "request_forwarding",
+    local_rank: int = 0,
+    replica_id: int = 0,
+) -> ConnectorSpec:
+    """Copy and resolve a connector spec for one stage worker."""
+    extra = dict(spec.extra or {})
+    extra["stage_id"] = stage_id
+    if role is not None:
+        extra["role"] = role
+    if spec.name not in TRANSFER_ENGINE_CONNECTOR_NAMES:
+        return ConnectorSpec(name=spec.name, extra=extra)
+
+    resolved_port = None
+    if role != "receiver" or "sender_zmq_port" not in extra:
+        resolved_port = compute_connector_zmq_port(
+            expand_env_int(extra.get("zmq_port", 50051), "zmq_port"),
+            purpose=purpose,
+            from_stage=int(extra.get("from_stage", stage_id)),
+            local_rank=local_rank,
+            replica_id=replica_id,
+        )
+    if role == "receiver":
+        extra.setdefault("sender_host", extra.get("host", "127.0.0.1"))
+        extra.setdefault("sender_zmq_port", resolved_port)
+        extra.pop("zmq_port", None)
+        outgoing = extra.get("outgoing")
+        if spec.name == "NixlConnector" and isinstance(outgoing, dict):
+            extra["zmq_port"] = compute_connector_zmq_port(
+                expand_env_int(outgoing.get("zmq_port", 50051), "zmq_port"),
+                purpose=purpose,
+                from_stage=int(outgoing.get("from_stage", stage_id)),
+                local_rank=local_rank,
+                replica_id=replica_id,
+            )
+            extra["host"] = outgoing.get("host", "auto")
+    else:
+        extra["zmq_port"] = resolved_port
+    return ConnectorSpec(name=spec.name, extra=extra)
+
 
 def initialize_connectors_from_config(
     config_path: str | Path | None = None,
@@ -82,52 +168,40 @@ 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:
             if connector_spec.name in TRANSFER_ENGINE_CONNECTOR_NAMES:
                 extra = dict(connector_spec.extra) if connector_spec.extra else {}
-                base_port = expand_env_int(extra.get("zmq_port", 50051), "zmq_port")
                 try:
-                    stage_offset = int(from_stage)
-                except (TypeError, ValueError):
-                    stage_offset = 0
-
-                local_rank = 0 if str(caller_stage_id) == "orchestrator" else get_connector_local_rank()
-                if str(caller_stage_id) == "orchestrator":
-                    rank0_port = base_port + orchestrator_port_offset + stage_offset
-                else:
-                    rank0_port = base_port + port_offset + stage_offset
-                adjusted_port = rank0_port + local_rank * KV_RANK_PORT_STRIDE
-                extra["zmq_port"] = adjusted_port
+                    from_stage_id = int(from_stage)
+                except (TypeError, ValueError) as exc:
+                    raise ValueError(f"Transfer-engine stage id must be an integer, got {from_stage!r}") from exc
 
+                is_orchestrator = str(caller_stage_id) == "orchestrator"
+                local_rank = 0 if is_orchestrator else get_connector_local_rank()
+                resolved_purpose: ConnectorPurpose = "orchestrator" if is_orchestrator else purpose  # type: ignore[assignment]
                 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"
+                    role = "sender" if is_sender else "receiver"
+                elif caller_stage_id is not None and str(caller_stage_id) == to_stage:
+                    role = "receiver"
+                elif caller_stage_id is not None and str(caller_stage_id) == from_stage:
+                    role = "sender"
                 else:
-                    extra["role"] = extra.get("role", "auto")
-
-                connector = OmniConnectorFactory.create_connector(ConnectorSpec(name=connector_spec.name, extra=extra))
+                    role = str(extra.get("role", "auto"))
+                extra["from_stage"] = from_stage_id
+                extra["to_stage"] = int(to_stage)
+                replica_id = max(int(os.environ.get("VLLM_OMNI_REPLICA_ID", "0")), 0)
+                resolved_spec = resolve_connector_spec(
+                    ConnectorSpec(name=connector_spec.name, extra=extra),
+                    stage_id=int(caller_stage_id) if str(caller_stage_id).isdigit() else from_stage_id,
+                    role=role,
+                    purpose=resolved_purpose,
+                    local_rank=local_rank,
+                    replica_id=replica_id,
+                )
+                connector = OmniConnectorFactory.create_connector(resolved_spec)
             else:
                 connector = OmniConnectorFactory.create_connector(connector_spec)
             connectors[edge_key] = connector
@@ -174,12 +248,16 @@ def get_connectors_config_for_stage(transfer_config: OmniTransferConfig | None,
             # Incoming edge → this stage is the receiver
             extra = dict(spec.extra) if spec.extra else {}
             extra.setdefault("role", "receiver")
+            extra["from_stage"] = int(from_stage)
+            extra["to_stage"] = int(to_stage)
             stage_connectors_config[f"from_stage_{from_stage}"] = {"spec": {"name": spec.name, "extra": extra}}
         elif from_stage == target_stage and target_stage == "0":
             # Outgoing edge for stage 0 — included for async_chunk spec
             # extraction (omni_stage.py), NOT for connector instantiation.
             extra = dict(spec.extra) if spec.extra else {}
             extra.setdefault("role", "sender")
+            extra["from_stage"] = int(from_stage)
+            extra["to_stage"] = int(to_stage)
             stage_connectors_config[f"to_stage_{to_stage}"] = {"spec": {"name": spec.name, "extra": extra}}
 
     return stage_connectors_config

vllm_omni/distributed/omni_connectors/utils/kv_utils.py

@@ -16,7 +16,7 @@ from vllm.distributed.parallel_state import (
 )
 from vllm.logger import init_logger
 
-from .initialization import KV_RANK_PORT_STRIDE, KV_REPLICA_PORT_STRIDE, KV_TRANSFER_PORT_OFFSET
+from .initialization import compute_connector_zmq_port
 
 logger = init_logger(__name__)
 
@@ -123,12 +123,12 @@ def kv_zmq_port(
 
     """
     replica = get_omni_replica_id() if replica_id is None else max(int(replica_id), 0)
-    return (
-        base_port
-        + KV_TRANSFER_PORT_OFFSET
-        + replica * KV_REPLICA_PORT_STRIDE
-        + local_rank * KV_RANK_PORT_STRIDE
-        + from_stage
+    return compute_connector_zmq_port(
+        base_port,
+        purpose="kv_transfer",
+        from_stage=from_stage,
+        local_rank=local_rank,
+        replica_id=replica,
     )
 
 

vllm_omni/engine/__init__.py

@@ -79,6 +79,7 @@ class OmniEngineCoreRequest(EngineCoreRequest):
     # GPUModelRunner.model_intermediate_buffer instead of using the deprecated
     # additional_information request transport.
     model_intermediate_buffer: dict[str, Any] | None = None
+    payload_sender_info: dict[str, Any] | None = None
 
     @classmethod
     def from_request(
@@ -88,6 +89,7 @@ class OmniEngineCoreRequest(EngineCoreRequest):
         prompt_embeds: torch.Tensor | None = None,
         additional_information: AdditionalInformationPayload | None = None,
         model_intermediate_buffer: dict[str, Any] | None = None,
+        payload_sender_info: dict[str, Any] | None = None,
     ) -> "OmniEngineCoreRequest":
         """Clone an EngineCoreRequest into an OmniEngineCoreRequest with optional payload overrides."""
 
@@ -97,6 +99,8 @@ class OmniEngineCoreRequest(EngineCoreRequest):
             additional_information = getattr(request, "additional_information", None)
         if model_intermediate_buffer is None:
             model_intermediate_buffer = getattr(request, "model_intermediate_buffer", None)
+        if payload_sender_info is None:
+            payload_sender_info = getattr(request, "payload_sender_info", None)
 
         return cls(
             request_id=request.request_id,
@@ -121,6 +125,7 @@ class OmniEngineCoreRequest(EngineCoreRequest):
             abort_immediately=request.abort_immediately,
             additional_information=additional_information,
             model_intermediate_buffer=model_intermediate_buffer,
+            payload_sender_info=payload_sender_info,
         )
 
 

vllm_omni/engine/orchestrator.py

@@ -2148,6 +2148,7 @@ class Orchestrator:
         *,
         mm_features: list | None = None,
         resumable: bool = False,
+        payload_sender_info: dict[str, Any] | None = None,
     ) -> Any:
         next_pool = self.stage_pools[next_stage_id]
         if self._next_stage_input_is_tokens(next_input):
@@ -2160,6 +2161,7 @@ class Orchestrator:
                 resumable=resumable,
             )
             request.external_req_id = request.request_id
+            request.payload_sender_info = payload_sender_info
             return request
 
         processor = self._get_stage_input_processor(next_stage_id)
@@ -2181,6 +2183,7 @@ class Orchestrator:
         )
         request = self._upgrade_processed_stage_request(request, next_input)
         request.external_req_id = req_id
+        request.payload_sender_info = payload_sender_info
         return request
 
     @staticmethod
@@ -2829,6 +2832,7 @@ class Orchestrator:
                 params=params,
                 mm_features=mm_features,
                 resumable=next_stage_resumable,
+                payload_sender_info=self._build_payload_sender_info(src_stage_id, request_id=req_id),
             )
 
             if already_submitted:
@@ -2967,6 +2971,7 @@ class Orchestrator:
                     model_config=next_pool.stage_vllm_config.model_config,
                     resumable=downstream_resumable,
                 )
+                request.payload_sender_info = self._build_payload_sender_info(next_stage_id - 1, request_id=request_id)
                 request.external_req_id = request.request_id
                 submitted = await self._dispatch_or_fail_request(
                     lambda: next_pool.submit_initial(
@@ -3043,6 +3048,23 @@ class Orchestrator:
 
         return sender_infos or None
 
+    def _build_payload_sender_info(
+        self,
+        sender_stage_id: int,
+        *,
+        request_id: str | None = None,
+    ) -> dict[str, Any] | None:
+        if sender_stage_id < 0 or sender_stage_id >= len(self.stage_pools):
+            return None
+        sender_pool = self.stage_pools[sender_stage_id]
+        sender_stage = sender_pool.get_bound_client(request_id) if request_id is not None else None
+        if sender_stage is None:
+            sender_stage = sender_pool.stage_client
+        get_sender_info = getattr(sender_stage, "get_payload_sender_info", None)
+        if not callable(get_sender_info):
+            return None
+        return get_sender_info()
+
     # ---- Shutdown / lifecycle ----
 
     def _shutdown_stages(self) -> None:

vllm_omni/engine/stage_client.py

@@ -94,6 +94,8 @@ class StagePoolLLMClient(StagePoolClient, Protocol):
         kv_transfer_port_offset: int = ...,
     ) -> dict[str, Any] | None: ...
 
+    def get_payload_sender_info(self) -> dict[str, Any] | None: ...
+
 
 class StagePoolDiffusionClient(StagePoolClient, Protocol):
     """Pool-facing API for diffusion stages."""

vllm_omni/engine/stage_engine_core_client.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """
 Stage Engine Core Client for vLLM-Omni multi-stage runtime.
 
@@ -210,6 +213,7 @@ class StageEngineCoreClientBase(StageClientBase):
                 )
             raise
 
+        self._payload_sender_info = self._build_payload_sender_info()
         self._initialize_kv_sender_endpoint()
 
         logger.info(
@@ -291,6 +295,39 @@ class StageEngineCoreClientBase(StageClientBase):
             return None
         return connector_config
 
+    def _build_payload_sender_info(self) -> dict[str, Any] | None:
+        model_config = getattr(self.vllm_config, "model_config", None)
+        connector_config = getattr(model_config, "stage_connector_config", None)
+        if not isinstance(connector_config, dict):
+            return None
+        extra = connector_config.get("extra")
+        if not isinstance(extra, dict):
+            return None
+        if isinstance(extra.get("outgoing"), dict):
+            extra = extra["outgoing"]
+        elif extra.get("role") != "sender":
+            return None
+        base_port = extra.get("zmq_port", 50051)
+        if base_port is None:
+            return None
+        sender_host = self._resolve_sender_host_from_config(extra)
+        if sender_host is None:
+            return None
+        from vllm_omni.distributed.omni_connectors.utils.initialization import compute_connector_zmq_port
+
+        return {
+            "host": sender_host,
+            "zmq_port": compute_connector_zmq_port(
+                int(os.path.expandvars(str(base_port))),
+                purpose="request_forwarding",
+                from_stage=int(extra.get("from_stage", self.stage_id)),
+                replica_id=self.replica_id,
+            ),
+        }
+
+    def get_payload_sender_info(self) -> dict[str, Any] | None:
+        return dict(self._payload_sender_info) if self._payload_sender_info is not None else None
+
     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", "::"}:

vllm_omni/engine/stage_engine_core_proc.py

@@ -88,6 +88,7 @@ class StageEngineCoreProc(EngineCoreProc):
         scheduler_request, current_wave = super().preprocess_add_request(request)
         scheduler_request.additional_information = request.additional_information
         scheduler_request.external_req_id = getattr(request, "external_req_id", request.request_id)
+        scheduler_request.payload_sender_info = request.payload_sender_info
         return scheduler_request, current_wave
 
     @staticmethod

vllm_omni/engine/stage_init_utils.py

@@ -50,6 +50,9 @@ from vllm_omni.config.omni_config import (
 )
 from vllm_omni.config.stage_config import StageType
 from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.distributed.omni_connectors.utils.config import (
+    TRANSFER_ENGINE_CONNECTOR_NAMES,
+)
 from vllm_omni.engine.arg_utils import OmniEngineArgs
 from vllm_omni.entrypoints.stage_utils import _to_dict, set_stage_devices
 from vllm_omni.entrypoints.utils import filter_dataclass_kwargs
@@ -1882,17 +1885,37 @@ def get_stage_connector_spec(
 
     stage_connectors_cfg = get_stage_connector_config(omni_transfer_config, stage_id)
     for cfg in stage_connectors_cfg.values():
-        return dict(cfg.get("spec", {}))
+        connector_spec = dict(cfg.get("spec", {}))
+        if connector_spec.get("name") == "NixlConnector":
+            # A middle worker uses one connector for both directions. Preserve
+            # the incoming edge and carry the outgoing bind config separately.
+            for (source, target), outgoing in getattr(omni_transfer_config, "connectors", {}).items():
+                if source == str(stage_id):
+                    if outgoing.name != "NixlConnector":
+                        raise ValueError("A NIXL middle stage requires NIXL on its outgoing edge")
+                    extra = dict(connector_spec.get("extra", {}))
+                    extra["outgoing"] = {**(outgoing.extra or {}), "from_stage": int(source), "to_stage": int(target)}
+                    connector_spec["extra"] = extra
+                    break
+        if connector_spec.get("name") not in TRANSFER_ENGINE_CONNECTOR_NAMES:
+            extra = dict(connector_spec.get("extra", {}))
+            extra.pop("from_stage", None)
+            extra.pop("to_stage", None)
+            connector_spec["extra"] = extra
+        return connector_spec
 
     # A producer does not consume connector data itself. Keep its connector
     # for both async-chunk and terminal full-payload sends, but mark it
     # sender-only so the scheduler does not park orchestrator-provided inputs
     # waiting for an upstream payload.
     target_stage = str(stage_id)
-    for (from_stage, _to_stage), spec in getattr(omni_transfer_config, "connectors", {}).items():
+    for (from_stage, to_stage), spec in getattr(omni_transfer_config, "connectors", {}).items():
         if from_stage == target_stage:
             extra = dict(spec.extra or {})
             extra.setdefault("role", "sender")
+            if spec.name in TRANSFER_ENGINE_CONNECTOR_NAMES:
+                extra["from_stage"] = int(from_stage)
+                extra["to_stage"] = int(to_stage)
             return {"name": spec.name, "extra": extra}
     return {}