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__) # 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]: """ 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( "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}. " f"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: """Shared Mooncake transfer engine for RDMA/transfer operations.""" 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://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 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, 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: ret = self.engine.transfer_sync_write( session_id, buffer, peer_buffer_address, length ) except Exception: ret = -1 if ret < 0: 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 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 def get_engine(self): return self.engine.get_engine() def get_ib_device(self): return self.ib_device 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