From 5850aa14c3e929cc86a4ba484ec00f9058e606a8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=8E=8B=E9=B9=A4=E7=94=B7?= Date: Fri, 29 May 2026 14:21:15 +0800 Subject: [PATCH] fix(mooncake): honour MOONCAKE_PROTOCOL so EFA hardware can select efa transport (#25083) Co-authored-by: whn09 --- .../mooncake_transfer_engine.py | 29 +++++++++---------- python/sglang/srt/environ.py | 2 +- .../sglang/srt/model_executor/model_runner.py | 5 +++- 3 files changed, 18 insertions(+), 18 deletions(-) diff --git a/python/sglang/srt/distributed/device_communicators/mooncake_transfer_engine.py b/python/sglang/srt/distributed/device_communicators/mooncake_transfer_engine.py index 9efefe177..0b5653bb0 100644 --- a/python/sglang/srt/distributed/device_communicators/mooncake_transfer_engine.py +++ b/python/sglang/srt/distributed/device_communicators/mooncake_transfer_engine.py @@ -183,23 +183,20 @@ class MooncakeTransferEngine: """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 "", - ) + suffix = self.gpu_id if npu_phy_id == -1 else npu_phy_id + hostname += f":{get_free_port()}:npu_{suffix}" + protocol = "ascend" else: - ret_value = self.engine.initialize( - hostname, - "P2PHANDSHAKE", - "rdma", - device_name if device_name is not None else "", - ) + # MOONCAKE_PROTOCOL selects the transport (rdma | efa | tcp | ...). + # Default is "rdma"; set MOONCAKE_PROTOCOL=efa on AWS EFA hardware. + protocol = envs.MOONCAKE_PROTOCOL.get() + + ret_value = self.engine.initialize( + hostname, + "P2PHANDSHAKE", + protocol, + 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.") diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 6e66ad6fe..1fd02be23 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -356,7 +356,7 @@ class Envs: MOONCAKE_LOCAL_HOSTNAME = EnvStr("localhost") MOONCAKE_TE_META_DATA_SERVER = EnvStr("P2PHANDSHAKE") MOONCAKE_GLOBAL_SEGMENT_SIZE = EnvStr("4gb") - MOONCAKE_PROTOCOL = EnvStr("tcp") + MOONCAKE_PROTOCOL = EnvStr("rdma") MOONCAKE_DEVICE = EnvStr("") MOONCAKE_MASTER_METRICS_PORT = EnvInt(9003) MOONCAKE_CHECK_SERVER = EnvBool(False) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 6d9db720a..54d8fafba 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -925,7 +925,10 @@ class ModelRunner(ModelRunnerKVCacheMixin): self.remote_instance_transfer_engine = TransferEngine() local_ip = get_local_ip_auto() self.remote_instance_transfer_engine.initialize( - local_ip, "P2PHANDSHAKE", "rdma", envs.MOONCAKE_DEVICE.get() + local_ip, + "P2PHANDSHAKE", + envs.MOONCAKE_PROTOCOL.get(), + envs.MOONCAKE_DEVICE.get(), ) self.remote_instance_transfer_engine_session_id = NetworkAddress( local_ip, self.remote_instance_transfer_engine.get_rpc_port()