From db58e76c334548416265c9c7a55b8ff6bf21b8f8 Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Sun, 7 Jun 2026 02:48:47 -0700 Subject: [PATCH] Add all_to_all_single to GroupCoordinator (#27492) Co-authored-by: Ming Yang --- .../sglang/srt/distributed/parallel_state.py | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index 810d752a6..8489e9f3f 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/python/sglang/srt/distributed/parallel_state.py @@ -196,6 +196,17 @@ def reg_reduce_scatter_tensor( group._reduce_scatter_tensor(output, input) +@register_custom_op(mutates_args=["output"]) +def reg_all_to_all_single( + output: torch.Tensor, input: torch.Tensor, group_name: str +) -> None: + assert group_name in _groups, f"Group {group_name} is not found." + group = _groups[group_name]() + if group is None: + raise ValueError(f"Group {group_name} is destroyed.") + group._all_to_all_single(output, input) + + class GroupCoordinator: """ PyTorch ProcessGroup wrapper for a group of processes. @@ -776,6 +787,15 @@ class GroupCoordinator: else: reg_reduce_scatter_tensor(output, input, group_name=self.unique_name) + def _all_to_all_single(self, output: torch.Tensor, input: torch.Tensor) -> None: + torch.distributed.all_to_all_single(output, input, group=self.device_group) + + def all_to_all_single(self, output: torch.Tensor, input: torch.Tensor): + if self.world_size == 1: + output.copy_(input) + return + reg_all_to_all_single(output, input, group_name=self.unique_name) + def reduce_scatter( self, output: torch.Tensor,