[EC][Mooncake] Default TransferEngine protocol to TCP with host buffers

tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh

@@ -13,7 +13,7 @@
 #   MODEL                    HF model id (default: Qwen/Qwen2.5-VL-3B-Instruct)
 #   GPU_SINGLE / GPU_E / GPU_PD   GPU ids (defaults 0 / 1 / 2)
 #   ENDPOINT_PORT, ENCODE_PORT, PREFILL_DECODE_PORT
-#   MOONCAKE_EC_PROTOCOL        rdma | tcp (default rdma)
+#   MOONCAKE_EC_PROTOCOL        rdma | tcp (default rdma; use tcp without verbs)
 #   USE_MM_PROMPTS              1 (default) or 0 for text-only quick sanity
 #   TIMEOUT_SECONDS             wait_for_server timeout (default 1200)
 #   SKIP_BASELINE               set to 1 to reuse existing BASELINE_FILE

tests/v1/ec_connector/unit/test_ec_mooncake_connector.py

@@ -1696,3 +1696,88 @@ def test_producer_scheduler_has_cache_item_false(
             )
             mm_hash = mock_request_with_3_mm.mm_features[0].identifier
             assert not scheduler.has_cache_item(mm_hash)
+
+
+class TestECMooncakeTCPTransport:
+    def test_default_protocol_is_rdma(self, mock_vllm_config_producer):
+        mock_vllm_config_producer.ec_transfer_config.ec_connector_extra_config = {}
+        with patch_ec_mooncake_deps():
+            connector = ECMooncakeConnector(
+                mock_vllm_config_producer, ECConnectorRole.WORKER
+            )
+            try:
+                assert connector._protocol == "rdma"
+                # Software RoCE / missing nvidia_peermem cannot GPUDirect.
+                assert connector._uses_host_transport()
+                assert connector._transport_device(torch.device("cuda")).type == "cpu"
+            finally:
+                connector.shutdown()
+
+    def test_rdma_keeps_requested_cuda_transport_when_host_buffers_disabled(
+        self, mock_vllm_config_producer
+    ):
+        mock_vllm_config_producer.ec_transfer_config.ec_connector_extra_config = {
+            "mooncake_protocol": "rdma",
+            "mooncake_host_buffers": False,
+        }
+        with patch_ec_mooncake_deps():
+            connector = ECMooncakeConnector(
+                mock_vllm_config_producer, ECConnectorRole.WORKER
+            )
+            try:
+                assert connector._protocol == "rdma"
+                assert not connector._uses_host_transport()
+                cuda = torch.device("cuda")
+                assert connector._transport_device(cuda).type == "cuda"
+            finally:
+                connector.shutdown()
+
+    def test_tcp_uses_host_transport(self, mock_vllm_config_producer):
+        mock_vllm_config_producer.ec_transfer_config.ec_connector_extra_config = {
+            "mooncake_protocol": "tcp",
+        }
+        with patch_ec_mooncake_deps():
+            connector = ECMooncakeConnector(
+                mock_vllm_config_producer, ECConnectorRole.WORKER
+            )
+            try:
+                assert connector._protocol == "tcp"
+                assert connector._uses_host_transport()
+                assert connector._transport_device(torch.device("cuda")).type == "cpu"
+            finally:
+                connector.shutdown()
+
+    def test_tcp_consumer_pool_stays_on_host_when_buffer_device_is_cuda(
+        self, mock_vllm_config_consumer
+    ):
+        mock_vllm_config_consumer.ec_transfer_config.ec_buffer_device = "cuda"
+        mock_vllm_config_consumer.ec_transfer_config.ec_buffer_size = 4096
+        mock_vllm_config_consumer.ec_transfer_config.ec_connector_extra_config[
+            "consumer_buffer_pool_size"
+        ] = 4096
+        with patch_ec_mooncake_deps():
+            consumer = ECMooncakeConnector(
+                mock_vllm_config_consumer, ECConnectorRole.WORKER
+            )
+            try:
+                consumer._ensure_consumer_pool(torch.device("cuda"))
+                assert consumer._consumer_pool is not None
+                assert consumer._consumer_pool.device.type == "cpu"
+            finally:
+                consumer.shutdown()
+
+    def test_tcp_materialize_keeps_cpu_when_buffer_device_is_cpu(
+        self, mock_vllm_config_consumer
+    ):
+        mock_vllm_config_consumer.ec_transfer_config.ec_buffer_device = "cpu"
+        host = torch.randn(4, 8)
+        with patch_ec_mooncake_deps():
+            consumer = ECMooncakeConnector(
+                mock_vllm_config_consumer, ECConnectorRole.WORKER
+            )
+            try:
+                out = consumer._materialize_for_encoder_cache(host)
+                assert out.device.type == "cpu"
+                assert torch.equal(out, host)
+            finally:
+                consumer.shutdown()

