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

Signed-off-by: Shangming Cai <csmthu@gmail.com>
This commit is contained in:
Shangming Cai
2026-02-09 10:53:11 +08:00
committed by GitHub
parent bf89cc3803
commit bffd765417
8 changed files with 121 additions and 37 deletions
@@ -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