minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Add mooncake_tcp transfer backend (mooncake over TCP) (#26346)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-05-26 18:15:55 -07:00
committed by GitHub
parent 737c6cd6d1
commit 6076066e38
3 changed files with 99 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/arg_groups/pd_disaggregation_hook.py
+80
View File
@@ -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}'."
)
python/sglang/srt/distributed/device_communicators/mooncake_transfer_engine.py
+6 -1
View File
@@ -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,
python/sglang/srt/server_args.py
+13 -58
View File
@@ -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(