vllm/config/ec_transfer.py

@@ -20,9 +20,10 @@ class ECTransferConfig:
     """The EC connector for vLLM to transmit EC caches between vLLM instances.
 
     Built-in options include ``ECExampleConnector`` (shared filesystem via
-    safetensors) and ``ECMooncakeConnector`` (Mooncake TransferEngine RDMA;
-    requires ``mooncake-transfer-engine`` and matching producer/consumer
-    ``ec_connector_extra_config``; see ``mooncake_ec_connector`` module docstring).
+    safetensors) and ``ECMooncakeConnector`` (Mooncake TransferEngine, RDMA
+    by default or TCP; requires ``mooncake-transfer-engine`` and matching
+    producer/consumer ``ec_connector_extra_config``; see
+    ``mooncake_ec_connector`` module docstring).
     """
 
     engine_id: str | None = None

vllm/distributed/ec_transfer/ec_connector/mooncake_ec_connector.py

@@ -4,14 +4,16 @@
 Encoder-cache (EC) connector backed by Mooncake TransferEngine.
 
 Used in disaggregated setups where an encoder / prefill instance produces
-multimodal encoder outputs and a decode instance loads them over RDMA-capable
-Mooncake transport instead of shared filesystem.
+multimodal encoder outputs and a decode instance loads them over Mooncake
+(RDMA by default; TCP when ``mooncake_protocol`` is ``tcp``) instead of a
+shared filesystem.
 """
 
 from __future__ import annotations
 
 import bisect
 import math
+import os
 import threading
 import time
 import uuid
@@ -621,9 +623,15 @@ class ECMooncakeConnector(ECConnectorBase):
     Extra config (``ec_connector_extra_config``):
 
     - ``mooncake_protocol`` (optional): Passed to ``TransferEngine.initialize``
-      (default ``"rdma"``).
+      (default ``"rdma"``; set ``"tcp"`` when no verbs device is available).
+    - ``mooncake_device`` (optional): RDMA device name passed as the
+      TransferEngine NIC list (empty string lets Mooncake auto-discover).
+    - ``mooncake_host_buffers`` (optional): Force pinned-host bounce buffers.
+      Defaults to on for TCP, and for RDMA when GPUDirect is unavailable
+      (software RoCE / missing ``nvidia_peermem``).
     - ``consumer_buffer_pool_size`` (consumer, optional): Bytes reserved for a
-      long-lived registered CUDA receive arena (default ``ec_buffer_size``).
+      long-lived registered receive arena (pinned host when GPUDirect is
+      unavailable, CUDA otherwise; default ``ec_buffer_size``).
     - ``reservation_zmq_port`` (consumer worker, required): Exposes registered
       receive addresses over ZMQ. Replica ``d`` of the first pipeline stage owns
       the block starting at ``port + d * tensor_parallel_size``; tensor-parallel
@@ -705,7 +713,9 @@ def __init__(self, vllm_config: VllmConfig, role: ECConnectorRole):
         assert ec_cfg is not None
         self._ec_cfg = ec_cfg
         self._extra = self._ec_cfg.ec_connector_extra_config
-        self._protocol: str = self._extra.get("mooncake_protocol", "rdma")
+        self._protocol: str = str(
+            self._extra.get("mooncake_protocol", "rdma")
+        ).lower()
         reservation_port = self._extra.get("reservation_zmq_port")
         self._reservation_zmq_port = (
             int(reservation_port) if reservation_port is not None else None
@@ -858,14 +868,18 @@ def _ensure_engine(self) -> TransferEngine:
             if self._engine is not None:
                 return self._engine
             eng = TransferEngine()
-            ret = eng.initialize(self._hostname, "P2PHANDSHAKE", self._protocol, "")
+            device = str(self._extra.get("mooncake_device", "") or "")
+            ret = eng.initialize(
+                self._hostname, "P2PHANDSHAKE", self._protocol, device
+            )
             if ret != 0:
                 raise RuntimeError("Mooncake TransferEngine initialization failed.")
             self._engine = eng
             logger.info(
-                "ECMooncakeConnector TransferEngine ready at %s:%d",
+                "ECMooncakeConnector TransferEngine ready at %s:%d protocol=%s",
                 self._hostname,
                 eng.get_rpc_port(),
+                self._protocol,
             )
         return self._engine
 
@@ -1040,19 +1054,78 @@ def _release_push_source_registrations(self, addresses: list[int]) -> bool:
                 self._pending_unregister.pop(address, None)
             return True
 
+    def _has_gpu_direct_rdma(self) -> bool:
+        """True when CUDA buffers can be registered for GPUDirect RDMA."""
+        if os.environ.get("WITH_NVIDIA_PEERMEM", "1") == "0":
+            return False
+        ib_root = "/sys/class/infiniband"
+        try:
+            names = os.listdir(ib_root)
+        except OSError:
+            return False
+        if not names or all(name.startswith("rxe") for name in names):
+            return False
+        return os.path.exists("/sys/module/nvidia_peermem")
+
+    def _uses_host_transport(self) -> bool:
+        """Use pinned host bounce buffers when GPUDirect is not usable.
+
+        TCP never GPUDirects. RDMA still needs host staging on software RoCE
+        or hosts without ``nvidia_peermem``; otherwise CUDA pointers hang
+        inside ``batch_transfer_sync_write``.
+        """
+        flag = self._extra.get("mooncake_host_buffers")
+        if flag is not None:
+            return bool(flag)
+        if self._protocol == "tcp":
+            return True
+        return not self._has_gpu_direct_rdma()
+
+    def _transport_device(self, requested: torch.device) -> torch.device:
+        if self._uses_host_transport():
+            return torch.device("cpu")
+        return requested
+
+    def _encoder_cache_device(self) -> torch.device:
+        raw = self._ec_cfg.ec_buffer_device
+        name = raw.lower() if isinstance(raw, str) and raw else "cuda"
+        if name == "cuda" and not torch.cuda.is_available():
+            return torch.device("cpu")
+        return torch.device(name)
+
+    def _alloc_transport_buffer(
+        self, nbytes: int, device: torch.device
+    ) -> torch.Tensor:
+        kwargs: dict[str, Any] = {}
+        if device.type == "cpu":
+            kwargs["pin_memory"] = torch.cuda.is_available()
+        return torch.empty(nbytes, dtype=torch.uint8, device=device, **kwargs)
+
+    def _materialize_for_encoder_cache(self, tensor: torch.Tensor) -> torch.Tensor:
+        """Move TCP host buffers onto the encoder-cache device."""
+        if tensor.device.type != "cpu" or not self._uses_host_transport():
+            return tensor
+        target = self._encoder_cache_device()
+        if tensor.device.type == target.type:
+            return tensor
+        return tensor.to(device=target, non_blocking=False)
+
     def _ensure_consumer_pool(
         self, device: torch.device, *, allow_host: bool = False
     ) -> None:
+        device = self._transport_device(device)
         if (
             self._consumer_pool is not None
             or self._consumer_pool_disabled
-            or (device.type != "cuda" and not allow_host)
+            or (
+                device.type != "cuda"
+                and not allow_host
+                and not self._uses_host_transport()
+            )
         ):
             return
         try:
-            pool = torch.empty(
-                self._consumer_pool_capacity, dtype=torch.uint8, device=device
-            )
+            pool = self._alloc_transport_buffer(self._consumer_pool_capacity, device)
             if self._is_receiving_rank:
                 # Producers write into this pool directly, so it needs a memory
                 # region. Later pipeline stages never receive and skip it.
@@ -1072,8 +1145,9 @@ def _ensure_consumer_pool(
         self._consumer_pool = pool
         self._consumer_pool_allocator = _ContiguousAllocator(pool.nbytes)
         logger.info(
-            "Prepared %d-byte CUDA receive pool for Mooncake EC (registered=%s)",
+            "Prepared %d-byte %s receive pool for Mooncake EC (registered=%s)",
             pool.nbytes,
+            device.type,
             self._is_receiving_rank,
         )
 
@@ -1082,16 +1156,18 @@ def _ensure_producer_pool(self, device: torch.device) -> None:
 
         Registering the encoder output itself costs more than the transfer
         (register+unregister dominated the push path); staging into a slab
-        that is registered once trades that for a device-to-device copy.
+        that is registered once trades that for a copy. TCP uses pinned host
+        memory because Mooncake cannot GPUDirect over TCP.
         """
