From f5c9f88ee284d05c55fcd1484e1af05c67b162a7 Mon Sep 17 00:00:00 2001 From: Yihao Wang <42559837+AgainstEntropy@users.noreply.github.com> Date: Thu, 11 Jun 2026 20:05:17 -0700 Subject: [PATCH] [plugin][distributed] use active platform's backend in `get_default_distributed_backend` (#23969) --- .../sglang/srt/distributed/parallel_state.py | 19 +++-- python/sglang/srt/platforms/__init__.py | 5 +- python/sglang/srt/platforms/device_mixin.py | 21 ++++- .../test_get_default_distributed_backend.py | 77 +++++++++++++++++++ 4 files changed, 109 insertions(+), 13 deletions(-) create mode 100644 test/registered/unit/distributed/test_get_default_distributed_backend.py diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index 95920bc1b..62dc01641 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/python/sglang/srt/distributed/parallel_state.py @@ -42,12 +42,14 @@ import torch import torch.distributed from torch.distributed import Backend, ProcessGroup +from sglang.srt import platforms from sglang.srt.compilation.compilation_config import register_split_op from sglang.srt.distributed.utils import set_global_tcp_store from sglang.srt.environ import envs from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( is_in_tc_piecewise_cuda_graph, ) +from sglang.srt.platforms.device_mixin import _DEVICE_TO_DISTRIBUTED_BACKEND from sglang.srt.utils import ( get_current_device_stream_fast, get_int_env_var, @@ -1690,17 +1692,14 @@ 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" if not envs.SGLANG_ZBAL_LOCAL_MEM_SIZE.get() > 0 else "zbal", - "musa": "mccl", -} - - +# TODO: refactor in-tree platforms to get rid of this wrapper def get_default_distributed_backend(device: str) -> str: + # We deliberately go through ``platforms.current_platform`` (rather than + # ``from ... import current_platform``) so each call resolves through the + # platforms package's lazy ``__getattr__`` and picks up runtime overrides + # of ``_current_platform`` (e.g. in tests). + if device == platforms.current_platform.device_type: + return platforms.current_platform.get_torch_distributed_backend_str() return _DEVICE_TO_DISTRIBUTED_BACKEND.get(device, "gloo") diff --git a/python/sglang/srt/platforms/__init__.py b/python/sglang/srt/platforms/__init__.py index c64a187a7..849f0c312 100644 --- a/python/sglang/srt/platforms/__init__.py +++ b/python/sglang/srt/platforms/__init__.py @@ -139,7 +139,10 @@ def _load_platform_class(qualname: str) -> type: return cls -def __getattr__(name: str) -> SRTPlatform: +current_platform: SRTPlatform + + +def __getattr__(name: str): """Lazy initialization of current_platform on first access.""" if name == "current_platform": global _current_platform diff --git a/python/sglang/srt/platforms/device_mixin.py b/python/sglang/srt/platforms/device_mixin.py index f5fa42ae6..a87523aa7 100644 --- a/python/sglang/srt/platforms/device_mixin.py +++ b/python/sglang/srt/platforms/device_mixin.py @@ -32,6 +32,8 @@ from typing import NamedTuple, Optional import numpy as np import torch +from sglang.srt.environ import envs + class PlatformEnum(enum.Enum): """Enumeration of known platform types. @@ -79,6 +81,16 @@ class DeviceCapability(NamedTuple): return self.major * 10 + self.minor +_DEVICE_TO_DISTRIBUTED_BACKEND: dict[str, str] = { + "cuda": "nccl", + "xpu": "xccl", + "hpu": "hccl", + "cpu": "gloo", + "npu": "hccl" if not envs.SGLANG_ZBAL_LOCAL_MEM_SIZE.get() > 0 else "zbal", + "musa": "mccl", +} + + class DeviceMixin: """Mixin providing device identity queries and basic device operations. @@ -192,8 +204,13 @@ class DeviceMixin: # ---- Distributed ---- def get_torch_distributed_backend_str(self) -> str: - """[Planned] Return the torch.distributed backend string (e.g. "nccl", "hccl").""" - raise NotImplementedError + """Return the torch.distributed backend string (e.g. "nccl", "hccl"). + + Default: lookup ``self.device_type`` in ``_DEVICE_TO_DISTRIBUTED_BACKEND``, + falling back to ``"gloo"``. Subclasses override only when they need a + non-default backend (e.g. mooncake, or a brand-new device). + """ + return _DEVICE_TO_DISTRIBUTED_BACKEND.get(self.device_type, "gloo") def get_communicator_class(self) -> type | None: """[Planned] Return platform-specific communicator class, or None for default.""" diff --git a/test/registered/unit/distributed/test_get_default_distributed_backend.py b/test/registered/unit/distributed/test_get_default_distributed_backend.py new file mode 100644 index 000000000..c6fa3917d --- /dev/null +++ b/test/registered/unit/distributed/test_get_default_distributed_backend.py @@ -0,0 +1,77 @@ +"""Unit tests for sglang.srt.distributed.parallel_state — no server, no model loading.""" + +import unittest + +import sglang.srt.platforms as platforms_mod +from sglang.srt.distributed.parallel_state import get_default_distributed_backend +from sglang.srt.platforms.interface import SRTPlatform +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=2, suite="base-a-test-cpu") + + +# Worked example: an out-of-tree platform that overrides the torch backend. +# Future plugin authors can pattern-match against this in their own tests. +class _OverridingPlatform(SRTPlatform): + device_type = "cuda" + + def get_torch_distributed_backend_str(self) -> str: + return "fake_backend" + + +# Mirrors all in-tree platforms today: no override, so DeviceMixin's default +# body returns _DEVICE_TO_DISTRIBUTED_BACKEND.get(device_type, "gloo"). +class _DefaultPlatform(SRTPlatform): + device_type = "cuda" + + +class TestGetDefaultDistributedBackend(CustomTestCase): + """Cover all paths of get_default_distributed_backend. + + When ``device == current_platform.device_type`` the dispatcher asks + ``current_platform.get_torch_distributed_backend_str()``; otherwise it + looks up ``_DEVICE_TO_DISTRIBUTED_BACKEND[device]`` for cross-device + queries (e.g. auxiliary "cpu"/gloo groups on a CUDA process). + + Tests use real SRTPlatform subclasses so the assertions stay close to + how an actual out-of-tree plugin would interact with this dispatcher. + + ``current_platform`` is exposed via ``__getattr__`` on the platforms + module backed by the ``_current_platform`` singleton, so the test + overrides that singleton directly and restores it in tearDown. + """ + + def setUp(self): + self._saved_platform = platforms_mod._current_platform + + def tearDown(self): + platforms_mod._current_platform = self._saved_platform + + def _install(self, platform: SRTPlatform) -> None: + platforms_mod._current_platform = platform + + def test_overriding_platform_supplies_backend(self): + self._install(_OverridingPlatform()) + self.assertEqual(get_default_distributed_backend("cuda"), "fake_backend") + + def test_overriding_platform_skipped_for_non_active_device(self): + # Even when current_platform overrides get_torch_distributed_backend_str, + # callers asking for a different device (e.g. an auxiliary "cpu" gloo + # group on a CUDA process) must keep going through the dict. + self._install(_OverridingPlatform()) + self.assertEqual(get_default_distributed_backend("cpu"), "gloo") + + def test_default_platform_uses_device_mixin_table(self): + # No override: DeviceMixin's default body looks up the device_type in + # _DEVICE_TO_DISTRIBUTED_BACKEND, so cuda still resolves to nccl. + self._install(_DefaultPlatform()) + self.assertEqual(get_default_distributed_backend("cuda"), "nccl") + + def test_unknown_device_returns_gloo_default(self): + self._install(_DefaultPlatform()) + self.assertEqual(get_default_distributed_backend("unobtanium"), "gloo") + + +if __name__ == "__main__": + unittest.main()