[diffusion] feat: use srt custom allreduce for tp groups (#28324)

python/sglang/multimodal_gen/runtime/distributed/group_coordinator.py

@@ -158,6 +158,7 @@ class GroupCoordinator:
         local_rank: int,
         torch_distributed_backend: Union[str, Backend],
         use_device_communicator: bool = True,
+        use_srt_custom_allreduce: bool = False,
         use_message_queue_broadcaster: bool = False,
         group_name: str | None = None,
     ):
@@ -213,11 +214,29 @@ class GroupCoordinator:
                 )
 
         self.mq_broadcaster = None
+        self.srt_custom_allreduce = None
+        if (
+            use_srt_custom_allreduce
+            and current_platform.is_cuda_alike()
+            and self.world_size > 1
+        ):
+            # srt owns topology, dtype, contiguity, and size dispatch for custom ar
+            self._init_srt_custom_allreduce()
 
         # TODO(will): check if this is needed
         # self.use_custom_op_call = current_platform.is_cuda_alike()
         self.use_custom_op_call = False
 
+    def _init_srt_custom_allreduce(self) -> None:
+        from sglang.srt.distributed.device_communicators.custom_all_reduce import (
+            CustomAllreduce,
+        )
+
+        self.srt_custom_allreduce = CustomAllreduce(
+            group=self.cpu_group,
+            device=self.device,
+        )
+
     @property
     def first_rank(self):
         """Return the global rank of the first process in the group"""
@@ -326,6 +345,18 @@ class GroupCoordinator:
         if self.world_size == 1:
             return input_
         else:
+            custom_ar = self.srt_custom_allreduce
+            if (
+                not async_op
+                and custom_ar is not None
+                and op == torch.distributed.ReduceOp.SUM
+                and not input_.is_cpu
+                and not custom_ar.disabled
+                and custom_ar.should_custom_ar(input_)
+            ):
+                if custom_ar._IS_CAPTURING:
+                    return custom_ar.custom_all_reduce(input_)
+                return custom_ar._all_reduce_impl(input_, registered=False)
             if (
                 current_platform.is_cpu()
                 and is_shm_available(input_.dtype, self.world_size, len(self.ranks))
@@ -769,6 +800,9 @@ class GroupCoordinator:
             self.cpu_group = None
         if self.device_communicator is not None:
             self.device_communicator.destroy()
+        if self.srt_custom_allreduce is not None:
+            self.srt_custom_allreduce.close()
+            self.srt_custom_allreduce = None
         if self.mq_broadcaster is not None:
             self.mq_broadcaster = None
 

python/sglang/multimodal_gen/runtime/distributed/parallel_state.py

@@ -139,6 +139,20 @@ def init_world_group(
     )
 
 
+def _sync_srt_world_group() -> None:
+    import sglang.srt.distributed.parallel_state as srt_parallel_state
+
+    if srt_parallel_state._WORLD is None:
+        srt_parallel_state._WORLD = _WORLD
+
+
+def _clear_srt_world_group() -> None:
+    import sglang.srt.distributed.parallel_state as srt_parallel_state
+
+    if srt_parallel_state._WORLD is _WORLD:
+        srt_parallel_state._WORLD = None
+
+
 def init_parallel_group_coordinator(
     group_ranks: List[List[int]],
     local_rank: int,
@@ -175,8 +189,14 @@ def init_parallel_group_coordinator(
             group_ranks=group_ranks,
             local_rank=local_rank,
             torch_distributed_backend=backend,
+            use_device_communicator=parallel_mode != "tensor",
+            use_srt_custom_allreduce=parallel_mode == "tensor",
             group_name=(
-                "vae_decode_group" if parallel_mode == "vae_decode" else "cfg_group"
+                "tp_group"
+                if parallel_mode == "tensor"
+                else (
+                    "vae_decode_group" if parallel_mode == "vae_decode" else "cfg_group"
+                )
             ),
         )
 
@@ -264,6 +284,7 @@ def init_distributed_environment(
         assert (
             _WORLD.world_size == torch.distributed.get_world_size()
         ), "world group already initialized with a different world size"
+    _sync_srt_world_group()
 
 
 def get_sp_group() -> SequenceParallelGroupCoordinator:
@@ -591,6 +612,7 @@ def get_tp_rank() -> int:
 
 def destroy_distributed_environment() -> None:
     global _WORLD
+    _clear_srt_world_group()
     if _WORLD:
         _WORLD.destroy()
     _WORLD = None