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

68
python/sglang/srt/disaggregation/mooncake/transfer_engine.py → python/sglang/srt/distributed/device_communicators/mooncake_transfer_engine.py
View File

@@ -8,6 +8,9 @@ from sglang.srt.utils import get_free_port, maybe_wrap_ipv6_address
logger = logging.getLogger(__name__)
# Module-level shared engine instance, set by init_mooncake_transfer_engine().
_mooncake_transfer_engine: Optional["MooncakeTransferEngine"] = None
def get_ib_devices_for_gpu(ib_device_str: Optional[str], gpu_id: int) -> Optional[str]:
"""
@@ -61,7 +64,8 @@ def get_ib_devices_for_gpu(ib_device_str: Optional[str], gpu_id: int) -> Optiona
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"
"Invalid format: keys must be integers (or string "
"representations of integers) and values must be strings"
)
if not gpu_mapping:
@@ -72,7 +76,8 @@ def get_ib_devices_for_gpu(ib_device_str: Optional[str], gpu_id: int) -> Optiona
return gpu_mapping[gpu_id]
else:
raise ValueError(
f"No IB devices configured for GPU {gpu_id}. Available GPUs: {list(gpu_mapping.keys())}"
f"No IB devices configured for GPU {gpu_id}. "
f"Available GPUs: {list(gpu_mapping.keys())}"
)
except json.JSONDecodeError:
@@ -86,21 +91,27 @@ def get_ib_devices_for_gpu(ib_device_str: Optional[str], gpu_id: int) -> Optiona
class MooncakeTransferEngine:
"""Shared Mooncake transfer engine for RDMA/transfer operations."""
def __init__(self, hostname: str, gpu_id: int, ib_device: Optional[str] = None):
def __init__(
self,
hostname: str,
gpu_id: Optional[int] = None,
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
"https://kvcache-ai.github.io/Mooncake/getting_started/build.html "
"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.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)
self.initialize(
hostname=self.hostname,
@@ -139,8 +150,8 @@ class MooncakeTransferEngine:
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."
"Mooncake's batch register requires a newer version of "
"mooncake-transfer-engine. Please upgrade Mooncake."
)
if ret_value != 0:
@@ -193,17 +204,13 @@ class MooncakeTransferEngine:
) -> 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,
@@ -227,16 +234,17 @@ class MooncakeTransferEngine:
)
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'"
"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",
"Failed to batch transfer data. Buffers: %s, Session: %s, "
"Peer addresses: %s",
buffers,
session_id,
peer_buffer_addresses,
@@ -245,3 +253,31 @@ class MooncakeTransferEngine:
def get_session_id(self):
return self.session_id
def get_engine(self):
return self.engine.get_engine()
def init_mooncake_transfer_engine(
hostname: str,
gpu_id: Optional[int] = None,
ib_device: Optional[str] = None,
) -> MooncakeTransferEngine:
"""
Initialize the shared MooncakeTransferEngine. Note: if already
initialized with the same (hostname, gpu_id, ib_device), returns existing
instance. Call from parallel_state when model parallel is set up and
mooncake transfer is needed.
"""
global _mooncake_transfer_engine
if _mooncake_transfer_engine is not None:
return _mooncake_transfer_engine
_mooncake_transfer_engine = MooncakeTransferEngine(
hostname=hostname, gpu_id=gpu_id, ib_device=ib_device
)
return _mooncake_transfer_engine
def get_mooncake_transfer_engine() -> Optional[MooncakeTransferEngine]:
"""Return the shared MooncakeTransferEngine if initialized, else None."""
return _mooncake_transfer_engine

12
python/sglang/srt/distributed/parallel_state.py
View File

@@ -1433,6 +1433,18 @@ def get_pp_group() -> GroupCoordinator:
get_pipeline_model_parallel_group = get_pp_group
def get_mooncake_transfer_engine():
"""
Return the shared MooncakeTransferEngine if initialized in device_communicators,
else None. Used by disaggregation mooncake backend and mem_cache mooncake_store.
"""
from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import (
get_mooncake_transfer_engine as _get_engine,
)
return _get_engine()
@contextmanager
def graph_capture(stream: Optional[torch.cuda.Stream] = None):
"""

2
python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py
View File

@@ -310,6 +310,8 @@ class MooncakeStore(HiCacheStorage):
self.config.client_server_address,
)
else:
# TODO(shangming): disable mooncake transfer engine reuse for hicache temporary
# Need to wait for the next mooncake release
ret_code = self.store.setup(
self.config.local_hostname,
self.config.metadata_server,

43
python/sglang/srt/model_executor/model_runner.py
View File

@@ -373,6 +373,9 @@ class ModelRunner(ModelRunnerKVCacheMixin):
if self.device == "cpu":
self.init_threads_binding()
# Initialize MooncakeTransferEngine
self.init_shared_mooncake_transfer_engine()
# Get memory before model loading
min_per_gpu_memory = self.init_torch_distributed()
@@ -828,6 +831,46 @@ class ModelRunner(ModelRunnerKVCacheMixin):
)
return min_per_gpu_memory
def init_shared_mooncake_transfer_engine(self):
"""
Need MooncakeTransferEngine when:
1) PD disaggregation uses mooncake for KV transfer (prefill/decode)
2) HiCache uses mooncake storage backend
3) Encoder disaggregation uses mooncake
"""
use_mooncake_te = (
(
self.server_args.disaggregation_mode != "null"
and self.server_args.disaggregation_transfer_backend == "mooncake"
)
or (
self.server_args.enable_hierarchical_cache
and self.server_args.hicache_storage_backend == "mooncake"
)
or (
self.server_args.encoder_only
and self.server_args.encoder_transfer_backend == "mooncake"
)
or (
self.server_args.language_only
and self.server_args.encoder_transfer_backend == "mooncake"
)
)
if use_mooncake_te:
from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import (
init_mooncake_transfer_engine,
)
init_mooncake_transfer_engine(
hostname=get_local_ip_auto(),
gpu_id=self.gpu_id,
ib_device=(
self.server_args.disaggregation_ib_device
or self.server_args.mooncake_ib_device
),
)
def load_model(self):
tic_total = time.perf_counter()
before_avail_memory = get_available_gpu_memory(self.device, self.gpu_id)