diff --git a/python/sglang/multimodal_gen/runtime/distributed/group_coordinator.py b/python/sglang/multimodal_gen/runtime/distributed/group_coordinator.py index 4b1e694fa..46e26cb87 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/group_coordinator.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py b/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py index 7f4a06d12..f18cb119b 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py +++ b/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