Take the parallel getters off the package's public surface (#40344)

python/sglang/srt/distributed/bootstrap.py

@@ -7,7 +7,9 @@ import torch
 import torch.distributed as dist
 
 from sglang.srt.configs.model_config import ModelConfig
-from sglang.srt.distributed import (
+from sglang.srt.distributed.gated_launch import maybe_wait_for_gated_launch
+from sglang.srt.distributed.parallel_state import (
+    _tag_groups_for_flashinfer_allreduce_only,
     get_default_distributed_backend,
     get_tp_group,
     get_world_group,
@@ -18,10 +20,6 @@ from sglang.srt.distributed import (
     set_mscclpp_all_reduce,
     set_torch_symm_mem_all_reduce,
 )
-from sglang.srt.distributed.gated_launch import maybe_wait_for_gated_launch
-from sglang.srt.distributed.parallel_state import (
-    _tag_groups_for_flashinfer_allreduce_only,
-)
 from sglang.srt.environ import envs
 from sglang.srt.layers.dp_attention import initialize_dp_attention
 from sglang.srt.layers.layernorm_sp import initialize_layernorm_sp

python/sglang/srt/distributed/device_communicators/pynccl_allocator.py

@@ -427,7 +427,7 @@ def prealloc_symmetric_memory_pool(
     ):
         return
 
-    from sglang.srt.distributed import get_tp_group
+    from sglang.srt.distributed.parallel_state import get_tp_group
 
     # Memory allocation is tied to a cuda stream, use the forward stream
     with torch.get_device_module(device).stream(forward_stream):

python/sglang/srt/distributed/device_communicators/triton_symm_mem_ag.py

@@ -469,8 +469,10 @@ class MultimemAllGatherer:
         self._state = self._UNINIT if enabled else None
         if self._state is self._UNINIT:
             # Lazy import avoids a module-load dependency on the distributed facade.
-            from sglang.srt.distributed import get_tp_group
-            from sglang.srt.distributed.parallel_state import in_the_same_node_as
+            from sglang.srt.distributed.parallel_state import (
+                get_tp_group,
+                in_the_same_node_as,
+            )
 
             tp_group = get_tp_group()
             # Only probe node topology when the deployment can actually span
@@ -527,7 +529,7 @@ class MultimemAllGatherer:
         if x.shape[-1] % _NUMEL_PER_THREAD != 0:
             return None
         try:
-            from sglang.srt.distributed import get_tp_group
+            from sglang.srt.distributed.parallel_state import get_tp_group
 
             tp_group = get_tp_group()
             if tp_group.world_size <= 1:

python/sglang/srt/distributed/gated_launch.py

@@ -9,7 +9,7 @@ import uvicorn
 from fastapi import FastAPI
 from fastapi.responses import PlainTextResponse
 
-from sglang.srt.distributed import get_world_group
+from sglang.srt.distributed.parallel_state import get_world_group
 
 logger = logging.getLogger(__name__)
 

python/sglang/srt/distributed/parallel_state.py

@@ -3534,3 +3534,14 @@ for _name, _replacement in _CONTEXT_NAME_OF.items():
     if _fn is not None:
         globals()[_name] = _warn_if_called_from_outside(_name, _replacement)(_fn)
 del _name, _replacement, _fn
+
+
+# What `from sglang.srt.distributed import *` re-exports: everything public
+# except the deprecated getters. Business code reaches them through
+# `get_parallel()`, and the package that defines them imports them from this
+# module by name, so nothing needs the package path to reach one.
+__all__ = [
+    _public
+    for _public in list(globals())
+    if not _public.startswith("_") and _public not in _CONTEXT_NAME_OF
+]

python/sglang/srt/layers/dp_attention.py

@@ -15,9 +15,6 @@ from sglang.srt.arg_groups.model_override_base import (
 )
 from sglang.srt.distributed import (
     GroupCoordinator,
-)
-from sglang.srt.distributed import get_moe_dp_group as _get_moe_dp_group
-from sglang.srt.distributed import (
     tensor_model_parallel_all_reduce,
 )
 from sglang.srt.distributed.device_communicators.pynccl_allocator import (
@@ -1069,15 +1066,15 @@ def attn_cp_all_gather_into_tensor(output: torch.Tensor, input: torch.Tensor):
 
 def get_moe_cp_group() -> GroupCoordinator:
     """Returns the MOE_DP group, which includes CP partners when attn_cp_size > moe_dp_size."""
-    return _get_moe_dp_group()
+    return get_parallel().moe_dp_group
 
 
 def get_moe_cp_rank() -> int:
-    return _get_moe_dp_group().rank_in_group
+    return get_parallel().moe_dp_group.rank_in_group
 
 
 def get_moe_cp_size() -> int:
-    return _get_moe_dp_group().world_size
+    return get_parallel().moe_dp_group.world_size
 
 
 def is_enable_moe_cp_allgather() -> bool:
@@ -1090,7 +1087,7 @@ def is_enable_moe_cp_allgather() -> bool:
 
 
 def moe_cp_all_gather_into_tensor(output: torch.Tensor, input: torch.Tensor):
-    return _get_moe_dp_group().all_gather_into_tensor(output, input)
+    return get_parallel().moe_dp_group.all_gather_into_tensor(output, input)
 
 
 def attn_tp_all_gather(output_list: List[torch.Tensor], input: torch.Tensor):