[Feature] Omni Connector + ray supported (#215)

Signed-off-by: wzliu <wzliu@connect.hku.hk>

.gitignore

@@ -229,3 +229,5 @@ configs/development.yaml
 # Docker
 .dockerignore
 Dockerfile.dev
+discussion
+tmp_test

docs/design/connectors/omni_connector_design.md

@@ -0,0 +1,200 @@
+# OmniConnector: a unified connector interface of disaggregated inference for omni-modality models
+
+This guide explains how to configure and use distributed connectors (vllm_omni/distributed/connectors) in vllm-omni for multi-stage pipelines.
+
+## 1. Overview
+
+Connectors enable data transfer between pipeline stages (e.g., Thinker -> Talker).
+Currently supported connectors operate in **D2H2D (Device-to-Host-to-Device)** mode:
+1. **SharedMemoryConnector**: Uses system shared memory.
+2. **MooncakeConnector**: Uses [Mooncake](https://github.com/kvcache-ai/Mooncake).
+
+*   **SharedMemoryConnector (Default)**: Zero-copy (on host), lowest latency. Best for **single-node** deployments. Auto-configured if no connectors are specified.
+*   **MooncakeConnector**: TCP/RDMA based. Best for **multi-node** distributed deployments. Requires a Mooncake Master service.
+
+## 2. API Design
+
+The connector system is built around the `OmniConnectorBase` abstraction, which decouples data transport from stage logic.
+
+### Core Interface
+
+```python
+class OmniConnectorBase(ABC):
+    @abstractmethod
+    def put(self, from_stage: str, to_stage: str, request_id: str, data: Any) -> tuple[bool, int, Optional[dict]]:
+        """
+        Store data.
+        Returns: (success, serialized_size, metadata)
+        """
+        pass
+
+    @abstractmethod
+    def get(self, from_stage: str, to_stage: str, request_id: str, metadata: Optional[dict] = None) -> Optional[tuple[Any, int]]:
+        """
+        Retrieve data.
+        Args: metadata - Transport-specific handles returned by put() (e.g., SHM name).
+        Returns: (object, serialized_size)
+        """
+        pass
+```
+
+### Key Concept: Metadata Passing
+Unlike a pure key-value store, some connectors (like `SharedMemoryConnector`) generate transient resources (e.g., a shared memory block name) during `put()`. This `metadata` **must be passed via the control plane** (e.g., HTTP headers, queue messages) from the producer stage to the consumer stage so `get()` can locate the data.
+
+## 3. Backends & Use Cases
+
+### 3.1 SharedMemoryConnector
+**Best for:** Single-node, high-performance IPC.
+
+*   **Mechanism:**
+    *   **Small Payloads (< Threshold)**: Data is serialized and passed directly "inline" within the `metadata` dictionary. This avoids the overhead of creating SHM blocks for tiny messages.
+    *   **Large Payloads (>= Threshold)**: Data is written to a named System Shared Memory block. The block name is returned in `metadata`.
+*   **Configuration:**
+    *   `shm_threshold_bytes`: Size in bytes to switch from inline to SHM (default: 64KB).
+
+### 3.2 MooncakeConnector
+**Best for:** Multi-node distributed inference.
+
+*   **Mechanism:** Uses Mooncake's distributed KVCache store.
+    *   **Data Plane**: TCP or RDMA for high-bandwidth transfer.
+    *   **Control Plane**: Uses a centralized Mooncake Master and Metadata Server.
+    *   **Keying**: Deterministic keys based on `request_id/from_stage_to_stage`.
+*   **Requirements**: Requires a running Mooncake Master service.
+
+## 4. Relationship with vLLM
+
+vLLM provides specialized distributed mechanisms for specific artifacts:
+*   **KV Transfer** (`vllm.distributed.kv_transfer`): Optimized for transferring KV caches between prefill and decode instances (using NCCL, Mooncake, etc.).
+*   **EC Transfer** (`vllm.distributed.ec_transfer`): Optimized for sharing encoder embeddings.
+*   **Device Communicators** (`vllm.distributed.device_communicators`): Low-level primitives (NCCL, SHM) for Tensor/Pipeline Parallelism.
+
+`vllm-omni` complements this by introducing a **Generalized Connector Abstraction** (`OmniConnector`) for multimodal pipelines. While vLLM's connectors are artifact-specific, `vllm-omni`:
+
+1.  **Unifies Transport**: Provides a single API (`put`/`get`) to transport *any* stage artifact (Input Embeddings, Hidden States, Audio/Image Tensors, KV Cache, Final Output) between arbitrary pipeline stages (e.g., AudioEncoder -> LLM -> AudioGenerator).
+2.  **Extends Connectivity**: Enables flexible construction of complex DAGs (Directed Acyclic Graphs) where stages can run in the same process, same node, or across nodes, using the most appropriate backend (SHM, Mooncake, etc.) for each edge.
+3.  **Wraps & Adapts**: Can internally utilize vLLM's specialized `kv_transfer` for KV paths while using generic transports (SHM/Mooncake) for other data types, presenting a consistent interface to the application layer.
+
+## 5. Installation (Mooncake)
+
+If using `MooncakeConnector`, install the library first:
+
+```bash
+# For CUDA-enabled systems (Recommended)
+pip install mooncake-transfer-engine
+
+# For non-CUDA systems
+pip install mooncake-transfer-engine-non-cuda
+```
+
+## 6. Using MooncakeConnector
+
+### 6.1 Start Mooncake Master
+
+Start the master service on your primary node:
+
+```bash
+# if you use mooncake SSD storage
+mkdir -p ./mc_storage
+
+mooncake_master \
+  --rpc_port=50051 \
+  --enable_http_metadata_server=true \
+  --http_metadata_server_host=0.0.0.0 \
+  --http_metadata_server_port=8080 \
+  --metrics_port=9003 \
+  --root_fs_dir=./mc_storage/ \
+  --cluster_id=mc-local-1 &
+```
+
+### 6.2 Configuration (YAML)
+
+Edit your stage config (e.g., `qwen2_5_omni.yaml`).
+
+**Step 1: Define Connector in Global Runtime**
+
+```yaml
+runtime:
+  connectors:
+    connector_of_mooncake:
+      name: MooncakeConnector
+      extra:
+        host: "127.0.0.1"           # Local Worker IP
+        metadata_server: "http://<MASTER_IP>:8080/metadata"
+        master: "<MASTER_IP>:50051"
+        segment: 512000000          # 512MB segment
+        localbuf: 64000000          # 64MB buffer
+        proto: "tcp"                # "tcp" or "rdma"
+    ```
+
+**Mooncake Configuration Parameters:**
+
+*   **host**: The hostname or IP address of the local machine (worker). Mooncake uses this to register itself in the metadata server so other nodes can find it.
+*   **metadata_server**: The URL of the metadata server. This is used for service discovery and connection establishment (e.g., exchanging QP information for RDMA).
+*   **master**: The address of the Mooncake Master Server (e.g., `<MASTER_IP>:50051`). This is used for global state management and control plane operations.
+*   **segment**: The size of the global memory segment in bytes (default: ~512MB). This defines the shared memory region accessible by Mooncake for data transfer.
+*   **localbuf**: The size of the local buffer in bytes (default: ~64MB). Used for local data buffering during transfer operations.
+*   **proto**: The transport protocol to use. Options:
+    *   `tcp`: Standard TCP/IP (easier setup, universal compatibility).
+    *   `rdma`: Remote Direct Memory Access (higher performance, requires RDMA-capable hardware).
+
+For more details, refer to the [Mooncake Repository](https://github.com/kvcache-ai/Mooncake).
+
+    **Step 2: Reference in Stages**
+
+Explicitly link stages using `input_connectors` and `output_connectors`:
+
+```yaml
+stage_args:
+  - stage_id: 0
+    # ...
+    output_connectors:
+      to_stage_1: connector_of_mooncake
+
+  - stage_id: 1
+    # ...
+    input_connectors:
+      from_stage_0: connector_of_mooncake
+```
+
+## 7. Using SharedMemoryConnector (Auto-Mode)
+
+**Best for single-node.**
+
+The system will automatically create `SharedMemoryConnector`s for any pipeline edge that does not have an explicit connector defined. This is inferred from:
+1.  `runtime.edges` list in the config.
+2.  `engine_input_source` dependencies defined in `stage_args`.
+
+### Threshold Configuration
+By default, payloads larger than **64KB** (default threshold) are transferred via shared memory, while smaller ones use the control queue (inline).
+
+To adjust this threshold (e.g., to 1GB), add the following to your `runtime.connectors`:
+
+```yaml
+runtime:
+  connectors:
+    connector_of_shared_memory:
+      name: SharedMemoryConnector
+      extra:
+        shm_threshold_bytes: 1024 # 1KB threshold
+```
+
+## 8. Summary
+
+| Use Case | Recommended Connector | Configuration |
+| :--- | :--- | :--- |
+| **Single Node** | `SharedMemoryConnector` | **None** (Automatic) or Custom Threshold |
+| **Multi Node** | `MooncakeConnector` | Explicit YAML + Mooncake Master |
+
+## 9. Operational Notes (important)
+
+- **Fail-fast config validation**: the loader raises if any expected edge is missing a connector. Define `input_connectors`/`output_connectors` or rely on auto-SHM filling; otherwise startup aborts.
+- **Missing payloads halt the stage**: workers expect connector payloads; if metadata or connector config is missing, the stage raises and stops. Verify connector wiring and metadata propagation before production.
+
+## 10. Future Roadmap: Device-to-Device (D2D) Transport
+
+The current implementations (`SharedMemoryConnector`, `MooncakeConnector`) utilize a **D2H2D (Device-to-Host-to-Device)** data path. Tensors are moved to CPU memory (Host) for transport, which incurs PCIe overhead.
+
+As outlined in the design RFC, future versions will introduce **D2D (Device-to-Device)** connectors:
+
+*   **Goal**: Direct GPU-to-GPU transfer (via NCCL, UCX, or IPC) to minimize latency for large tensor payloads.
+*   **Mechanism**: The `OmniConnector` API allows `put()` to initiate a transfer and return a lightweight handle (metadata) via the control plane, while the heavy payload flows directly between devices.

docs/design/connectors/ray_based_execution.md

@@ -0,0 +1,61 @@
+# VLLM-Omni Distributed Utilities
+
+This directory (vllm_omni/distributed/ray_utils) contains utilities for distributed execution in vllm-omni, supporting both **Ray** and **Multiprocessing** backends.
+
+## 1. Ray Utils
+
+The `ray_utils` module provides helper functions for managing Ray clusters and actors, which is essential for:
+*   **Multi-node deployment**: Running pipeline stages across different physical machines.
+*   **Resource management**: Efficient GPU/CPU allocation.
+
+### 1.1 Basic Usage
+
+To use the Ray backend, specify `worker_backend="ray"` when initializing the engine.
+
+**Command Line Example:**
+```bash
+vllm serve Qwen/Qwen2.5-Omni-7B \
+  --omni \
+  --port 8091 \
+  --worker-backend ray \
+  --ray-address auto
+```
+
+### 1.2 Cluster Setup
+
+**Step 1: Start Head Node**
+Run this on your primary machine:
+```bash
+ray start --head --port=6399
+```
+
+**Step 2: Connect Worker Nodes**
+Run this on each worker machine:
+```bash
+ray start --address=<HEAD_NODE_IP>:6399
+```
+
+> **Tip**: For a complete cluster setup script, refer to the vLLM example:
+> [run_cluster.sh](https://github.com/vllm-project/vllm/blob/main/examples/online_serving/run_cluster.sh)
+
+### 1.3 Distributed Connector Support
+
+When running on Ray, the system automatically adapts its communication strategy:
+
+*   **Cross-Node**: Recommended to use `MooncakeConnector` (requires separate configuration).
+*   **Same-Node**: Can still use `SharedMemoryConnector` for efficiency, or Ray's native object store (plasma).
+*   **SHM threshold default differs**: when `worker_backend="ray"`, the SharedMemoryConnector default threshold is set to `sys.maxsize`, which forces payloads to go inline (no SHM). Override `shm_threshold_bytes` in the connector config if you want SHM for Ray runs.
+
+### 1.4 Internal Helpers
+
+*   **`initialize_ray_cluster`**: Connects to an existing Ray cluster or starts a local one.
+
+## 2. Troubleshooting
+
+*   **Connection Issues**: Ensure the Ray head node is accessible and ports (default 6399 in this example) are open.
+*   **Version Mismatch**: Ensure all nodes run the same version of Ray and Python.
+
+### Installation
+```bash
+pip install "ray[default]"
+```

examples/offline_inference/qwen2_5_omni/end2end.py

@@ -496,7 +496,15 @@ def parse_args():
         default=16000,
         help="Sampling rate for audio loading (default: 16000).",
     )
-
+    parser.add_argument(
+        "--worker-backend", type=str, default="multi_process", choices=["multi_process", "ray"], help="backend"
+    )
+    parser.add_argument(
+        "--ray-address",
+        type=str,
+        default=None,
+        help="Address of the Ray cluster.",
+    )
     return parser.parse_args()
 
 

tests/distributed/omni_connectors/test_adapter_and_flow.py

@@ -0,0 +1,236 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from unittest.mock import MagicMock
+
+import pytest
+
+from vllm_omni.distributed.omni_connectors.adapter import try_recv_via_connector, try_send_via_connector
+from vllm_omni.distributed.omni_connectors.connectors.shm_connector import SharedMemoryConnector
+from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec, OmniTransferConfig
+from vllm_omni.distributed.omni_connectors.utils.initialization import get_connectors_config_for_stage
+
+
+@pytest.fixture
+def mock_objects():
+    return {"connector": MagicMock(), "metrics": MagicMock(), "queue_fn": MagicMock()}
+
+
+def test_send_success(mock_objects):
+    """Test try_send_via_connector success path."""
+    # Setup
+    mock_connector = mock_objects["connector"]
+    mock_metrics = mock_objects["metrics"]
+    mock_queue_fn = mock_objects["queue_fn"]
+
+    stage_id = 0
+    next_stage_id = 1
+    req_id = "req_123"
+    inputs = {"input_ids": [1, 2, 3]}
+    sampling_params = {"temperature": 0.7}
+    prompt = "test prompt"
+
+    # Mock connector.put return
+    # Returns: (success, size, metadata)
+    mock_metadata = {"handle": "xyz"}
+    mock_connector.put.return_value = (True, 100, mock_metadata)
+
+    # Execute
+    result = try_send_via_connector(
+        connector=mock_connector,
+        stage_id=stage_id,
+        next_stage_id=next_stage_id,
+        req_id=req_id,
+        next_inputs=inputs,
+        sampling_params=sampling_params,
+        original_prompt=prompt,
+        next_stage_queue_submit_fn=mock_queue_fn,
+        metrics=mock_metrics,
+    )
+
+    # Verify
+    assert result is True
+
+    # 1. Verify connector.put called correctly
+    mock_connector.put.assert_called_once()
+    args, _ = mock_connector.put.call_args
+    assert args[0] == "0"  # from_stage
+    assert args[1] == "1"  # to_stage
+    assert args[2] == req_id
+    # Verify payload structure in put
+    payload = args[3]
+    assert payload["engine_inputs"] == inputs
+    assert payload["sampling_params"] == sampling_params
+
+    # 2. Verify queue notification submitted
+    mock_queue_fn.assert_called_once()
+    notify_payload = mock_queue_fn.call_args[0][0]
+    assert notify_payload["request_id"] == req_id
+    assert notify_payload["from_connector"] is True
+    assert notify_payload["connector_metadata"] == mock_metadata
+
+    # 3. Verify metrics recorded
+    mock_metrics.on_forward.assert_called_once()
+
+
+def test_send_fail(mock_objects):
+    """Test try_send_via_connector when connector fails."""
+    mock_connector = mock_objects["connector"]
+    mock_metrics = mock_objects["metrics"]
+    mock_queue_fn = mock_objects["queue_fn"]
+
+    mock_connector.put.return_value = (False, 0, None)
+
+    result = try_send_via_connector(
+        connector=mock_connector,
+        stage_id=0,
+        next_stage_id=1,
+        req_id="req_fail",
+        next_inputs={},
+        sampling_params={},
+        original_prompt="",
+        next_stage_queue_submit_fn=mock_queue_fn,
+        metrics=mock_metrics,
+    )
+
+    assert result is False
+    mock_queue_fn.assert_not_called()
+
+
+def test_recv_success(mock_objects):
+    """Test try_recv_via_connector success path."""
+    mock_connector = mock_objects["connector"]
+
+    # Setup task received from queue
+    task = {
+        "request_id": "req_recv",
+        "from_connector": True,
+        "from_stage": "0",
+        "connector_metadata": {"handle": "xyz"},
+    }
+
+    # Setup connectors dict
+    connectors = {("0", "1"): mock_connector}
+
+    # Mock connector.get return
+    expected_data = {"engine_inputs": {"ids": [1]}}
+    # get returns: (data_obj, size)
+    mock_connector.get.return_value = (expected_data, 50)
+    # serialize_obj needed for metrics calculation if size not returned directly
+    mock_connector.serialize_obj.return_value = b"bytes"
+
+    # Execute
+    # We are stage 1 receiving from stage 0
+    inputs, rx_metrics = try_recv_via_connector(task, connectors, stage_id=1)
+
+    # Verify
+    assert inputs == expected_data["engine_inputs"]
+    assert rx_metrics is not None
+    mock_connector.get.assert_called_once_with("0", "1", "req_recv", metadata={"handle": "xyz"})
+
+
+def test_recv_no_connector():
+    """Test recv fails when no connector exists for edge."""
+    task = {"request_id": "req_missing", "from_connector": True, "from_stage": "0"}
+    connectors = {}  # Empty connectors
+
+    inputs, _ = try_recv_via_connector(task, connectors, stage_id=1)
+    assert inputs is None
+
+
+def test_shm_connector_flow():
+    """
+    Verify the full flow: Send -> Adapter -> Connector -> Adapter -> Recv.
+    Using real SharedMemoryConnector (inline mode for simplicity).
+    """
+    # 1. Setup Connector
+    config = {"shm_threshold_bytes": 1024}  # Large threshold to use inline
+    connector = SharedMemoryConnector(config)
+    connectors_map = {("0", "1"): connector}
+
+    # 2. Setup Data
+    stage_id = 0
+    next_stage_id = 1
+    req_id = "flow_req"
+    inputs = {"tokens": [10, 20, 30]}
+    sampling_params = {"n": 1}
+
+    # Queue capture mechanism
+    queue_capture = []
+
+    def mock_submit(payload):
+        queue_capture.append(payload)
+
+    mock_metrics = MagicMock()
+
+    # 3. Send
+    success = try_send_via_connector(
+        connector=connector,
+        stage_id=stage_id,
+        next_stage_id=next_stage_id,
+        req_id=req_id,
+        next_inputs=inputs,
+        sampling_params=sampling_params,
+        original_prompt="prompt",
+        next_stage_queue_submit_fn=mock_submit,
+        metrics=mock_metrics,
+    )
+    assert success is True
+    assert len(queue_capture) == 1
+
+    # 4. Recv
+    # The 'task' is what would be popped from the queue
+    received_task = queue_capture[0]
+
+    # Verify queue payload contains what we expect
+    assert received_task["from_connector"] is True
+    assert received_task["from_stage"] == "0"
+
+    # Decode
+    decoded_inputs, _ = try_recv_via_connector(received_task, connectors_map, stage_id=1)
+
+    # 5. Verify Data Integrity
+    assert decoded_inputs == inputs
+
+
+def test_get_connectors_for_stage():
+    """Test filtering logic for stage config."""
+    # Config has edges: 0->1, 1->2
+    config = OmniTransferConfig(connectors={("0", "1"): ConnectorSpec(name="C1"), ("1", "2"): ConnectorSpec(name="C2")})
+
+    # Get config for Stage 1
+    # Stage 1 receives from 0 (input) and sends to 2 (output)
+    # get_connectors_config_for_stage ONLY returns INPUT connectors for the worker to initialize
+
+    stage_config = get_connectors_config_for_stage(config, stage_id=1)
+
+    # Should contain "from_stage_0"
+    assert "from_stage_0" in stage_config
+    assert stage_config["from_stage_0"]["spec"]["name"] == "C1"
+
+    # Should NOT contain "from_stage_1" or related to output
+    assert "from_stage_1" not in stage_config
+
+    # Verify Stage 2
+    stage_2_config = get_connectors_config_for_stage(config, stage_id=2)
+    assert "from_stage_1" in stage_2_config
+    assert stage_2_config["from_stage_1"]["spec"]["name"] == "C2"
+
+
+def test_recv_with_missing_metadata():
+    """Test recv when queue payload is malformed (missing metadata)."""
+    # Connector expects metadata but task doesn't have it
+    task = {
+        "request_id": "req_bad",
+        "from_connector": True,
+        "from_stage": "0",
+        # Missing "connector_metadata"
+    }
+    mock_conn = MagicMock()
+    # If get is called with None metadata, connector usually handles it or adapter handles exception
+    mock_conn.get.side_effect = Exception("Get failed")
+
+    connectors = {("0", "1"): mock_conn}
+
+    inputs, _ = try_recv_via_connector(task, connectors, stage_id=1)
+    assert inputs is None

tests/distributed/omni_connectors/test_basic_connectors.py

@@ -0,0 +1,99 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from unittest.mock import MagicMock
+
+import pytest
+
+from vllm_omni.distributed.omni_connectors.connectors.shm_connector import SharedMemoryConnector
+from vllm_omni.distributed.omni_connectors.factory import OmniConnectorFactory
+from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec
+from vllm_omni.distributed.omni_connectors.utils.serialization import OmniSerializer
+
+
+def test_pickle_serialization():
+    """Test basic pickle serialization."""
+    data = {"key": "value", "list": [1, 2, 3]}
+    serialized = OmniSerializer.serialize(data, method="cloudpickle")
+    assert isinstance(serialized, bytes)
+
+    deserialized = OmniSerializer.deserialize(serialized, method="cloudpickle")
+    assert data == deserialized
+
+
+def test_create_shm_connector():
+    """Test creating SharedMemoryConnector via Factory."""
+    spec = ConnectorSpec(name="SharedMemoryConnector", extra={"shm_threshold_bytes": 1024})
+    connector = OmniConnectorFactory.create_connector(spec)
+    assert isinstance(connector, SharedMemoryConnector)
+    assert connector.threshold == 1024
+
+
+def test_create_unknown_connector():
+    """Test error when creating unknown connector."""
+    spec = ConnectorSpec(name="UnknownConnector")
+    with pytest.raises(ValueError):
+        OmniConnectorFactory.create_connector(spec)
+
+
+@pytest.fixture
+def shm_connector():
+    config = {"shm_threshold_bytes": 100}  # Small threshold for testing
+    return SharedMemoryConnector(config)
+
+
+def test_put_get_inline(shm_connector):
+    """Test inline transfer for small data."""
+    data = {"small": "data"}
+    # Ensure data is smaller than threshold (100 bytes)
+
+    success, size, metadata = shm_connector.put("stage_0", "stage_1", "req_1", data)
+    assert success is True
+    assert "inline_bytes" in metadata
+    assert "shm" not in metadata
+
+    # Retrieve
+    retrieved_data, ret_size = shm_connector.get("stage_0", "stage_1", "req_1", metadata)
+    assert data == retrieved_data
+    assert size == ret_size
+
+
+def test_put_get_shm(shm_connector, monkeypatch):
+    """Test SHM transfer logic for large data (Mocked)."""
+    # Create data larger than 100 bytes
+    data = {"large": "x" * 200}
+
+    # Mock SHM return values
+    mock_handle = {"name": "test_shm", "size": 200}
+    mock_write = MagicMock(return_value=mock_handle)
+    monkeypatch.setattr("vllm_omni.distributed.omni_connectors.connectors.shm_connector.shm_write_bytes", mock_write)
+
+    # When reading, return the serialized bytes of the data
+    serialized_data = shm_connector.serialize_obj(data)
+    mock_read = MagicMock(return_value=serialized_data)
+    monkeypatch.setattr("vllm_omni.distributed.omni_connectors.connectors.shm_connector.shm_read_bytes", mock_read)
+
+    # Put
+    success, size, metadata = shm_connector.put("stage_0", "stage_1", "req_2", data)
+
+    assert success is True
+    # Should use SHM because data > threshold
+    assert "shm" in metadata
+    assert metadata["shm"] == mock_handle
+    assert "inline_bytes" not in metadata
+
+    mock_write.assert_called_once()
+
+    # Get
+    retrieved_data, ret_size = shm_connector.get("stage_0", "stage_1", "req_2", metadata)
+
+    assert data == retrieved_data
+    mock_read.assert_called_once_with(mock_handle)
+
+
+def test_get_invalid_metadata(shm_connector):
+    """Test get with invalid metadata."""
+    result = shm_connector.get("stage_0", "stage_1", "req_3", {})
+    assert result is None
+
+    result = shm_connector.get("stage_0", "stage_1", "req_3", {"unknown": "format"})
+    assert result is None

tests/distributed/omni_connectors/test_omni_connector_configs.py

@@ -0,0 +1,64 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from pathlib import Path
+
+import pytest
+
+# Use the new import path for initialization utilities
+from vllm_omni.distributed.omni_connectors.utils.initialization import load_omni_transfer_config
+
+
+def get_config_files():
+    """Helper to find config files."""
+    # Go up two levels from 'tests/distributed/omni_connectors' (approx) to 'vllm-omni' root
+    # Adjust based on file location: vllm-omni/tests/distributed/omni_connectors/test_omni_connector_configs.py
+    # This file is 4 levels deep from root if we count from tests?
+    # vllm-omni/tests/distributed/omni_connectors -> parent -> distributed -> parent -> tests -> parent -> vllm-omni
+    # Let's use resolve to be safe.
+
+    # Path(__file__) = .../vllm-omni/tests/distributed/omni_connectors/test_omni_connector_configs.py
+    # .parent = omni_connectors
+    # .parent = distributed
+    # .parent = tests
+    # .parent = vllm-omni
+
+    base_dir = Path(__file__).resolve().parent.parent.parent.parent
+    config_dir = base_dir / "vllm_omni" / "model_executor" / "stage_configs"
+
+    if not config_dir.exists():
+        return []
+
+    return list(config_dir.glob("qwen*.yaml"))
+
+
+# Collect files at module level for parametrization
+config_files = get_config_files()
+
+
+@pytest.mark.skipif(len(config_files) == 0, reason="No config files found or directory missing")
+@pytest.mark.parametrize("yaml_file", config_files, ids=lambda p: p.name)
+def test_load_qwen_yaml_configs(yaml_file):
+    """
+    Scan and test loading of all qwen*.yaml config files.
+    This ensures that existing stage configs are compatible with the OmniConnector system.
+    """
+    print(f"Testing config load: {yaml_file.name}")
+    try:
+        # Attempt to load the config
+        # default_shm_threshold doesn't matter much for loading correctness, using default
+        config = load_omni_transfer_config(yaml_file)
+
+        assert config is not None, "Config should not be None"
+
+        # Basic validation
+        # Note: Some configs might not have 'runtime' or 'connectors' section if they rely on auto-shm
+        # but the load function should succeed regardless.
+
+        # If the config defines stages, we expect connectors to be populated (either explicit or auto SHM)
+        # We can't strictly assert len(config.connectors) > 0 because a single stage pipeline might have 0 edges.
+
+        print(f"  -> Successfully loaded. Connectors: {len(config.connectors)}")
+
+    except Exception as e:
+        pytest.fail(f"Failed to load config {yaml_file.name}: {e}")

tests/test_omni_llm.py

@@ -104,6 +104,7 @@ class _FakeStage:
         shm_threshold_bytes: int = 65536,
         ctx=None,
         batch_timeout: int = 10,
+        **kwargs,
     ):
         """Mock init_stage_worker: don't start real process, just send stage_ready message."""
         # Create a mock process object

vllm_omni/distributed/__init__.py

@@ -0,0 +1,25 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from .omni_connectors import (
+    ConnectorSpec,
+    MooncakeConnector,
+    OmniConnectorBase,
+    OmniConnectorFactory,
+    OmniTransferConfig,
+    SharedMemoryConnector,
+    load_omni_transfer_config,
+)
+
+__all__ = [
+    # Config
+    "ConnectorSpec",
+    "OmniTransferConfig",
+    # Connectors
+    "OmniConnectorBase",
+    "OmniConnectorFactory",
+    "MooncakeConnector",
+    "SharedMemoryConnector",
+    # Utilities
+    "load_omni_transfer_config",
+]

vllm_omni/distributed/omni_connectors/__init__.py

@@ -0,0 +1,37 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from .connectors.base import OmniConnectorBase
+from .connectors.mooncake_connector import MooncakeConnector
+from .connectors.shm_connector import SharedMemoryConnector
+from .factory import OmniConnectorFactory
+from .utils.config import ConnectorSpec, OmniTransferConfig
+from .utils.initialization import (
+    build_stage_connectors,
+    get_connectors_config_for_stage,
+    get_stage_connector_config,
+    initialize_connectors_from_config,
+    initialize_orchestrator_connectors,
+    load_omni_transfer_config,
+)
+
+__all__ = [
+    # Config
+    "ConnectorSpec",
+    "OmniTransferConfig",
+    # Base classes and implementations
+    "OmniConnectorBase",
+    # Factory
+    "OmniConnectorFactory",
+    # Specific implementations
+    "MooncakeConnector",
+    "SharedMemoryConnector",
+    # Utilities
+    "load_omni_transfer_config",
+    "initialize_connectors_from_config",
+    "get_connectors_config_for_stage",
+    # Manager helpers
+    "initialize_orchestrator_connectors",
+    "get_stage_connector_config",
+    "build_stage_connectors",
+]

vllm_omni/distributed/omni_connectors/adapter.py

@@ -0,0 +1,168 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# temporary for compatibility with vllm_omni.entrypoints.omni_stage.py
+# and vllm_omni.entrypoints.omni_llm.py
+
+import time
+from typing import Any, Callable, Optional
+
+from .utils.logging import get_connector_logger
+
+logger = get_connector_logger(__name__)
+
+
+def try_send_via_connector(
+    connector: Any,
+    stage_id: int,
+    next_stage_id: int,
+    req_id: str,
+    next_inputs: Any,
+    sampling_params: Any,
+    original_prompt: Any,
+    next_stage_queue_submit_fn: Callable[[dict[str, Any]], None],
+    metrics: Any,
+) -> bool:
+    """
+    Attempts to send data via OmniConnector.
+    Returns True if successful, False otherwise.
+    Encapsulates the logic of preparing payload, sending via connector,
+    sending notification, and recording metrics.
+    """
+    try:
+        t0 = time.time()
+
+        # Prepare data for connector
+        payload_data = {
+            "engine_inputs": next_inputs,
+            "sampling_params": sampling_params,
+            "metadata": {
+                "original_prompt": original_prompt,
+                "stage_transition": f"{stage_id}->{next_stage_id}",
+                "timestamp": time.time(),
+            },
+        }
+
+        # Send data via connector
+        success, serialized_size, metadata = connector.put(str(stage_id), str(next_stage_id), str(req_id), payload_data)
+
+        if success:
+            # Send lightweight notification via queue
+            notify_payload = {
+                "request_id": req_id,
+                "sampling_params": sampling_params,
+                "from_connector": True,
+                "from_stage": str(stage_id),
+                "to_stage": str(next_stage_id),
+                "sent_ts": time.time(),
+            }
+            # Merge connector metadata (e.g. shm handle or inline data) into queue payload
+            if metadata:
+                notify_payload["connector_metadata"] = metadata
+
+            next_stage_queue_submit_fn(notify_payload)
+
+            t1 = time.time()
+            tx_ms = (t1 - t0) * 1000.0
+
+            metrics.on_forward(
+                stage_id,
+                next_stage_id,
+                req_id,
+                serialized_size,  # Use size from connector
+                float(tx_ms),
+                True,  # Mark as using connector
+            )
+            return True
+        else:
+            # If put returned False, we let the caller handle fallback
+            return False
+
+    except Exception as e:
+        logger.warning(
+            "[Orchestrator] OmniConnector failed for req %s: %s; falling back to queue",
+            req_id,
+            e,
+        )
+        return False
+
+
+def try_recv_via_connector(
+    task: dict[str, Any],
+    connectors: dict[Any, Any],
+    stage_id: int,
+) -> tuple[Any, Optional[dict[str, Any]]]:
+    """
+    Attempts to resolve input data from either connector or IPC.
+    Returns (engine_inputs, rx_metrics) or (None, None) if failed/skipped.
+    """
+    rid = task["request_id"]
+
+    if task.get("from_connector"):
+        from_stage = task.get("from_stage")
+        to_stage = str(stage_id)
+
+        if not from_stage:
+            logger.error(
+                "[Stage-%s] 'from_connector' is true but 'from_stage' is missing for request %s", stage_id, rid
+            )
+            return None, None
+
+        # Get connector for this edge
+        connector_key = (from_stage, to_stage)
+        connector = connectors.get(connector_key)
+
+        if connector:
+            try:
+                # Get data from connector with timeout
+                _t_start = time.time()
+                connector_metadata = task.get("connector_metadata")
+                payload = connector.get(from_stage, to_stage, str(rid), metadata=connector_metadata)
+                _t_end = time.time()
+
+                if payload:
+                    if isinstance(payload, tuple):
+                        payload_data, serialized_size = payload
+                    else:
+                        payload_data = payload
+                        serialized_size = len(connector.serialize_obj(payload_data))
+                else:
+                    payload_data = None
+                    serialized_size = 0
+
+                if payload_data and isinstance(payload_data, dict):
+                    ein = payload_data.get("engine_inputs")
+                    decode_ms = (_t_end - _t_start) * 1000.0
+
+                    rx_metrics = {"rx_decode_time_ms": decode_ms, "rx_transfer_bytes": serialized_size}
+                    return ein, rx_metrics
+                else:
+                    logger.error(
+                        "[Stage-%s] Failed to get data from connector for request %s or payload is empty", stage_id, rid
+                    )
+                    return None, None
+            except Exception as e:
+                logger.error("[Stage-%s] Error retrieving data from connector for request %s: %s", stage_id, rid, e)
+                return None, None
+        else:
+            logger.error(
+                "[Stage-%s] No connector found for edge %s -> %s for request %s", stage_id, from_stage, to_stage, rid
+            )
+            return None, None
+    else:
+        # Data comes from queue as usual (e.g. seed request for Stage-0)
+        # Since fallback logic is deprecated, we assume this is a direct inputs payload.
+        # We still need to decode it if it used SHM (via legacy stage_utils logic, or new shm_connector format)
+        # For Stage-0 specifically, 'engine_inputs' is often directly in the task dict.
+
+        # Try to use the new stage_utils which uses OmniSerializer
+        from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc_with_metrics
+
+        try:
+            ein, metrics = maybe_load_from_ipc_with_metrics(task, "engine_inputs", "engine_inputs_shm")
+            # If metrics are empty or zero, we might want to populate dummy metrics
+            return ein, metrics
+        except Exception:
+            # If engine_inputs is missing, it might be a different kind of payload,
+            # but for Stage-0 seed it should be there.
+            # We'll return None to let caller handle error if strictly required.
+            return None, None

vllm_omni/distributed/omni_connectors/connectors/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm_omni/distributed/omni_connectors/connectors/base.py

@@ -0,0 +1,72 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from abc import ABC, abstractmethod
+from typing import Any, Optional
+
+from ..utils.logging import get_connector_logger
+
+logger = get_connector_logger(__name__)
+
+
+class OmniConnectorBase(ABC):
+    """Base class for all OmniConnectors."""
+
+    @abstractmethod
+    def put(
+        self, from_stage: str, to_stage: str, request_id: str, data: Any
+    ) -> tuple[bool, int, Optional[dict[str, Any]]]:
+        """Store Python object, internal serialization handled by connector.
+
+        Args:
+            from_stage: Source stage identifier
+            to_stage: Destination stage identifier
+            request_id: Unique request identifier
+            data: Python object to store
+
+        Returns:
+            tuple: (success: bool, serialized_size: int, metadata: Optional[dict])
+                   Metadata may contain transport-specific handles or inline data.
+        """
+        pass
+
+    @abstractmethod
+    def get(
+        self, from_stage: str, to_stage: str, request_id: str, metadata: Optional[dict[str, Any]] = None
+    ) -> Optional[tuple[Any, int]]:
+        """Retrieve Python object and payload size (bytes).
+
+        Args:
+            from_stage: Source stage identifier
+            to_stage: Destination stage identifier
+            request_id: Unique request identifier
+            metadata: Optional transport-specific metadata from the put operation
+
+        Returns:
+            Tuple of (Python object, serialized byte size) if found, None otherwise
+        """
+        pass
+
+    @abstractmethod
+    def cleanup(self, request_id: str) -> None:
+        """Clean up resources for a request."""
+        pass
+
+    @abstractmethod
+    def health(self) -> dict[str, Any]:
+        """Return health status and metrics."""
+        pass
+
+    @staticmethod
+    def serialize_obj(obj: Any) -> bytes:
+        """Serialize a Python object to bytes using centralized serializer."""
+        from ..utils.serialization import OmniSerializer
+
+        return OmniSerializer.serialize(obj)
+
+    @staticmethod
+    def deserialize_obj(data: bytes) -> Any:
+        """Deserialize bytes to Python object using centralized serializer."""
+        from ..utils.serialization import OmniSerializer
+
+        return OmniSerializer.deserialize(data)

vllm_omni/distributed/omni_connectors/connectors/mooncake_connector.py

@@ -0,0 +1,170 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import time
+from typing import Any, Optional
+
+from ..utils.logging import get_connector_logger
+from .base import OmniConnectorBase
+
+logger = get_connector_logger(__name__)
+
+try:
+    from mooncake.store import MooncakeDistributedStore, ReplicateConfig
+except ImportError:
+    try:
+        from mooncake import MooncakeDistributedStore, ReplicateConfig
+    except ImportError:
+        logger.warning("Mooncake not available, MooncakeOmniConnector will not work")
+        MooncakeDistributedStore = None
+        ReplicateConfig = None
+
+
+class MooncakeConnector(OmniConnectorBase):
+    """Mooncake-based distributed connector for OmniConnector."""
+
+    def __init__(self, config: dict[str, Any]):
+        if MooncakeDistributedStore is None:
+            raise ImportError("Mooncake not available")
+
+        self.config = config
+        self.host = config.get("host", "127.0.0.1")
+        self.metadata = config.get("metadata_server", "http://127.0.0.1:8080/metadata")
+        self.master = config.get("master", "127.0.0.1:50051")
+        self.segment = config.get("segment", 512 * 1024 * 1024)  # 512MB
+        self.localbuf = config.get("localbuf", 64 * 1024 * 1024)  # 64MB
+        self.proto = config.get("proto", "tcp")
+        self.rdma = config.get("rdma", "")
+
+        self.store: Optional[MooncakeDistributedStore] = None
+        self.pin: Optional[ReplicateConfig] = None
+
+        self._metrics = {
+            "puts": 0,
+            "gets": 0,
+            "bytes_transferred": 0,
+            "errors": 0,
+            "timeouts": 0,
+        }
+
+        self._init_store()
+
+    def _make_key(self, rid: str, from_stage: str, to_stage: str) -> str:
+        """Generate store key for request between stages."""
+        return f"{rid}/{from_stage}_to_{to_stage}"
+
+    def _init_store(self):
+        """Initialize Mooncake store."""
+        try:
+            self.store = MooncakeDistributedStore()
+            rc = self.store.setup(
+                self.host, self.metadata, self.segment, self.localbuf, self.proto, self.rdma, self.master
+            )
+            if rc != 0:
+                raise RuntimeError(f"Mooncake setup failed: {rc}")
+
+            self.pin = ReplicateConfig()
+            self.pin.with_soft_pin = True
+            logger.info("MooncakeConnector initialized successfully")
+        except Exception as e:
+            logger.error("Failed to initialize Mooncake store: %s", e)
+            raise
+
+    # Use base class serialization methods for consistency
+
+    def put(
+        self, from_stage: str, to_stage: str, request_id: str, data: Any
+    ) -> tuple[bool, int, Optional[dict[str, Any]]]:
+        if not self.store:
+            logger.error("Store not initialized")
+            return False, 0, None
+
+        try:
+            serialized_data = self.serialize_obj(data)
+            key = self._make_key(request_id, from_stage, to_stage)
+            self.store.put(key, serialized_data, self.pin)
+
+            self._metrics["puts"] += 1
+            self._metrics["bytes_transferred"] += len(serialized_data)
+
+            logger.debug(
+                "MooncakeConnector: stored %s (%s -> %s) %d bytes",
+                key,
+                from_stage,
+                to_stage,
+                len(serialized_data),
+            )
+            return True, len(serialized_data), None
+
+        except Exception as e:
+            self._metrics["errors"] += 1
+            logger.error("MooncakeConnector put failed: %s", e)
+            return False, 0, None
+
+    def get(
+        self, from_stage: str, to_stage: str, request_id: str, metadata: Optional[dict[str, Any]] = None
+    ) -> Optional[tuple[Any, int]]:
+        if not self.store:
+            logger.error("Store not initialized")
+            return None
+
+        retries = 20
+        sleep_s = 0.05
+        key = self._make_key(request_id, from_stage, to_stage)
+
+        for attempt in range(retries):
+            try:
+                raw_data = self.store.get(key)
+
+                if raw_data:
+                    data = self.deserialize_obj(raw_data)
+                    self._metrics["gets"] += 1
+                    payload_size = len(raw_data)
+                    logger.debug(
+                        "MooncakeConnector: retrieved %s (%s -> %s) %d bytes",
+                        key,
+                        from_stage,
+                        to_stage,
+                        payload_size,
+                    )
+                    return data, payload_size
+
+            except Exception as e:
+                logger.debug("MooncakeConnector get attempt %s failed: %s", attempt, e)
+
+            if attempt < retries - 1:
+                time.sleep(sleep_s)
+
+        self._metrics["timeouts"] += 1
+        logger.warning("MooncakeConnector: timeout waiting for %s", key)
+        return None
+
+    def cleanup(self, request_id: str) -> None:
+        if not self.store:
+            return
+
+        # Note: Mooncake doesn't have explicit delete, data will be garbage collected
+        # We could implement a cleanup mechanism by storing deletion markers
+        logger.debug("MooncakeConnector: cleanup requested for %s (no-op)", request_id)
+
+    def health(self) -> dict[str, Any]:
+        if not self.store:
+            return {"status": "unhealthy", "error": "Store not initialized"}
+
+        return {
+            "status": "healthy",
+            "host": self.host,
+            "metadata_server": self.metadata,
+            "master": self.master,
+            **self._metrics,
+        }
+
+    def close(self):
+        """Clean shutdown."""
+        if self.store:
+            try:
+                self.store.close()
+                self.store = None
+                logger.info("MooncakeConnector closed")
+            except Exception as e:
+                logger.error("Error closing Mooncake store: %s", e)

vllm_omni/distributed/omni_connectors/connectors/shm_connector.py

@@ -0,0 +1,118 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from typing import Any, Optional
+
+from vllm_omni.entrypoints.stage_utils import shm_read_bytes, shm_write_bytes
+
+from ..utils.logging import get_connector_logger
+from .base import OmniConnectorBase
+
+logger = get_connector_logger(__name__)
+
+
+class SharedMemoryConnector(OmniConnectorBase):
+    """
+    Connector that uses SharedMemory for large objects and inline data for small objects.
+    Acts as a unified replacement for the legacy IPC fallback logic.
+    """
+
+    def __init__(self, config: dict[str, Any]):
+        self.config = config
+        # Default threshold matches legacy behavior (64KB)
+        self.threshold = int(config.get("shm_threshold_bytes", 65536))
+        self._metrics = {
+            "puts": 0,
+            "gets": 0,
+            "bytes_transferred": 0,
+            "shm_writes": 0,
+            "inline_writes": 0,
+        }
+
+    def put(
+        self, from_stage: str, to_stage: str, request_id: str, data: Any
+    ) -> tuple[bool, int, Optional[dict[str, Any]]]:
+        try:
+            # Always serialize first to check size (and for SHM writing)
+            # Note: For extremely large objects in "inline" mode (e.g. Ray),
+            # we might double-serialize if we're not careful, but here we assume
+            # if it's huge we use SHM, or if Ray, threshold is maxsize.
+            payload = self.serialize_obj(data)
+            size = len(payload)
+
+            if size > self.threshold:
+                # Use Shared Memory
+                meta = shm_write_bytes(payload)
+                # meta contains {'name': ..., 'size': ...}
+                metadata = {"shm": meta, "size": size}
+                self._metrics["shm_writes"] += 1
+            else:
+                # Inline - pass bytes directly to avoid double serialization of the object
+                # We already serialized it to check size, so we pass the bytes.
+                # The Queue will pickle these bytes (fast), avoiding re-serializing the complex object.
+                metadata = {"inline_bytes": payload, "size": size}
+                self._metrics["inline_writes"] += 1
+
+            self._metrics["puts"] += 1
+            self._metrics["bytes_transferred"] += size
+
+            return True, size, metadata
+
+        except Exception as e:
+            logger.error(f"SharedMemoryConnector put failed for req {request_id}: {e}")
+            return False, 0, None
+
+    def get(
+        self, from_stage: str, to_stage: str, request_id: str, metadata: Optional[dict[str, Any]] = None
+    ) -> Optional[tuple[Any, int]]:
+        if not metadata:
+            logger.error(f"SharedMemoryConnector get called without metadata for req {request_id}")
+            return None
+
+        try:
+            obj = None
+            size = 0
+
+            if "shm" in metadata:
+                meta = metadata["shm"]
+                # shm_read_bytes handles reading and unlinking
+                data_bytes = shm_read_bytes(meta)
+                obj = self.deserialize_obj(data_bytes)
+                size = metadata.get("size", len(data_bytes))
+            elif "inline_bytes" in metadata:
+                # Deserialize bytes back to object
+                payload = metadata["inline_bytes"]
+                obj = self.deserialize_obj(payload)
+                size = metadata.get("size", len(payload))
+            elif "inline" in metadata:
+                obj = metadata["inline"]
+                size = metadata.get("size", 0)
+                if size == 0:
+                    # Fallback if size wasn't recorded
+                    try:
+                        size = len(self.serialize_obj(obj))
+                    except Exception:
+                        pass
+            else:
+                logger.error(
+                    f"Unknown metadata format in SharedMemoryConnector for req {request_id}: {list(metadata.keys())}"
+                )
+                return None
+
+            self._metrics["gets"] += 1
+            return obj, size
+
+        except Exception as e:
+            logger.error(f"SharedMemoryConnector get failed for req {request_id}: {e}")
+            return None
+
+    def cleanup(self, request_id: str) -> None:
+        # SHM segments are automatically unlinked during 'get' (shm_read_bytes).
+        # If 'get' is never called (e.g. error flow), the SHM segment might leak.
+        # A robust implementation might track created segments and unlink them here
+        # if they haven't been consumed.
+        # For now, we rely on the consumer to read and unlink.
+        pass
+
+    def health(self) -> dict[str, Any]:
+        return {"status": "healthy", "threshold": self.threshold, **self._metrics}

vllm_omni/distributed/omni_connectors/factory.py

@@ -0,0 +1,87 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections.abc import Callable
+from typing import Any
+
+from .utils.logging import get_connector_logger
+
+try:
+    from .connectors.base import OmniConnectorBase
+    from .utils.config import ConnectorSpec
+except ImportError:
+    # Fallback for direct execution
+    import os
+    import sys
+
+    sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
+    from omni_connectors.connectors.base import OmniConnectorBase
+    from omni_connectors.utils.config import ConnectorSpec
+
+logger = get_connector_logger(__name__)
+
+
+class OmniConnectorFactory:
+    """Factory for creating OmniConnectors."""
+
+    _registry: dict[str, Callable[[dict[str, Any]], OmniConnectorBase]] = {}
+
+    @classmethod
+    def register_connector(cls, name: str, constructor: Callable[[dict[str, Any]], OmniConnectorBase]) -> None:
+        """Register a connector constructor."""
+        if name in cls._registry:
+            raise ValueError(f"Connector '{name}' is already registered.")
+        cls._registry[name] = constructor
+        logger.debug(f"Registered connector: {name}")
+
+    @classmethod
+    def create_connector(cls, spec: ConnectorSpec) -> OmniConnectorBase:
+        """Create a connector from specification."""
+        if spec.name not in cls._registry:
+            raise ValueError(f"Unknown connector: {spec.name}. Available: {list(cls._registry.keys())}")
+
+        constructor = cls._registry[spec.name]
+        try:
+            connector = constructor(spec.extra)
+            logger.info(f"Created connector: {spec.name}")
+            return connector
+        except Exception as e:
+            logger.error(f"Failed to create connector {spec.name}: {e}")
+            raise
+
+    @classmethod
+    def list_registered_connectors(cls) -> list[str]:
+        """List all registered connector names."""
+        return list(cls._registry.keys())
+
+
+# Register built-in connectors with lazy imports
+def _create_mooncake_connector(config: dict[str, Any]) -> OmniConnectorBase:
+    try:
+        from .connectors.mooncake_connector import MooncakeConnector
+    except ImportError:
+        # Fallback import
+        import os
+        import sys
+
+        sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
+        from omni_connectors.connectors.mooncake_connector import MooncakeConnector
+    return MooncakeConnector(config)
+
+
+def _create_shm_connector(config: dict[str, Any]) -> OmniConnectorBase:
+    try:
+        from .connectors.shm_connector import SharedMemoryConnector
+    except ImportError:
+        # Fallback import
+        import os
+        import sys
+
+        sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
+        from omni_connectors.connectors.shm_connector import SharedMemoryConnector
+    return SharedMemoryConnector(config)
+
+
+# Register connectors
+OmniConnectorFactory.register_connector("MooncakeConnector", _create_mooncake_connector)
+OmniConnectorFactory.register_connector("SharedMemoryConnector", _create_shm_connector)

vllm_omni/distributed/omni_connectors/utils/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm_omni/distributed/omni_connectors/utils/config.py

@@ -0,0 +1,40 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from dataclasses import dataclass, field
+from typing import Any, Optional
+
+from .logging import get_connector_logger
+
+logger = get_connector_logger(__name__)
+
+
+@dataclass
+class ConnectorSpec:
+    """Specification for a connector instance."""
+
+    name: str  # e.g., "MooncakeConnector", "SharedMemoryConnector"
+    extra: dict[str, Any] = field(default_factory=dict)  # backend-specific config
+
+
+@dataclass
+class OmniTransferConfig:
+    """
+    Top-level configuration for OmniConnector system.
+    Members:
+        connectors: A dictionary of connectors, keyed by (from_stage, to_stage).
+        default_connector: The default connector to use if no connector is specified for an edge.
+    """
+
+    # Direct mapping: (from_stage, to_stage) -> connector
+    connectors: dict[tuple[str, str], ConnectorSpec] = field(default_factory=dict)
+    default_connector: Optional[ConnectorSpec] = None
+
+    def get_connector_for_edge(self, from_stage: str, to_stage: str) -> Optional[ConnectorSpec]:
+        """Get connector spec for a specific edge."""
+        edge_key = (from_stage, to_stage)
+        return self.connectors.get(edge_key, self.default_connector)
+
+    def has_connector_for_edge(self, from_stage: str, to_stage: str) -> bool:
+        """Check if there's a connector configured for the edge."""
+        return self.get_connector_for_edge(from_stage, to_stage) is not None

vllm_omni/distributed/omni_connectors/utils/initialization.py

@@ -0,0 +1,322 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Utilities for OmniConnector configuration and validation."""
+
+import json
+import sys
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Optional, Union
+
+from ..factory import OmniConnectorFactory
+from .config import ConnectorSpec, OmniTransferConfig
+from .logging import get_connector_logger
+
+if TYPE_CHECKING:
+    from ..connectors.base import OmniConnectorBase
+else:
+    OmniConnectorBase = Any
+
+logger = get_connector_logger(__name__)
+
+
+def initialize_connectors_from_config(
+    config_path: Optional[Union[str, Path]] = None, default_shm_threshold: int = 65536
+) -> tuple[Optional[OmniTransferConfig], dict[tuple[str, str], OmniConnectorBase]]:
+    """
+    Initialize connectors from configuration file.
+
+    Returns:
+        tuple: (OmniTransferConfig, dict of {(from, to): connector_instance})
+    """
+    transfer_config = load_omni_transfer_config(config_path, default_shm_threshold=default_shm_threshold)
+
+    if not transfer_config:
+        logger.info("No OmniTransferConfig provided")
+        return None, {}
+
+    # create connectors from config
+    connectors = create_connectors_from_config(transfer_config.connectors)
+    return transfer_config, connectors
+
+
+def create_connectors_from_config(
+    connectors_config: dict[tuple[str, str], ConnectorSpec],
+) -> dict[tuple[str, str], OmniConnectorBase]:
+    """
+    Create connectors from config.
+
+    Args:
+        connectors_config: A dictionary of connector configurations.
+
+    Returns:
+        A dictionary of connectors.
+    """
+    connectors = {}
+    for edge_key, connector_spec in connectors_config.items():
+        try:
+            connector = OmniConnectorFactory.create_connector(connector_spec)
+            connectors[edge_key] = connector
+            logger.info(f"Created connector for {edge_key[0]} -> {edge_key[1]}: {type(connector).__name__}")
+        except Exception as e:
+            raise RuntimeError(f"Failed to initialize connector for edge {edge_key}: {e}") from e
+
+    return connectors
+
+
+def get_connectors_config_for_stage(
+    transfer_config: Optional[OmniTransferConfig], stage_id: Union[str, int]
+) -> dict[str, Any]:
+    """
+    Extract connector configurations relevant for a specific stage worker.
+
+    Returns a dict compatible with worker initialization:
+    {
+        "from_stage_X": {
+            "spec": {
+                "name": "ConnectorName",
+                "extra": {...}
+            }
+        },
+        ...
+    }
+    """
+    if not transfer_config:
+        return {}
+
+    stage_connectors_config = {}
+    target_stage = str(stage_id)
+
+    # Iterate through all configured edges
+    for (from_stage, to_stage), spec in transfer_config.connectors.items():
+        # We only care about incoming edges for the worker process
+        # (Worker needs to create connectors to receive data)
+        if to_stage == target_stage:
+            stage_connectors_config[f"from_stage_{from_stage}"] = {"spec": {"name": spec.name, "extra": spec.extra}}
+
+    return stage_connectors_config
+
+
+def load_omni_transfer_config(
+    config_path: Optional[Union[str, Path]] = None,
+    config_dict: Optional[dict[str, Any]] = None,
+    default_shm_threshold: int = 65536,
+) -> Optional[OmniTransferConfig]:
+    """Load OmniTransferConfig from file or dict."""
+    if config_path is None and config_dict is None:
+        # Even if no config provided, we might want to return a default config with SHM connectors
+        # But without stage info we can't do much.
+        return None
+
+    if config_path is not None:
+        config_path = Path(config_path)
+        if not config_path.exists():
+            raise FileNotFoundError(f"Config file not found: {config_path}")
+
+        with open(config_path, encoding="utf-8") as f:
+            if config_path.suffix.lower() == ".json":
+                config_dict = json.load(f)
+            elif config_path.suffix.lower() in [".yaml", ".yml"]:
+                try:
+                    import yaml
+
+                    config_dict = yaml.safe_load(f)
+                except ImportError:
+                    raise ImportError("PyYAML required for YAML config files")
+            else:
+                raise ValueError(f"Unsupported config file format: {config_path.suffix}")
+
+    if config_dict is None:
+        return None
+
+    # Parse connectors
+    connectors = {}
+    runtime_config = config_dict.get("runtime", {})
+
+    # Parse global connectors (from runtime.connectors)
+    global_connectors = runtime_config.get("connectors", {})
+
+    # Parse stage-level connectors
+    stage_args = config_dict.get("stage_args", [])
+    expected_edges: set[tuple[str, str]] = set()
+    for stage_config in stage_args:
+        stage_id = str(stage_config["stage_id"])
+
+        # Input connectors
+        for input_key, conn_ref in stage_config.get("input_connectors", {}).items():
+            if isinstance(conn_ref, str):
+                # Reference to global connector
+                if conn_ref in global_connectors:
+                    conn_config = global_connectors[conn_ref]
+                    connector = ConnectorSpec(name=conn_config["name"], extra=conn_config.get("extra", {}))
+                else:
+                    raise ValueError(f"Undefined connector reference: {conn_ref}")
+            else:
+                # Inline connector definition
+                connector = ConnectorSpec(name=conn_ref["name"], extra=conn_ref.get("extra", {}))
+
+            # Parse from_stage from key (e.g., "from_stage_0" -> "0")
+            from_stage = input_key.replace("from_stage_", "")
+            edge_key = (from_stage, stage_id)
+            connectors[edge_key] = connector
+            expected_edges.add(edge_key)
+
+        # Output connectors
+        for output_key, conn_ref in stage_config.get("output_connectors", {}).items():
+            if isinstance(conn_ref, str):
+                # Reference to global connector
+                if conn_ref in global_connectors:
+                    conn_config = global_connectors[conn_ref]
+                    connector = ConnectorSpec(name=conn_config["name"], extra=conn_config.get("extra", {}))
+                else:
+                    raise ValueError(f"Undefined connector reference: {conn_ref}")
+            else:
+                # Inline connector definition
+                connector = ConnectorSpec(name=conn_ref["name"], extra=conn_ref.get("extra", {}))
+
+            # Parse to_stage from key (e.g., "to_stage_1" -> "1")
+            to_stage = output_key.replace("to_stage_", "")
+            edge_key = (stage_id, to_stage)
+            connectors[edge_key] = connector
+            expected_edges.add(edge_key)
+
+    # Auto-configure SharedMemoryConnector for missing edges based on runtime edges / engine_input_source
+    if stage_args:
+        try:
+            # Prefer explicit runtime edges if provided
+            runtime_edges = runtime_config.get("edges", [])
+            if isinstance(runtime_edges, list) and runtime_edges:
+                for edge in runtime_edges:
+                    from_stage = edge.get("from")
+                    to_stage = edge.get("to")
+                    if from_stage is None or to_stage is None:
+                        continue
+                    edge_key = (str(from_stage), str(to_stage))
+                    expected_edges.add(edge_key)
+                    if edge_key not in connectors:
+                        logger.info(f"Auto-configuring SharedMemoryConnector for edge {edge_key}")
+                        connectors[edge_key] = ConnectorSpec(
+                            name="SharedMemoryConnector",
+                            extra={"shm_threshold_bytes": default_shm_threshold},
+                        )
+
+            # Fallback: infer edges from engine_input_source for each stage
+            for stage_config in stage_args:
+                to_stage = str(stage_config["stage_id"])
+                # Check explicit input sources
+                sources = stage_config.get("engine_input_source", [])
+
+                for from_stage in sources:
+                    from_stage_str = str(from_stage)
+                    edge_key = (from_stage_str, to_stage)
+                    expected_edges.add(edge_key)
+
+                    if edge_key not in connectors:
+                        logger.info(f"Auto-configuring SharedMemoryConnector for edge {edge_key}")
+                        connectors[edge_key] = ConnectorSpec(
+                            name="SharedMemoryConnector", extra={"shm_threshold_bytes": default_shm_threshold}
+                        )
+
+        except Exception as e:
+            logger.warning(f"Failed to auto-configure SHM connectors: {e}")
+
+    # Fail fast if any expected edge is still missing a connector
+    missing_edges = [edge for edge in expected_edges if edge not in connectors]
+    if missing_edges:
+        missing_str = ", ".join([f"{f}->{t}" for f, t in missing_edges])
+        raise ValueError(
+            "Connector configuration missing for edges: "
+            f"{missing_str}. Define connectors or allow auto SHM creation for these edges."
+        )
+
+    config = OmniTransferConfig(connectors=connectors)
+
+    logger.info(f"Loaded OmniTransferConfig with {len(connectors)} connector configurations")
+    return config
+
+
+# High-level management functions
+
+
+def initialize_orchestrator_connectors(
+    config_path: Optional[str], worker_backend: Optional[str] = "multi_process", shm_threshold_bytes: int = 65536
+) -> tuple[Optional[OmniTransferConfig], dict[tuple[str, str], OmniConnectorBase]]:
+    """Initialize connectors shared at orchestrator level.
+    Args:
+        config_path: The path to the configuration file.
+        worker_backend: The backend to use for the worker.
+    Returns:
+        A tuple containing the OmniTransferConfig and a dictionary of connectors.
+    """
+    if worker_backend == "ray":
+        default_shm_threshold = sys.maxsize
+    else:
+        default_shm_threshold = max(0, shm_threshold_bytes)
+    transfer_config, connectors = initialize_connectors_from_config(
+        config_path, default_shm_threshold=default_shm_threshold
+    )
+    return transfer_config, connectors
+
+
+def get_stage_connector_config(
+    transfer_config: Optional[OmniTransferConfig],
+    stage_id: int,
+) -> dict[str, Any]:
+    """Return the serialized connector config payload for a specific stage."""
+    if transfer_config is None:
+        return {}
+
+    try:
+        return get_connectors_config_for_stage(transfer_config, stage_id)
+    except Exception as exc:  # pragma: no cover - defensive logging
+        logger.warning(
+            "Failed to build connector config for stage %s: %s. Using IPC fallback.",
+            stage_id,
+            exc,
+        )
+        return {}
+
+
+def build_stage_connectors(
+    stage_id: int,
+    connectors_config: dict[str, Any],
+) -> Optional[dict[tuple[str, str], Any]]:
+    """Instantiate OmniConnectors for a stage based on config."""
+    if not connectors_config:
+        return {}
+
+    logger.info(
+        "[Stage-%s] Initializing OmniConnectors with config keys: %s",
+        stage_id,
+        list(connectors_config.keys()),
+    )
+
+    from .config import ConnectorSpec
+
+    connectors: dict[tuple[str, str], Any] = {}
+    # 将字典格式的配置转换为ConnectorSpec对象
+    stage_connector_specs = {}
+    for input_key, config in connectors_config.items():
+        if not input_key.startswith("from_stage_"):
+            continue
+
+        from_stage = input_key.replace("from_stage_", "")
+        spec_dict = config.get("spec", {})
+        if not spec_dict:
+            continue
+
+        connector_spec = ConnectorSpec(
+            name=spec_dict.get("name", "SharedMemoryConnector"),
+            extra=spec_dict.get("extra", {}),
+        )
+        stage_connector_specs[(str(from_stage), str(stage_id))] = connector_spec
+
+    try:
+        # 使用统一的连接器创建逻辑
+        connectors = create_connectors_from_config(stage_connector_specs)
+    except Exception as exc:  # pragma: no cover - defensive logging
+        # Fail fast so the stage does not start with missing connectors.
+        logger.exception("[Stage-%s] Failed to initialize connectors: %s", stage_id, exc)
+        raise
+
+    return connectors

vllm_omni/distributed/omni_connectors/utils/logging.py

@@ -0,0 +1,14 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import logging
+
+try:
+    from vllm.logger import init_logger as _vllm_init_logger
+except Exception:  # pragma: no cover - optional dependency
+    _vllm_init_logger = None
+
+
+def get_connector_logger(name: str) -> logging.Logger:
+    """Return a logger preferring vLLM's init_logger when available."""
+    return _vllm_init_logger(name) if _vllm_init_logger else logging.getLogger(name)

vllm_omni/distributed/omni_connectors/utils/serialization.py

@@ -0,0 +1,79 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import io
+from typing import Any
+
+try:
+    import torch
+
+    _has_torch = True
+except ImportError:
+    _has_torch = False
+
+try:
+    import cloudpickle
+
+    _has_cloudpickle = True
+except ImportError:
+    _has_cloudpickle = False
+
+from .logging import get_connector_logger
+
+logger = get_connector_logger(__name__)
+
+
+class OmniSerializer:
+    """
+    Centralized serialization handler for OmniConnectors.
+    Supports multiple backends (torch, cloudpickle) to ensure consistency across connectors.
+    """
+
+    @staticmethod
+    def serialize(obj: Any, method: str = "cloudpickle") -> bytes:
+        """
+        Serialize an object to bytes.
+
+        Args:
+            obj: The object to serialize.
+            method: Serialization method ("cloudpickle" or "torch").
+                    Defaults to "cloudpickle" (consistent with legacy stage_utils).
+        """
+        if method == "torch" and _has_torch:
+            buf = io.BytesIO()
+            torch.save(obj, buf)
+            return buf.getvalue()
+        elif method == "cloudpickle" or (method == "torch" and not _has_torch):
+            if not _has_cloudpickle:
+                # Fallback to standard pickle if cloudpickle is missing
+                import pickle
+
+                return pickle.dumps(obj)
+            return cloudpickle.dumps(obj)
+        else:
+            # Default fallback
+            import pickle
+
+            return pickle.dumps(obj)
+
+    @staticmethod
+    def deserialize(data: bytes, method: str = "cloudpickle") -> Any:
+        """
+        Deserialize bytes to an object.
+
+        Args:
+            data: The bytes to deserialize.
+            method: Serialization method used ("cloudpickle" or "torch").
+        """
+        if method == "torch" and _has_torch:
+            return torch.load(io.BytesIO(data), map_location="cpu", weights_only=False)
+        elif method == "cloudpickle" or (method == "torch" and not _has_torch):
+            if not _has_cloudpickle:
+                import pickle
+
+                return pickle.loads(data)
+            return cloudpickle.loads(data)
+        else:
+            import pickle
+
+            return pickle.loads(data)

vllm_omni/distributed/ray_utils/__init__.py

@@ -0,0 +1,10 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from .utils import calculate_total_bytes, is_ray_initialized, maybe_disable_pin_memory_for_ray
+
+__all__ = [
+    "calculate_total_bytes",
+    "is_ray_initialized",
+    "maybe_disable_pin_memory_for_ray",
+]

vllm_omni/distributed/ray_utils/utils.py

@@ -0,0 +1,183 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import logging
+import os
+from contextlib import contextmanager
+from typing import Any, Optional
+
+import torch
+
+try:
+    import ray
+    from ray.util.queue import Queue as RayQueue
+    from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+
+    RAY_AVAILABLE = True
+    from ray.util.placement_group import PlacementGroup
+except ImportError:
+    ray = None
+    RayQueue = None
+    PlacementGroupSchedulingStrategy = None
+    RAY_AVAILABLE = False
+    PlacementGroup = Any
+
+logger = logging.getLogger(__name__)
+
+
+def is_ray_initialized():
+    """Check if Ray is initialized without hard dependency on Ray."""
+    # 1. Try standard API
+    if RAY_AVAILABLE:
+        if ray.is_initialized():
+            return True
+    # 2. Fallback: Check environment variables typical for Ray Workers
+    # RAY_RAYLET_PID is always set in Ray workers
+    if "RAY_RAYLET_PID" in os.environ:
+        return True
+    return False
+
+
+def calculate_total_bytes(size_args, dtype):
+    """
+    Calculate total bytes for a tensor allocation, handling nested tuples in size args.
+    """
+    num_elements = 1
+    for s in size_args:
+        if isinstance(s, (tuple, list)):
+            for inner in s:
+                num_elements *= inner
+        else:
+            num_elements *= s
+
+    element_size = torch.tensor([], dtype=dtype).element_size()
+    return num_elements * element_size
+
+
+@contextmanager
+def maybe_disable_pin_memory_for_ray(obj, size_bytes, threshold=32 * 1024 * 1024):
+    """
+    Context manager to temporarily disable pin_memory if running in Ray and
+    the allocation size exceeds the threshold.
+
+    This is a workaround for Ray workers often having low ulimit -l (locked memory),
+    causing OS call failed errors when allocating large pinned buffers.
+    """
+    should_disable = False
+    old_pin = False
+
+    # Check 1: Are we in a Ray-like environment?
+    in_ray = is_ray_initialized()
+
+    # Check 2: Is the size large enough to worry?
+    is_large = size_bytes > threshold
+
+    # Check 3: Is pinning currently enabled?
+    is_pinned = getattr(obj, "pin_memory", False)
+
+    if in_ray and is_large and is_pinned:
+        should_disable = True
+        old_pin = obj.pin_memory
+        obj.pin_memory = False
+
+    try:
+        yield
+    finally:
+        if should_disable:
+            obj.pin_memory = old_pin
+
+
+# --- Ray specific utilities ---
+
+
+def get_ray_queue_class():
+    if not RAY_AVAILABLE:
+        raise ImportError("ray is required for worker_backend='ray'")
+    return lambda: RayQueue(maxsize=0)
+
+
+def initialize_ray_cluster(address: Optional[str] = None):
+    if not RAY_AVAILABLE:
+        logger.warning("Ray is not available, skipping initialization.")
+        return
+
+    if not ray.is_initialized():
+        # Pass current PYTHONPATH to workers to ensure they can find vllm_omni
+        runtime_env = {"env_vars": {"PYTHONPATH": os.environ.get("PYTHONPATH", "")}}
+        ray.init(address=address, ignore_reinit_error=True, runtime_env=runtime_env)
+
+
+def create_placement_group(
+    number_of_stages: int, address: Optional[str] = None, strategy: str = "PACK"
+) -> PlacementGroup:
+    """Create a placement group for the given number of stages.
+    Args:
+        number_of_stages: The number of stages to create the placement group for.
+        strategy: The strategy to use for the placement group.
+    Returns:
+        The placement group.
+    """
+    if not RAY_AVAILABLE:
+        raise ImportError("ray is required for creating placement group")
+
+    # Initialize Ray if not already initialized (using default args if needed)
+    if not ray.is_initialized():
+        logger.warning("[Orchestrator] Ray is not initialized. Initializing with default settings.")
+        initialize_ray_cluster(address)
+
+    bundles = [{"GPU": 1.0, "CPU": 1.0} for _ in range(number_of_stages)]
+    pg = ray.util.placement_group(bundles, strategy=strategy)
+    ray.get(pg.ready())
+    logger.info("[Orchestrator] Ray Placement Group created")
+    return pg
+
+
+def remove_placement_group(pg):
+    if pg and RAY_AVAILABLE:
+        try:
+            ray.util.remove_placement_group(pg)
+        except Exception as e:
+            logger.warning(f"Failed to remove placement group: {e}")
+
+
+def try_close_ray(pg=None):
+    """Try to clean up Ray resources including placement group and shutdown."""
+    if pg:
+        remove_placement_group(pg)
+    # Note: We typically don't shutdown ray.init() here as it might be used by other components
+    # or the user might want it to persist. If full shutdown is needed, ray.shutdown() can be called.
+
+
+def kill_ray_actor(actor):
+    if actor and RAY_AVAILABLE:
+        try:
+            ray.kill(actor)
+        except Exception as e:
+            logger.warning(f"Failed to kill actor: {e}")
+
+
+def start_ray_actor(
+    worker_entry_fn,
+    placement_group,
+    placement_group_bundle_index: int,
+    *args,
+    **kwargs,
+):
+    if not RAY_AVAILABLE:
+        raise ImportError("ray is required for starting ray actor")
+
+    @ray.remote(num_gpus=1)
+    class OmniStageRayWorker:
+        def run(self, func, *args, **kwargs):
+            return func(*args, **kwargs)
+
+    worker_actor = OmniStageRayWorker.options(
+        scheduling_strategy=PlacementGroupSchedulingStrategy(
+            placement_group=placement_group, placement_group_bundle_index=placement_group_bundle_index
+        ),
+        runtime_env={"env_vars": {"PYTHONPATH": os.environ.get("PYTHONPATH", "")}, "CUDA_LAUNCH_BLOCKING": "1"},
+    ).remote()
+
+    worker_actor.run.remote(worker_entry_fn, *args, **kwargs)
+
+    return worker_actor

vllm_omni/entrypoints/async_omni.py

@@ -33,6 +33,16 @@ from vllm.v1.metrics.loggers import StatLoggerFactory, StatLoggerManager
 
 # Internal imports (our code)
 from vllm_omni.config import OmniModelConfig
+from vllm_omni.distributed.omni_connectors import (
+    get_stage_connector_config,
+    initialize_orchestrator_connectors,
+)
+from vllm_omni.distributed.omni_connectors.adapter import try_send_via_connector
+from vllm_omni.distributed.ray_utils.utils import (
+    create_placement_group,
+    get_ray_queue_class,
+    try_close_ray,
+)
 from vllm_omni.engine.arg_utils import AsyncOmniEngineArgs
 from vllm_omni.engine.output_processor import MultimodalOutputProcessor
 from vllm_omni.engine.processor import OmniProcessor
@@ -43,10 +53,12 @@ from vllm_omni.entrypoints.log_utils import (
     remove_old_logs,
 )
 from vllm_omni.entrypoints.omni_stage import OmniStage
-from vllm_omni.entrypoints.stage_utils import encode_for_ipc as _encode
 from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load
-from vllm_omni.entrypoints.stage_utils import serialize_obj as _set
-from vllm_omni.entrypoints.utils import load_stage_configs_from_model, load_stage_configs_from_yaml
+from vllm_omni.entrypoints.utils import (
+    load_stage_configs_from_model,
+    load_stage_configs_from_yaml,
+    resolve_model_config_path,
+)
 from vllm_omni.outputs import OmniRequestOutput
 
 logger = init_logger(__name__)
@@ -93,16 +105,29 @@ class AsyncOmni(EngineClient):
         cli_args: Namespace,
         **kwargs: Any,
     ):
+        self.worker_backend = cli_args.worker_backend
+        self.ray_address = cli_args.ray_address
+        self._ray_pg = None
+
         self.batch_timeout = cli_args.batch_timeout
         self._enable_stats: bool = bool(cli_args.log_stats)
 
         base_engine_args = AsyncOmniEngineArgs.from_cli_args(cli_args).__dict__.copy()
 
         if cli_args.stage_configs_path is None:
+            self.config_path = resolve_model_config_path(model)
             self.stage_configs = load_stage_configs_from_model(model, base_engine_args)
         else:
+            self.config_path = cli_args.stage_configs_path
             self.stage_configs = load_stage_configs_from_yaml(cli_args.stage_configs_path, base_engine_args)
 
+        shm_threshold_bytes = cli_args.shm_threshold_bytes
+
+        # Initialize connectors
+        self.omni_transfer_config, self.connectors = initialize_orchestrator_connectors(
+            self.config_path, worker_backend=self.worker_backend, shm_threshold_bytes=shm_threshold_bytes
+        )
+
         self.stage_list: list[OmniStage] = []
         self.default_sampling_params_list: list[SamplingParams] = []
         # Optional file handler for orchestrator
@@ -138,34 +163,56 @@ class AsyncOmni(EngineClient):
         self.default_sampling_params_list = [st.default_sampling_params for st in self.stage_list]
         logger.debug("[Orchestrator] Loaded %d stages", len(self.stage_list))
 
-        self._ctx = mp.get_context("spawn")
+        if self.worker_backend == "ray":
+            self._queue_cls = get_ray_queue_class()
+        else:
+            self._ctx = mp.get_context("spawn")
+            self._queue_cls = lambda: self._ctx.Queue(maxsize=0)
+
         self._stage_in_queues: list[mp.Queue] = []
         self._stage_out_queues: list[mp.Queue] = []
         self._init_sleep_seconds = max(0, int(init_sleep_seconds))
         self._shm_threshold_bytes = max(0, int(shm_threshold_bytes))
-        self._start_stage_processes(model)
+        self._start_stages(model)
         # Wait for all stages to report readiness before seeding
         self._stages_ready: set[int] = set()
         self._wait_for_stages_ready(timeout=init_timeout)
 
-    def _start_stage_processes(self, model: str) -> None:
+    def _start_stages(self, model: str) -> None:
+        if self.worker_backend == "ray":
+            # Initialize Ray Cluster
+            self._ray_pg = create_placement_group(
+                number_of_stages=len(self.stage_list), address=self.ray_address, strategy="PACK"
+            )
+
         for stage_id, stage in enumerate(self.stage_list):
             # Use unbounded queues to avoid deadlock when seeding many requests
-            in_q: mp.Queue = self._ctx.Queue(maxsize=0)
-            out_q: mp.Queue = self._ctx.Queue(maxsize=0)
+            in_q = self._queue_cls()
+            out_q = self._queue_cls()
             self._stage_in_queues.append(in_q)
             self._stage_out_queues.append(out_q)
 
             # Attach queues and start Stage-owned worker process
             stage.attach_queues(in_q, out_q)
+
+            # Build connectors config for this stage
+            stage_connectors_config = get_stage_connector_config(
+                self.omni_transfer_config,
+                stage_id,
+            )
+
             stage.init_stage_worker(
                 model,
                 is_async=True,
                 log_file=self._log_file,
                 shm_threshold_bytes=self._shm_threshold_bytes,
-                ctx=self._ctx,
+                ctx=self._ctx if self.worker_backend != "ray" else None,
                 batch_timeout=self.batch_timeout,
+                connectors_config=stage_connectors_config,
+                worker_backend=self.worker_backend,
+                ray_placement_group=self._ray_pg,
             )
+
             logger.debug("[Orchestrator] Stage-%s process started", stage_id)
             time.sleep(self._init_sleep_seconds)
 
@@ -191,6 +238,8 @@ class AsyncOmni(EngineClient):
             except Exception as e:
                 logger.warning("[Orchestrator] Failed to stop stage worker: %s", e)
 
+        try_close_ray(self._ray_pg)
+
     def __del__(self) -> None:  # best-effort
         try:
             self.close()
@@ -360,8 +409,8 @@ class AsyncOmni(EngineClient):
                     # End-to-end timing and time-per-token for final output
                     # (only once per request at the designated final stage)
                     try:
-                        rid_int = int(req_id) if isinstance(req_id, (int, str)) and str(req_id).isdigit() else req_id
-                        if stage_id == final_stage_id_for_e2e and rid_int not in metrics.e2e_done:
+                        rid_key = str(req_id)
+                        if stage_id == final_stage_id_for_e2e and rid_key not in metrics.e2e_done:
                             metrics.on_finalize_request(
                                 stage_id,
                                 req_id,
@@ -390,52 +439,36 @@ class AsyncOmni(EngineClient):
                     next_stage: OmniStage = self.stage_list[next_stage_id]
                     next_inputs = next_stage.process_engine_inputs(self.stage_list, prompt)
                     sp_next: SamplingParams = sampling_params_list[next_stage_id]
-                    try:
-                        # Measure transfer size and time (encode + enqueue)
-                        size_bytes = 0
-                        try:
-                            size_bytes = len(_set(next_inputs))
-                        except Exception:
-                            size_bytes = 0
-                        t0 = time.time()
-                        ipc_payload = _encode(
-                            next_inputs,
-                            getattr(self, "_shm_threshold_bytes", 65536),
-                            obj_key="engine_inputs",
-                            shm_key="engine_inputs_shm",
+
+                    # Check if we have a connector for this edge
+                    connector_key = (str(stage_id), str(next_stage_id))
+                    connector = self.connectors.get(connector_key)
+
+                    sent_via_connector = False
+                    if connector:
+                        sent_via_connector = try_send_via_connector(
+                            connector=connector,
+                            stage_id=stage_id,
+                            next_stage_id=next_stage_id,
+                            req_id=req_id,
+                            next_inputs=next_inputs,
+                            sampling_params=sp_next,
+                            original_prompt=prompt,
+                            next_stage_queue_submit_fn=self.stage_list[next_stage_id].submit,
+                            metrics=metrics,
                         )
-                        ipc_payload.update(
-                            {
-                                "request_id": req_id,
-                                "sampling_params": sp_next,
-                                "sent_ts": time.time(),
-                            }
-                        )
-                        self.stage_list[next_stage_id].submit(ipc_payload)
-                        t1 = time.time()
-                        tx_ms = (t1 - t0) * 1000.0
-                        metrics.on_forward(
-                            stage_id,
-                            next_stage_id,
-                            req_id,
-                            int(size_bytes),
-                            float(tx_ms),
-                            bool("engine_inputs_shm" in ipc_payload),
-                        )
-                    except Exception as e:
-                        logger.warning(
-                            "[Orchestrator] IPC encode failed for req %s: %s; \
-                                falling back to inline payload",
-                            req_id,
-                            e,
-                        )
-                        self.stage_list[next_stage_id].submit(
-                            {
-                                "request_id": req_id,
-                                "engine_inputs": next_inputs,
-                                "sampling_params": sp_next,
-                            }
+
+                    if not sent_via_connector:
+                        # Fallback logic removed as we now enforce connector usage.
+                        # If no connector is found or send fails, we log an error and raise,
+                        # because continuing would cause the request to be silently dropped
+                        # and the orchestrator to hang waiting for completion.
+                        error_msg = (
+                            f"[Orchestrator] Failed to send request {req_id} to stage-{next_stage_id} via connector. "
+                            "Configure a connector for this edge or inspect connector logs for details."
                         )
+                        logger.error(error_msg)
+                        raise RuntimeError(error_msg)
                     logger.debug(
                         "[Orchestrator] Forwarded request %s to stage-%s",
                         req_id,

vllm_omni/entrypoints/cli/serve.py

@@ -95,6 +95,19 @@ class OmniServeCommand(CLISubcommand):
             default=10,
             help="The timeout for the batch.",
         )
+        serve_parser.add_argument(
+            "--worker-backend",
+            type=str,
+            default="multi_process",
+            choices=["multi_process", "ray"],
+            help="The backend to use for stage workers.",
+        )
+        serve_parser.add_argument(
+            "--ray-address",
+            type=str,
+            default=None,
+            help="The address of the Ray cluster to connect to.",
+        )
         return serve_parser
 
 

vllm_omni/entrypoints/log_utils.py

@@ -301,7 +301,7 @@ def compute_and_log_stage_request_stats(
 
 
 def record_stage_metrics(
-    per_request: dict[int, dict[str, Any]],
+    per_request: dict[str, dict[str, Any]],
     stage_req_counts: list[int],
     stage_total_time_ms: list[float],
     stage_total_tokens: list[int],
@@ -312,8 +312,8 @@ def record_stage_metrics(
     try:
         stage_req_counts[stage_id] += 1
         stage_total_tokens[stage_id] += int(metrics.get("num_tokens_out", 0))
-        rid_int = int(req_id)
-        pr = per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
+        rid_key = str(req_id)
+        pr = per_request.setdefault(rid_key, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
         pr_stages = pr["stages"]  # type: ignore[index]
         pr_stages[stage_id] = {
             "stage_gen_time_ms": float(metrics.get("stage_gen_time_ms", 0.0)),
@@ -324,9 +324,9 @@ def record_stage_metrics(
 
 
 def aggregate_rx_and_maybe_total(
-    transfer_edge_req: dict[tuple[int, int, int], dict[str, float]],
+    transfer_edge_req: dict[tuple[int, int, str], dict[str, float]],
     transfer_agg: dict[tuple[int, int], dict[str, float]],
-    per_request: dict[int, dict[str, Any]],
+    per_request: dict[str, dict[str, Any]],
     stage_id: int,
     req_id: Any,
     rx_bytes: float,
@@ -355,11 +355,8 @@ def aggregate_rx_and_maybe_total(
             agg["rx_count"] += 1.0
 
             # Try combine with sender-side timing if present
-            try:
-                rid_int = int(req_id)
-            except Exception:
-                rid_int = req_id  # best effort
-            s = transfer_edge_req.get((stage_id - 1, stage_id, rid_int))
+            rid_key = str(req_id)
+            s = transfer_edge_req.get((stage_id - 1, stage_id, rid_key))
             if s is None:
                 return None
             tx_ms = float(s.get("tx_ms", 0.0))
@@ -369,7 +366,7 @@ def aggregate_rx_and_maybe_total(
             agg["total_count"] += 1.0
             # accumulate per-request transfer totals
             try:
-                pr = per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
+                pr = per_request.setdefault(rid_key, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
                 pr["transfers_ms"] = float(pr.get("transfers_ms", 0.0)) + total_ms  # type: ignore[index]
                 pr["transfers_bytes"] = int(pr.get("transfers_bytes", 0)) + int(rx_bytes)  # type: ignore[index]
             except Exception:
@@ -382,7 +379,7 @@ def aggregate_rx_and_maybe_total(
 
 def record_sender_transfer_agg(
     transfer_agg: dict[tuple[int, int], dict[str, float]],
-    transfer_edge_req: dict[tuple[int, int, int], dict[str, float]],
+    transfer_edge_req: dict[tuple[int, int, str], dict[str, float]],
     from_stage: int,
     to_stage: int,
     req_id: Any,
@@ -408,11 +405,8 @@ def record_sender_transfer_agg(
         agg["sum_ms"] += float(tx_ms)
         agg["count"] += 1.0
         # Store sender-side timing for per-request combination
-        try:
-            rid_int = int(req_id)
-        except Exception:
-            rid_int = req_id
-        transfer_edge_req[(from_stage, to_stage, rid_int)] = {
+        rid_key = str(req_id)
+        transfer_edge_req[(from_stage, to_stage, rid_key)] = {
             "tx_ms": float(tx_ms),
             "size_bytes": float(size_bytes),
         }
@@ -512,12 +506,12 @@ class OrchestratorMetrics:
         self.stage_total_tokens: list[int] = [0 for _ in range(self.num_stages)]
         self.stage_req_counts: list[int] = [0 for _ in range(self.num_stages)]
         self.transfer_agg: dict[tuple[int, int], dict[str, float]] = {}
-        self.transfer_edge_req: dict[tuple[int, int, int], dict[str, float]] = {}
+        self.transfer_edge_req: dict[tuple[int, int, str], dict[str, float]] = {}
         self.e2e_total_ms: float = 0.0
         self.e2e_total_tokens: int = 0
         self.e2e_count: int = 0
-        self.e2e_done: set[int] = set()
-        self.per_request: dict[int, dict[str, Any]] = {}
+        self.e2e_done: set[str] = set()
+        self.per_request: dict[str, dict[str, Any]] = {}
         self.sum_per_request_transfer_ms: float = 0.0
         self.wall_start_ts: float = float(wall_start_ts)
         self.last_finish_ts: float = float(wall_start_ts)
@@ -614,10 +608,7 @@ class OrchestratorMetrics:
         )
 
     def on_finalize_request(self, stage_id: int, req_id: Any, engine_outputs: list[Any], req_start_ts: float) -> None:
-        try:
-            rid_int = int(req_id)
-        except Exception:
-            rid_int = req_id
+        rid_key = str(req_id)
         _t0 = float(req_start_ts)
         _t1 = time.time()
         # Update last output time for this stage
@@ -629,11 +620,11 @@ class OrchestratorMetrics:
         self.e2e_total_ms += e2e_ms
         self.e2e_total_tokens += int(num_tokens)
         self.e2e_count += 1
-        self.e2e_done.add(rid_int)
-        pr = self.per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
+        self.e2e_done.add(rid_key)
+        pr = self.per_request.setdefault(rid_key, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
         per_req_record = {
             "type": "overall_request",
-            "request_id": rid_int,
+            "request_id": rid_key,
             "e2e_time_ms": e2e_ms,
             "num_tokens_out": int(num_tokens),
             "transfers_total_time_ms": float(pr.get("transfers_ms", 0.0)),

vllm_omni/entrypoints/omni_llm.py

@@ -1,6 +1,7 @@
 import multiprocessing as mp
 import os
 import time
+import uuid
 from collections.abc import Sequence
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from typing import Any, Optional, Union
@@ -19,7 +20,18 @@ from vllm.usage.usage_lib import UsageContext
 from vllm.utils import Counter
 from vllm.v1.engine.llm_engine import LLMEngine
 
+from vllm_omni.distributed.omni_connectors import (
+    get_stage_connector_config,
+    initialize_orchestrator_connectors,
+)
+
 # Internal imports (our code)
+from vllm_omni.distributed.omni_connectors.adapter import try_send_via_connector
+from vllm_omni.distributed.ray_utils.utils import (
+    create_placement_group,
+    get_ray_queue_class,
+    try_close_ray,
+)
 from vllm_omni.engine.arg_utils import OmniEngineArgs
 from vllm_omni.engine.output_processor import MultimodalOutputProcessor
 from vllm_omni.engine.processor import OmniProcessor
@@ -30,10 +42,12 @@ from vllm_omni.entrypoints.log_utils import (
     remove_old_logs,
 )
 from vllm_omni.entrypoints.omni_stage import OmniStage
-from vllm_omni.entrypoints.stage_utils import encode_for_ipc as _encode
 from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load
-from vllm_omni.entrypoints.stage_utils import serialize_obj as _set
-from vllm_omni.entrypoints.utils import load_stage_configs_from_model, load_stage_configs_from_yaml
+from vllm_omni.entrypoints.utils import (
+    load_stage_configs_from_model,
+    load_stage_configs_from_yaml,
+    resolve_model_config_path,
+)
 from vllm_omni.outputs import OmniRequestOutput
 
 logger = init_logger(__name__)
@@ -84,14 +98,25 @@ class OmniLLM:
         init_timeout: int = 300,
         **kwargs: Any,
     ):
+        self.worker_backend = kwargs.get("worker_backend", "multi_process")
+        self.ray_address = kwargs.get("ray_address", None)
+        self._ray_pg = None
         self.batch_timeout = batch_timeout
         self._enable_stats: bool = bool(log_stats)
+
         # Do NOT call super().__init__ to avoid creating OmniStageLLM instances in parent.
         if stage_configs_path is None:
+            self.config_path = resolve_model_config_path(model)
             self.stage_configs = load_stage_configs_from_model(model)
         else:
+            self.config_path = stage_configs_path
             self.stage_configs = load_stage_configs_from_yaml(stage_configs_path)
 
+        # Initialize connectors
+        self.omni_transfer_config, self.connectors = initialize_orchestrator_connectors(
+            self.config_path, worker_backend=self.worker_backend, shm_threshold_bytes=shm_threshold_bytes
+        )
+
         # Optional file handler for orchestrator
         self._log_file = log_file
 
@@ -124,33 +149,51 @@ class OmniLLM:
         self.stage_list = [st for _, st in results]
         logger.debug("[Orchestrator] Loaded %d stages", len(self.stage_list))
 
-        self._ctx = mp.get_context("spawn")
+        if self.worker_backend == "ray":
+            self._queue_cls = get_ray_queue_class()
+        else:
+            self._ctx = mp.get_context("spawn")
+            self._queue_cls = lambda: self._ctx.Queue(maxsize=0)
+
         self._stage_in_queues: list[mp.Queue] = []
         self._stage_out_queues: list[mp.Queue] = []
         self._init_sleep_seconds = max(0, int(init_sleep_seconds))
         self._shm_threshold_bytes = max(0, int(shm_threshold_bytes))
-        self._start_stage_processes(model)
+        self._start_stages(model)
         # Wait for all stages to report readiness before seeding
         self._stages_ready: set[int] = set()
         self._wait_for_stages_ready(timeout=init_timeout)
 
-    def _start_stage_processes(self, model: str) -> None:
+    def _start_stages(self, model: str) -> None:
+        if self.worker_backend == "ray":
+            # Initialize Ray Cluster
+            self._ray_pg = create_placement_group(
+                number_of_stages=len(self.stage_list), address=self.ray_address, strategy="PACK"
+            )
+
         for stage_id, stage in enumerate(self.stage_list):
-            # Use unbounded queues to avoid deadlock when seeding many requests
-            in_q: mp.Queue = self._ctx.Queue(maxsize=0)
-            out_q: mp.Queue = self._ctx.Queue(maxsize=0)
+            in_q = self._queue_cls()
+            out_q = self._queue_cls()
             self._stage_in_queues.append(in_q)
             self._stage_out_queues.append(out_q)
-
-            # Attach queues and start Stage-owned worker process
             stage.attach_queues(in_q, out_q)
+
+            stage_connectors_config = get_stage_connector_config(
+                self.omni_transfer_config,
+                stage_id,
+            )
+
             stage.init_stage_worker(
                 model,
                 log_file=self._log_file,
                 shm_threshold_bytes=self._shm_threshold_bytes,
-                ctx=self._ctx,
+                ctx=self._ctx if self.worker_backend != "ray" else None,
                 batch_timeout=self.batch_timeout,
+                connectors_config=stage_connectors_config,
+                worker_backend=self.worker_backend,
+                ray_placement_group=self._ray_pg,
             )
+
             logger.debug("[Orchestrator] Stage-%s process started", stage_id)
             time.sleep(self._init_sleep_seconds)
 
@@ -175,6 +218,8 @@ class OmniLLM:
             except Exception as e:
                 logger.warning("[Orchestrator] Failed to stop stage worker: %s", e)
 
+        try_close_ray(self._ray_pg)
+
     def __del__(self) -> None:  # best-effort
         try:
             self.close()
@@ -237,11 +282,12 @@ class OmniLLM:
         # Orchestrator keeps stage objects for input derivation
         num_stages = len(self.stage_list)
 
-        # Map from request_id to original prompt
-        request_id_to_prompt: dict[int, PromptType] = {i: p for i, p in enumerate(request_prompts)}
+        # Generate globally unique request IDs and map them to original prompts
+        request_ids: list[str] = [f"{i}_{uuid.uuid4()}" for i in range(len(request_prompts))]
+        request_id_to_prompt: dict[str, PromptType] = {rid: p for rid, p in zip(request_ids, request_prompts)}
 
         # Track per-request start time for end-to-end timing
-        _req_start_ts: dict[int, float] = {}
+        _req_start_ts: dict[str, float] = {}
         _wall_start_ts: float = time.time()
 
         # Determine the final stage for E2E stats (highest stage_id with final_output=True; fallback to last stage)
@@ -357,8 +403,8 @@ class OmniLLM:
                     # End-to-end timing and time-per-token for final output
                     # (only once per request at the designated final stage)
                     try:
-                        rid_int = int(req_id) if isinstance(req_id, (int, str)) and str(req_id).isdigit() else req_id
-                        if stage_id == final_stage_id_for_e2e and rid_int not in metrics.e2e_done:
+                        rid_key = str(req_id)
+                        if stage_id == final_stage_id_for_e2e and rid_key not in metrics.e2e_done:
                             metrics.on_finalize_request(
                                 stage_id,
                                 req_id,
@@ -387,50 +433,28 @@ class OmniLLM:
                         )
                         continue
                     sp_next: SamplingParams = sampling_params_list[next_stage_id]  # type: ignore[index]
-                    try:
-                        # Measure transfer size and time (encode + enqueue)
-                        size_bytes = 0
-                        try:
-                            size_bytes = len(_set(next_inputs))
-                        except Exception:
-                            size_bytes = 0
-                        t0 = time.time()
-                        ipc_payload = _encode(
-                            next_inputs,
-                            getattr(self, "_shm_threshold_bytes", 65536),
-                            obj_key="engine_inputs",
-                            shm_key="engine_inputs_shm",
+
+                    # Check if we have a connector for this edge
+                    connector_key = (str(stage_id), str(next_stage_id))
+                    connector = self.connectors.get(connector_key)
+                    sent_via_connector = False
+                    if connector:
+                        sent_via_connector = try_send_via_connector(
+                            connector=connector,
+                            stage_id=stage_id,
+                            next_stage_id=next_stage_id,
+                            req_id=req_id,
+                            next_inputs=next_inputs,
+                            sampling_params=sp_next,
+                            original_prompt=request_id_to_prompt[req_id],
+                            next_stage_queue_submit_fn=self.stage_list[next_stage_id].submit,
+                            metrics=metrics,
                         )
-                        ipc_payload.update(
-                            {
-                                "request_id": req_id,
-                                "sampling_params": sp_next,
-                                "sent_ts": time.time(),
-                            }
-                        )
-                        self.stage_list[next_stage_id].submit(ipc_payload)
-                        t1 = time.time()
-                        tx_ms = (t1 - t0) * 1000.0
-                        metrics.on_forward(
-                            stage_id,
-                            next_stage_id,
-                            req_id,
-                            int(size_bytes),
-                            float(tx_ms),
-                            bool("engine_inputs_shm" in ipc_payload),
-                        )
-                    except Exception as e:
-                        logger.warning(
-                            "[Orchestrator] IPC encode failed for req %s: %s; falling back to inline payload",
-                            req_id,
-                            e,
-                        )
-                        self.stage_list[next_stage_id].submit(
-                            {
-                                "request_id": req_id,
-                                "engine_inputs": next_inputs,
-                                "sampling_params": sp_next,
-                            }
+
+                    if not sent_via_connector:
+                        raise RuntimeError(
+                            f"[Orchestrator] Failed to send request {req_id} to stage-{next_stage_id} via connector. "
+                            "Configure a connector for this edge or inspect connector logs for details."
                         )
                     logger.debug(
                         "[Orchestrator] Forwarded request %s to stage-%s",

vllm_omni/entrypoints/omni_stage.py

@@ -15,6 +15,7 @@ import importlib
 import logging
 import multiprocessing as mp
 import os
+import sys
 from typing import Any, Optional, Union
 
 from vllm.inputs import TextPrompt
@@ -27,11 +28,11 @@ from vllm.v1.engine import EngineCoreOutput
 from vllm.v1.engine.async_llm import AsyncLLM
 from vllm.v1.engine.llm_engine import LLMEngine
 
+from vllm_omni.distributed.ray_utils.utils import kill_ray_actor, start_ray_actor
 from vllm_omni.engine.arg_utils import AsyncOmniEngineArgs
 from vllm_omni.entrypoints.stage_utils import (
     _to_dict,
     maybe_dump_to_shm,
-    maybe_load_from_ipc_with_metrics,
     set_stage_devices,
 )
 from vllm_omni.inputs.data import OmniTokensPrompt
@@ -163,6 +164,9 @@ class OmniStage:
         shm_threshold_bytes: int = 65536,
         ctx: Optional[mp.context.BaseContext] = None,
         batch_timeout: int = 10,
+        connectors_config: Optional[dict] = None,
+        worker_backend: str = "multi_process",
+        **kwargs: Any,
     ) -> None:
         """Initialize and start the stage worker process.
 
@@ -176,13 +180,23 @@ class OmniStage:
             shm_threshold_bytes: Threshold for using shared memory for IPC
             ctx: Optional multiprocessing context (default: spawn)
             batch_timeout: Timeout in seconds for batching requests
+            connectors_config: Configuration for stage connectors
+            worker_backend: Backend type ("multi_process" or "ray")
+            **kwargs: Additional arguments (e.g. ray_placement_group)
 
         Raises:
             AssertionError: If queues are not attached before calling this method
         """
         assert self._in_q is not None and self._out_q is not None, "Queues must be attached before start_process"
         self._log_file = log_file
-        self._shm_threshold_bytes = shm_threshold_bytes
+
+        if worker_backend == "ray":
+            ray_placement_group = kwargs.get("ray_placement_group", None)
+            assert ray_placement_group is not None, "Ray placement group must be provided"
+            self._shm_threshold_bytes = sys.maxsize
+        else:
+            self._shm_threshold_bytes = shm_threshold_bytes
+
         ctx = ctx or mp.get_context("spawn")
         # Prepare lightweight dict config for worker
         engine_args = _to_dict(self.engine_args)
@@ -192,43 +206,74 @@ class OmniStage:
             "engine_args": engine_args,
             "runtime": runtime_cfg,
             "shm_threshold_bytes": self._shm_threshold_bytes,
+            "connectors_config": connectors_config or {},
         }
-        if is_async:
-            self._proc = ctx.Process(
-                target=_stage_worker_async_entry,
-                args=(
+
+        if worker_backend == "ray":
+            if is_async:
+                self._ray_actor = start_ray_actor(
+                    _stage_worker_async_entry,
+                    ray_placement_group,
+                    self.stage_id,
                     self,
-                    model,
-                    stage_payload,
-                    batch_timeout,
-                ),
-            )
+                    model=model,
+                    stage_payload=stage_payload,
+                    batch_timeout=batch_timeout,
+                )
+            else:
+                self._ray_actor = start_ray_actor(
+                    _stage_worker,
+                    ray_placement_group,
+                    self.stage_id,
+                    model=model,
+                    stage_payload=stage_payload,
+                    in_q=self._in_q,
+                    out_q=self._out_q,
+                    log_file=self._log_file,
+                    batch_timeout=batch_timeout,
+                )
         else:
-            self._proc = ctx.Process(
-                target=_stage_worker,
-                args=(
-                    model,
-                    stage_payload,
-                    self._in_q,
-                    self._out_q,
-                    self._log_file,
-                    batch_timeout,
-                ),
-            )
-        self._proc.start()
+            if is_async:
+                self._proc = ctx.Process(
+                    target=_stage_worker_async_entry,
+                    args=(
+                        self,
+                        model,
+                        stage_payload,
+                        batch_timeout,
+                    ),
+                )
+            else:
+                self._proc = ctx.Process(
+                    target=_stage_worker,
+                    args=(
+                        model,
+                        stage_payload,
+                        self._in_q,
+                        self._out_q,
+                        self._log_file,
+                        batch_timeout,
+                    ),
+                )
+            self._proc.start()
 
     def stop_stage_worker(self) -> None:
         """Stop the stage worker process gracefully.
 
         Sends shutdown signal to the worker and waits for it to terminate.
         If graceful shutdown fails, forcefully terminates the process.
+        Handles both multiprocessing Process and Ray Actor.
         """
         if self._in_q is not None:
             try:
                 self._in_q.put_nowait(None)
             except Exception as e:
                 self._logger.warning("[Stage-%s] Failed to send shutdown to in_q: %s", self.stage_id, e)
-        if self._proc is not None:
+
+        if hasattr(self, "_ray_actor") and self._ray_actor:
+            kill_ray_actor(self._ray_actor)
+            self._ray_actor = None
+        elif self._proc is not None:
             try:
                 self._proc.join(timeout=5)
             except Exception as e:
@@ -238,6 +283,7 @@ class OmniStage:
                     self._proc.terminate()
                 except Exception as e:
                     self._logger.warning("[Stage-%s] terminate() failed: %s", self.stage_id, e)
+
         # Cleanup temporary stage log if we created one (only when no log_file provided)
         try:
             if not self._log_file:
@@ -334,6 +380,8 @@ def _stage_worker(
     import os as _os
     import time as _time
 
+    from vllm_omni.distributed.omni_connectors import build_stage_connectors
+    from vllm_omni.distributed.omni_connectors.adapter import try_recv_via_connector
     from vllm_omni.entrypoints.log_utils import (
         compute_and_log_stage_request_stats,
         count_tokens_from_outputs,
@@ -347,6 +395,7 @@ def _stage_worker(
     engine_args = stage_payload.get("engine_args", {})
     runtime_cfg = stage_payload.get("runtime", {})
     shm_threshold_bytes = int(stage_payload.get("shm_threshold_bytes", 65536))
+    connectors_config = stage_payload.get("connectors_config", {})
 
     # Per-stage logger: clear inherited handlers to avoid broken parent streams
     try:
@@ -440,6 +489,18 @@ def _stage_worker(
     )
     stage_engine = OmniStageLLM(model=model, **engine_args)
     _logging.getLogger(__name__).debug("[Stage-%s] Engine initialized", stage_id)
+
+    # Initialize OmniConnectors if configured
+    connectors = {}
+    if connectors_config:
+        built_connectors = build_stage_connectors(
+            stage_id=stage_id,
+            connectors_config=connectors_config,
+        )
+        if built_connectors is None:
+            return
+        connectors = built_connectors
+
     # Signal readiness to orchestrator
     try:
         out_q.put({"type": "stage_ready", "stage_id": stage_id})
@@ -496,9 +557,25 @@ def _stage_worker(
                     _in_flight_ms_by_rid[rid] = 0.0
             except Exception:
                 _in_flight_ms_by_rid[rid] = 0.0
-            ein, _rx_metrics = maybe_load_from_ipc_with_metrics(t, obj_key="engine_inputs", shm_key="engine_inputs_shm")
-            _rx_decode_ms_by_rid[rid] = float(_rx_metrics.get("rx_decode_time_ms", 0.0))
-            _rx_bytes_by_rid[rid] = int(_rx_metrics.get("rx_transfer_bytes", 0))
+
+            # Resolve input data strictly via connectors if payload
+            # is larger than shm_threshold_bytes or using other connectors
+            ein, _rx_metrics = try_recv_via_connector(
+                task=t,
+                connectors=connectors,
+                stage_id=stage_id,
+            )
+
+            if ein is None or _rx_metrics is None:
+                raise RuntimeError(
+                    f"[Stage-{stage_id}] Missing connector payload for request {rid}. "
+                    "Ensure connectors are configured for all incoming edges."
+                )
+
+            if _rx_metrics:
+                _rx_decode_ms_by_rid[rid] = float(_rx_metrics.get("rx_decode_time_ms", 0.0))
+                _rx_bytes_by_rid[rid] = int(_rx_metrics.get("rx_transfer_bytes", 0))
+
             batch_request_ids.append(rid)
             if isinstance(ein, list):
                 batch_engine_inputs.extend(ein)
@@ -671,6 +748,8 @@ async def _stage_worker_async(
     import logging as _logging
     import time as _time
 
+    from vllm_omni.distributed.omni_connectors import build_stage_connectors
+    from vllm_omni.distributed.omni_connectors.adapter import try_recv_via_connector
     from vllm_omni.entrypoints.async_omni import AsyncOmniStageLLM
     from vllm_omni.entrypoints.log_utils import (
         compute_and_log_stage_request_stats,
@@ -685,6 +764,7 @@ async def _stage_worker_async(
     engine_args = stage_payload.get("engine_args", {})
     runtime_cfg = stage_payload.get("runtime", {})
     shm_threshold_bytes = int(stage_payload.get("shm_threshold_bytes", 65536))
+    connectors_config = stage_payload.get("connectors_config", {})
 
     log_file = omni_stage._log_file
     in_q = omni_stage._in_q
@@ -729,6 +809,17 @@ async def _stage_worker_async(
     except Exception as e:
         _logging.getLogger(__name__).warning("[Stage-%s] Device setup failed: %s", stage_id, e)
 
+    # Initialize OmniConnectors if configured to match sync worker behavior
+    connectors: dict[Any, Any] = {}
+    if connectors_config:
+        built_connectors = build_stage_connectors(
+            stage_id=stage_id,
+            connectors_config=connectors_config,
+        )
+        if built_connectors is None:
+            return
+        connectors = built_connectors
+
     # Init LLM
     _logging.getLogger(__name__).debug(
         "[Stage-%s] Initializing engine with args keys=%s",
@@ -788,7 +879,16 @@ async def _stage_worker_async(
                 _in_flight_ms_by_rid[rid] = 0.0
         except Exception:
             _in_flight_ms_by_rid[rid] = 0.0
-        ein, _rx_metrics = maybe_load_from_ipc_with_metrics(task, obj_key="engine_inputs", shm_key="engine_inputs_shm")
+        ein, _rx_metrics = try_recv_via_connector(
+            task=task,
+            connectors=connectors,
+            stage_id=stage_id,
+        )
+        if ein is None or _rx_metrics is None:
+            raise RuntimeError(
+                f"[Stage-{stage_id}] Missing connector payload for request {rid}. "
+                "Ensure connectors are configured for all incoming edges."
+            )
         _rx_decode_ms_by_rid[rid] = float(_rx_metrics.get("rx_decode_time_ms", 0.0))
         _rx_bytes_by_rid[rid] = int(_rx_metrics.get("rx_transfer_bytes", 0))
 

vllm_omni/entrypoints/stage_utils.py

@@ -3,11 +3,9 @@ from __future__ import annotations
 import json
 import logging
 import os
-import pickle
 from multiprocessing import shared_memory as _shm
 from typing import Any
 
-import cloudpickle
 from omegaconf import OmegaConf
 
 logger = logging.getLogger(__name__)
@@ -177,8 +175,10 @@ def set_stage_devices(
 
 
 def serialize_obj(obj: Any) -> bytes:
-    """Serialize a Python object to bytes using cloudpickle."""
-    return cloudpickle.dumps(obj)
+    """Serialize a Python object to bytes using centralized serializer (defaults to cloudpickle)."""
+    from vllm_omni.distributed.omni_connectors.utils.serialization import OmniSerializer
+
+    return OmniSerializer.serialize(obj)
 
 
 def shm_write_bytes(payload: bytes) -> dict[str, Any]:
@@ -260,7 +260,9 @@ def maybe_load_from_ipc(container: dict[str, Any], obj_key: str, shm_key: str) -
     decode-time and size metrics.
     """
     if shm_key in container:
-        return pickle.loads(shm_read_bytes(container[shm_key]))
+        from vllm_omni.distributed.omni_connectors.utils.serialization import OmniSerializer
+
+        return OmniSerializer.deserialize(shm_read_bytes(container[shm_key]))
     return container[obj_key]
 
 
@@ -275,11 +277,13 @@ def maybe_load_from_ipc_with_metrics(
     """
     import time as _time  # local import to avoid overhead at module import
 
+    from vllm_omni.distributed.omni_connectors.utils.serialization import OmniSerializer
+
     t0 = _time.time()
     if shm_key in container:
         meta = container[shm_key]  # type: ignore[index]
         payload = shm_read_bytes(meta)
-        obj = pickle.loads(payload)
+        obj = OmniSerializer.deserialize(payload)
         try:
             rx_bytes = int(meta.get("size", len(payload)))  # type: ignore[call-arg]
         except Exception:

vllm_omni/entrypoints/utils.py

@@ -2,7 +2,7 @@ import os
 from collections import Counter
 from dataclasses import asdict, is_dataclass
 from pathlib import Path
-from typing import Any
+from typing import Any, Optional
 
 from omegaconf import OmegaConf
 from vllm.transformers_utils.config import get_config
@@ -62,7 +62,42 @@ def _convert_dataclasses_to_dict(obj: Any) -> Any:
     return obj
 
 
-def load_stage_configs_from_model(model: str, base_engine_args: dict = {}) -> list:
+def resolve_model_config_path(model: str) -> str:
+    """Resolve the stage config file path from the model name.
+
+    Resolves stage configuration path based on the model type and device type.
+    First tries to find a device-specific YAML file from stage_configs/{device_type}/
+    directory. If not found, falls back to the default config file.
+
+    Args:
+        model: Model name or path (used to determine model_type)
+
+    Returns:
+        String path to the stage configuration file
+
+    Raises:
+        FileNotFoundError: If no stage config file exists for the model type
+    """
+    hf_config = get_config(model, trust_remote_code=True)
+    model_type = hf_config.model_type
+    device_type = detect_device_type()
+
+    # Try device-specific config first
+    if device_type != "cuda":
+        device_config_file = f"vllm_omni/model_executor/stage_configs/{device_type}/{model_type}.yaml"
+        device_config_path = PROJECT_ROOT / device_config_file
+        if os.path.exists(device_config_path):
+            return str(device_config_path)
+
+    # Fall back to default config
+    stage_config_file = f"vllm_omni/model_executor/stage_configs/{model_type}.yaml"
+    stage_config_path = PROJECT_ROOT / stage_config_file
+    if not os.path.exists(stage_config_path):
+        raise FileNotFoundError(f"Stage config file {stage_config_path} not found")
+    return str(stage_config_path)
+
+
+def load_stage_configs_from_model(model: str, base_engine_args: Optional[dict] = None) -> list:
     """Load stage configurations from model's default config file.
 
     Loads stage configurations based on the model type and device type.
@@ -78,30 +113,14 @@ def load_stage_configs_from_model(model: str, base_engine_args: dict = {}) -> li
     Raises:
         FileNotFoundError: If no stage config file exists for the model type
     """
-    hf_config = get_config(model, trust_remote_code=True)
-    model_type = hf_config.model_type
-    device_type = detect_device_type()
-
-    # Try device-specific config first
-    if device_type != "cuda":
-        device_config_file = f"vllm_omni/model_executor/stage_configs/{device_type}/{model_type}.yaml"
-        device_config_path = PROJECT_ROOT / device_config_file
-        if os.path.exists(device_config_path):
-            stage_configs = load_stage_configs_from_yaml(
-                config_path=str(device_config_path), base_engine_args=base_engine_args
-            )
-            return stage_configs
-
-    # Fall back to default config
-    stage_config_file = f"vllm_omni/model_executor/stage_configs/{model_type}.yaml"
-    stage_config_path = PROJECT_ROOT / stage_config_file
-    if not os.path.exists(stage_config_path):
-        raise FileNotFoundError(f"Stage config file {stage_config_path} not found")
-    stage_configs = load_stage_configs_from_yaml(config_path=str(stage_config_path), base_engine_args=base_engine_args)
+    if base_engine_args is None:
+        base_engine_args = {}
+    stage_config_path = resolve_model_config_path(model)
+    stage_configs = load_stage_configs_from_yaml(config_path=stage_config_path, base_engine_args=base_engine_args)
     return stage_configs
 
 
-def load_stage_configs_from_yaml(config_path: str, base_engine_args: dict = {}) -> list:
+def load_stage_configs_from_yaml(config_path: str, base_engine_args: Optional[dict] = None) -> list:
     """Load stage configurations from a YAML file.
 
     Args:
@@ -110,6 +129,8 @@ def load_stage_configs_from_yaml(config_path: str, base_engine_args: dict = {})
     Returns:
         List of stage configuration dictionaries from the file's stage_args
     """
+    if base_engine_args is None:
+        base_engine_args = {}
     config_data = OmegaConf.load(config_path)
     stage_args = config_data.stage_args
     # Convert any nested dataclass objects to dicts before creating OmegaConf

vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml

@@ -29,6 +29,7 @@ stage_args:
       seed: 42
       detokenize: True
       repetition_penalty: 1.1
+
   - stage_id: 1
     runtime:
       process: true
@@ -56,10 +57,11 @@ stage_args:
       detokenize: True
       repetition_penalty: 1.05
       stop_token_ids: [8294]
+
   - stage_id: 2
     runtime:
       process: true
-      devices: "0"            # Example: use a different GPU than the previous stage; use "0" if single GPU
+      devices: "2"            # Example: use a different GPU than the previous stage; use "0" if single GPU
       max_batch_size: 1
     engine_args:
       model_stage: code2wav
@@ -90,6 +92,7 @@ runtime:
   defaults:
     window_size: -1             # Simplified: trigger downstream only after full upstream completion
     max_inflight: 1             # Simplified: process serially within each stage
+
   edges:
     - from: 0                   # thinker → talker: trigger only after receiving full input (-1)
       to: 1

vllm_omni/model_executor/stage_configs/qwen2_5_omni_multiconnector.yaml

@@ -0,0 +1,130 @@
+# stage config for running qwen2.5-omni with architecture of OmniLLM.
+
+# The following config has been verified on 1x H100-80G GPU.
+stage_args:
+  - stage_id: 0
+    runtime:
+      process: true            # Run this stage in a separate process
+      devices: "0"            # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
+      max_batch_size: 1
+    engine_args:
+      model_stage: thinker
+      model_arch: Qwen2_5OmniForConditionalGeneration
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.8
+      enforce_eager: true  # Now we only support eager mode
+      trust_remote_code: true
+      engine_output_type: latent
+      enable_prefix_caching: false
+    is_comprehension: true
+    final_output: true
+    final_output_type: text
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.1
+    # Distributed connector configuration (optional)
+    output_connectors:
+      to_stage_1: mooncake_connector
+  - stage_id: 1
+    runtime:
+      process: true
+      devices: "1"
+      max_batch_size: 1
+    engine_args:
+      model_stage: talker
+      model_arch: Qwen2_5OmniForConditionalGeneration
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.8
+      enforce_eager: true
+      trust_remote_code: true
+      enable_prefix_caching: false
+      engine_output_type: latent
+    engine_input_source: [0]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen2_5_omni.thinker2talker
+    default_sampling_params:
+      temperature: 0.9
+      top_p: 0.8
+      top_k: 40
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.05
+      stop_token_ids: [8294]
+    # Distributed connector configuration (optional)
+    input_connectors:
+      from_stage_0: mooncake_connector
+    output_connectors:
+      to_stage_2: mooncake_connector
+  - stage_id: 2
+    runtime:
+      process: true
+      devices: "2"            # Example: use a different GPU than the previous stage; use "0" if single GPU
+      max_batch_size: 1
+    engine_args:
+      model_stage: code2wav
+      model_arch: Qwen2_5OmniForConditionalGeneration
+      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      gpu_memory_utilization: 0.3
+      enforce_eager: true
+      trust_remote_code: true
+      enable_prefix_caching: false
+      max_num_batched_tokens: 32768
+      engine_output_type: audio
+    engine_input_source: [1]
+    final_output: true
+    final_output_type: audio
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.1
+    # Distributed connector configuration (optional)
+    input_connectors:
+      from_stage_1: mooncake_connector
+
+# Top-level runtime config (concise): default windows and stage edges
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1             # Simplified: trigger downstream only after full upstream completion
+    max_inflight: 1             # Simplified: process serially within each stage
+
+  # Distributed connectors configuration (optional)
+  # More connectors will be supported in the future.
+  connectors:
+    # Mooncake connector for cross-node/intra-node communication
+    mooncake_connector:
+      name: MooncakeConnector
+      extra:
+        host: "127.0.0.1"
+        metadata_server: "http://10.90.67.86:8080/metadata"
+        master: "10.90.67.86:50051"
+        segment: 512000000    # 512MB
+        localbuf: 64000000     # 64MB
+        proto: "tcp"
+
+    # SharedMemory connector for intra-node communication
+    #Alternative SHM connector with different threshold
+    shared_memory_connector:
+      name: SharedMemoryConnector
+      extra:
+        shm_threshold_bytes: 65536 # 64KB threshold
+
+  edges:
+    - from: 0                   # thinker → talker: trigger only after receiving full input (-1)
+      to: 1
+      window_size: -1
+    - from: 1                   # talker → code2wav: trigger only after receiving full input (-1)
+      to: 2
+      window_size: -1

vllm_omni/model_executor/stage_configs/qwen3_omni_moe_multiconnector.yaml

@@ -0,0 +1,140 @@
+# Stage config for running Qwen3-Omni-MoE with 3-stage architecture
+# Stage 0: Thinker (multimodal understanding + text generation)
+# Stage 1: Talker (text embeddings -> 8-layer RVQ codec codes)
+# Stage 2: Code2Wav (8-layer RVQ codes -> audio waveform)
+
+# The following config has been verified on 2x H100-80G GPUs.
+stage_args:
+  - stage_id: 0
+    runtime:
+      devices: "0,1"
+      max_batch_size: 1
+    engine_args:
+      model_stage: thinker
+      model_arch: Qwen3OmniMoeForConditionalGeneration
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
+      gpu_memory_utilization: 0.6
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: latent  # Output hidden states for talker
+      distributed_executor_backend: "mp"
+      enable_prefix_caching: false
+      hf_config_name: thinker_config
+      tensor_parallel_size: 2
+    final_output: true
+    final_output_type: text
+    is_comprehension: true
+    default_sampling_params:
+      temperature: 0.4
+      top_p: 0.9
+      top_k: 1
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.05
+    # Distributed connector configuration
+    output_connectors:
+      to_stage_1: connector_of_mooncake
+
+  - stage_id: 1
+    runtime:
+       devices: "1"
+       max_batch_size: 1
+    engine_args:
+       model_stage: talker
+       model_arch: Qwen3OmniMoeForConditionalGeneration
+       worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+       scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
+       gpu_memory_utilization: 0.3
+       enforce_eager: true
+       trust_remote_code: true
+       engine_output_type: latent  # Output codec codes for code2wav
+      #  tensor_parallel_size: 2
+       enable_prefix_caching: false
+       distributed_executor_backend: "mp"
+       hf_config_name: talker_config
+    engine_input_source: [0]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
+    # final_output: true
+    # final_output_type: text
+    default_sampling_params:
+      temperature: 0.9
+      top_k: 50
+      max_tokens: 4096
+      seed: 42
+      detokenize: False
+      repetition_penalty: 1.05
+      stop_token_ids: [2150]
+    # Distributed connector configuration
+    input_connectors:
+      from_stage_0: connector_of_mooncake
+    output_connectors:
+      to_stage_2: connector_of_mooncake
+
+  - stage_id: 2
+    runtime:
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: code2wav
+      model_arch: Qwen3OmniMoeForConditionalGeneration
+      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      scheduler_cls: vllm_omni.core.sched.generation_scheduler.GenerationScheduler
+      enforce_eager: true
+      trust_remote_code: true
+      enable_prefix_caching: false
+      engine_output_type: audio  # Final output: audio waveform
+      gpu_memory_utilization: 0.1
+      distributed_executor_backend: "mp"
+      max_num_batched_tokens: 1000000
+      hf_config_name: thinker_config
+    engine_input_source: [1]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
+    final_output: true
+    final_output_type: audio
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 65536
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.1
+    # Distributed connector configuration
+    input_connectors:
+      from_stage_1: connector_of_mooncake
+
+# Top-level runtime config: default windows and stage edges
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 1
+
+  # Distributed connectors configuration
+  connectors:
+    # Mooncake connector for cross-node/intra-node communication
+    connector_of_mooncake:
+      name: MooncakeConnector
+      extra:
+        host: "127.0.0.1"
+        metadata_server: "http://10.90.67.86:8080/metadata"
+        master: "10.90.67.86:50051"
+        segment: 512000000    # 512MB
+        localbuf: 64000000     # 64MB
+        proto: "tcp"
+
+    # SharedMemory connector for intra-node communication
+    connector_of_shared_memory:
+      name: SharedMemoryConnector
+      extra:
+        shm_threshold_bytes: 65536 # 64KB threshold
+
+  edges:
+    - from: 0
+      to: 1
+      window_size: -1
+    - from: 1
+      to: 2
+      window_size: -1

vllm_omni/worker/gpu_ar_model_runner.py

@@ -47,6 +47,19 @@ class GPUARModelRunner(OmniGPUModelRunner):
         self.hidden_size = self.model_config.hf_text_config.hidden_size
         self.inputs_embeds = self._make_buffer(self.max_num_tokens, self.hidden_size, dtype=self.dtype, numpy=False)
 
+    def _make_buffer(self, *size, dtype, numpy=True):
+        # Prevent ray from pinning the buffer due to large size
+        from vllm_omni.distributed.ray_utils.utils import (
+            calculate_total_bytes,
+            maybe_disable_pin_memory_for_ray,
+        )
+
+        total_bytes = calculate_total_bytes(size, dtype)
+
+        # Use the context manager to temporarily disable pinning if needed
+        with maybe_disable_pin_memory_for_ray(self, total_bytes):
+            return super()._make_buffer(*size, dtype=dtype, numpy=numpy)
+
     @torch.inference_mode()
     def execute_model(
         self,