Yuanrong TransferEngine Connector for NPU (#3180)

Signed-off-by: yangsonglin13 <yangsonglin566@gmail.com>
2026-05-20 14:51:33 +08:00
1 parent 9dd36e3afd
commit 1c530d984d
@@ -9,6 +9,7 @@ Backend-specific setup lives in separate docs:
 - [MooncakeStoreConnector](omni_connectors/mooncake_store_connector.md)
 - [MooncakeTransferEngineConnector](omni_connectors/mooncake_transfer_engine_connector.md)
 - [YuanrongConnector](omni_connectors/yuanrong_connector.md)
+- [YuanrongTransferEngineConnector](omni_connectors/yuanrong_transfer_engine_connector.md)

 ## Overview

@@ -23,6 +24,7 @@ Current connectors operate in D2H2D (device to host to device) mode.
 | Multi node (Mooncake Store) | MooncakeStoreConnector | TCP-based, requires Mooncake Master + metadata server. |
 | Multi node (Mooncake RDMA) | MooncakeTransferEngineConnector | RDMA/TCP direct transfer with managed memory pool. Fastest. |
 | Multi node (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"`. |

 ## Core API

@@ -0,0 +1,211 @@
+# YuanrongTransferEngineConnector
+
+## When to Use
+
+Use `YuanrongTransferEngineConnector` for high-performance peer-to-peer KV cache
+transfer on Ascend NPU. It is different from `YuanrongConnector`:
+
+- `YuanrongConnector` uses Yuanrong Datasystem as a distributed KV store and
+  requires Datasystem workers plus etcd.
+- `YuanrongTransferEngineConnector` uses Yuanrong TransferEngine directly. It
+  registers local memory in each stage worker and lets the receiver pull data
+  from the sender with TransferEngine.
+
+The implementation currently lives under
+`vllm_omni/platforms/npu/omni_connectors/` because the supported backend is
+Ascend NPU only. The connector is still registered with the generic
+OmniConnector factory under the name `YuanrongTransferEngineConnector`.
+
+For Ascend P2P transfer, the connector should use NPU memory as its transfer
+pool. The Mooncake-style CPU pool is not supported for this connector.
+
+## Prerequisites
+
+Install Yuanrong Datasystem/TransferEngine Python bindings in the runtime
+environment. The packaged install is:
+
+```bash
+pip install openyuanrong-datasystem
+```
+
+To build from source, follow the upstream project instructions:
+
+- Source repository: [openeuler/yuanrong-datasystem](https://atomgit.com/openeuler/yuanrong-datasystem)
+- Installation guide: [Yuanrong Datasystem Linux installation](https://pages.openeuler.openatom.cn/openyuanrong-datasystem/docs/zh-cn/latest/installation/installation_linux.html)
+
+The Ascend driver, CANN runtime, HCCL/HCCP, and NPU device network must also be
+configured on every machine participating in the transfer.
+
+## Check NPU Device IPv4
+
+Yuanrong TransferEngine with `protocol: "ascend"` uses the NPU device network,
+not only the host management IP. Each participating NPU must have an IPv4 HCCN
+device IP.
+
+Check the device IP, link, and health state:
+
+```bash
+for i in {0..7}; do
+  echo "===== NPU $i ====="
+  hccn_tool -i "$i" -ip -g
+  hccn_tool -i "$i" -link -g
+  hccn_tool -i "$i" -net_health -g
+done
+```
+
+Expected output for each NPU should include an IPv4 address and healthy link:
+
+```text
+ipaddr:10.86.1.80
+netmask:255.255.255.0
+link status: UP
+net health status: Success
+```
+
+If the output says `Get ipconf failed, because no ip was preset there!`, the NPU
+device IP is not configured and Ascend P2P will fail. If the address is IPv6,
+configure IPv4 for this connector path or validate that the underlying Yuanrong
+TransferEngine build supports the IPv6 HCCN path before using it.
+
+## Check Device-to-Device Connectivity
+
+After collecting the target-stage NPU IPv4 addresses, ping them from each
+source-stage NPU:
+
+```bash
+# Example: AR uses NPU 0..3 and DiT uses NPU 4..7.
+# Replace the destination IPs with the values reported by hccn_tool -i 4..7 -ip -g.
+for src in 0 1 2 3; do
+  for dst_ip in 10.86.1.78 10.86.1.79 10.86.1.80 10.86.1.81; do
+    echo "===== NPU $src -> $dst_ip ====="
+    hccn_tool -i "$src" -ping -g address "$dst_ip" || true
+  done
+done
+```
+
+The expected result is `0.00% packet loss` for every required source/target pair.
+This verifies basic HCCN IPv4 connectivity. TransferEngine can still fail later
+if RDMA/QP/HCCL setup is broken, but failed ping means the device network must be
+fixed first.
+
+## YAML Configuration
+
+For Ascend P2P, configure the connector with `protocol: "ascend"` and an NPU
+memory pool:
+
+```yaml
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 1
+  connectors:
+    yuanrong_te_connector:
+      name: YuanrongTransferEngineConnector
+      extra:
+        host: "auto"
+        zmq_port: 50051
+        rpc_port: "auto"
+        protocol: "ascend"
+        device_name: "auto"
+        memory_pool_size: 1073741824
+        memory_pool_device: "npu"
+  edges:
+    - from: 0
+      to: 1
+      window_size: -1
+```
+
+Important fields:
+
+| Parameter | Recommended Value | Notes |
+|---|---|---|
+| `host` | `"auto"` | Advertises a routable local host IP for ZMQ metadata exchange. |
+| `zmq_port` | `50051` | Base ZMQ port. The runtime applies stage/rank offsets. |
+| `rpc_port` | `"auto"` | Lets TransferEngine choose an available RPC port. |
+| `protocol` | `"ascend"` | Uses the Ascend TransferEngine backend. |
+| `device_name` | `"auto"` | Each TP worker resolves to its local logical NPU, for example `npu:0`, `npu:1`, etc. |
+| `memory_pool_device` | `"npu"` | Required for the Ascend fast path. This registers NPU memory with TransferEngine. |
+| `memory_pool_size` | `1073741824` | 1 GiB per worker is a practical starting point for KV-cache transfer validation. Increase if the transfer pool is exhausted. |
+
+`memory_pool_device: "cpu"` is rejected for
+`YuanrongTransferEngineConnector` with `protocol: "ascend"`. CPU pool registration
+can lead to Ascend P2P failures such as HCCL RA init errors, QP timeouts, or
+FFTS/SDMA runtime errors. `memory_pool_device: "cpu"` is common for
+`MooncakeTransferEngineConnector`, but Yuanrong TE on Ascend only supports
+`"npu"`.
+
+## 4-Card AR-to-DiT TP Example
+
+For AR 4-way TP sending KV cache to DiT 4-way TP, set the stage TP topology in
+both stage `omni_kv_config` blocks:
+
+```yaml
+stage_args:
+  - stage_id: 0
+    runtime:
+      devices: "0,1,2,3"
+    engine_args:
+      tensor_parallel_size: 4
+      omni_kv_config:
+        need_send_cache: true
+        rank_mapping:
+          from_tp: 4
+          to_tp: 4
+    output_connectors:
+      to_stage_1: yuanrong_te_connector
+
+  - stage_id: 1
+    runtime:
+      devices: "4,5,6,7"
+    engine_args:
+      parallel_config:
+        tensor_parallel_size: 4
+      omni_kv_config:
+        need_recv_cache: true
+        rank_mapping:
+          from_tp: 4
+          to_tp: 4
+    input_connectors:
+      from_stage_0: yuanrong_te_connector
+```
+
+In this topology each DiT rank receives the corresponding AR rank shard:
+
+```text
+DiT rank 0 <- AR rank 0
+DiT rank 1 <- AR rank 1
+DiT rank 2 <- AR rank 2
+DiT rank 3 <- AR rank 3
+```
+
+## Validate a Successful Run
+
+Search the inference log for Yuanrong TE receive-side messages:
+
+```bash
+grep -E "\[YR GET\]|Successfully received KV cache|transfer_engine get failed" inference.log
+```
+
+A successful Ascend fast-path transfer contains lines like:
+
+```text
+[YR GET] ... fast_path, zero-copy
+Successfully received KV cache ... across 1 key(s)
+Applied CFG KV caches
+[KV Reuse] ...
+```
+
+Sender-side `KV transfer OK` only means that the sender registered the payload
+and metadata. The actual P2P data transfer happens on the receiver side during
+`[YR GET]`.
+
+## Troubleshooting
+
+| Symptom | Likely Cause | Action |
+|---|---|---|
+| `IPv4 device IP not found` | NPU HCCN IPv4 is missing. | Run `hccn_tool -i <id> -ip -g` and configure IPv4 for every participating NPU. |
+| `P2PCommInitRootInfo failed`, `ra init failed`, `RA_QP_STATUS_TIMEOUT` | Ascend P2P/RDMA/QP setup failed. | First verify device IPv4, link health, and device-to-device ping. Then test a minimal Yuanrong TE pair outside vLLM. |
+| `fftsplus sdma error`, `context is abort` | Runtime SDMA/FFTS failure, often after an invalid or failed P2P transfer. | Use `memory_pool_device: "npu"` for Ascend TE, reduce `memory_pool_size` if OOM occurs, and isolate the failing NPU pair with a Yuanrong TE demo. |
+| Sender logs `KV transfer OK` but receiver never logs `[YR GET]` | Sender `put()` succeeded, receiver P2P read failed or timed out. | Check receiver logs and search for `transfer_engine get failed`. |
+| Pool allocation fails or OOM | Transfer pool is too small or too large for available NPU memory. | Start with 1 GiB (`1073741824`), reduce to 512 MiB for validation, or increase for larger KV payloads. |
@@ -15,8 +15,9 @@ from vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_c
    BufferAllocator,
    ManagedBuffer,
 )
+from vllm_omni.distributed.omni_connectors.kv_transfer_manager import KVCacheTransferData

-# All tests in this file are pure-CPU unit tests for the memory allocator.
+# Most tests in this file are pure-CPU unit tests; CUDA smoke tests are skip-guarded.
 pytestmark = [pytest.mark.cpu, pytest.mark.parallel, pytest.mark.core_model]


@@ -225,3 +226,70 @@ class TestManagedBuffer:
        # Space should be reusable
        new_offset = self.allocator.alloc(128)
        assert new_offset == offset
+
+    def test_to_bytes_cpu_pool(self):
+        """Smoke-test Mooncake's ManagedBuffer byte extraction on CPU pools."""
+        expected = bytes(range(64))
+        offset = self.allocator.alloc(len(expected))
+        buf = ManagedBuffer(self.allocator, offset, len(expected), self.pool)
+        self.pool[offset : offset + len(expected)] = torch.tensor(list(expected), dtype=torch.uint8)
+
+        try:
+            assert buf.to_bytes() == expected
+        finally:
+            buf.release()
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for GPU pool smoke test")
+    def test_to_bytes_cuda_pool(self):
+        """Smoke-test Mooncake's existing CUDA pool D2H byte path."""
+        expected = bytes(range(64))
+        allocator = BufferAllocator(total_size=4096, alignment=64)
+        pool = torch.zeros(4096, dtype=torch.uint8, device="cuda:0")
+        offset = allocator.alloc(len(expected))
+        buf = ManagedBuffer(allocator, offset, len(expected), pool)
+        pool[offset : offset + len(expected)].copy_(torch.tensor(list(expected), dtype=torch.uint8, device="cuda:0"))
+
+        try:
+            assert buf.to_bytes() == expected
+        finally:
+            buf.release()
+
+
+def _make_kv_payload(device: str | torch.device = "cpu") -> tuple[KVCacheTransferData, torch.Tensor, torch.Tensor]:
+    key_tensor = torch.arange(12, dtype=torch.float32, device=device).reshape(3, 4)
+    value_tensor = key_tensor + 100
+    payload = KVCacheTransferData(
+        request_id="moon-smoke",
+        layer_blocks={"key_cache": [key_tensor], "value_cache": [value_tensor]},
+        block_ids=[1],
+        metadata={"seq_len": 3},
+    )
+    return payload, key_tensor, value_tensor
+
+
+class TestMooncakePackedPayloadSmoke:
+    """Smoke tests for Mooncake raw payload helpers without real Mooncake/RDMA."""
+
+    def test_load_header_from_cpu_tensor(self):
+        payload, _, _ = _make_kv_payload()
+        raw = payload.to_bytes()
+        packed = torch.tensor(list(raw), dtype=torch.uint8)
+        header, data_start = KVCacheTransferData._load_header_from_tensor(packed)
+
+        assert header["rid"] == "moon-smoke"
+        assert data_start == 4 + int.from_bytes(raw[:4], "big")
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for packed CUDA smoke test")
+    def test_from_bytes_device_cuda_round_trip(self):
+        payload, key_tensor, value_tensor = _make_kv_payload(device="cuda:0")
+        packed = payload.to_gpu_tensor()
+
+        header, data_start = KVCacheTransferData._load_header_from_tensor(packed)
+        decoded = KVCacheTransferData.from_bytes_device(packed)
+
+        assert header["rid"] == "moon-smoke"
+        assert data_start == 4 + int.from_bytes(payload.to_bytes()[:4], "big")
+        assert decoded["layer_blocks"]["key_cache"][0].is_cuda
+        assert decoded["layer_blocks"]["value_cache"][0].is_cuda
+        assert torch.equal(decoded["layer_blocks"]["key_cache"][0].cpu(), key_tensor.cpu())
+        assert torch.equal(decoded["layer_blocks"]["value_cache"][0].cpu(), value_tensor.cpu())
@@ -15,6 +15,16 @@ connectors:
      device_name: ""
      memory_pool_size: 4294967296
      memory_pool_device: "cpu"
+  yuanrong_te_connector:
+    name: YuanrongTransferEngineConnector
+    extra:
+      host: "auto"
+      zmq_port: 50051
+      rpc_port: "auto"
+      protocol: "ascend"
+      device_name: "auto"
+      memory_pool_size: 1073741824
+      memory_pool_device: "npu"
  shared_memory_connector:
    name: SharedMemoryConnector
    extra:
@@ -11,6 +11,7 @@ from .omni_connectors import (
    OmniTransferConfig,
    SharedMemoryConnector,
    YuanrongConnector,
+    YuanrongTransferEngineConnector,
    load_omni_transfer_config,
 )

@@ -26,6 +27,7 @@ __all__ = [
    "MooncakeTransferEngineConnector",
    "SharedMemoryConnector",
    "YuanrongConnector",
+    "YuanrongTransferEngineConnector",
    # Utilities
    "load_omni_transfer_config",
 ]
@@ -6,6 +6,13 @@ from .connectors.mooncake_store_connector import MooncakeStoreConnector
 from .connectors.shm_connector import SharedMemoryConnector
 from .connectors.yuanrong_connector import YuanrongConnector

+try:
+    from vllm_omni.platforms.npu.omni_connectors.yuanrong_transfer_engine_connector import (
+        YuanrongTransferEngineConnector,
+    )
+except ImportError:
+    YuanrongTransferEngineConnector = None
+
 try:
    from .connectors.mooncake_transfer_engine_connector import MooncakeTransferEngineConnector
 except ImportError:
@@ -39,6 +46,7 @@ __all__ = [
    "MooncakeTransferEngineConnector",
    "SharedMemoryConnector",
    "YuanrongConnector",
+    "YuanrongTransferEngineConnector",
    # Utilities
    "load_omni_transfer_config",
    "initialize_connectors_from_config",
@@ -16,6 +16,7 @@ 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

@@ -66,160 +67,6 @@ class MooncakeAgentMetadata:
    lengths: list[int]


-class BufferAllocator:
-    """
-    Manages the allocation of memory segments within the registered pool.
-    Thread-safe implementation using a simple free list.
-    """
-
-    def __init__(self, total_size: int, alignment: int = 4096):
-        self.total_size = total_size
-        self.alignment = alignment
-        self.lock = threading.Lock()
-        # Free list: [(start, size), ...] sorted by start
-        self.free_blocks = [(0, total_size)]
-
-    def alloc(self, size: int) -> int:
-        """
-        Allocates a block of 'size' bytes.
-        Returns the starting offset.
-        """
-        # Align size upwards
-        aligned_size = (size + self.alignment - 1) // self.alignment * self.alignment
-
-        with self.lock:
-            for i, (start, block_size) in enumerate(self.free_blocks):
-                if block_size >= aligned_size:
-                    # Found a block
-                    new_start = start + aligned_size
-                    new_size = block_size - aligned_size
-
-                    if new_size > 0:
-                        self.free_blocks[i] = (new_start, new_size)
-                    else:
-                        self.free_blocks.pop(i)
-                    return start
-
-        raise MemoryError(f"Out of memory in buffer pool. Requested {size} bytes (aligned {aligned_size}).")
-
-    def free(self, offset: int, size: int):
-        """
-        Frees a previously allocated block.
-        """
-        aligned_size = (size + self.alignment - 1) // self.alignment * self.alignment
-
-        with self.lock:
-            # Check for double-free and corruption
-            for start, length in self.free_blocks:
-                # Case 1: Exact match = double free, safe to ignore
-                if offset == start and aligned_size == length:
-                    logger.warning(f"Double free detected at offset {offset}, size {aligned_size}. Ignoring.")
-                    return
-                # Case 2: Block is fully contained within an existing free block = also double free
-                # This happens when the block was freed and then merged with adjacent blocks
-                if offset >= start and offset + aligned_size <= start + length:
-                    logger.warning(
-                        f"Double free detected: block {offset}-{offset + aligned_size} "
-                        f"is already within free block {start}-{start + length}. Ignoring."
-                    )
-                    return
-                # Case 3: Partial overlap (but not fully contained) = memory corruption
-                if not (offset + aligned_size <= start or start + length <= offset):
-                    raise RuntimeError(
-                        f"Memory corruption detected: freeing {offset}-{offset + aligned_size} "
-                        f"partially overlaps with free block {start}-{start + length}"
-                    )
-
-            self.free_blocks.append((offset, aligned_size))
-            self.free_blocks.sort()  # Sort by offset
-
-            # Merge adjacent blocks
-            i = 0
-            while i < len(self.free_blocks) - 1:
-                curr_start, curr_size = self.free_blocks[i]
-                next_start, next_size = self.free_blocks[i + 1]
-
-                if curr_start + curr_size == next_start:
-                    self.free_blocks[i] = (curr_start, curr_size + next_size)
-                    self.free_blocks.pop(i + 1)
-                else:
-                    i += 1
-
-
-class ManagedBuffer:
-    """
-    A temporary view into the global memory pool.
-    Must be kept alive while the data view is being used.
-    """
-
-    def __init__(self, allocator: BufferAllocator, offset: int, size: int, pool_tensor: torch.Tensor):
-        self.allocator = allocator
-        self.offset = offset
-        self.size = size
-        self.pool_tensor = pool_tensor
-        self._released = False
-
-    def release(self):
-        """Explicitly release the buffer back to the pool."""
-        if not self._released:
-            self.allocator.free(self.offset, self.size)
-            self._released = True
-
-    def __del__(self):
-        self.release()
-
-    def __enter__(self):
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        self.release()
-
-    @property
-    def tensor(self) -> torch.Tensor:
-        """
-        Returns a 1D uint8 zero-copy view of the buffer.
-        """
-        return self.pool_tensor[self.offset : self.offset + self.size]
-
-    def as_tensor(self, dtype: torch.dtype, shape: tuple) -> torch.Tensor:
-        """
-        Returns a typed, shaped zero-copy view.
-        Validates size, shape, and alignment.
-        """
-        itemsize = torch.tensor([], dtype=dtype).element_size()
-
-        # Calculate expected size
-        expected_bytes = itemsize
-        for dim in shape:
-            if dim < 0:
-                raise ValueError("Dynamic dimension (-1) is not supported in as_tensor")
-            expected_bytes *= dim
-
-        if expected_bytes != self.size:
-            raise ValueError(
-                f"Shape {shape} with dtype {dtype} requires {expected_bytes} bytes, but buffer size is {self.size}"
-            )
-
-        # Check alignment (offset must be divisible by itemsize)
-        if self.offset % itemsize != 0:
-            raise RuntimeError(f"Buffer offset {self.offset} is not aligned for dtype {dtype} (itemsize {itemsize})")
-
-        raw_view = self.tensor
-        # view() requires contiguous memory, slice of contiguous tensor is contiguous
-        typed_view = raw_view.view(dtype)
-        return typed_view.reshape(shape)
-
-    def to_bytes(self) -> bytes:
-        """
-        Returns a copy of the data as python bytes.
-        Performs D2H copy if pool is on GPU.
-        """
-        t = self.tensor
-        if t.is_cuda:
-            t = t.cpu()
-        return t.numpy().tobytes()
-
-
 class MooncakeTransferEngineConnector(OmniConnectorBase):
    """
    OmniConnector implementation using Mooncake Transfer Engine with a managed memory pool.
@@ -0,0 +1,10 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Compatibility import for the NPU-only Yuanrong TransferEngine connector."""
+
+from vllm_omni.platforms.npu.omni_connectors.yuanrong_transfer_engine_connector import (
+    YuanrongTransferEngineConnector,
+)
+
+__all__ = ["YuanrongTransferEngineConnector"]
@@ -91,6 +91,18 @@ def _create_yuanrong_connector(config: dict[str, Any]) -> OmniConnectorBase:
    return YuanrongConnector(config)


+def _create_yuanrong_transfer_engine_connector(config: dict[str, Any]) -> OmniConnectorBase:
+    try:
+        from vllm_omni.platforms.npu.omni_connectors import YuanrongTransferEngineConnector
+    except ImportError as exc:
+        raise ImportError(
+            "YuanrongTransferEngineConnector is only available in the NPU platform "
+            "environment. Install the Ascend/Yuanrong runtime dependencies before "
+            "using this connector."
+        ) from exc
+    return YuanrongTransferEngineConnector(config)
+
+
 def _create_mooncake_transfer_engine_connector(config: dict[str, Any]) -> OmniConnectorBase:
    try:
        from .connectors.mooncake_transfer_engine_connector import MooncakeTransferEngineConnector
@@ -107,5 +119,6 @@ OmniConnectorFactory.register_connector("MooncakeStoreConnector", _create_moonca
 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)
 # Backward-compatible aliases – will be removed in the future
 OmniConnectorFactory.register_connector("MooncakeConnector", _create_mooncake_store_connector)
@@ -13,7 +13,8 @@ import torch
 from vllm.logger import init_logger

 from .factory import OmniConnectorFactory
-from .utils.config import ConnectorSpec
+from .utils.config import TRANSFER_ENGINE_CONNECTOR_NAMES, ConnectorSpec
+from .utils.env import expand_env_int
 from .utils.initialization import KV_RANK_PORT_STRIDE
 from .utils.kv_utils import (
    KVTPTopology,
@@ -106,7 +107,7 @@ class KVCacheTransferData:
                t = tensor.detach().contiguous()
                if cpu:
                    t = t.cpu()
-                elif device is None and t.is_cuda:
+                elif device is None and getattr(t.device, "type", "cpu") != "cpu":
                    device = t.device
                nbytes = t.numel() * t.element_size()
                tensors_desc.append(
@@ -143,10 +144,10 @@ class KVCacheTransferData:
        return b"".join([self._build_header_bytes(tensors_desc)] + chunks)

    def to_gpu_tensor(self) -> torch.Tensor:
-        """Convert to a packed GPU tensor for raw-data connectors."""
+        """Convert to a packed device tensor for raw-data connectors."""
        tensors_desc, chunks, data_offset, device = self._build_tensors_desc(cpu=False)
        if device is None:
-            raise RuntimeError("No CUDA tensors found, use to_bytes() instead")
+            raise RuntimeError("No device tensors found, use to_bytes() instead")
        header_prefix = self._build_header_bytes(tensors_desc)
        output = torch.empty(len(header_prefix) + data_offset, dtype=torch.uint8, device=device)
        output[: len(header_prefix)].copy_(torch.frombuffer(bytearray(header_prefix), dtype=torch.uint8))
@@ -169,19 +170,19 @@ class KVCacheTransferData:
        return json.loads(bytes(raw_mv[4 : 4 + header_len])), raw_mv[4 + header_len :]

    @staticmethod
-    def _load_header_from_tensor(gpu_tensor: torch.Tensor) -> tuple[dict[str, Any], int]:
-        if gpu_tensor.dtype != torch.uint8 or gpu_tensor.dim() != 1:
-            raise ValueError("Packed GPU KV payload must be a 1-D uint8 tensor")
+    def _load_header_from_tensor(tensor: torch.Tensor) -> tuple[dict[str, Any], int]:
+        if tensor.dtype != torch.uint8 or tensor.dim() != 1:
+            raise ValueError("Packed device KV payload must be a 1-D uint8 tensor")

-        total_bytes = int(gpu_tensor.numel())
+        total_bytes = int(tensor.numel())
        if total_bytes < 4:
            raise ValueError("Corrupted KV payload: missing 4-byte header length")

-        header_len = struct.unpack(">I", gpu_tensor[:4].cpu().numpy().tobytes())[0]
+        header_len = struct.unpack(">I", tensor[:4].cpu().numpy().tobytes())[0]
        if header_len > total_bytes - 4:
            raise ValueError(f"Corrupted KV payload: header_len={header_len} exceeds buffer size={total_bytes}")

-        header_bytes = gpu_tensor[4 : 4 + header_len].cpu().numpy().tobytes()
+        header_bytes = tensor[4 : 4 + header_len].cpu().numpy().tobytes()
        return json.loads(header_bytes), 4 + header_len

    @staticmethod
@@ -260,17 +261,22 @@ class KVCacheTransferData:
        return KVCacheTransferData._populate_caches(header, _get)

    @staticmethod
-    def from_bytes_gpu(gpu_tensor: torch.Tensor) -> dict[str, Any]:
-        """Reconstruct KV cache data from a packed GPU tensor."""
-        header, data_start = KVCacheTransferData._load_header_from_tensor(gpu_tensor)
-        data_len = int(gpu_tensor.numel()) - data_start
+    def from_bytes_device(tensor: torch.Tensor) -> dict[str, Any]:
+        """Reconstruct KV cache data from a packed device tensor."""
+        header, data_start = KVCacheTransferData._load_header_from_tensor(tensor)
+        data_len = int(tensor.numel()) - data_start

        def _get(info: dict) -> torch.Tensor:
            offset, nbytes = KVCacheTransferData._validate_tensor_span(info["n"], info, data_len)
-            return gpu_tensor[data_start + offset : data_start + offset + nbytes].clone()
+            return tensor[data_start + offset : data_start + offset + nbytes].clone()

        return KVCacheTransferData._populate_caches(header, _get)

+    @staticmethod
+    def from_bytes_gpu(tensor: torch.Tensor) -> dict[str, Any]:
+        """Compatibility alias for callers using the old GPU-specific name."""
+        return KVCacheTransferData.from_bytes_device(tensor)
+

 class OmniKVTransferManager:
    """Unified management for OmniConnector and KV cache transfer.
@@ -405,8 +411,8 @@ class OmniKVTransferManager:
            if cfg and (c_type := cfg.get("type")):
                try:
                    c_extra = {k: v for k, v in cfg.items() if k != "type"}
-                    if c_type == "MooncakeTransferEngineConnector":
-                        base_port = c_extra.get("zmq_port", 50051)
+                    if c_type in TRANSFER_ENGINE_CONNECTOR_NAMES:
+                        base_port = expand_env_int(c_extra.get("zmq_port", 50051), "zmq_port")
                        c_extra["from_stage"] = (
                            str(self.config.from_stage) if self.config.from_stage is not None else "0"
                        )
@@ -429,8 +435,13 @@ class OmniKVTransferManager:
                            c_extra["zmq_port"] = zmq_port
                        elif self.config.need_recv_cache:
                            c_extra["role"] = "receiver"
-                            c_extra.setdefault("sender_host", c_extra.get("host", "127.0.0.1"))
-                            c_extra.setdefault("sender_zmq_port", zmq_port)
+                            # Receiver-side sender endpoints are request-scoped.
+                            # They are attached by the orchestrator as
+                            # kv_sender_info and applied in update_sender_info().
+                            # Do not derive sender_host from this receiver's
+                            # local YAML host; on multi-node runs that points at
+                            # the wrong process. Explicit sender_* YAML values
+                            # are still preserved for standalone connector use.

                    logger.info(
                        "Initializing OmniConnector type=%s role=%s",
@@ -1058,8 +1069,8 @@ class OmniKVTransferManager:
                        managed_buffer = raw_data
                        try:
                            buf_tensor = raw_data.tensor
-                            if buf_tensor.is_cuda:
-                                data = KVCacheTransferData.from_bytes_gpu(buf_tensor)
+                            if getattr(buf_tensor.device, "type", "cpu") != "cpu":
+                                data = KVCacheTransferData.from_bytes_device(buf_tensor)
                                raw_data.release()
                                managed_buffer = None
                            else:
@@ -1075,7 +1086,10 @@ class OmniKVTransferManager:
                    elif isinstance(raw_data, (bytes, bytearray)):
                        data = KVCacheTransferData.from_bytes(raw_data)
                    elif isinstance(raw_data, torch.Tensor) and raw_data.dtype == torch.uint8 and raw_data.dim() == 1:
-                        data = KVCacheTransferData.from_bytes(raw_data.cpu().numpy().tobytes())
+                        if getattr(raw_data.device, "type", "cpu") != "cpu":
+                            data = KVCacheTransferData.from_bytes_device(raw_data)
+                        else:
+                            data = KVCacheTransferData.from_bytes(raw_data.numpy().tobytes())
                    else:
                        data = raw_data

@@ -8,6 +8,13 @@ from .logging import get_connector_logger

 logger = get_connector_logger(__name__)

+TRANSFER_ENGINE_CONNECTOR_NAMES = frozenset(
+    {
+        "MooncakeTransferEngineConnector",
+        "YuanrongTransferEngineConnector",
+    }
+)
+

@dataclass
 class ConnectorSpec:
@@ -0,0 +1,47 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import os
+import re
+from typing import Any
+
+AUTO_HOST_VALUES = {"", "auto", "*", "0.0.0.0", "::"}
+AUTO_DEVICE_VALUES = {"", "auto"}
+
+_UNEXPANDED_ENV_VAR_PATTERN = re.compile(r"\$\{([^}]+)\}|\$([A-Za-z_][A-Za-z0-9_]*)|%([^%]+)%")
+
+
+class EnvVarExpansionError(ValueError):
+    pass
+
+
+def _unexpanded_env_vars(value: str) -> list[str]:
+    names: list[str] = []
+    for match in _UNEXPANDED_ENV_VAR_PATTERN.finditer(value):
+        name = next(group for group in match.groups() if group)
+        if name not in names:
+            names.append(name)
+    return names
+
+
+def expand_env_value(value: Any, *, strict: bool = False, field_name: str = "value") -> Any:
+    if not isinstance(value, str):
+        return value
+
+    expanded = os.path.expandvars(value)
+    if strict:
+        unresolved = _unexpanded_env_vars(expanded)
+        if unresolved:
+            names = ", ".join(unresolved)
+            raise EnvVarExpansionError(
+                f"{field_name} references unset environment variable(s): {names}. Original value: {value!r}"
+            )
+    return expanded
+
+
+def expand_env_int(value: Any, field_name: str) -> int:
+    expanded = expand_env_value(value, strict=True, field_name=field_name)
+    try:
+        return int(expanded)
+    except (TypeError, ValueError) as exc:
+        raise ValueError(f"{field_name} must be an integer after environment expansion, got {expanded!r}") from exc
@@ -10,7 +10,9 @@ from pathlib import Path
 from typing import TYPE_CHECKING, Any

 from ..factory import OmniConnectorFactory
-from .config import ConnectorSpec, OmniTransferConfig
+from .config import TRANSFER_ENGINE_CONNECTOR_NAMES, ConnectorSpec, OmniTransferConfig
+from .env import expand_env_int
+from .local_rank import get_connector_local_rank
 from .logging import get_connector_logger

 if TYPE_CHECKING:
@@ -93,18 +95,20 @@ def create_connectors_from_config(
    for edge_key, connector_spec in connectors_config.items():
        from_stage, to_stage = edge_key
        try:
-            if connector_spec.name == "MooncakeTransferEngineConnector":
+            if connector_spec.name in TRANSFER_ENGINE_CONNECTOR_NAMES:
                extra = dict(connector_spec.extra) if connector_spec.extra else {}
-                base_port = extra.get("zmq_port", 50051)
+                base_port = expand_env_int(extra.get("zmq_port", 50051), "zmq_port")
                try:
                    stage_offset = int(from_stage)
                except (TypeError, ValueError):
                    stage_offset = 0

+                local_rank = 0 if str(caller_stage_id) == "orchestrator" else get_connector_local_rank()
                if str(caller_stage_id) == "orchestrator":
-                    adjusted_port = base_port + orchestrator_port_offset + stage_offset
+                    rank0_port = base_port + orchestrator_port_offset + stage_offset
                else:
-                    adjusted_port = base_port + port_offset + stage_offset
+                    rank0_port = base_port + port_offset + stage_offset
+                adjusted_port = rank0_port + local_rank * KV_RANK_PORT_STRIDE
                extra["zmq_port"] = adjusted_port

                if is_sender is not None:
@@ -0,0 +1,49 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import os
+
+_LOCAL_RANK_ENV_NAMES = (
+    "LOCAL_RANK",
+    "VLLM_LOCAL_RANK",
+    "OMPI_COMM_WORLD_LOCAL_RANK",
+    "MV2_COMM_WORLD_LOCAL_RANK",
+)
+_GLOBAL_RANK_ENV_NAMES = ("RANK", "RANK_ID")
+
+
+def _visible_device_count() -> int | None:
+    visible = os.environ.get("ASCEND_RT_VISIBLE_DEVICES") or os.environ.get("CUDA_VISIBLE_DEVICES")
+    if not visible:
+        return None
+    devices = [item.strip() for item in visible.split(",") if item.strip()]
+    return len(devices) or None
+
+
+def _get_tensor_model_parallel_rank() -> int:
+    from vllm.distributed.parallel_state import get_tensor_model_parallel_rank
+
+    return int(get_tensor_model_parallel_rank())
+
+
+def get_connector_local_rank() -> int:
+    try:
+        return _get_tensor_model_parallel_rank()
+    except Exception:
+        pass
+
+    for env_name in _LOCAL_RANK_ENV_NAMES:
+        try:
+            return int(os.environ[env_name])
+        except (KeyError, TypeError, ValueError):
+            pass
+
+    visible_count = _visible_device_count()
+    for env_name in _GLOBAL_RANK_ENV_NAMES:
+        try:
+            rank = int(os.environ[env_name])
+        except (KeyError, TypeError, ValueError):
+            continue
+        return rank % visible_count if visible_count else rank
+
+    return 0
@@ -0,0 +1,153 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import threading
+
+import torch
+
+from .logging import get_connector_logger
+
+logger = get_connector_logger(__name__)
+
+
+class BufferAllocator:
+    """
+    Manages the allocation of memory segments within a registered pool.
+    Thread-safe implementation using a simple free list.
+    """
+
+    def __init__(self, total_size: int, alignment: int = 4096):
+        self.total_size = total_size
+        self.alignment = alignment
+        self.lock = threading.Lock()
+        self.free_blocks = [(0, total_size)]
+
+    def alloc(self, size: int) -> int:
+        """
+        Allocates a block of *size* bytes and returns the starting offset.
+        """
+        aligned_size = (size + self.alignment - 1) // self.alignment * self.alignment
+
+        with self.lock:
+            for i, (start, block_size) in enumerate(self.free_blocks):
+                if block_size >= aligned_size:
+                    new_start = start + aligned_size
+                    new_size = block_size - aligned_size
+
+                    if new_size > 0:
+                        self.free_blocks[i] = (new_start, new_size)
+                    else:
+                        self.free_blocks.pop(i)
+                    return start
+
+        raise MemoryError(f"Out of memory in buffer pool. Requested {size} bytes (aligned {aligned_size}).")
+
+    def free(self, offset: int, size: int) -> None:
+        """
+        Frees a previously allocated block.
+        """
+        aligned_size = (size + self.alignment - 1) // self.alignment * self.alignment
+
+        with self.lock:
+            for start, length in self.free_blocks:
+                if offset == start and aligned_size == length:
+                    logger.warning("Double free detected at offset %s, size %s. Ignoring.", offset, aligned_size)
+                    return
+                if offset >= start and offset + aligned_size <= start + length:
+                    logger.warning(
+                        "Double free detected: block %s-%s is already within free block %s-%s. Ignoring.",
+                        offset,
+                        offset + aligned_size,
+                        start,
+                        start + length,
+                    )
+                    return
+                if not (offset + aligned_size <= start or start + length <= offset):
+                    raise RuntimeError(
+                        f"Memory corruption detected: freeing {offset}-{offset + aligned_size} "
+                        f"partially overlaps with free block {start}-{start + length}"
+                    )
+
+            self.free_blocks.append((offset, aligned_size))
+            self.free_blocks.sort()
+
+            i = 0
+            while i < len(self.free_blocks) - 1:
+                curr_start, curr_size = self.free_blocks[i]
+                next_start, next_size = self.free_blocks[i + 1]
+
+                if curr_start + curr_size == next_start:
+                    self.free_blocks[i] = (curr_start, curr_size + next_size)
+                    self.free_blocks.pop(i + 1)
+                else:
+                    i += 1
+
+
+class ManagedBuffer:
+    """
+    A temporary view into a global memory pool.
+    Must be kept alive while the data view is being used.
+    """
+
+    def __init__(self, allocator: BufferAllocator, offset: int, size: int, pool_tensor: torch.Tensor):
+        self.allocator = allocator
+        self.offset = offset
+        self.size = size
+        self.pool_tensor = pool_tensor
+        self._released = False
+
+    def release(self) -> None:
+        """Explicitly release the buffer back to the pool."""
+        if not self._released:
+            self.allocator.free(self.offset, self.size)
+            self._released = True
+
+    def __del__(self):
+        self.release()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.release()
+
+    @property
+    def tensor(self) -> torch.Tensor:
+        """
+        Returns a 1D uint8 zero-copy view of the buffer.
+        """
+        return self.pool_tensor[self.offset : self.offset + self.size]
+
+    def as_tensor(self, dtype: torch.dtype, shape: tuple) -> torch.Tensor:
+        """
+        Returns a typed, shaped zero-copy view.
+        Validates size, shape, and alignment.
+        """
+        itemsize = torch.tensor([], dtype=dtype).element_size()
+
+        expected_bytes = itemsize
+        for dim in shape:
+            if dim < 0:
+                raise ValueError("Dynamic dimension (-1) is not supported in as_tensor")
+            expected_bytes *= dim
+
+        if expected_bytes != self.size:
+            raise ValueError(
+                f"Shape {shape} with dtype {dtype} requires {expected_bytes} bytes, but buffer size is {self.size}"
+            )
+
+        if self.offset % itemsize != 0:
+            raise RuntimeError(f"Buffer offset {self.offset} is not aligned for dtype {dtype} (itemsize {itemsize})")
+
+        typed_view = self.tensor.view(dtype)
+        return typed_view.reshape(shape)
+
+    def to_bytes(self) -> bytes:
+        """
+        Returns a copy of the data as Python bytes.
+        Performs D2H copy if the pool is on a device.
+        """
+        t = self.tensor
+        if getattr(t.device, "type", "cpu") != "cpu":
+            t = t.cpu()
+        return t.numpy().tobytes()
@@ -8,6 +8,7 @@ from __future__ import annotations

 import inspect
 import multiprocessing.connection
+import os
 import socket
 import threading
 import weakref
@@ -20,6 +21,9 @@ from vllm.v1.engine import EngineCoreRequest
 from vllm.v1.engine.core_client import AsyncMPClient, DPLBAsyncMPClient
 from vllm.v1.engine.exceptions import EngineDeadError

+from vllm_omni.distributed.omni_connectors.utils.config import (
+    TRANSFER_ENGINE_CONNECTOR_NAMES,
+)
 from vllm_omni.distributed.omni_connectors.utils.initialization import (
    KV_TRANSFER_PORT_OFFSET,
 )
@@ -331,11 +335,13 @@ class StageEngineCoreClientBase(StageClientBase):
        if sender_host is not None:
            self._kv_sender_host = sender_host

+        connector_type = connector_config.get("type")
        sender_port = connector_config.get("sender_zmq_port")
-        if sender_port is None:
+        if connector_type in TRANSFER_ENGINE_CONNECTOR_NAMES or sender_port is None:
            base_port = connector_config.get("zmq_port")
            if base_port is None:
                return
+            base_port = os.path.expandvars(str(base_port))

            omni_kv_config = getattr(self, "_omni_kv_config", None)
            from_stage = self.stage_id
@@ -0,0 +1,11 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+try:
+    from vllm_omni.platforms.npu.omni_connectors.yuanrong_transfer_engine_connector import (
+        YuanrongTransferEngineConnector,
+    )
+except ImportError:
+    YuanrongTransferEngineConnector = None
+
+__all__ = ["YuanrongTransferEngineConnector"]
@@ -0,0 +1,862 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import queue
+import socket
+import threading
+import time as _time_mod
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass
+from typing import Any
+
+import msgspec
+import torch
+import zmq
+
+from vllm_omni.distributed.omni_connectors.connectors.base import OmniConnectorBase
+from vllm_omni.distributed.omni_connectors.utils.env import (
+    AUTO_DEVICE_VALUES,
+    AUTO_HOST_VALUES,
+    expand_env_int,
+    expand_env_value,
+)
+from vllm_omni.distributed.omni_connectors.utils.local_rank import get_connector_local_rank
+from vllm_omni.distributed.omni_connectors.utils.logging import get_connector_logger
+from vllm_omni.distributed.omni_connectors.utils.memory_pool import BufferAllocator, ManagedBuffer
+from vllm_omni.distributed.omni_connectors.utils.serialization import OmniSerializer
+
+logger = get_connector_logger(__name__)
+
+try:
+    from yr.datasystem import TransferEngine
+except ImportError:
+    TransferEngine = None
+
+_BUFFER_TTL_SECONDS = 300
+
+QUERY_INFO = b"query_info"
+CLEANUP_KEY = b"cleanup_key"
+ACK_OK = b"ack_ok"
+ACK_ERROR = b"ack_error"
+INFO_NOT_FOUND = b"info_not_found"
+
+
+@dataclass
+class QueryRequest:
+    """Request to query metadata for a specific key."""
+
+    request_id: str
+
+
+@dataclass
+class CleanupRequest:
+    """Response containing metadata for a request."""
+
+    request_id: str
+
+
+@dataclass
+class QueryResponse:
+    request_id: str
+    rpc_port: int
+    data_size: int
+    is_fast_path: bool
+    src_addrs: list[int]
+    lengths: list[int]
+
+
+def _resolve_device_name(configured_device: Any, protocol: str) -> str:
+    device_name = str(expand_env_value(configured_device or "auto")).strip()
+    if device_name.lower() not in AUTO_DEVICE_VALUES:
+        return device_name
+    if protocol.lower() in {"ascend", "hccl", "npu"}:
+        return f"npu:{get_connector_local_rank()}"
+    return ""
+
+
+def _resolve_pool_device(configured_device: Any) -> str:
+    pool_device = str(expand_env_value(configured_device or "npu")).strip().lower()
+    if pool_device in AUTO_DEVICE_VALUES:
+        return "npu"
+    if not (pool_device == "npu" or pool_device.startswith("npu:")):
+        raise ValueError(
+            f"YuanrongTransferEngineConnector requires an NPU memory pool. Got memory_pool_device={pool_device!r}."
+        )
+    return pool_device
+
+
+class YuanrongTransferEngineConnector(OmniConnectorBase):
+    """Pull-based connector backed by Yuanrong transfer_engine."""
+
+    supports_raw_data: bool = True
+
+    def __init__(self, config: dict[str, Any]):
+        if TransferEngine is None:
+            raise ImportError("Yuanrong is not available.")
+
+        self._closed = False
+        self._bind_error: Exception | None = None
+        self._stop_event = threading.Event()
+        self._listener_thread: threading.Thread | None = None
+        self._listener_ready = threading.Event()
+        self._control_executor: ThreadPoolExecutor | None = None
+        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()
+
+        self._metrics = {
+            "puts": 0,
+            "gets": 0,
+            "bytes_transferred": 0,
+            "errors": 0,
+            "timeouts": 0,
+        }
+
+        self.config = {key: expand_env_value(value) for key, value in config.items()}
+        config = self.config
+
+        raw_host_config = config.get("host")
+        host_config = "auto" if raw_host_config is None else str(raw_host_config).strip()
+        self.host = self._get_local_ip() if host_config.lower() in AUTO_HOST_VALUES else host_config
+        self.zmq_port = self._resolve_port(config.get("zmq_port"), self.host, "zmq_port")
+        self.rpc_port = self._resolve_port(config.get("rpc_port"), self.host, "rpc_port")
+        self.protocol = str(config.get("protocol", "ascend"))
+        self.device_name = _resolve_device_name(config.get("device_name", "auto"), self.protocol)
+        self.pool_size = int(config.get("memory_pool_size", 1024**3))
+        self.pool_device = _resolve_pool_device(config.get("memory_pool_device", "npu"))
+        self.sender_host = config.get("sender_host")
+        sender_zmq_port = config.get("sender_zmq_port")
+        self.sender_zmq_port = self._resolve_optional_port(sender_zmq_port, "sender_zmq_port")
+
+        role = str(config.get("role", "sender")).lower()
+        if role not in {"sender", "receiver"}:
+            raise ValueError(
+                f"Invalid role={role!r} for YuanrongTransferEngineConnector. Expected 'sender' or 'receiver'."
+            )
+        self.can_put = role == "sender"
+
+        self.engine = TransferEngine()
+        local_endpoint = f"{self.host}:{self.rpc_port}"
+        self._ensure_result_ok(
+            self.engine.initialize(local_endpoint, self.protocol, self.device_name),
+            f"initialize transfer engine on {local_endpoint}",
+        )
+        resolved_rpc_port = int(self.engine.get_rpc_port())
+        if resolved_rpc_port > 0:
+            self.rpc_port = resolved_rpc_port
+
+        logger.info(
+            "YuanrongTransferEngineConnector initialized at %s:%s (protocol=%s, device=%s)",
+            self.host,
+            self.rpc_port,
+            self.protocol,
+            self.device_name,
+        )
+
+        try:
+            self.pool = torch.empty(self.pool_size, dtype=torch.uint8, device=self.pool_device)
+            self.base_ptr = int(self.pool.data_ptr())
+            self._ensure_result_ok(
+                self.engine.register_memory(self.base_ptr, self.pool_size),
+                f"register memory pool ({self.pool_size} bytes)",
+            )
+        except Exception as exc:
+            logger.error("Failed to allocate/register transfer_engine pool: %s", exc)
+            raise
+
+        self.allocator = BufferAllocator(self.pool_size, alignment=4096)
+        self.zmq_ctx = zmq.Context()
+        self._control_executor = ThreadPoolExecutor(max_workers=4, thread_name_prefix="yuanrong-transfer-control")
+
+        if self.can_put:
+            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"YuanrongTransferEngineConnector failed to bind ZMQ on "
+                    f"{self.host}:{self.zmq_port}: {self._bind_error}"
+                ) from self._bind_error
+        else:
+            logger.info(
+                "YuanrongTransferEngineConnector started in receiver mode (sender_host=%r, sender_zmq_port=%r)",
+                self.sender_host,
+                self.sender_zmq_port,
+            )
+
+    @staticmethod
+    def _ensure_result_ok(result: Any, action: str) -> None:
+        if result.is_error():
+            raise RuntimeError(f"{action} failed: {result.to_string()}")
+
+    @staticmethod
+    def _resolve_optional_port(configured_port: Any, field_name: str) -> int | None:
+        configured_port = expand_env_value(configured_port, strict=True, field_name=field_name)
+        if configured_port is None or str(configured_port).strip().lower() in {"", "auto"}:
+            return None
+        return expand_env_int(configured_port, field_name)
+
+    @classmethod
+    def _resolve_port(cls, configured_port: Any, host: str, field_name: str) -> int:
+        resolved_port = cls._resolve_optional_port(configured_port, field_name)
+        if resolved_port is not None:
+            return resolved_port
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+            sock.bind((host, 0))
+            return int(sock.getsockname()[1])
+
+    @staticmethod
+    def _get_local_ip() -> str:
+        try:
+            with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock:
+                sock.connect(("8.8.8.8", 80))
+                return sock.getsockname()[0]
+        except Exception:
+            try:
+                return socket.gethostbyname(socket.gethostname())
+            except Exception:
+                return "127.0.0.1"
+
+    @staticmethod
+    def _sync_tensor_device(tensor: torch.Tensor | None) -> None:
+        if tensor is None:
+            return
+        device = getattr(tensor, "device", None)
+        if device is None or getattr(device, "type", "cpu") == "cpu":
+            return
+
+        backend = getattr(torch, device.type, None)
+        if backend is None:
+            return
+
+        current_stream = getattr(backend, "current_stream", None)
+        if callable(current_stream):
+            try:
+                current_stream(device).synchronize()
+                return
+            except TypeError:
+                try:
+                    current_stream().synchronize()
+                    return
+                except Exception:
+                    pass
+            except Exception:
+                pass
+
+        synchronize = getattr(backend, "synchronize", None)
+        if callable(synchronize):
+            try:
+                synchronize(device)
+            except TypeError:
+                synchronize()
+
+    def get_connection_info(self) -> dict[str, Any]:
+        return {
+            "host": self.host,
+            "zmq_port": self.zmq_port,
+            "rpc_port": self.rpc_port,
+            "can_put": self.can_put,
+        }
+
+    def update_sender_info(
+        self,
+        sender_host: str,
+        sender_zmq_port: int,
+    ) -> None:
+        self.sender_host = sender_host
+        self.sender_zmq_port = int(sender_zmq_port)
+
+    def _get_req_socket(self, zmq_addr: str, timeout_ms: int) -> zmq.Socket:
+        cache: dict[str, zmq.Socket] = getattr(self._req_local, "cache", None)  # type: ignore[assignment]
+        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] = getattr(self._req_local, "cache", None)  # type: ignore[assignment]
+        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 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 data: YuanrongTransferEngineConnector is closed")
+
+        if not self.can_put:
+            self._metrics["errors"] += 1
+            logger.warning("Rejecting put for %s: connector is in receiver mode", put_key)
+            return False, 0, None
+
+        put_key = self._make_key(put_key, from_stage, to_stage)
+
+        try:
+            src_addr = 0
+            size = 0
+            holder = None
+            is_fast_path = True
+            should_release_holder = False
+
+            if isinstance(data, bytes) and len(data) == 0:
+                self._metrics["errors"] += 1
+                logger.warning("Rejecting put for %s: empty bytes payload", put_key)
+                return False, 0, None
+            if isinstance(data, torch.Tensor) and data.nbytes == 0:
+                self._metrics["errors"] += 1
+                logger.warning("Rejecting put for %s: zero-size tensor", put_key)
+                return False, 0, None
+
+            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():
+                    logger.warning("ManagedBuffer from different pool detected. Falling back to copy path.")
+                    data = data.tensor.contiguous()
+                else:
+                    src_addr = self.base_ptr + data.offset
+                    size = data.size
+                    holder = data
+                    should_release_holder = 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_holder = True
+                except MemoryError:
+                    self._metrics["errors"] += 1
+                    logger.error("Pool exhausted, cannot put data size %s", size)
+                    return False, 0, None
+
+                try:
+                    dst_tensor = holder.tensor
+                    if isinstance(data, torch.Tensor):
+                        if not data.is_contiguous():
+                            data = data.contiguous()
+                        src_view = data.view(torch.uint8).flatten()
+                        dst_tensor.copy_(src_view, non_blocking=src_view.device != dst_tensor.device)
+                        self._sync_tensor_device(src_view)
+                        self._sync_tensor_device(dst_tensor)
+                    else:
+                        dst_tensor.copy_(tensor_data)
+                        self._sync_tensor_device(dst_tensor)
+                except Exception as exc:
+                    holder.release()
+                    self._metrics["errors"] += 1
+                    logger.error("Failed to copy data to pool: %s", exc)
+                    return False, 0, None
+
+                src_addr = self.base_ptr + offset
+
+            if size <= 0:
+                if should_release_holder and isinstance(holder, ManagedBuffer):
+                    holder.release()
+                self._metrics["errors"] += 1
+                logger.warning("Rejecting put for %s: final payload size is %s", put_key, size)
+                return False, 0, None
+
+            with self._local_buffers_lock:
+                old_item = self._local_buffers.pop(put_key, None)
+                if old_item:
+                    _, _, old_holder, old_should_release, _, _ = old_item
+                    if old_should_release and isinstance(old_holder, ManagedBuffer):
+                        old_holder.release()
+                self._local_buffers[put_key] = (
+                    [src_addr],
+                    [size],
+                    holder,
+                    should_release_holder,
+                    is_fast_path,
+                    _time_mod.monotonic(),
+                )
+
+            metadata = {
+                "source_host": self.host,
+                "source_port": self.zmq_port,
+                "source_rpc_port": self.rpc_port,
+                "source_addrs": [src_addr],
+                "lengths": [size],
+                "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 exc:
+            self._metrics["errors"] += 1
+            logger.error("transfer_engine put failed for %s: %s", put_key, exc, exc_info=True)
+            return False, 0, None
+
+    def _query_metadata_at(self, get_key: str, host: str, port: int) -> dict[str, Any] | None:
+        zmq_addr = f"tcp://{host}:{port}"
+        req_socket = self._get_req_socket(zmq_addr, timeout_ms=5000)
+
+        try:
+            query = QueryRequest(request_id=get_key)
+            req_socket.send(QUERY_INFO + msgspec.msgpack.encode(query))
+            resp = req_socket.recv()
+            if resp == INFO_NOT_FOUND:
+                return None
+
+            query_resp = msgspec.msgpack.decode(resp, type=QueryResponse)
+            return {
+                "source_host": host,
+                "source_port": port,
+                "source_rpc_port": query_resp.rpc_port,
+                "source_addrs": query_resp.src_addrs,
+                "lengths": query_resp.lengths,
+                "data_size": query_resp.data_size,
+                "is_fast_path": query_resp.is_fast_path,
+            }
+        except Exception as exc:
+            self._invalidate_req_socket(zmq_addr)
+            logger.debug("Failed to query metadata at %s for %s: %s", zmq_addr, get_key, exc)
+            return None
+
+    @classmethod
+    def _buffer_to_bytes(cls, buffer: ManagedBuffer) -> bytes:
+        tensor = buffer.tensor
+        cls._sync_tensor_device(tensor)
+        if getattr(tensor.device, "type", "cpu") != "cpu":
+            tensor = tensor.cpu()
+        return tensor.numpy().tobytes()
+
+    def _request_sender_cleanup(self, request_id: str, source_host: str | None, source_port: int | None) -> bool:
+        if not source_host or not source_port:
+            return False
+
+        zmq_addr = f"tcp://{source_host}:{source_port}"
+        req_socket = self._get_req_socket(zmq_addr, timeout_ms=3000)
+
+        try:
+            payload = CleanupRequest(request_id=request_id)
+            req_socket.send(CLEANUP_KEY + msgspec.msgpack.encode(payload))
+            return req_socket.recv() == ACK_OK
+        except Exception as exc:
+            self._invalidate_req_socket(zmq_addr)
+            logger.warning("Failed to cleanup sender buffer for %s: %s", request_id, exc)
+            return False
+
+    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 data: YuanrongTransferEngineConnector is closed")
+
+        get_key = self._make_key(get_key, from_stage, to_stage)
+        t0 = _time_mod.perf_counter()
+
+        if not metadata:
+            if not self.sender_host or self.sender_zmq_port is None or str(self.sender_host).lower() == "auto":
+                self._metrics["errors"] += 1
+                raise RuntimeError(
+                    "get(metadata=None) requires sender_host and sender_zmq_port to be resolved. "
+                    "Call update_sender_info(host, port) before polling."
+                )
+            metadata = self._query_metadata_at(get_key, str(self.sender_host), int(self.sender_zmq_port))
+            if not metadata:
+                return None
+        elif not (metadata.get("source_host") and metadata.get("source_rpc_port") and metadata.get("source_addrs")):
+            partial_host = metadata.get("source_host")
+            partial_port = metadata.get("source_port")
+            if not partial_host or not partial_port:
+                self._metrics["errors"] += 1
+                logger.warning(
+                    "get(%s): partial metadata missing source_host/source_port. metadata=%s",
+                    get_key,
+                    metadata,
+                )
+                return None
+            queried = self._query_metadata_at(get_key, str(partial_host), int(partial_port))
+            if not queried:
+                return None
+            metadata = queried
+
+        t1 = _time_mod.perf_counter()
+        query_ms = (t1 - t0) * 1000
+
+        source_host = metadata.get("source_host")
+        source_port = metadata.get("source_port")
+        source_rpc_port = metadata.get("source_rpc_port")
+        source_addrs = list(metadata.get("source_addrs") or [])
+        lengths = list(metadata.get("lengths") or [])
+        data_size = int(metadata.get("data_size", 0))
+        is_fast_path = bool(metadata.get("is_fast_path", False))
+
+        if not source_host or not source_rpc_port or str(source_host).lower() == "auto":
+            self._metrics["errors"] += 1
+            logger.error(
+                "Invalid metadata for %s: source_host=%r, source_rpc_port=%r",
+                get_key,
+                source_host,
+                source_rpc_port,
+            )
+            return None
+        if not source_addrs:
+            self._metrics["errors"] += 1
+            logger.error("Invalid metadata for %s: missing source_addrs", get_key)
+            return None
+        if not lengths:
+            lengths = [data_size]
+        if len(source_addrs) != len(lengths):
+            self._metrics["errors"] += 1
+            logger.error(
+                "Invalid metadata for %s: source_addrs/lengths mismatch (%s vs %s)",
+                get_key,
+                len(source_addrs),
+                len(lengths),
+            )
+            return None
+
+        total_size = sum(int(length) for length in lengths)
+        if data_size <= 0:
+            data_size = total_size
+        if data_size <= 0:
+            self._metrics["errors"] += 1
+            logger.warning("Skipping get for %s: data_size is 0", get_key)
+            return None
+
+        try:
+            offset = self.allocator.alloc(data_size)
+            recv_buffer = ManagedBuffer(self.allocator, offset, data_size, self.pool)
+            base_dst_ptr = self.base_ptr + offset
+        except MemoryError:
+            self._metrics["errors"] += 1
+            logger.error("Failed to allocate %s bytes in receive pool", data_size)
+            return None
+
+        dst_addrs: list[int] = []
+        cursor = 0
+        for length in lengths:
+            dst_addrs.append(base_dst_ptr + cursor)
+            cursor += int(length)
+        if cursor != data_size:
+            recv_buffer.release()
+            self._metrics["errors"] += 1
+            logger.error("Length sum mismatch for %s: lengths_sum=%s, data_size=%s", get_key, cursor, data_size)
+            return None
+
+        t2 = _time_mod.perf_counter()
+        alloc_ms = (t2 - t1) * 1000
+        target_hostname = f"{source_host}:{int(source_rpc_port)}"
+
+        try:
+            result = self.engine.batch_transfer_sync_read(target_hostname, dst_addrs, source_addrs, lengths)
+            self._ensure_result_ok(result, f"read from {target_hostname}")
+            self._sync_tensor_device(recv_buffer.tensor)
+
+            t3 = _time_mod.perf_counter()
+            read_ms = (t3 - t2) * 1000
+
+            cleanup_ok = self._request_sender_cleanup(get_key, source_host, int(source_port) if source_port else None)
+            if not cleanup_ok:
+                logger.warning("Sender-side cleanup not confirmed for %s", get_key)
+
+            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(
+                    "[YR GET] %s: query=%.1fms, alloc=%.1fms, read=%.1fms, total=%.1fms, %.1f MB/s "
+                    "(fast_path, zero-copy)",
+                    get_key,
+                    query_ms,
+                    alloc_ms,
+                    read_ms,
+                    total_ms,
+                    mbps,
+                )
+                self._metrics["gets"] += 1
+                self._metrics["bytes_transferred"] += data_size
+                return recv_buffer, data_size
+
+            try:
+                copy_start = _time_mod.perf_counter()
+                raw_bytes = self._buffer_to_bytes(recv_buffer)
+                copy_end = _time_mod.perf_counter()
+                value = OmniSerializer.deserialize(raw_bytes)
+                end = _time_mod.perf_counter()
+                total_ms = (end - t0) * 1000
+                mbps = (data_size / 1024 / 1024) / (total_ms / 1000) if total_ms > 0 else 0
+                logger.info(
+                    "[YR GET] %s: query=%.1fms, alloc=%.1fms, read=%.1fms, copy=%.1fms, total=%.1fms, %.1f MB/s",
+                    get_key,
+                    query_ms,
+                    alloc_ms,
+                    read_ms,
+                    (copy_end - copy_start) * 1000,
+                    total_ms,
+                    mbps,
+                )
+                self._metrics["gets"] += 1
+                self._metrics["bytes_transferred"] += data_size
+                return value, data_size
+            finally:
+                recv_buffer.release()
+        except Exception as exc:
+            self._metrics["errors"] += 1
+            self._metrics["timeouts"] += 1
+            logger.error("transfer_engine get failed for %s: %s", get_key, exc, exc_info=True)
+            recv_buffer.release()
+            return None
+
+    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. "
+                f"Got from_stage={from_stage!r}, to_stage={to_stage!r}"
+            )
+        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 not item:
+                return
+            _, _, holder, should_release, _, _ = item
+            if should_release and isinstance(holder, ManagedBuffer):
+                holder.release()
+
+    def health(self) -> dict[str, Any]:
+        if self._closed:
+            return {"status": "unhealthy", "error": "Connector is closed"}
+
+        return {
+            "status": "healthy",
+            "host": self.host,
+            "zmq_port": self.zmq_port,
+            "rpc_port": self.rpc_port,
+            "protocol": self.protocol,
+            "pool_device": self.pool_device,
+            "pool_size": self.pool_size,
+            **self._metrics,
+        }
+
+    def close(self) -> None:
+        if getattr(self, "_closed", True):
+            return
+        self._closed = True
+
+        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._control_executor is not None:
+            self._control_executor.shutdown(wait=True, cancel_futures=False)
+
+        with self._local_buffers_lock:
+            for _, item in list(self._local_buffers.items()):
+                _, _, holder, should_release, _, _ = item
+                if should_release and isinstance(holder, ManagedBuffer):
+                    holder.release()
+            self._local_buffers.clear()
+
+        cache: dict[str, zmq.Socket] = getattr(self._req_local, "cache", None)  # type: ignore[assignment]
+        if cache:
+            for sock in cache.values():
+                try:
+                    sock.close(linger=0)
+                except Exception:
+                    pass
+            cache.clear()
+
+        try:
+            if hasattr(self, "engine") and hasattr(self.engine, "unregister_memory"):
+                result = self.engine.unregister_memory(self.base_ptr)
+                if result.is_error():
+                    logger.warning("Failed to unregister memory: %s", result.to_string())
+        except Exception as exc:
+            logger.warning("Failed to unregister memory: %s", exc)
+
+        try:
+            if hasattr(self, "engine") and hasattr(self.engine, "finalize"):
+                result = self.engine.finalize()
+                if result.is_error():
+                    logger.warning("Failed to finalize transfer engine: %s", result.to_string())
+        except Exception as exc:
+            logger.warning("Failed to finalize transfer engine: %s", exc)
+
+        try:
+            if hasattr(self, "zmq_ctx"):
+                self.zmq_ctx.term()
+        except Exception as exc:
+            logger.warning("Failed to terminate ZMQ context: %s", exc)
+
+        self.pool = None
+
+    def _cleanup_stale_buffers(self) -> None:
+        now = _time_mod.monotonic()
+        with self._local_buffers_lock:
+            stale_keys = [k for k, v in self._local_buffers.items() if now - v[5] > _BUFFER_TTL_SECONDS]
+            for key in stale_keys:
+                item = self._local_buffers.pop(key)
+                _, _, holder, should_release, _, _ = item
+                if should_release and isinstance(holder, ManagedBuffer):
+                    holder.release()
+                logger.warning("TTL expired (%ss): cleaned up stale buffer for %s", _BUFFER_TTL_SECONDS, key)
+
+    def _zmq_listener_loop(self) -> None:
+        socket_obj = self.zmq_ctx.socket(zmq.ROUTER)
+        try:
+            socket_obj.bind(f"tcp://{self.host}:{self.zmq_port}")
+        except zmq.ZMQError as exc:
+            logger.error("ZMQ bind failed on %s:%s: %s", self.host, self.zmq_port, exc)
+            self.can_put = False
+            self._bind_error = exc
+            self._listener_ready.set()
+            return
+
+        self._listener_ready.set()
+
+        notify_addr = f"inproc://notify-{id(self)}"
+        notify_recv = self.zmq_ctx.socket(zmq.PULL)
+        notify_recv.bind(notify_addr)
+
+        poller = zmq.Poller()
+        poller.register(socket_obj, 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()
+                            socket_obj.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 socket_obj in events:
+                        frames = socket_obj.recv_multipart()
+                        if len(frames) >= 2:
+                            self._control_executor.submit(
+                                self._handle_control_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)
+                socket_obj.close(linger=0)
+            except Exception:
+                pass
+
+    def _handle_control_request(self, response_queue: queue.Queue, notify_addr: str, identity, payload: bytes) -> None:
+        try:
+            if payload.startswith(QUERY_INFO):
+                self._handle_query_request(response_queue, notify_addr, identity, payload[len(QUERY_INFO) :])
+                return
+            if payload.startswith(CLEANUP_KEY):
+                self._handle_cleanup_request(response_queue, notify_addr, identity, payload[len(CLEANUP_KEY) :])
+                return
+
+            response_queue.put((identity, ACK_ERROR))
+        except Exception as exc:
+            logger.error("Control request failed: %s", exc)
+            response_queue.put((identity, ACK_ERROR))
+
+        self._notify_listener(notify_addr)
+
+    def _handle_query_request(self, response_queue: queue.Queue, notify_addr: str, identity, payload: bytes) -> None:
+        try:
+            query = msgspec.msgpack.decode(payload, type=QueryRequest)
+            with self._local_buffers_lock:
+                item = self._local_buffers.get(query.request_id)
+
+            if not item:
+                response_queue.put((identity, INFO_NOT_FOUND))
+            else:
+                src_addrs, src_lengths, _, _, is_fast_path, _ = item
+                resp = QueryResponse(
+                    request_id=query.request_id,
+                    rpc_port=self.rpc_port,
+                    data_size=sum(int(length) for length in src_lengths),
+                    is_fast_path=is_fast_path,
+                    src_addrs=src_addrs,
+                    lengths=src_lengths,
+                )
+                response_queue.put((identity, msgspec.msgpack.encode(resp)))
+        except Exception as exc:
+            logger.error("Query request failed: %s", exc)
+            response_queue.put((identity, INFO_NOT_FOUND))
+
+        self._notify_listener(notify_addr)
+
+    def _handle_cleanup_request(self, response_queue: queue.Queue, notify_addr: str, identity, payload: bytes) -> None:
+        try:
+            request = msgspec.msgpack.decode(payload, type=CleanupRequest)
+            self.cleanup(request.request_id)
+            response_queue.put((identity, ACK_OK))
+        except Exception as exc:
+            logger.error("Cleanup request failed: %s", exc)
+            response_queue.put((identity, ACK_ERROR))
+
+        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