From 6076066e380a6abae257167ff4c764e624137111 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Tue, 26 May 2026 18:15:55 -0700 Subject: [PATCH] Add mooncake_tcp transfer backend (mooncake over TCP) (#26346) --- .../srt/arg_groups/pd_disaggregation_hook.py | 80 +++++++++++++++++++ .../mooncake_transfer_engine.py | 7 +- python/sglang/srt/server_args.py | 71 +++------------- 3 files changed, 99 insertions(+), 59 deletions(-) create mode 100644 python/sglang/srt/arg_groups/pd_disaggregation_hook.py diff --git a/python/sglang/srt/arg_groups/pd_disaggregation_hook.py b/python/sglang/srt/arg_groups/pd_disaggregation_hook.py new file mode 100644 index 000000000..7b46d68f7 --- /dev/null +++ b/python/sglang/srt/arg_groups/pd_disaggregation_hook.py @@ -0,0 +1,80 @@ +import logging +import os +from typing import TYPE_CHECKING + +from sglang.srt.environ import envs + +if TYPE_CHECKING: + from sglang.srt.server_args import ServerArgs + +logger = logging.getLogger(__name__) + + +def handle_pd_disaggregation(server_args: "ServerArgs") -> None: + """Validate and normalize PD-disaggregation server args.""" + # "mooncake_tcp" is mooncake with the TCP transport forced: set MC_FORCE_TCP + # so mooncake installs TcpTransport instead of RDMA, rewrite the backend to + # mooncake, and skip RDMA HCA selection. Must run before backend-name checks. + if server_args.disaggregation_transfer_backend == "mooncake_tcp": + os.environ.setdefault("MC_FORCE_TCP", "1") + server_args.disaggregation_transfer_backend = "mooncake" + server_args.disaggregation_ib_device = None + logger.info( + "disaggregation transfer backend 'mooncake_tcp' -> mooncake " + "with MC_FORCE_TCP=1 (TCP transport, no RDMA)" + ) + + if server_args.disaggregation_mode == "decode": + if server_args.disaggregation_decode_enable_radix_cache: + if server_args.enable_hisparse: + raise ValueError( + "--disaggregation-decode-enable-radix-cache is incompatible " + "with --enable-hisparse" + ) + if server_args.disaggregation_transfer_backend not in ("nixl", "mooncake"): + raise ValueError( + "--disaggregation-decode-enable-radix-cache currently " + "requires --disaggregation-transfer-backend in " + "('nixl', 'mooncake'), but got " + f"{server_args.disaggregation_transfer_backend!r}" + ) + if server_args.speculative_algorithm is not None: + raise ValueError( + "--disaggregation-decode-enable-radix-cache is incompatible " + "with speculative decoding " + f"(--speculative-algorithm {server_args.speculative_algorithm})" + ) + if server_args.enable_dp_attention: + logger.warning( + "EXPERIMENTAL: Decode radix cache with DP attention. " + "Requires prefix-aware DP rank routing for optimal cache hits." + ) + server_args.disable_radix_cache = False + logger.warning("EXPERIMENTAL: Radix cache is enabled for decode server") + else: + server_args.disable_radix_cache = True + logger.warning("KV cache is forced as chunk cache for decode server") + if server_args.enable_mamba_extra_buffer(): + logger.warning( + "Mamba extra_buffer is disabled because decode disaggregation " + "currently forces chunk cache. Falling back to no_buffer." + ) + server_args.mamba_scheduler_strategy = "no_buffer" + + elif server_args.disaggregation_mode == "prefill": + assert ( + server_args.disaggregation_transfer_backend != "fake" + ), "Prefill server does not support 'fake' as the transfer backend" + + server_args.disable_cuda_graph = True + + if server_args.disaggregation_mode in ("prefill", "decode"): + if ( + envs.SGLANG_DISAGG_STAGING_BUFFER.get() + and server_args.disaggregation_transfer_backend not in ("mooncake", "nixl") + ): + raise ValueError( + f"SGLANG_DISAGG_STAGING_BUFFER requires " + f"disaggregation_transfer_backend='mooncake' or 'nixl', " + f"got '{server_args.disaggregation_transfer_backend}'." + ) diff --git a/python/sglang/srt/distributed/device_communicators/mooncake_transfer_engine.py b/python/sglang/srt/distributed/device_communicators/mooncake_transfer_engine.py index 372919b1c..9efefe177 100644 --- a/python/sglang/srt/distributed/device_communicators/mooncake_transfer_engine.py +++ b/python/sglang/srt/distributed/device_communicators/mooncake_transfer_engine.py @@ -111,7 +111,12 @@ class MooncakeTransferEngine: self.engine = TransferEngine() self.hostname = hostname self.gpu_id = gpu_id if gpu_id is not None else 0 - self.ib_device = get_ib_devices_for_gpu(ib_device, self.gpu_id) + # MC_FORCE_TCP=1 makes mooncake install TcpTransport instead of RDMA, + # in which case RDMA HCA selection is irrelevant; pass empty device. + if os.environ.get("MC_FORCE_TCP") == "1": + self.ib_device = "" + else: + self.ib_device = get_ib_devices_for_gpu(ib_device, self.gpu_id) self.initialize( hostname=self.hostname, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 660735bc8..350683be8 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -183,7 +183,14 @@ DETERMINISTIC_ATTENTION_BACKEND_CHOICES = ["flashinfer", "fa3", "triton"] RADIX_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND = ["fa3", "triton"] -DISAGG_TRANSFER_BACKEND_CHOICES = ["mooncake", "nixl", "ascend", "fake", "mori"] +DISAGG_TRANSFER_BACKEND_CHOICES = [ + "mooncake", + "nixl", + "ascend", + "fake", + "mori", + "mooncake_tcp", +] GRAMMAR_BACKEND_CHOICES = ["xgrammar", "outlines", "llguidance", "none"] @@ -859,7 +866,11 @@ class ServerArgs: self._handle_ssl_validation() # Validate PD disaggregation flags early (before dummy-model short-circuit). - self._handle_pd_disaggregation() + from sglang.srt.arg_groups.pd_disaggregation_hook import ( + handle_pd_disaggregation, + ) + + handle_pd_disaggregation(self) # Validate --prefill-only-disable-kv-cache args early (before dummy-model # short-circuit). The backend check is run later after backends settle. @@ -3748,62 +3759,6 @@ class ServerArgs: except Exception: return False - def _handle_pd_disaggregation(self): - if self.disaggregation_mode == "decode": - if self.disaggregation_decode_enable_radix_cache: - if self.enable_hisparse: - raise ValueError( - "--disaggregation-decode-enable-radix-cache is incompatible " - "with --enable-hisparse" - ) - if self.disaggregation_transfer_backend not in ("nixl", "mooncake"): - raise ValueError( - "--disaggregation-decode-enable-radix-cache currently " - "requires --disaggregation-transfer-backend in " - "('nixl', 'mooncake'), but got " - f"{self.disaggregation_transfer_backend!r}" - ) - if self.speculative_algorithm is not None: - raise ValueError( - "--disaggregation-decode-enable-radix-cache is incompatible " - "with speculative decoding " - f"(--speculative-algorithm {self.speculative_algorithm})" - ) - if self.enable_dp_attention: - logger.warning( - "EXPERIMENTAL: Decode radix cache with DP attention. " - "Requires prefix-aware DP rank routing for optimal cache hits." - ) - self.disable_radix_cache = False - logger.warning("EXPERIMENTAL: Radix cache is enabled for decode server") - else: - self.disable_radix_cache = True - logger.warning("KV cache is forced as chunk cache for decode server") - if self.enable_mamba_extra_buffer(): - logger.warning( - "Mamba extra_buffer is disabled because decode disaggregation " - "currently forces chunk cache. Falling back to no_buffer." - ) - self.mamba_scheduler_strategy = "no_buffer" - - elif self.disaggregation_mode == "prefill": - assert ( - self.disaggregation_transfer_backend != "fake" - ), "Prefill server does not support 'fake' as the transfer backend" - - self.disable_cuda_graph = True - - if self.disaggregation_mode in ("prefill", "decode"): - if ( - envs.SGLANG_DISAGG_STAGING_BUFFER.get() - and self.disaggregation_transfer_backend not in ("mooncake", "nixl") - ): - raise ValueError( - f"SGLANG_DISAGG_STAGING_BUFFER requires " - f"disaggregation_transfer_backend='mooncake' or 'nixl', " - f"got '{self.disaggregation_transfer_backend}'." - ) - def _handle_encoder_disaggregation(self): if self.enable_prefix_mm_cache and not self.encoder_only: raise ValueError(