[plugin][distributed] use active platform's backend in get_default_distributed_backend (#23969)

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")
 
 

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

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."""

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()