diff --git a/tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh b/tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh index 04e28c70483d..704bfaead5ed 100755 --- a/tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh +++ b/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 diff --git a/tests/v1/ec_connector/unit/test_ec_mooncake_connector.py b/tests/v1/ec_connector/unit/test_ec_mooncake_connector.py index 20ea692032bc..5814a09e4cf8 100644 --- a/tests/v1/ec_connector/unit/test_ec_mooncake_connector.py +++ b/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() diff --git a/vllm/config/ec_transfer.py b/vllm/config/ec_transfer.py index 3948f3205934..7a524d3806b0 100644 --- a/vllm/config/ec_transfer.py +++ b/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 diff --git a/vllm/distributed/ec_transfer/ec_connector/mooncake_ec_connector.py b/vllm/distributed/ec_transfer/ec_connector/mooncake_ec_connector.py index b1103d1aa90c..9b4e939f504c 100644 --- a/vllm/distributed/ec_transfer/ec_connector/mooncake_ec_connector.py +++ b/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(