add moriio transfer engine (#1742)

Signed-off-by: inkcherry <mingzhi.liu@amd.com> Signed-off-by: zejwang <zejwang@amd.com> Signed-off-by: junkang1991 <junkangchow@gmail.com> Signed-off-by: Zejian Wang <zejianwang@sjtu.edu.cn> Co-authored-by: zejwang <zejwang@amd.com> Co-authored-by: junkang1991 <junkangchow@gmail.com> Co-authored-by: Zejian Wang <zejianwang@sjtu.edu.cn> Co-authored-by: Zejian Wang <Zejian.Wang@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
2026-06-10 08:37:33 +08:00
1 parent 57227dc7f3
commit e1b0b83238
@@ -8,6 +8,7 @@ Backend-specific setup lives in separate docs:
 - [SharedMemoryConnector](omni_connectors/shared_memory_connector.md)
 - [MooncakeStoreConnector](omni_connectors/mooncake_store_connector.md)
 - [MooncakeTransferEngineConnector](omni_connectors/mooncake_transfer_engine_connector.md)
+- [MoriTransferEngineConnector](omni_connectors/mori_transfer_engine_connector.md)
 - [YuanrongConnector](omni_connectors/yuanrong_connector.md)
 - [YuanrongTransferEngineConnector](omni_connectors/yuanrong_transfer_engine_connector.md)

@@ -23,6 +24,7 @@ Current connectors operate in D2H2D (device to host to device) mode.
 | Single node | SharedMemoryConnector | Auto-configured if no connector is specified. |
 | Multi node (Mooncake Store) | MooncakeStoreConnector | TCP-based, requires Mooncake Master + metadata server. |
 | Multi node (Mooncake RDMA) | MooncakeTransferEngineConnector | RDMA/TCP direct transfer with managed memory pool. Fastest. |
+| Multi node (Mori RDMA) | MoriTransferEngineConnector | RDMA direct transfer via Mori IOEngine. |
 | Multi node (Yuanrong) | YuanrongConnector | Requires Yuanrong Datasystem + etcd. |
 | Ascend NPU P2P (Yuanrong TE) | YuanrongTransferEngineConnector | Uses Yuanrong TransferEngine directly. Configure NPU device IPv4 and `memory_pool_device: "npu"`. |

@@ -0,0 +1,76 @@
+# MoriTransferEngineConnector
+
+## When to Use
+
+Currently supports intra-node deployment with Mori.
+
+As noted in #1742, inter-node support will be added back in a future
+refactor.
+
+## Mechanism
+
+Uses Mori's `IOEngine` / `MemoryDesc` API for zero-copy RDMA transfers.
+
+- Data Plane: RDMA (InfiniBand/RoCE) with managed memory pool.
+- Control Plane: ZMQ for pull-request handshake and async completion.
+
+## Installation
+
+See the [Mori repository](https://github.com/ROCm/mori) for installation instructions.
+
+## Configuration
+
+Mori is configured through the new deploy-config schema (see
+[`docs/configuration/stage_configs.md`](../../../configuration/stage_configs.md)).
+Define the connector at the top level of the deploy YAML and reference it
+by name from each stage's `input_connectors` / `output_connectors`:
+
+```yaml
+connectors:
+  mori_connector:
+    name: MoriTransferEngineConnector
+    extra:
+      host: "auto"
+      zmq_port: 50051
+      device_name: ""
+      memory_pool_size: 536870912
+      memory_pool_device: "cuda"
+
+stages:
+  - stage_id: 0
+    output_connectors:
+      to_stage_1: mori_connector
+
+  - stage_id: 1
+    input_connectors:
+      from_stage_0: mori_connector
+```
+
+A ready-to-run intra-node example for Qwen3-Omni-MoE on AMD MI300X lives
+at `vllm_omni/deploy/qwen3_omni_moe_mori_intranode.yaml` and can be loaded
+with:
+
+```bash
+vllm-omni serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --log-stats \
+    --deploy-config vllm_omni/deploy/qwen3_omni_moe_mori_intranode.yaml
+```
+
+The yaml wires `MoriTransferEngineConnector` (with `backend_type: xgmi`) to
+the chunk_transfer_adapter path (`async_chunk: true`) so stage-to-stage
+hidden-state and codec-frame streams ship GPU-to-GPU over AMD Infinity
+Fabric instead of SHM.  Qwen2.5-Omni + Mori is not yet functional on
+the chunk path: it needs `thinker2talker_async_chunk` /
+`talker2code2wav_async_chunk` input processors that do not exist yet
+(the orchestrator-level path the upstream PR originally targeted was
+lost during the entrypoints → engine refactor and removed in #1742).
+
+Parameters:
+
+- host: local RDMA IP (`"auto"` for auto-detect).
+- zmq_port: ZMQ base port for control-plane communication.
+- device_name: RDMA device (e.g., `"mlx5_0"`), empty for auto-detect.
+- memory_pool_size: RDMA memory pool size in bytes.
+- memory_pool_device: `"cpu"` (pinned) or `"cuda"` (GPUDirect / XGMI RDMA).
+
+For more details, refer to the
+[Mori repository](https://github.com/ROCm/mori).
@@ -0,0 +1,107 @@
+# Qwen3-Omni-MoE intra-node deploy using Mori XGMI transfer.
+#
+# Targets a single AMD Instinct MI300X OAM node (8x GPU, 192GB HBM3,
+# gfx942).  Each stage occupies one dedicated GPU (3 GPUs total) and
+# stage-to-stage streaming chunks are shipped via Mori's RDMA IOEngine
+# (with ``backend_type: xgmi`` selecting the AMD Infinity Fabric
+# GPU-to-GPU direct path for intranode transfers).
+#
+# Pipeline wiring (stage_id, execution_type, model_stage, input_sources,
+# custom processors, async-chunk builders) lives in
+# ``vllm_omni/model_executor/models/qwen3_omni/pipeline.py`` and is
+# merged into the stages below by the deploy loader.
+#
+# Why this yaml exists:
+#   * On CUDA/H100 Qwen3-Omni-MoE ships via ``SharedMemoryConnector``
+#     (see ``qwen3_omni_moe.yaml``).  On MI300X the SHM path still works
+#     but XGMI RDMA can substantially reduce stage-transition latency
+#     for the thinker->talker hidden-state stream and the talker->code2wav
+#     codec-frame stream, both of which run in async-chunk mode.
+#   * This config swaps the connector definition to
+#     ``MoriTransferEngineConnector`` while keeping the pipeline topology,
+#     sampling params, and async_chunk=true gating identical to the CUDA
+#     deploy, so the chunk_transfer_adapter path (scheduler-owned
+#     background put/get threads) is fully exercised.
+#
+# Usage:
+#   vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --log-stats \
+#     --deploy-config vllm_omni/deploy/qwen3_omni_moe_mori_intranode.yaml
+#
+# Dependencies: mori (``pip install mori``), torch-rocm>=2.9, vllm-omni
+# with #2383 deploy schema, ibverbs/RDMA NICs exposed via /dev/infiniband.
+#
+# Limitation: TP=1-only.  Connector raises NotImplementedError on
+# tp_world_size > 1 (chunk-path endpoint derivation isn't rank-aware yet).
+async_chunk: true
+
+connectors:
+  mori_connector:
+    name: MoriTransferEngineConnector
+    extra:
+      host: "auto"
+      zmq_port: 50051
+      backend_type: "xgmi"            # AMD Infinity Fabric GPU-to-GPU direct
+      device_name: ""                 # leave empty to honour $MORI_RDMA_DEVICES
+      memory_pool_size: 536870912     # 512 MB, matches Qwen2.5-Omni Mori deploy
+      memory_pool_device: "cuda"      # GPU memory required for XGMI RDMA transfers
+      xgmi_num_streams: 64
+      xgmi_num_events: 64
+      qp_per_transfer: 1
+      num_worker_threads: 1
+      post_batch_size: -1
+      # Chunk-transfer knobs consumed by talker2code2wav_async_chunk:
+      # 25 decoded codec frames per connector.put (~40 ms at 625 Hz
+      # codec rate) with a 25-frame left context to satisfy the
+      # code2wav receptive field.  Kept identical to the SHM deploy so
+      # accuracy stays comparable when swapping backends.
+      codec_chunk_frames: 25
+      codec_left_context_frames: 25
+
+stages:
+  - stage_id: 0
+    gpu_memory_utilization: 0.9
+    enforce_eager: true               # ROCm override: avoid flashinfer autotune path
+    devices: "0"
+    output_connectors:
+      to_stage_1: mori_connector
+    default_sampling_params:
+      temperature: 0.4
+      top_p: 0.9
+      top_k: 1
+      max_tokens: 2048
+      seed: 42
+      repetition_penalty: 1.05
+
+  - stage_id: 1
+    gpu_memory_utilization: 0.9
+    enforce_eager: true
+    devices: "1"
+    input_connectors:
+      from_stage_0: mori_connector
+    output_connectors:
+      to_stage_2: mori_connector
+    default_sampling_params:
+      temperature: 0.9
+      top_k: 50
+      max_tokens: 4096
+      seed: 42
+      repetition_penalty: 1.05
+
+  - stage_id: 2
+    gpu_memory_utilization: 0.3
+    max_num_seqs: 1
+    enforce_eager: true
+    async_scheduling: false
+    # Codec prefill length (Q * num_frames) exceeds 32k default;
+    # matches qwen3_omni_moe.yaml code2wav sizing.
+    max_num_batched_tokens: 51200
+    devices: "2"
+    input_connectors:
+      from_stage_1: mori_connector
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 65536
+      seed: 42
+      repetition_penalty: 1.1
@@ -6,6 +6,7 @@ from .omni_connectors import (
    MooncakeConnector,
    MooncakeStoreConnector,
    MooncakeTransferEngineConnector,
+    MoriTransferEngineConnector,
    OmniConnectorBase,
    OmniConnectorFactory,
    OmniTransferConfig,
@@ -25,6 +26,7 @@ __all__ = [
    "MooncakeConnector",  # compat alias
    "MooncakeStoreConnector",
    "MooncakeTransferEngineConnector",
+    "MoriTransferEngineConnector",
    "SharedMemoryConnector",
    "YuanrongConnector",
    "YuanrongTransferEngineConnector",
@@ -17,6 +17,11 @@ try:
    from .connectors.mooncake_transfer_engine_connector import MooncakeTransferEngineConnector
 except ImportError:
    MooncakeTransferEngineConnector = None  # RDMA deps (msgspec/zmq/mooncake) not installed
+
+try:
+    from .connectors.mori_transfer_engine_connector import MoriTransferEngineConnector
+except ImportError:
+    MoriTransferEngineConnector = None  # RDMA deps (msgspec/zmq/mori) not installed
 from .factory import OmniConnectorFactory
 from .utils.config import ConnectorSpec, OmniTransferConfig
 from .utils.initialization import (
@@ -44,6 +49,7 @@ __all__ = [
    "MooncakeConnector",  # compat alias → MooncakeStoreConnector
    "MooncakeStoreConnector",
    "MooncakeTransferEngineConnector",
+    "MoriTransferEngineConnector",
    "SharedMemoryConnector",
    "YuanrongConnector",
    "YuanrongTransferEngineConnector",
@@ -0,0 +1,962 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""OmniConnector backed by the Mori RDMA transfer engine.
+
+Implements the ``OmniConnector`` interface using Mori's ``IOEngine`` /
+``MemoryDesc`` API to perform zero-copy RDMA transfers between
+disaggregated prefill and decode workers.
+
+Notable design points:
+  * Remote peers must be registered via ``register_remote_engine()``
+    with the peer's ``EngineDesc`` before any data transfer can proceed.
+  * Memory regions are described by ``MemoryDesc`` objects (serialisable
+    via pack/unpack) rather than raw virtual addresses.
+  * Transfers are dispatched through ``IOEngine.batch_write()`` and
+    tracked asynchronously via ``TransferStatus`` objects.
+
+ZMQ handshake protocol:
+  * **Pull request** – receiver sends ``MoriPullRequest`` (msgspec-encoded)
+    containing its ``EngineDesc`` and pool ``MemoryDesc`` so the sender can
+    register the remote engine and RDMA-write directly into the receiver's
+    pool at the specified offset.
+  * **Query request** – receiver sends ``QUERY_INFO`` prefix + ``QueryRequest``
+    to check whether data is ready on the sender side (metadata-less get path).
+"""
+
+import os
+import queue
+import socket
+import threading
+import time as _time_mod
+import uuid
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any
+
+import msgspec
+import torch
+import zmq
+
+from ..utils.logging import get_connector_logger
+from ..utils.memory_pool import BufferAllocator, ManagedBuffer
+from ..utils.serialization import OmniSerializer
+from .base import OmniConnectorBase
+
+logger = get_connector_logger(__name__)
+
+try:
+    from mori.cpp import TransferStatus  # noqa: F401
+    from mori.io import (
+        BackendType,
+        EngineDesc,
+        IOEngine,
+        IOEngineConfig,
+        MemoryDesc,
+        PollCqMode,
+        RdmaBackendConfig,
+        XgmiBackendConfig,
+    )
+except ImportError:
+    IOEngine = None
+
+# Supported backend types for Mori. Kept as string constants so configuration
+# (YAML / CLI / dict) stays transport-agnostic. ``rdma`` uses NIC-based RDMA
+# (RoCE / IB, GDR-capable); ``xgmi`` uses AMD Infinity Fabric GPU-to-GPU
+# direct links and therefore requires a CUDA pool.
+_SUPPORTED_BACKENDS = ("rdma", "xgmi")
+
+_BUFFER_TTL_SECONDS = 300
+
+TRANS_DONE = b"trans_done"
+TRANS_ERROR = b"trans_error"
+QUERY_INFO = b"query_info"
+INFO_NOT_FOUND = b"info_not_found"
+
+
+# ---------------------------------------------------------------------------
+# ZMQ message types
+# ---------------------------------------------------------------------------
+
+
+class MoriPullRequest(msgspec.Struct):
+    """Receiver → sender: request an RDMA write into the receiver's pool."""
+
+    request_id: str
+    engine_desc_packed: bytes
+    mem_desc_packed: bytes
+    dst_offset: int
+    length: int
+
+
+class QueryRequest(msgspec.Struct):
+    """Receiver → sender: query whether data for *request_id* is ready."""
+
+    request_id: str
+
+
+class QueryResponse(msgspec.Struct):
+    """Sender → receiver: metadata about a ready buffer."""
+
+    request_id: str
+    data_size: int
+    is_fast_path: bool
+
+
+# ---------------------------------------------------------------------------
+# Connector
+# ---------------------------------------------------------------------------
+
+
+class MoriTransferEngineConnector(OmniConnectorBase):
+    """OmniConnector backed by the Mori RDMA IOEngine.
+
+    Topology: 1 sender ↔ 1 receiver per key (same as
+    ``MooncakeTransferEngineConnector``).
+    """
+
+    supports_raw_data: bool = True
+
+    # ------------------------------------------------------------------ init
+    def __init__(self, config: dict[str, Any]):
+        if IOEngine is None:
+            raise ImportError("Mori is not available. Install via: pip install mori")
+
+        self._closed = False
+        self._bind_error: Exception | None = None
+
+        # Teardown-safe defaults
+        self._stop_event = threading.Event()
+        self._sender_executor: ThreadPoolExecutor | None = None
+        self._listener_thread: threading.Thread | None = None
+        self._listener_ready = threading.Event()
+        self._local_buffers: dict[str, Any] = {}
+        self._local_buffers_lock = threading.Lock()
+        self._req_local = threading.local()
+        self._worker_local = threading.local()
+        self._last_ttl_check: float = _time_mod.monotonic()
+
+        # Track remote engines already registered with IOEngine.
+        # ``mori.io.IOEngine.deregister_remote_engine`` requires the original
+        # ``EngineDesc`` object (its pybind signature is typed
+        # ``deregister_remote_engine(self, engine_desc: EngineDesc)``), not
+        # just the key string -- so we keep the desc alongside the key so
+        # ``close()`` can pass a valid object back at teardown.
+        self._registered_engines: dict[str, EngineDesc] = {}
+        self._registered_engines_lock = threading.Lock()
+
+        self._metrics = {
+            "puts": 0,
+            "gets": 0,
+            "bytes_transferred": 0,
+            "errors": 0,
+            "timeouts": 0,
+        }
+
+        self.config = config
+
+        # Stage id is read by OmniChunkTransferAdapter (process_pending_chunks /
+        # _send_single_request / _poll_single_request / ...) to decide whether
+        # this connector sits at the sender or receiver end of a stage pair.
+        # Kept parallel to ``SharedMemoryConnector`` and other OmniConnector
+        # implementations. Populated by ``build_stage_connectors`` via the
+        # ``stage_id`` key that ``get_connectors_config_for_stage`` injects.
+        self.stage_id = config.get("stage_id", -1)
+
+        # ---- Host / ZMQ ----
+        host_cfg = config.get("host", "127.0.0.1")
+        if host_cfg.lower() == "auto":
+            self.host = self._get_local_ip()
+            logger.info(f"Auto-detected local IP: {self.host}")
+        else:
+            self.host = host_cfg
+        self.zmq_port = config.get("zmq_port", 50051)
+
+        # ---- Backend selection ----
+        # Defaults to "rdma" so existing deployments (pre-XGMI support) keep
+        # working unchanged. ``xgmi`` selects AMD Infinity Fabric GPU-to-GPU
+        # direct links; see ``mori.io.XgmiBackendConfig``.
+        backend_type_cfg = str(config.get("backend_type", "rdma")).lower()
+        if backend_type_cfg not in _SUPPORTED_BACKENDS:
+            raise ValueError(f"Invalid backend_type={backend_type_cfg!r}. Supported: {list(_SUPPORTED_BACKENDS)}.")
+        self.backend_type = backend_type_cfg
+
+        # ---- RDMA device (RDMA backend only) ----
+        self.device_name = ""
+        if self.backend_type == "rdma":
+            self.device_name = config.get("device_name", "")
+            if not self.device_name:
+                env_dev = os.environ.get("MORI_RDMA_DEVICES", "")
+                if env_dev:
+                    self.device_name = env_dev
+                    logger.info(f"Using MORI_RDMA_DEVICES from env: {self.device_name}")
+        elif config.get("device_name"):
+            logger.warning(
+                "device_name=%r is ignored for backend_type='xgmi' (XGMI does not use an RDMA NIC).",
+                config.get("device_name"),
+            )
+
+        # ---- Pool config ----
+        self.pool_size = config.get("memory_pool_size", 1024**3)
+        self.pool_device = config.get("memory_pool_device", "cpu")
+        if self.backend_type == "xgmi" and self.pool_device == "cpu":
+            raise ValueError(
+                "backend_type='xgmi' requires memory_pool_device='cuda': "
+                "XGMI is a GPU-to-GPU fabric and cannot address CPU memory."
+            )
+
+        # ---- Sender info (receiver uses this when metadata=None) ----
+        self.sender_host = config.get("sender_host", None)
+        self.sender_zmq_port = config.get("sender_zmq_port", None)
+
+        # ---- Role ----
+        # ``sender``   : binds the ZMQ listener, accepts ``put()``.
+        # ``receiver`` : no listener, only ``get()`` by querying an
+        #                upstream sender whose endpoint lives in
+        #                ``sender_host`` / ``sender_zmq_port``.
+        role = str(config.get("role", "sender")).lower()
+        if role not in {"sender", "receiver"}:
+            raise ValueError(f"Invalid role={role!r} for MoriTransferEngineConnector. Expected 'sender' or 'receiver'.")
+        self.role = role
+        self.can_put = role == "sender"
+
+        # ---- Mori IOEngine ----
+        if self.device_name:
+            os.environ["MORI_RDMA_DEVICES"] = self.device_name
+
+        engine_config = IOEngineConfig(host=self.host, port=0)
+        self.engine_key = f"omni-{role}-{uuid.uuid4().hex[:8]}-pid{os.getpid()}-{self.host}"
+        self.engine = IOEngine(self.engine_key, engine_config)
+
+        # ---- Backend creation (per backend_type) ----
+        # ``IOEngine.batch_write`` is backend-agnostic, so only construction
+        # differs; the data-plane / ZMQ handshake paths below are identical.
+        if self.backend_type == "xgmi":
+            xgmi_cfg = XgmiBackendConfig()
+            xgmi_cfg.num_streams = config.get("xgmi_num_streams", 64)
+            xgmi_cfg.num_events = config.get("xgmi_num_events", 64)
+            self.engine.create_backend(BackendType.XGMI, xgmi_cfg)
+            logger.info(f"Mori backend: XGMI (num_streams={xgmi_cfg.num_streams}, num_events={xgmi_cfg.num_events})")
+        else:  # rdma
+            qp_per_transfer = config.get("qp_per_transfer", 1)
+            post_batch_size = config.get("post_batch_size", -1)
+            num_workers = config.get("num_worker_threads", 1)
+            rdma_cfg = RdmaBackendConfig(
+                qp_per_transfer,
+                post_batch_size,
+                num_workers,
+                PollCqMode.POLLING,
+                False,
+            )
+            self.engine.create_backend(BackendType.RDMA, rdma_cfg)
+            logger.info(
+                f"Mori backend: RDMA (qp_per_transfer={qp_per_transfer}, "
+                f"num_workers={num_workers}, device={self.device_name or 'auto'})"
+            )
+
+        self.engine_desc: EngineDesc = self.engine.get_engine_desc()
+        self.engine_desc_packed: bytes = self.engine_desc.pack()
+        logger.info(f"Mori IOEngine ready: key={self.engine_key} at {self.engine_desc.host}:{self.engine_desc.port}")
+
+        # ---- Pool allocation & Mori memory registration ----
+        logger.info(
+            f"Allocating {self.backend_type.upper()} pool: {self.pool_size / 1024**2:.2f} MB on {self.pool_device}"
+        )
+        try:
+            if self.pool_device == "cpu":
+                self.pool = torch.empty(self.pool_size, dtype=torch.uint8).pin_memory()
+            else:
+                self.pool = torch.empty(
+                    self.pool_size,
+                    dtype=torch.uint8,
+                    device=self.pool_device,
+                )
+            self.pool_mem_desc: MemoryDesc = self.engine.register_torch_tensor(self.pool)
+            self.pool_mem_desc_packed: bytes = self.pool_mem_desc.pack()
+            self.base_ptr = self.pool.data_ptr()
+        except Exception as e:
+            logger.error(f"Failed to allocate/register pool: {e}")
+            raise
+
+        self.allocator = BufferAllocator(self.pool_size, alignment=4096)
+
+        # ---- ZMQ / threading ----
+        self.zmq_ctx = zmq.Context()
+        self._sender_executor = ThreadPoolExecutor(max_workers=4, thread_name_prefix="mori-sender")
+
+        logger.info(
+            f"MoriTransferEngineConnector config:\n"
+            f"  Local: host={self.host}, zmq_port={self.zmq_port}\n"
+            f"  Remote: sender_host={self.sender_host}, "
+            f"sender_zmq_port={self.sender_zmq_port}\n"
+            f"  Role: can_put={self.can_put}, "
+            f"configured_role={config.get('role', 'sender')}"
+        )
+
+        if self.can_put:
+            self._last_ttl_check = _time_mod.monotonic()
+            self._listener_thread = threading.Thread(target=self._zmq_listener_loop, daemon=True)
+            self._listener_thread.start()
+            self._listener_ready.wait(timeout=1.0)
+            if self._bind_error is not None:
+                raise RuntimeError(
+                    f"MoriTransferEngineConnector failed to bind ZMQ on {self.host}:{self.zmq_port}: {self._bind_error}"
+                ) from self._bind_error
+            logger.info(f"MoriTransferEngineConnector SENDER ready (ZMQ on {self.host}:{self.zmq_port})")
+        else:
+            if not self.sender_host or self.sender_host.lower() == "auto":
+                logger.info("MoriTransferEngineConnector RECEIVER: awaiting sender info via update_sender_info().")
+            else:
+                logger.info(
+                    f"MoriTransferEngineConnector RECEIVER ready (sender at {self.sender_host}:{self.sender_zmq_port})"
+                )
+
+    # -------------------------------------------------------- public helpers
+    def get_connection_info(self) -> dict[str, Any]:
+        return {
+            "host": self.host,
+            "zmq_port": self.zmq_port,
+            "engine_key": self.engine_key,
+            "can_put": self.can_put,
+        }
+
+    def update_sender_info(self, sender_host: str, sender_zmq_port: int) -> None:
+        """Inject the sender's ZMQ endpoint into the receiver connector."""
+        self.sender_host = sender_host
+        self.sender_zmq_port = sender_zmq_port
+        logger.info(f"Sender info updated: host={sender_host!r}, zmq_port={sender_zmq_port}")
+
+    # -------------------------------------------------- internal helpers
+    @staticmethod
+    def _get_local_ip() -> str:
+        try:
+            with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
+                s.connect(("8.8.8.8", 80))
+                return s.getsockname()[0]
+        except Exception:
+            try:
+                return socket.gethostbyname(socket.gethostname())
+            except Exception:
+                return "127.0.0.1"
+
+    def _get_req_socket(self, zmq_addr: str, timeout_ms: int) -> zmq.Socket:
+        cache: dict[str, zmq.Socket] | None = getattr(self._req_local, "cache", None)
+        if cache is None:
+            cache = {}
+            self._req_local.cache = cache
+
+        sock = cache.get(zmq_addr)
+        if sock is None:
+            sock = self.zmq_ctx.socket(zmq.REQ)
+            sock.connect(zmq_addr)
+            cache[zmq_addr] = sock
+        sock.setsockopt(zmq.SNDTIMEO, timeout_ms)
+        sock.setsockopt(zmq.RCVTIMEO, timeout_ms)
+        return sock
+
+    def _invalidate_req_socket(self, zmq_addr: str) -> None:
+        cache: dict[str, zmq.Socket] | None = getattr(self._req_local, "cache", None)
+        if cache is None:
+            return
+        sock = cache.pop(zmq_addr, None)
+        if sock is not None:
+            try:
+                sock.close(linger=0)
+            except Exception:
+                pass
+
+    def _ensure_remote_registered(self, engine_desc_packed: bytes) -> str:
+        """Register a remote engine with the local IOEngine (idempotent)."""
+        desc = EngineDesc.unpack(engine_desc_packed)
+        key = desc.key
+        with self._registered_engines_lock:
+            if key not in self._registered_engines:
+                self.engine.register_remote_engine(desc)
+                self._registered_engines[key] = desc
+                logger.debug(f"Registered remote Mori engine: {key}")
+        return key
+
+    # ----------------------------------------------------------------- put()
+    def put(
+        self,
+        from_stage: str,
+        to_stage: str,
+        put_key: str,
+        data: Any,
+    ) -> tuple[bool, int, dict[str, Any] | None]:
+        if self._closed:
+            raise RuntimeError("Cannot put: MoriTransferEngineConnector is closed")
+        if not self.can_put:
+            logger.warning(f"Rejecting put for {put_key}: connector is receiver-only")
+            return False, 0, None
+
+        put_key = self._make_key(put_key, from_stage, to_stage)
+        try:
+            src_offset = 0
+            size = 0
+            holder = None
+            is_fast_path = True
+            should_release = False
+
+            # Reject empty payloads early
+            if isinstance(data, bytes) and len(data) == 0:
+                return False, 0, None
+            if isinstance(data, torch.Tensor) and data.nbytes == 0:
+                return False, 0, None
+
+            # Serialize non-raw types
+            if not isinstance(data, (ManagedBuffer, torch.Tensor, bytes)):
+                data = OmniSerializer.serialize(data)
+                is_fast_path = False
+
+            if isinstance(data, ManagedBuffer):
+                if data.pool_tensor.data_ptr() != self.pool.data_ptr():
+                    data = data.tensor.contiguous()
+                else:
+                    src_offset = data.offset
+                    size = data.size
+                    holder = data
+                    should_release = False
+
+            if isinstance(data, (torch.Tensor, bytes)):
+                if isinstance(data, torch.Tensor):
+                    size = data.nbytes
+                    tensor_data = data
+                else:
+                    size = len(data)
+                    tensor_data = torch.frombuffer(data, dtype=torch.uint8)
+
+                try:
+                    offset = self.allocator.alloc(size)
+                    holder = ManagedBuffer(self.allocator, offset, size, self.pool)
+                    should_release = True
+                except MemoryError:
+                    logger.error(f"Pool exhausted for {size} bytes")
+                    return False, 0, None
+
+                try:
+                    dst_t = holder.tensor
+                    if isinstance(data, torch.Tensor):
+                        if not data.is_contiguous():
+                            data = data.contiguous()
+                        src_view = data.view(torch.uint8).flatten()
+                        if src_view.device != dst_t.device:
+                            dst_t.copy_(src_view, non_blocking=True)
+                            if src_view.is_cuda:
+                                with torch.cuda.device(src_view.device):
+                                    torch.cuda.current_stream().synchronize()
+                            elif dst_t.is_cuda:
+                                with torch.cuda.device(dst_t.device):
+                                    torch.cuda.current_stream().synchronize()
+                        else:
+                            dst_t.copy_(src_view)
+                            if dst_t.is_cuda:
+                                with torch.cuda.device(dst_t.device):
+                                    torch.cuda.current_stream().synchronize()
+                    else:
+                        dst_t.copy_(tensor_data)
+                        if dst_t.is_cuda:
+                            with torch.cuda.device(dst_t.device):
+                                torch.cuda.current_stream().synchronize()
+                except Exception as e:
+                    holder.release()
+                    logger.error(f"Failed to copy data to pool: {e}")
+                    return False, 0, None
+
+                src_offset = offset
+
+            if size <= 0:
+                if should_release and isinstance(holder, ManagedBuffer):
+                    holder.release()
+                return False, 0, None
+
+            with self._local_buffers_lock:
+                old = self._local_buffers.pop(put_key, None)
+                if old:
+                    _, _, oh, osr, _, _ = old
+                    if osr and isinstance(oh, ManagedBuffer):
+                        oh.release()
+                        logger.warning(f"Released stale buffer for duplicate key: {put_key}")
+                # (offset, size, holder, should_release, is_fast_path, ts)
+                self._local_buffers[put_key] = (
+                    src_offset,
+                    size,
+                    holder,
+                    should_release,
+                    is_fast_path,
+                    _time_mod.monotonic(),
+                )
+
+            metadata = {
+                "source_host": self.host,
+                "source_port": self.zmq_port,
+                "data_size": size,
+                "is_fast_path": is_fast_path,
+            }
+            self._metrics["puts"] += 1
+            self._metrics["bytes_transferred"] += size
+            return True, size, metadata
+
+        except Exception as e:
+            self._metrics["errors"] += 1
+            logger.error(f"Put failed for {put_key}: {e}", exc_info=True)
+            return False, 0, None
+
+    # ------------------------------------------------ query (no-metadata get)
+    def _query_metadata_from_sender(self, get_key: str) -> dict[str, Any] | None:
+        zmq_addr = f"tcp://{self.sender_host}:{self.sender_zmq_port}"
+        req_socket = self._get_req_socket(zmq_addr, timeout_ms=5000)
+        try:
+            q = QueryRequest(request_id=get_key)
+            req_socket.send(QUERY_INFO + msgspec.msgpack.encode(q))
+            resp = req_socket.recv()
+            if resp == INFO_NOT_FOUND:
+                return None
+            qr = msgspec.msgpack.decode(resp, type=QueryResponse)
+            return {
+                "source_host": self.sender_host,
+                "source_port": self.sender_zmq_port,
+                "data_size": qr.data_size,
+                "is_fast_path": qr.is_fast_path,
+            }
+        except Exception as e:
+            self._invalidate_req_socket(zmq_addr)
+            logger.debug(f"Metadata query failed for {get_key}: {e}")
+            return None
+
+    # ----------------------------------------------------------------- get()
+    def get(
+        self,
+        from_stage: str,
+        to_stage: str,
+        get_key: str,
+        metadata: dict[str, Any] | None = None,
+    ) -> tuple[Any, int] | None:
+        if self._closed:
+            raise RuntimeError("Cannot get: MoriTransferEngineConnector is closed")
+
+        get_key = self._make_key(get_key, from_stage, to_stage)
+        _t0 = _time_mod.perf_counter()
+
+        # Resolve metadata
+        if not metadata:
+            if not self.sender_host or not self.sender_zmq_port or str(self.sender_host).lower() == "auto":
+                raise RuntimeError("get(metadata=None) requires sender info. Call update_sender_info() first.")
+            metadata = self._query_metadata_from_sender(get_key)
+            if not metadata:
+                return None
+
+        _t1 = _time_mod.perf_counter()
+        _query_ms = (_t1 - _t0) * 1000
+
+        src_host = metadata.get("source_host")
+        src_port = metadata.get("source_port")
+        data_size = metadata.get("data_size", 0)
+        is_fast_path = metadata.get("is_fast_path", False)
+
+        if not src_host or not src_port or str(src_host).lower() == "auto":
+            logger.error(f"Invalid metadata for {get_key}")
+            return None
+        if data_size == 0:
+            logger.warning(f"Skipping get for {get_key}: data_size is 0")
+            return None
+
+        # Allocate destination buffer
+        try:
+            offset = self.allocator.alloc(data_size)
+            recv_buf = ManagedBuffer(self.allocator, offset, data_size, self.pool)
+        except MemoryError:
+            logger.error(f"Failed to allocate {data_size} bytes for get")
+            return None
+
+        _t2 = _time_mod.perf_counter()
+        _alloc_ms = (_t2 - _t1) * 1000
+
+        # Build pull request with Mori-specific EngineDesc + MemoryDesc
+        pull_req = MoriPullRequest(
+            request_id=get_key,
+            engine_desc_packed=self.engine_desc_packed,
+            mem_desc_packed=self.pool_mem_desc_packed,
+            dst_offset=offset,
+            length=data_size,
+        )
+
+        _base_timeout_ms = 30000
+        _size_timeout_ms = max(0, (data_size // (100 * 1024 * 1024))) * 5000
+        _total_timeout_ms = _base_timeout_ms + _size_timeout_ms
+        zmq_addr = f"tcp://{src_host}:{src_port}"
+        req_socket = self._get_req_socket(zmq_addr, timeout_ms=_total_timeout_ms)
+
+        try:
+            req_socket.send(msgspec.msgpack.encode(pull_req))
+            resp = req_socket.recv()
+
+            _t3 = _time_mod.perf_counter()
+            _rdma_ms = (_t3 - _t2) * 1000
+
+            if resp == TRANS_DONE:
+                if self.pool.is_cuda:
+                    with torch.cuda.device(self.pool.device):
+                        torch.cuda.current_stream().synchronize()
+
+                _t4 = _time_mod.perf_counter()
+                _sync_ms = (_t4 - _t3) * 1000
+
+                if is_fast_path:
+                    _total_ms = (_time_mod.perf_counter() - _t0) * 1000
+                    _mbps = (data_size / 1024 / 1024) / (_total_ms / 1000) if _total_ms > 0 else 0
+                    logger.info(
+                        f"[MORI GET] {get_key}: query={_query_ms:.1f}ms, "
+                        f"alloc={_alloc_ms:.1f}ms, rdma={_rdma_ms:.1f}ms, "
+                        f"sync={_sync_ms:.1f}ms, total={_total_ms:.1f}ms, "
+                        f"{_mbps:.1f} MB/s (fast_path)"
+                    )
+                    self._metrics["gets"] += 1
+                    self._metrics["bytes_transferred"] += data_size
+                    return recv_buf, data_size
+                else:
+                    try:
+                        raw = recv_buf.to_bytes()
+                        val = OmniSerializer.deserialize(raw)
+                        _total_ms = (_time_mod.perf_counter() - _t0) * 1000
+                        _mbps = (data_size / 1024 / 1024) / (_total_ms / 1000) if _total_ms > 0 else 0
+                        logger.info(f"[MORI GET] {get_key}: total={_total_ms:.1f}ms, {_mbps:.1f} MB/s (deserialized)")
+                        self._metrics["gets"] += 1
+                        self._metrics["bytes_transferred"] += data_size
+                        return val, data_size
+                    finally:
+                        recv_buf.release()
+            else:
+                self._metrics["errors"] += 1
+                logger.error(f"MORI get failed: received {resp} instead of TRANS_DONE")
+                recv_buf.release()
+                return None
+        except Exception as e:
+            self._invalidate_req_socket(zmq_addr)
+            self._metrics["timeouts"] += 1
+            logger.error(f"MORI get error: {e}", exc_info=True)
+            recv_buf.release()
+            return None
+
+    # ------------------------------------------------------------- cleanup()
+    def cleanup(
+        self,
+        request_id: str,
+        from_stage: str | None = None,
+        to_stage: str | None = None,
+    ) -> None:
+        if (from_stage is None) != (to_stage is None):
+            raise ValueError("cleanup() requires both from_stage and to_stage, or neither.")
+        if from_stage is not None and to_stage is not None:
+            request_id = self._make_key(request_id, from_stage, to_stage)
+        with self._local_buffers_lock:
+            item = self._local_buffers.pop(request_id, None)
+            if item:
+                _, _, holder, sr, _, _ = item
+                if sr and isinstance(holder, ManagedBuffer):
+                    holder.release()
+
+    # -------------------------------------------------------------- health()
+    def health(self) -> dict[str, Any]:
+        if self._closed:
+            return {"status": "unhealthy", "error": "Connector is closed"}
+        return {
+            "status": "healthy",
+            "host": self.host,
+            "pool_device": self.pool_device,
+            "pool_size": self.pool_size,
+            "engine_key": self.engine_key,
+            **self._metrics,
+        }
+
+    # --------------------------------------------------------------- close()
+    def close(self) -> None:
+        if getattr(self, "_closed", True):
+            return
+        self._closed = True
+        logger.info("Closing MoriTransferEngineConnector...")
+
+        self._stop_event.set()
+
+        if self._listener_thread is not None and self._listener_thread.is_alive():
+            self._listener_thread.join(timeout=2.0)
+            if self._listener_thread.is_alive():
+                logger.warning("Listener thread did not stop gracefully")
+
+        if self._sender_executor is not None:
+            self._sender_executor.shutdown(wait=True, cancel_futures=False)
+
+        with self._local_buffers_lock:
+            for _k, item in list(self._local_buffers.items()):
+                _, _, holder, sr, _, _ = item
+                if sr and isinstance(holder, ManagedBuffer):
+                    holder.release()
+            self._local_buffers.clear()
+
+        cache: dict[str, zmq.Socket] | None = getattr(self._req_local, "cache", None)
+        if cache:
+            for _addr, sock in cache.items():
+                try:
+                    sock.close(linger=0)
+                except Exception:
+                    pass
+            cache.clear()
+
+        try:
+            if hasattr(self, "zmq_ctx"):
+                self.zmq_ctx.term()
+        except Exception as e:
+            logger.warning(f"Failed to terminate ZMQ context: {e}")
+
+        # ---- Release Mori IOEngine resources ----
+        # ``mori.io.IOEngine`` does not expose an explicit ``shutdown()`` /
+        # ``close()`` entry-point, so we best-effort deregister every resource
+        # we registered with it and then drop the only reference we hold so
+        # Python GC triggers the C++ destructor.
+        engine = getattr(self, "engine", None)
+        if engine is not None:
+            pool_desc = getattr(self, "pool_mem_desc", None)
+            if pool_desc is not None:
+                try:
+                    engine.deregister_memory(pool_desc)
+                except Exception as e:
+                    logger.warning(f"Mori deregister_memory failed: {e}")
+                self.pool_mem_desc = None  # type: ignore[assignment]
+
+            with self._registered_engines_lock:
+                remote_descs = list(self._registered_engines.values())
+                self._registered_engines.clear()
+            for desc in remote_descs:
+                try:
+                    engine.deregister_remote_engine(desc)
+                except Exception as e:
+                    logger.debug(f"Mori deregister_remote_engine({desc.key!r}) failed: {e}")
+
+            # Drop the only reference to IOEngine; the C++ destructor runs
+            # under Python GC and unwinds backend / RDMA fabric state.
+            self.engine = None  # type: ignore[assignment]
+
+        self.pool = None  # type: ignore[assignment]
+        logger.info("MoriTransferEngineConnector closed.")
+
+    # ============================================================== LISTENER
+    def _cleanup_stale_buffers(self) -> None:
+        """Reclaim buffers older than ``_BUFFER_TTL_SECONDS``.
+
+        Prevents permanent memory leaks when a receiver crashes or times out
+        without ever pulling the data.
+
+        TODO(zejwang): In extreme rare case, long transfer time, there might
+        exist TTL cleanup vs in-flight RDMA transfer conflict, which will be
+        handled in a follow-up PR. Same race is acknowledged in
+        ``MooncakeTransferEngineConnector._cleanup_stale_buffers``.
+        """
+        now = _time_mod.monotonic()
+        with self._local_buffers_lock:
+            stale = [k for k, v in self._local_buffers.items() if now - v[5] > _BUFFER_TTL_SECONDS]
+            for k in stale:
+                item = self._local_buffers.pop(k)
+                _, _, holder, sr, _, _ = item
+                if sr and isinstance(holder, ManagedBuffer):
+                    holder.release()
+                logger.warning(f"TTL expired ({_BUFFER_TTL_SECONDS}s): reclaimed buffer for {k}")
+
+    def _zmq_listener_loop(self) -> None:
+        sock = self.zmq_ctx.socket(zmq.ROUTER)
+
+        try:
+            sock.bind(f"tcp://{self.host}:{self.zmq_port}")
+        except zmq.ZMQError as exc:
+            # Any bind failure (EADDRINUSE, EADDRNOTAVAIL, EACCES, etc.) is
+            # fatal for a sender — fail fast so __init__ propagates the error.
+            # There is no silent receiver fallback; roles are explicitly
+            # assigned (matches MooncakeTransferEngineConnector).
+            logger.error(f"ZMQ bind failed on {self.host}:{self.zmq_port}: {exc} (errno={exc.errno})")
+            self.can_put = False
+            self._bind_error = exc
+            self._listener_ready.set()
+            return
+
+        self._listener_ready.set()
+
+        notify_addr = f"inproc://mori-notify-{id(self)}"
+        notify_recv = self.zmq_ctx.socket(zmq.PULL)
+        notify_recv.bind(notify_addr)
+
+        poller = zmq.Poller()
+        poller.register(sock, zmq.POLLIN)
+        poller.register(notify_recv, zmq.POLLIN)
+
+        response_queue: queue.Queue = queue.Queue()
+
+        try:
+            while not self._stop_event.is_set():
+                try:
+                    events = dict(poller.poll(1000))
+
+                    if notify_recv in events:
+                        while True:
+                            try:
+                                notify_recv.recv(zmq.NOBLOCK)
+                            except zmq.Again:
+                                break
+
+                    while True:
+                        try:
+                            identity, response = response_queue.get_nowait()
+                            sock.send_multipart([identity, b"", response])
+                        except queue.Empty:
+                            break
+
+                    now = _time_mod.monotonic()
+                    if now - self._last_ttl_check >= 10.0:
+                        self._last_ttl_check = now
+                        self._cleanup_stale_buffers()
+
+                    if sock in events:
+                        frames = sock.recv_multipart()
+                        if len(frames) >= 2:
+                            self._sender_executor.submit(
+                                self._handle_pull_request,
+                                response_queue,
+                                notify_addr,
+                                frames[0],
+                                frames[-1],
+                            )
+                except zmq.ContextTerminated:
+                    break
+                except Exception:
+                    logger.debug("Listener loop error", exc_info=True)
+        finally:
+            try:
+                notify_recv.close(linger=0)
+                sock.close(linger=0)
+            except Exception:
+                pass
+
+    def _handle_pull_request(
+        self,
+        response_queue: queue.Queue,
+        notify_addr: str,
+        identity: bytes,
+        payload: bytes,
+    ) -> None:
+        try:
+            if payload.startswith(QUERY_INFO):
+                self._handle_query_request(
+                    response_queue,
+                    notify_addr,
+                    identity,
+                    payload[len(QUERY_INFO) :],
+                )
+                return
+
+            pull = msgspec.msgpack.decode(payload, type=MoriPullRequest)
+
+            with self._local_buffers_lock:
+                item = self._local_buffers.get(pull.request_id)
+
+            if not item:
+                response_queue.put((identity, TRANS_ERROR))
+                self._notify_listener(notify_addr)
+                return
+
+            src_offset, src_size, _, _, _, _ = item
+
+            # Validate against sender's own src_size before RDMA write:
+            # length < src_size silently truncates the payload; length >
+            # src_size reads past this allocation into adjacent in-flight
+            # buffers because pool_mem_desc covers the whole pool, not just
+            # this slot, so Mori does not bound-check. Refuse on mismatch
+            # rather than corrupt or leak data silently.
+            if pull.length != src_size:
+                logger.error(
+                    "Length mismatch for %s: sender src_size=%d != receiver length=%d; refusing transfer",
+                    pull.request_id,
+                    src_size,
+                    pull.length,
+                )
+                response_queue.put((identity, TRANS_ERROR))
+                self._notify_listener(notify_addr)
+                return
+
+            # Register the receiver's IOEngine (idempotent)
+            self._ensure_remote_registered(pull.engine_desc_packed)
+
+            # Reconstruct the receiver's pool MemoryDesc
+            remote_mem = MemoryDesc.unpack(pull.mem_desc_packed)
+
+            # RDMA write from local pool → remote pool. Use src_size (sender's
+            # own record); it has just been validated to equal pull.length.
+            transfer_uid = self.engine.allocate_transfer_uid()
+            statuses = self.engine.batch_write(
+                [self.pool_mem_desc],
+                [[src_offset]],
+                [remote_mem],
+                [[pull.dst_offset]],
+                [[src_size]],
+                [transfer_uid],
+            )
+
+            success = True
+            for st in statuses:
+                st.Wait()
+                if st.Failed():
+                    logger.error(f"RDMA write failed for {pull.request_id}: {st.Message()}")
+                    success = False
+
+            if success:
+                self.cleanup(pull.request_id)
+                response_queue.put((identity, TRANS_DONE))
+            else:
+                logger.warning(f"RDMA write failed for {pull.request_id}. Buffer retained for retry.")
+                response_queue.put((identity, TRANS_ERROR))
+
+        except Exception as e:
+            logger.error(f"Pull request handler error: {e}", exc_info=True)
+            response_queue.put((identity, TRANS_ERROR))
+
+        self._notify_listener(notify_addr)
+
+    def _handle_query_request(
+        self,
+        response_queue: queue.Queue,
+        notify_addr: str,
+        identity: bytes,
+        payload: bytes,
+    ) -> None:
+        try:
+            q = msgspec.msgpack.decode(payload, type=QueryRequest)
+            with self._local_buffers_lock:
+                item = self._local_buffers.get(q.request_id)
+            if not item:
+                response_queue.put((identity, INFO_NOT_FOUND))
+            else:
+                _, data_size, _, _, is_fast, _ = item
+                resp = QueryResponse(
+                    request_id=q.request_id,
+                    data_size=data_size,
+                    is_fast_path=is_fast,
+                )
+                response_queue.put((identity, msgspec.msgpack.encode(resp)))
+        except Exception as e:
+            logger.error(f"Query handler error: {e}")
+            response_queue.put((identity, INFO_NOT_FOUND))
+
+        self._notify_listener(notify_addr)
+
+    def _notify_listener(self, notify_addr: str) -> None:
+        try:
+            local = self._worker_local
+            sock = getattr(local, "notify_socket", None)
+            cached_addr = getattr(local, "notify_addr", None)
+            if sock is None or cached_addr != notify_addr:
+                if sock is not None:
+                    sock.close(linger=0)
+                sock = self.zmq_ctx.socket(zmq.PUSH)
+                sock.connect(notify_addr)
+                local.notify_socket = sock
+                local.notify_addr = notify_addr
+            sock.send(b"", zmq.NOBLOCK)
+        except Exception:
+            local.notify_socket = None
+            local.notify_addr = None
@@ -114,11 +114,23 @@ def _create_mooncake_transfer_engine_connector(config: dict[str, Any]) -> OmniCo
    return MooncakeTransferEngineConnector(config)


+def _create_mori_transfer_engine_connector(config: dict[str, Any]) -> OmniConnectorBase:
+    try:
+        from .connectors.mori_transfer_engine_connector import MoriTransferEngineConnector
+    except ImportError:
+        import sys
+
+        sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
+        from omni_connectors.connectors.mori_transfer_engine_connector import MoriTransferEngineConnector
+    return MoriTransferEngineConnector(config)
+
+
 # Register connectors
 OmniConnectorFactory.register_connector("MooncakeStoreConnector", _create_mooncake_store_connector)
 OmniConnectorFactory.register_connector("MooncakeTransferEngineConnector", _create_mooncake_transfer_engine_connector)
 OmniConnectorFactory.register_connector("SharedMemoryConnector", _create_shm_connector)
 OmniConnectorFactory.register_connector("YuanrongConnector", _create_yuanrong_connector)
 OmniConnectorFactory.register_connector("YuanrongTransferEngineConnector", _create_yuanrong_transfer_engine_connector)
+OmniConnectorFactory.register_connector("MoriTransferEngineConnector", _create_mori_transfer_engine_connector)
 # Backward-compatible aliases – will be removed in the future
 OmniConnectorFactory.register_connector("MooncakeConnector", _create_mooncake_store_connector)
@@ -11,6 +11,7 @@ logger = get_connector_logger(__name__)
 TRANSFER_ENGINE_CONNECTOR_NAMES = frozenset(
    {
        "MooncakeTransferEngineConnector",
+        "MoriTransferEngineConnector",
        "YuanrongTransferEngineConnector",
    }
 )
@@ -1363,6 +1363,7 @@ class AsyncOmniEngine:

        async def _run_orchestrator() -> None:
            self._initialize_stages(stage_init_timeout)
+
            pd_config = self._detect_pd_config()
            coordinator_pub_address: str | None = None
            load_balancer_factory: Callable[[], LoadBalancer] | None = None