diff --git a/docs/docs/advanced_features/pd_disaggregation.mdx b/docs/docs/advanced_features/pd_disaggregation.mdx index 1cb97e88d752..7fd878bc94c6 100644 --- a/docs/docs/advanced_features/pd_disaggregation.mdx +++ b/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: diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index de96b45e6fd2..d76257d2f1c9 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/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" diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index d6bc4aed9521..558505eb0e45 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/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() diff --git a/python/sglang/srt/disaggregation/mooncake/utils.py b/python/sglang/srt/disaggregation/mooncake/utils.py index 279cf194de87..816622ced1e0 100644 --- a/python/sglang/srt/disaggregation/mooncake/utils.py +++ b/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]]: diff --git a/test/registered/unit/disaggregation/test_mooncake_efa_allocator.py b/test/registered/unit/disaggregation/test_mooncake_efa_allocator.py new file mode 100644 index 000000000000..2ed4793bad56 --- /dev/null +++ b/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()