chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Reuse initialized transfer engine in mooncake store (#18460)

Browse Source 
Signed-off-by: Shangming Cai <csmthu@gmail.com>
This commit is contained in:
Shangming Cai
2026-02-13 01:21:35 +08:00
committed by GitHub
parent 0236522214
commit 2a8a48c0ca
2 changed files with 37 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
python/sglang/srt/distributed/device_communicators/mooncake_transfer_engine.py
View File

@@ -257,6 +257,9 @@ class MooncakeTransferEngine:
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,

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

@@ -310,16 +310,47 @@ 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
try:
from sglang.srt.distributed.parallel_state import (
get_mooncake_transfer_engine,
)
self._shared_mooncake_transfer_engine = (
get_mooncake_transfer_engine()
)
except Exception:
self._shared_mooncake_transfer_engine = None
logger.debug("Failed to reuse initialized mooncake transfer engine")
# Only reuse the shared MooncakeTransferEngine when its
# configuration matches the one used by MooncakeStore.
if (
self._shared_mooncake_transfer_engine is not None
and device_name
== self._shared_mooncake_transfer_engine.get_ib_device()
and self.config.metadata_server == "P2PHANDSHAKE"
and self.config.protocol == "rdma"
):
client_hostname = (
self._shared_mooncake_transfer_engine.get_session_id()
)
transfer_engine = self._shared_mooncake_transfer_engine.get_engine()
logger.info(
f"Reuse initialized mooncake transfer engine: {self._shared_mooncake_transfer_engine}"
)
else:
client_hostname = self.config.local_hostname
transfer_engine = None
ret_code = self.store.setup(
self.config.local_hostname,
client_hostname,
self.config.metadata_server,
per_tp_global_segment_size,
DEFAULT_LOCAL_BUFFER_SIZE, # Zero copy interface does not need local buffer
self.config.protocol,
device_name,
self.config.master_server_address,
transfer_engine,
)
if ret_code:
raise RuntimeError(