+        device = self._transport_device(device)
         if self._producer_pool is not None or self._producer_pool_disabled:
             return
         with self._producer_pool_lock:
             if self._producer_pool is not None or self._producer_pool_disabled:
                 return
             try:
-                pool = torch.empty(
-                    self._producer_pool_capacity, dtype=torch.uint8, device=device
+                pool = self._alloc_transport_buffer(
+                    self._producer_pool_capacity, device
                 )
                 ret = self._ensure_engine().batch_register_memory(
                     [pool.data_ptr()], [pool.nbytes]
@@ -1119,7 +1195,7 @@ def _stage_push_sources(
         """Copy the batch into the staging pool; None if it does not fit."""
         if not tensors:
             return [], []
-        self._ensure_producer_pool(tensors[0].device)
+        self._ensure_producer_pool(self._transport_device(tensors[0].device))
         pool = self._producer_pool
         allocator = self._producer_pool_allocator
         if pool is None or allocator is None:
@@ -1211,7 +1287,7 @@ def _take_resident_tensor(self, spec: ECMooncakeLoadSpec) -> torch.Tensor | None
             self._consumer_residents.pin(spec.mm_hash)
             self._consumer_retire_events.pop(spec.mm_hash, None)
             self._consumer_worker_metrics["residents_promoted"] += 1
-            return tensor
+            return self._materialize_for_encoder_cache(tensor)
 
     def _release_stale_consumer_allocations(
         self, encoder_cache: dict[str, torch.Tensor]
@@ -1234,13 +1310,17 @@ def _release_stale_consumer_allocations(
                 if id(allocation) in reserved_allocations:
                     continue
                 # Retire rather than free: the bytes stay valid and serve the
-                # next request that needs this item. The event orders the
-                # eventual reuse behind whatever still reads the tensor.
-                event = torch.Event()
-                event.record(
-                    torch.accelerator.current_stream(self._consumer_pool.device)
-                )
-                self._consumer_retire_events[mm_hash] = event
+                # next request that needs this item. CUDA events order reuse
+                # behind in-flight GPU reads; host TCP pools are already
+                # synchronized by the materialize copy.
+                if self._consumer_pool.device.type == "cuda":
+                    event = torch.Event()
+                    event.record(
+                        torch.accelerator.current_stream(
+                            self._consumer_pool.device
+                        )
+                    )
+                    self._consumer_retire_events[mm_hash] = event
                 self._consumer_residents.retire(mm_hash)
                 self._consumer_worker_metrics["residents_retired"] += 1
         self._poll_consumer_pool_frees()
@@ -1625,7 +1705,10 @@ def _take_pushed_tensor(
                 spec.mm_hash, reservation.allocation, reservation.allocation.size
             )
             self._consumer_worker_metrics["reservations_taken"] += 1
-            return reservation.allocation.tensor, reservation.allocation
+            return (
+                self._materialize_for_encoder_cache(reservation.allocation.tensor),
+                reservation.allocation,
+            )
 
     def _send_control(self, addr: str, request: dict[str, Any]) -> Any:
         return self._control_channel.request(addr, request)
@@ -1897,10 +1980,16 @@ def _push_batch(self, pushes: list[_PendingPush]) -> None:
                     sources, staged_regions = staged
                     # The NIC reads outside the CUDA stream, so the staging
                     # copies have to have landed before the transfer starts.
+                    # TCP stages onto host, so wait on the original CUDA
+                    # tensors rather than the CPU views.
                     if sources and sources[0].device.type == "cuda":
                         torch.accelerator.current_stream(
                             sources[0].device
                         ).synchronize()
+                    elif tensors and tensors[0].device.type == "cuda":
+                        torch.accelerator.current_stream(
+                            tensors[0].device
+                        ).synchronize()
                 else:
                     sources = tensors
                     registered_sources = self._acquire_push_source_registrations(