[PD] Validate Mooncake EFA allocator compatibility

docs/docs/advanced_features/pd_disaggregation.mdx

@@ -55,6 +55,39 @@ Serving-side capability enforcement is tracked in [#39122](https://github.com/sg
 uv pip install mooncake-transfer-engine
 ```
 
+### AWS EFA
+
+On an x86_64 AWS EFA instance, replace the standard CUDA 13 Mooncake package with the EFA build. The two packages install the same `mooncake` Python module, so do not install them side by side.
+
+The following commands replace the package in an SGLang CUDA 13 environment while preserving its existing Python dependencies:
+
+```bash
+python3 -m pip uninstall -y \
+  mooncake-transfer-engine-cuda13 \
+  mooncake-transfer-engine-efa-cuda13
+python3 -m pip install --no-deps \
+  mooncake-transfer-engine-efa-cuda13==0.3.13.post1
+
+export LD_LIBRARY_PATH="/opt/amazon/efa/lib:/opt/amazon/efa/lib64:${LD_LIBRARY_PATH:-}"
+export MOONCAKE_PROTOCOL=efa
+
+python3 -c 'from mooncake import engine; assert engine.SUPPORT_EFA'
+fi_info -p efa -t FI_EP_RDM
+```
+
+The EFA wheel expects a compatible AWS EFA userspace installation on the host and does not bundle `libfabric` or `libefa`. The current wheel is available only for x86_64. It is built without Mooncake EP/PG and the NVLink/MNNVL components, so use the standard Mooncake CUDA package for those features.
+
+The current libfabric EFA provider cannot transfer CUDA VMM-backed buffers. Until [libfabric issue #12853](https://github.com/ofiwg/libfabric/issues/12853) is resolved, apply these settings to both prefill and decode workers:
+
+- Do not set `SGLANG_MOONCAKE_CUSTOM_MEM_POOL` to `NVLINK` or `BAREX`.
+  `INTRA_NODE_NVLINK` is supported because it uses the default CUDA allocator.
+- Do not enable `expandable_segments:True` in `PYTORCH_CUDA_ALLOC_CONF` or `PYTORCH_ALLOC_CONF`.
+- Post-capture KV sizing is automatically disabled for the EFA protocol because
+  its KV cache uses CUDA VMM allocations.
+- Do not set `MC_FORCE_MNNVL` or `NCCL_MNNVL_ENABLE`; EFA is not an MNNVL transport.
+
+SGLang validates the first two constraints at startup to avoid a later `fi_read` or `fi_write` failure from libfabric.
+
 ### IB Device Configuration
 
 `--disaggregation-ib-device` supports the following formats when using the Mooncake backend:

python/sglang/srt/arg_groups/overrides.py

@@ -1821,6 +1821,13 @@ def post_capture_kv_sizing_planned(server_args: Any) -> bool:
         return False
     if envs.SGLANG_MOONCAKE_CUSTOM_MEM_POOL.get() is not None:
         return False
+    # Mooncake over EFA cannot register the CUDA VMM allocation used by
+    # post-capture KV sizing. Fall back to the regular cudaMalloc-backed pool.
+    if (
+        cfg.disaggregation_transfer_backend == "mooncake"
+        and envs.MOONCAKE_PROTOCOL.get().lower() == "efa"
+    ):
+        return False
 
     if (
         cfg.disaggregation_mode != "prefill"

python/sglang/srt/disaggregation/mooncake/conn.py

@@ -43,6 +43,7 @@
     unpack_int_lists,
 )
 from sglang.srt.disaggregation.mooncake.utils import (
+    _validate_efa_allocator_compatibility,
     check_mooncake_custom_mem_pool_enabled,
 )
 from sglang.srt.disaggregation.utils import (
@@ -216,6 +217,12 @@ def __init__(
         server_args: ServerArgs,
         is_mla_backend: Optional[bool] = False,
     ):
+        enable_custom_mem_pool, custom_mem_pool_type = (
+            check_mooncake_custom_mem_pool_enabled()
+        )
+        _validate_efa_allocator_compatibility(
+            enable_custom_mem_pool, custom_mem_pool_type
+        )
         super().__init__(args, disaggregation_mode, server_args, is_mla_backend)
         self.init_engine()
         self.register_buffer_to_engine()

python/sglang/srt/disaggregation/mooncake/utils.py

@@ -14,6 +14,7 @@
 """Mooncake-specific utilities for custom memory pool management."""
 
 import logging
+import os
 from typing import Any, Optional, Tuple
 
 import torch
@@ -26,6 +27,41 @@
 SUPPORTED_MOONCAKE_CUSTOM_MEM_POOL_TYPES = ["NVLINK", "BAREX", "INTRA_NODE_NVLINK"]
 
 
+def _cuda_expandable_segments_enabled() -> Optional[str]:
+    """Return the allocator env var that enables CUDA expandable segments."""
+    for var in ("PYTORCH_CUDA_ALLOC_CONF", "PYTORCH_ALLOC_CONF"):
+        for field in os.environ.get(var, "").split(","):
+            key, _, value = field.partition(":")
+            if key.strip() == "expandable_segments" and value.strip().lower() == "true":
+                return var
+    return None
+
+
+def _validate_efa_allocator_compatibility(
+    enable_custom_mem_pool: bool, custom_mem_pool_type: Optional[str]
+) -> None:
+    """Reject CUDA VMM allocators unsupported by the current libfabric EFA path."""
+    if envs.MOONCAKE_PROTOCOL.get().lower() != "efa":
+        return
+
+    if enable_custom_mem_pool and custom_mem_pool_type in ("NVLINK", "BAREX"):
+        raise ValueError(
+            f"SGLANG_MOONCAKE_CUSTOM_MEM_POOL={custom_mem_pool_type} is "
+            "incompatible with MOONCAKE_PROTOCOL=efa. Mooncake custom memory "
+            "pools use CUDA VMM allocations, which the current libfabric EFA "
+            "provider cannot transfer. Unset SGLANG_MOONCAKE_CUSTOM_MEM_POOL."
+        )
+
+    expandable_segments_var = _cuda_expandable_segments_enabled()
+    if expandable_segments_var is not None:
+        raise ValueError(
+            f"{expandable_segments_var} enables expandable_segments, which is "
+            "incompatible with MOONCAKE_PROTOCOL=efa because the current "
+            "libfabric EFA provider cannot transfer CUDA VMM allocations. "
+            "Disable expandable_segments."
+        )
+
+
 def init_mooncake_custom_mem_pool(
     device: str,
 ) -> Tuple[bool, Optional[Any], Optional[str]]:

test/registered/unit/disaggregation/test_mooncake_efa_allocator.py

@@ -0,0 +1,103 @@
+import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
+
+from sglang.srt.arg_groups.overrides import post_capture_kv_sizing_planned
+from sglang.srt.disaggregation.mooncake.utils import (
+    _validate_efa_allocator_compatibility,
+    check_mooncake_custom_mem_pool_enabled,
+)
+from sglang.srt.environ import envs
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=1, suite="base-a-test-cpu")
+
+
+class TestMooncakeEfaAllocator(unittest.TestCase):
+    def _check(self, protocol, custom_mem_pool=None, allocator_env=None):
+        allocator_env = allocator_env or {}
+        with (
+            patch.object(envs.MOONCAKE_PROTOCOL, "get", return_value=protocol),
+            patch.object(
+                envs.SGLANG_MOONCAKE_CUSTOM_MEM_POOL,
+                "get",
+                return_value=custom_mem_pool,
+            ),
+            patch(
+                "sglang.srt.disaggregation.mooncake.utils.os.environ.get",
+                side_effect=lambda key, default="": allocator_env.get(key, default),
+            ),
+        ):
+            result = check_mooncake_custom_mem_pool_enabled()
+            _validate_efa_allocator_compatibility(*result)
+            return result
+
+    def test_efa_uses_default_allocator(self):
+        self.assertEqual(self._check("efa"), (False, None))
+
+    def test_efa_rejects_custom_memory_pools(self):
+        for pool_type in ("true", "NVLINK", "BAREX"):
+            with self.subTest(pool_type=pool_type):
+                with self.assertRaisesRegex(
+                    ValueError, "incompatible with MOONCAKE_PROTOCOL=efa"
+                ):
+                    self._check("efa", custom_mem_pool=pool_type)
+
+    def test_efa_allows_intra_node_nvlink(self):
+        self.assertEqual(
+            self._check("efa", custom_mem_pool="INTRA_NODE_NVLINK"),
+            (True, "INTRA_NODE_NVLINK"),
+        )
+
+    def test_efa_rejects_expandable_segments(self):
+        for var in ("PYTORCH_CUDA_ALLOC_CONF", "PYTORCH_ALLOC_CONF"):
+            with self.subTest(var=var):
+                with self.assertRaisesRegex(ValueError, var):
+                    self._check(
+                        "efa",
+                        allocator_env={
+                            var: "garbage_collection_threshold:0.8,"
+                            "expandable_segments:True"
+                        },
+                    )
+
+    def test_efa_allows_disabled_expandable_segments(self):
+        self.assertEqual(
+            self._check(
+                "EFA",
+                allocator_env={"PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:False"},
+            ),
+            (False, None),
+        )
+
+    def test_non_efa_keeps_custom_memory_pool_behavior(self):
+        self.assertEqual(self._check("rdma", custom_mem_pool="true"), (True, "NVLINK"))
+
+    def test_efa_disables_post_capture_kv_sizing(self):
+        cfg = SimpleNamespace(
+            enable_unified_memory=False,
+            device="cuda",
+            dcp_size=1,
+            kv_cache_dtype="auto",
+            prefill_only_disable_kv_cache=False,
+            enable_memory_saver=False,
+            disaggregation_transfer_backend="mooncake",
+        )
+        with (
+            patch("sglang.srt.arg_groups.overrides.resolving_view", return_value=cfg),
+            patch(
+                "sglang.srt.arg_groups.overrides.use_mla_backend", return_value=False
+            ),
+            patch.object(
+                envs.SGLANG_ENABLE_POST_CAPTURE_KV_SIZING, "get", return_value=True
+            ),
+            patch.object(
+                envs.SGLANG_MOONCAKE_CUSTOM_MEM_POOL, "get", return_value=None
+            ),
+            patch.object(envs.MOONCAKE_PROTOCOL, "get", return_value="efa"),
+        ):
+            self.assertFalse(post_capture_kv_sizing_planned(object()))
+
+
+if __name__ == "__main__":
+    unittest.main()