diff --git a/python/sglang/srt/disaggregation/encode_server.py b/python/sglang/srt/disaggregation/encode_server.py index 12e36fb16..e1eabe0bd 100644 --- a/python/sglang/srt/disaggregation/encode_server.py +++ b/python/sglang/srt/disaggregation/encode_server.py @@ -25,6 +25,7 @@ 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.distributed.parallel_state import ( + get_default_distributed_backend, get_mooncake_transfer_engine, get_tp_group, init_distributed_environment, @@ -176,6 +177,7 @@ class MMEncoder: self.use_image_processor_gpu = use_image_processor_gpu init_distributed_environment( + backend=get_default_distributed_backend(self.device), world_size=server_args.tp_size, rank=rank, distributed_init_method=dist_init_method, diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index cd1577c7d..5b8a82232 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/python/sglang/srt/distributed/parallel_state.py @@ -1578,6 +1578,20 @@ def set_torch_symm_mem_all_reduce(enable: bool): _ENABLE_TORCH_SYMM_MEM_ALL_REDUCE = enable +_DEVICE_TO_DISTRIBUTED_BACKEND = { + "cuda": "nccl", + "xpu": "xccl", + "hpu": "hccl", + "cpu": "gloo", + "npu": "hccl", + "musa": "mccl", +} + + +def get_default_distributed_backend(device: str) -> str: + return _DEVICE_TO_DISTRIBUTED_BACKEND.get(device, "gloo") + + def init_distributed_environment( world_size: int = -1, rank: int = -1, diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 058030768..607766703 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -55,6 +55,7 @@ from sglang.srt.debug_utils.tensor_dump_forward_hook import ( register_forward_hook_for_model, ) from sglang.srt.distributed import ( + get_default_distributed_backend, get_pp_group, get_tp_group, get_world_group, @@ -734,29 +735,17 @@ class ModelRunner(ModelRunnerKVCacheMixin): ) raise - if self.device == "cuda": - if self.server_args.elastic_ep_backend == "mooncake": - backend = "mooncake" - if self.server_args.mooncake_ib_device: - mooncake_ib_device = self.server_args.mooncake_ib_device.split(",") - try: - from mooncake import ep as mooncake_ep + backend = get_default_distributed_backend(self.device) + if self.device == "cuda" and self.server_args.elastic_ep_backend == "mooncake": + backend = "mooncake" + if self.server_args.mooncake_ib_device: + mooncake_ib_device = self.server_args.mooncake_ib_device.split(",") + try: + from mooncake import ep as mooncake_ep - mooncake_ep.set_device_filter(mooncake_ib_device) - except: - pass # A warning will be raised in `init_distributed_environment` - else: - backend = "nccl" - elif self.device == "xpu": - backend = "xccl" - elif self.device == "hpu": - backend = "hccl" - elif self.device == "cpu": - backend = "gloo" - elif self.device == "npu": - backend = "hccl" - elif self.device == "musa": - backend = "mccl" + mooncake_ep.set_device_filter(mooncake_ib_device) + except: + pass # A warning will be raised in `init_distributed_environment` before_avail_memory = get_available_gpu_memory(self.device, self.gpu_id) if not self.server_args.enable_p2p_check: