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

Signed-off-by: Shangming Cai <csmthu@gmail.com>

python/sglang/srt/disaggregation/ascend/transfer_engine.py

@@ -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

python/sglang/srt/disaggregation/encode_receiver.py

@@ -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 (
-from sglang.srt.distributed.parallel_state import GroupCoordinator
+    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(
+            self.embeddings_engine = get_mooncake_transfer_engine()
-                hostname=get_local_ip_auto(),
-                gpu_id=None,
-                ib_device=server_args.disaggregation_ib_device,
-            )
             self.embeddings_buffer = dict()
         elif self.encoder_transfer_backend == "zmq_to_scheduler":
             self.pp_rank = pp_rank

python/sglang/srt/disaggregation/encode_server.py

@@ -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 = get_mooncake_transfer_engine()
-                self.engine = MooncakeTransferEngine(
-                    hostname=self.local_ip,
-                    gpu_id=None,
-                    ib_device=server_args.disaggregation_ib_device,
-                )
 
             self.embedding_to_send = dict()
             self.background_tasks: Set[asyncio.Task] = set()

python/sglang/srt/disaggregation/mooncake/conn.py

@@ -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(
+        self.engine = get_mooncake_transfer_engine()
-            hostname=self.local_ip,
-            gpu_id=self.kv_args.gpu_id,
-            ib_device=self.kv_args.ib_device,
-        )
 
     def register_buffer_to_engine(self):
         # Batch register KV data buffers

python/sglang/srt/distributed/device_communicators/mooncake_transfer_engine.py

@@ -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.gpu_id = gpu_id if gpu_id is not None else 0
-        self.ib_device = get_ib_devices_for_gpu(ib_device, gpu_id)
+        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. "
+                    "Mooncake's batch register requires a newer version of "
-                    "Please upgrade Mooncake."
+                    "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. "
+                    "Mooncake's batch transfer requires mooncake-transfer-engine "
-                    "Please upgrade Mooncake by 'pip install mooncake-transfer-engine --upgrade'"
+                    ">= 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

python/sglang/srt/distributed/parallel_state.py

@@ -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):
     """

python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py

@@ -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,

python/sglang/srt/model_executor/model_runner.py

@@ -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)