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,