diff --git a/python/sglang/srt/disaggregation/ascend/transfer_engine.py b/python/sglang/srt/disaggregation/ascend/transfer_engine.py index b7f0cbd92..0a80ee752 100644 --- a/python/sglang/srt/disaggregation/ascend/transfer_engine.py +++ b/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 diff --git a/python/sglang/srt/disaggregation/encode_receiver.py b/python/sglang/srt/disaggregation/encode_receiver.py index 25da76226..e33084c09 100644 --- a/python/sglang/srt/disaggregation/encode_receiver.py +++ b/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 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 diff --git a/python/sglang/srt/disaggregation/encode_server.py b/python/sglang/srt/disaggregation/encode_server.py index 47727465c..ba8663525 100644 --- a/python/sglang/srt/disaggregation/encode_server.py +++ b/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 = 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() diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 6b9ffc64d..0c5b87533 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/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( - 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 diff --git a/python/sglang/srt/disaggregation/mooncake/transfer_engine.py b/python/sglang/srt/distributed/device_communicators/mooncake_transfer_engine.py similarity index 79% rename from python/sglang/srt/disaggregation/mooncake/transfer_engine.py rename to python/sglang/srt/distributed/device_communicators/mooncake_transfer_engine.py index 67e4d4fc6..447d9cf58 100644 --- a/python/sglang/srt/disaggregation/mooncake/transfer_engine.py +++ b/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.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 diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index d45c87d72..df3499a65 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/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): """ diff --git a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py index f661e73fa..479b39ad5 100644 --- a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py +++ b/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, diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index ac34bf171..3fe264646 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/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)