fix(mooncake): keep TCP transport off RDMA devices

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",

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)
 # ---------------------------------------------------------------------------