chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Refactoring Mooncake TE as a shared distributed component (#17810)

Browse Source 
Signed-off-by: Shangming Cai <csmthu@gmail.com>
This commit is contained in:
Shangming Cai
2026-02-09 10:53:11 +08:00
committed by GitHub
parent bf89cc3803
commit bffd765417
8 changed files with 121 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/sglang/srt/disaggregation/ascend/transfer_engine.py
View File

@@ -4,8 +4,10 @@ from typing import List
import torch
from sglang.srt.disaggregation.mooncake.transfer_engine import MooncakeTransferEngine
from sglang.srt.disaggregation.utils import DisaggregationMode
from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import (
MooncakeTransferEngine,
)
try:
from memfabric_hybrid import TransferEngine

12
python/sglang/srt/disaggregation/encode_receiver.py
View File

@@ -14,8 +14,10 @@ import zmq
import zmq.asyncio
from transformers import PretrainedConfig
from sglang.srt.disaggregation.mooncake.transfer_engine import MooncakeTransferEngine
from sglang.srt.distributed.parallel_state import GroupCoordinator
from sglang.srt.distributed.parallel_state import (
GroupCoordinator,
get_mooncake_transfer_engine,
)
from sglang.srt.managers.io_struct import TokenizedGenerateReqInput
from sglang.srt.managers.multimodal_processor import get_mm_processor, import_processors
from sglang.srt.managers.schedule_batch import Req
@@ -287,11 +289,7 @@ class MMReceiverHTTP(MMReceiverBase):
self.host = get_local_ip_auto(server_args.host)
if self.encoder_transfer_backend == "mooncake":
self.dtype = dtype
self.embeddings_engine = MooncakeTransferEngine(
hostname=get_local_ip_auto(),
gpu_id=None,
ib_device=server_args.disaggregation_ib_device,
)
self.embeddings_engine = get_mooncake_transfer_engine()
self.embeddings_buffer = dict()
elif self.encoder_transfer_backend == "zmq_to_scheduler":
self.pp_rank = pp_rank

9
python/sglang/srt/disaggregation/encode_server.py
View File

@@ -24,8 +24,8 @@ from sglang.srt.configs.device_config import DeviceConfig
from sglang.srt.configs.load_config import LoadConfig
from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.disaggregation.encode_receiver import EmbeddingData
from sglang.srt.disaggregation.mooncake.transfer_engine import MooncakeTransferEngine
from sglang.srt.distributed.parallel_state import (
get_mooncake_transfer_engine,
init_distributed_environment,
initialize_model_parallel,
)
@@ -210,12 +210,7 @@ class MMEncoder:
if self.server_args.encoder_transfer_backend == "mooncake":
self.local_ip = get_local_ip_auto()
self.engine = MooncakeTransferEngine(
hostname=self.local_ip,
gpu_id=None,
ib_device=server_args.disaggregation_ib_device,
)
self.engine = get_mooncake_transfer_engine()
self.embedding_to_send = dict()
self.background_tasks: Set[asyncio.Task] = set()

8
python/sglang/srt/disaggregation/mooncake/conn.py
View File

@@ -27,11 +27,11 @@ from sglang.srt.disaggregation.common.utils import (
FastQueue,
group_concurrent_contiguous,
)
from sglang.srt.disaggregation.mooncake.transfer_engine import MooncakeTransferEngine
from sglang.srt.disaggregation.mooncake.utils import (
check_mooncake_custom_mem_pool_enabled,
)
from sglang.srt.disaggregation.utils import DisaggregationMode
from sglang.srt.distributed.parallel_state import get_mooncake_transfer_engine
from sglang.srt.environ import envs
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils import format_tcp_address, is_valid_ipv6_address
@@ -239,11 +239,7 @@ class MooncakeKVManager(CommonKVManager):
self.failure_lock = threading.Lock()
def init_engine(self):
self.engine = MooncakeTransferEngine(
hostname=self.local_ip,
gpu_id=self.kv_args.gpu_id,
ib_device=self.kv_args.ib_device,
)
self.engine = get_mooncake_transfer_engine()
def register_buffer_to_engine(self):
# Batch register KV data buffers

247
python/sglang/srt/disaggregation/mooncake/transfer_engine.py
View File

@@ -1,247 +0,0 @@
import json
import logging
import os
from typing import List, Optional
from sglang.srt.environ import envs
from sglang.srt.utils import get_free_port, maybe_wrap_ipv6_address
logger = logging.getLogger(__name__)
def get_ib_devices_for_gpu(ib_device_str: Optional[str], gpu_id: int) -> Optional[str]:
"""
Parse IB device string and get IB devices for a specific GPU ID.
Supports all the following formats:
1. Old format: "ib0, ib1, ib2"
2. New format: {0: "ib0, ib1", 1: "ib2, ib3", 2: "ib4"}
3. JSON file: path to a JSON file containing the mapping
Args:
ib_device_str: The original IB device string or path to JSON file
gpu_id: The GPU ID to get devices for
Returns:
IB devices string for the GPU, or None if not available
"""
if ib_device_str is None or not ib_device_str.strip():
return None
ib_device_str = ib_device_str.strip()
# Check if it's a JSON file first and load its content
is_json_file = ib_device_str.endswith(".json")
if is_json_file:
try:
if os.path.isfile(ib_device_str):
with open(ib_device_str, "r") as f:
ib_device_str = f.read()
else:
# File doesn't exist, treat as old format
raise RuntimeError(f"File {ib_device_str} does not exist.")
except (IOError, OSError) as e:
# File reading failed, raise exception
raise RuntimeError(f"Failed to read JSON file {ib_device_str}: {e}") from e
# Check if it's JSON format (new format)
try:
parsed_json = json.loads(ib_device_str)
if isinstance(parsed_json, dict):
# Validate format - keys should be integers (or string rep), values should be strings
gpu_mapping = {}
for gpu_key, ib_devices in parsed_json.items():
if (
isinstance(gpu_key, str)
and gpu_key.isdigit()
and isinstance(ib_devices, str)
):
gpu_mapping[int(gpu_key)] = ib_devices.strip()
elif isinstance(gpu_key, int) and isinstance(ib_devices, str):
gpu_mapping[gpu_key] = ib_devices.strip()
else:
raise ValueError(
f"Invalid format: keys must be integers (or string representations of integers) and values must be strings"
)
if not gpu_mapping:
raise ValueError("No valid GPU mappings found in JSON")
# Return devices for specific GPU
if gpu_id in gpu_mapping:
return gpu_mapping[gpu_id]
else:
raise ValueError(
f"No IB devices configured for GPU {gpu_id}. Available GPUs: {list(gpu_mapping.keys())}"
)
except json.JSONDecodeError:
if is_json_file:
# It was supposed to be a JSON file but failed to parse
raise RuntimeError(
f"Failed to parse JSON content from file {ib_device_str}"
)
# Not JSON format, treat as old format - return same devices for all GPUs
return ib_device_str
class MooncakeTransferEngine:
def __init__(self, hostname: str, gpu_id: int, ib_device: Optional[str] = None):
try:
from mooncake.engine import TransferEngine
except ImportError as e:
raise ImportError(
"Please install mooncake by following the instructions at "
"https://github.com/kvcache-ai/Mooncake/blob/main/doc/en/build.md " # noqa: E501
"to run SGLang with MooncakeTransferEngine."
) from e
self.engine = TransferEngine()
self.hostname = hostname
self.gpu_id = gpu_id
self.ib_device = get_ib_devices_for_gpu(ib_device, gpu_id)
self.initialize(
hostname=self.hostname,
device_name=self.ib_device,
)
self.session_id = (
f"{maybe_wrap_ipv6_address(self.hostname)}:{self.engine.get_rpc_port()}"
)
def register(self, ptr, length):
try:
ret_value = self.engine.register_memory(ptr, length)
except Exception:
# Mark register as failed
ret_value = -1
if ret_value != 0:
logger.debug("Mooncake memory registration %s failed.", ptr)
def deregister(self, ptr):
try:
ret_value = self.engine.unregister_memory(ptr)
except Exception:
# Mark deregister as failed
ret_value = -1
if ret_value != 0:
logger.debug("Mooncake memory deregistration %s failed.", ptr)
def batch_register(self, ptrs: List[int], lengths: List[int]) -> int:
"""Batch register multiple memory regions."""
try:
ret_value = self.engine.batch_register_memory(ptrs, lengths)
except Exception:
# Mark batch register as failed
ret_value = -1
if not hasattr(self.engine, "batch_register_memory"):
raise RuntimeError(
"Mooncake's batch register requires a newer version of mooncake-transfer-engine. "
"Please upgrade Mooncake."
)
if ret_value != 0:
logger.debug("Mooncake batch memory registration failed.")
return ret_value
def batch_deregister(self, ptrs: List[int]) -> int:
"""Batch deregister multiple memory regions."""
try:
ret_value = self.engine.batch_unregister_memory(ptrs)
except Exception:
# Mark batch deregister as failed
ret_value = -1
if ret_value != 0:
logger.debug("Mooncake batch memory deregistration failed.")
return ret_value
def initialize(
self,
hostname: str,
device_name: Optional[str],
) -> None:
"""Initialize the mooncake instance."""
if envs.ENABLE_ASCEND_TRANSFER_WITH_MOONCAKE.get():
npu_phy_id = envs.ASCEND_NPU_PHY_ID.get()
if npu_phy_id == -1:
hostname += f":{get_free_port()}:npu_{self.gpu_id}"
else:
hostname += f":{get_free_port()}:npu_{npu_phy_id}"
ret_value = self.engine.initialize(
hostname,
"P2PHANDSHAKE",
"ascend",
device_name if device_name is not None else "",
)
else:
ret_value = self.engine.initialize(
hostname,
"P2PHANDSHAKE",
"rdma",
device_name if device_name is not None else "",
)
if ret_value != 0:
logger.error("Mooncake Transfer Engine initialization failed.")
raise RuntimeError("Mooncake Transfer Engine initialization failed.")
def transfer_sync(
self, session_id: str, buffer: int, peer_buffer_address: int, length: int
) -> int:
"""Synchronously transfer data to the specified address."""
try:
# the first time: based on session_id (which contains remote_ip) to construct a queue pair, and cache the queue pair
# later: based on the cached queue pair to send data
ret = self.engine.transfer_sync_write(
session_id, buffer, peer_buffer_address, length
)
except Exception:
# Mark transfer request as failed
ret = -1
if ret < 0:
# Do not raise an exception here, since some transfer requests fail should be accepted and the execution thread should not be stopped.
logger.debug(
"Failed to transfer data from %s to %s - %s.",
buffer,
session_id,
peer_buffer_address,
)
return ret
def batch_transfer_sync(
self,
session_id: str,
buffers: List[int],
peer_buffer_addresses: List[int],
lengths: List[int],
) -> int:
"""Synchronously transfer data to the specified addresses in batches."""
try:
ret = self.engine.batch_transfer_sync_write(
session_id, buffers, peer_buffer_addresses, lengths
)
except Exception:
ret = -1
# Inform user to upgrade mooncake-transfer-engine >= 0.3.4.post2
if not hasattr(self.engine, "batch_transfer_sync_write"):
raise RuntimeError(
"Mooncake's batch transfer requires mooncake-transfer-engine >= 0.3.4.post2. "
"Please upgrade Mooncake by 'pip install mooncake-transfer-engine --upgrade'"
)
if ret < 0:
logger.debug(
"Failed to batch transfer data. Buffers: %s, Session: %s, Peer addresses: %s",
buffers,
session_id,
peer_buffer_addresses,
)
return ret
def get_session_id(self):
return self.session_id