diff --git a/atom/kv_transfer/disaggregation/mooncake/mooncake_connector.py b/atom/kv_transfer/disaggregation/mooncake/mooncake_connector.py index 0e341f7268..27bd66a85f 100644 --- a/atom/kv_transfer/disaggregation/mooncake/mooncake_connector.py +++ b/atom/kv_transfer/disaggregation/mooncake/mooncake_connector.py @@ -4,7 +4,7 @@ """ Worker-side and scheduler-side KV cache connectors for disaggregated P/D. -Uses Mooncake TransferEngine for RDMA-based push (WRITE) transfers of +Uses Mooncake TransferEngine for TCP- or RDMA-based push (WRITE) transfers of KV cache data from producer (prefill) to consumer (decode) nodes. """ @@ -112,6 +112,37 @@ def _auto_select_ib_device(phys_idx: int) -> str: return rdma_device +def _select_ib_device( + protocol: str, configured_device: str, phys_idx: int | None +) -> str: + """Resolve the Mooncake device filter without enabling RDMA for TCP. + + Mooncake's TCP transport requires an empty device list. Passing a usable + HCA alongside ``protocol=tcp`` allows the transfer engine to activate RDMA + as an alternate path, which violates the caller's explicit transport + choice. RDMA-family transports retain the existing configured/automatic + device selection. + """ + if protocol.strip().lower() == "tcp": + return "" + if configured_device: + return configured_device + if phys_idx is None: + raise ValueError("physical GPU index is required for RDMA device selection") + return _auto_select_ib_device(phys_idx) + + +def _configure_mooncake_transport(protocol: str) -> None: + """Make Mooncake honor ATOM's explicit transport selection. + + Legacy TransferEngine builds auto-discover installed HCAs independently + of the device filter. ``MC_FORCE_TCP`` is Mooncake's supported override + for preventing that implicit RDMA transport from being installed. + """ + if protocol.strip().lower() == "tcp": + os.environ["MC_FORCE_TCP"] = "true" + + # ZMQ side-channel message types MSG_WRITE_REQUEST = b"write_request" MSG_WRITE_DONE = b"write_done" @@ -461,19 +492,23 @@ def __init__(self, config: Config) -> None: if not _MOONCAKE_AVAILABLE: raise RuntimeError( "Mooncake is not installed but kv_connector='mooncake' was requested. " - "Install the mooncake package to use push-mode RDMA transfers." + "Install the mooncake package to use push-mode transfers." ) - # Determine which RDMA device this TP rank should use. + # Determine which RDMA device this TP rank should use. TCP is + # intentionally initialized with an empty device filter so Mooncake + # cannot activate an available HCA as an alternate path. # AMD GPU nodes pair GPU N with NIC N, but the HCA name is cluster # dependent: Spur MI350 exposes ionic_N while older setups used rdmaN. # Registering GPU memory with a non-local RDMA NIC fails with # EINVAL. Pass the device name as a filter so Mooncake only # creates a context for the local NIC. - ib_device = kv_transfer_config.get("ib_device", "") - if not ib_device: - ib_device = os.environ.get("ATOM_MOONCAKE_IB_DEVICE", "") - if not ib_device: + _configure_mooncake_transport(self.protocol) + configured_ib_device = kv_transfer_config.get( + "ib_device", "" + ) or os.environ.get("ATOM_MOONCAKE_IB_DEVICE", "") + phys_idx: int | None = None + if self.protocol.strip().lower() != "tcp" and not configured_ib_device: visible_idx = torch.cuda.current_device() visible_env = os.environ.get("HIP_VISIBLE_DEVICES") or os.environ.get( "CUDA_VISIBLE_DEVICES" @@ -483,7 +518,10 @@ def __init__(self, config: Config) -> None: phys_idx = int(visible_list[visible_idx]) else: phys_idx = visible_idx - ib_device = _auto_select_ib_device(phys_idx) + ib_device = _select_ib_device(self.protocol, configured_ib_device, phys_idx) + if self.protocol.strip().lower() == "tcp": + logger.info("Mooncake TCP selected; RDMA device selection is disabled") + elif not configured_ib_device: logger.info( "Auto-selecting RDMA device %s for physical GPU %d " "(visible_idx=%d, tp_rank=%d)", @@ -493,7 +531,11 @@ def __init__(self, config: Config) -> None: self.tp_rank, ) - rdma_local_ip = _ip_for_ib_device(ib_device, default_local_ip) + rdma_local_ip = ( + _ip_for_ib_device(ib_device, default_local_ip) + if ib_device + else default_local_ip + ) if rdma_local_ip != default_local_ip: logger.info( "Using RDMA-local IP %s for ib_device=%s instead of default IP %s", diff --git a/tests/test_pd_pp.py b/tests/test_pd_pp.py index 6a27bc40a4..fa7ec08514 100644 --- a/tests/test_pd_pp.py +++ b/tests/test_pd_pp.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: MIT # PD-disaggregation + pipeline-parallel unit tests (GPU-free). +import os import sys import threading import types @@ -165,6 +166,51 @@ def test_producer_advertises_remote_pp_size(): assert seq.kv_transfer_params_output["remote_block_ids"] == [1, 2, 3] +# --------------------------------------------------------------------------- +# Mooncake transport selection +# --------------------------------------------------------------------------- + + +def test_mooncake_tcp_disables_rdma_device_even_when_configured(): + mc = pytest.importorskip( + "atom.kv_transfer.disaggregation.mooncake.mooncake_connector" + ) + assert mc._select_ib_device("tcp", "rdma0", None) == "" + assert mc._select_ib_device(" TCP ", "ionic_0", None) == "" + + +def test_mooncake_tcp_forces_transfer_engine_transport(monkeypatch): + mc = pytest.importorskip( + "atom.kv_transfer.disaggregation.mooncake.mooncake_connector" + ) + monkeypatch.delenv("MC_FORCE_TCP", raising=False) + mc._configure_mooncake_transport(" TCP ") + assert os.environ["MC_FORCE_TCP"] == "true" + + +def test_mooncake_rdma_preserves_explicit_device(): + mc = pytest.importorskip( + "atom.kv_transfer.disaggregation.mooncake.mooncake_connector" + ) + assert mc._select_ib_device("rdma", "ionic_3", None) == "ionic_3" + + +def test_mooncake_rdma_auto_selects_from_physical_gpu(monkeypatch): + mc = pytest.importorskip( + "atom.kv_transfer.disaggregation.mooncake.mooncake_connector" + ) + monkeypatch.setattr(mc, "_auto_select_ib_device", lambda idx: f"auto{idx}") + assert mc._select_ib_device("rdma", "", 5) == "auto5" + + +def test_mooncake_rdma_requires_gpu_index_without_explicit_device(): + mc = pytest.importorskip( + "atom.kv_transfer.disaggregation.mooncake.mooncake_connector" + ) + with pytest.raises(ValueError, match="physical GPU index"): + mc._select_ib_device("rdma", "", None) + + # --------------------------------------------------------------------------- # Producer per-layer region mapping (consumer_region_indices) # ---------------------------------------------------------------------------