diff --git a/python/sglang/srt/distributed/device_communicators/torch_symm_mem.py b/python/sglang/srt/distributed/device_communicators/torch_symm_mem.py index cc9256280..fd38fb8f5 100644 --- a/python/sglang/srt/distributed/device_communicators/torch_symm_mem.py +++ b/python/sglang/srt/distributed/device_communicators/torch_symm_mem.py @@ -59,6 +59,8 @@ class TorchSymmMemCommunicator: """ self.disabled = True + self.buffer = None + self.max_size = 0 if not torch_symm_mem_available: return @@ -73,26 +75,24 @@ class TorchSymmMemCommunicator: self.group = group self.world_size = dist.get_world_size(self.group) self.device_capability = torch.cuda.get_device_capability(device)[0] - if self.device_capability < 9: + supported_max_sizes = TORCH_SYMM_MEM_ALL_REDUCE_MAX_SIZES.get( + self.device_capability + ) + if supported_max_sizes is None: logger.warning( "TorchSymmMemCommunicator: Device capability %s not supported, " "communicator is not available.", self.device_capability, ) return - if ( - self.world_size - not in TORCH_SYMM_MEM_ALL_REDUCE_MAX_SIZES[self.device_capability] - ): + if self.world_size not in supported_max_sizes: logger.warning( "TorchSymmMemCommunicator: World size %d not supported, " "communicator is not available.", self.world_size, ) return - self.max_size = TORCH_SYMM_MEM_ALL_REDUCE_MAX_SIZES[self.device_capability][ - self.world_size - ] + self.max_size = supported_max_sizes[self.world_size] self.buffer = torch_symm_mem.empty( self.max_size // self.dtype.itemsize, device=self.device, @@ -124,6 +124,8 @@ class TorchSymmMemCommunicator: """ if self.disabled: return False + if inp.device != self.device: + return False if inp.dtype != self.dtype: return False inp_size = inp.numel() * inp.element_size() @@ -150,10 +152,14 @@ class TorchSymmMemCommunicator: - Selects 'multimem' or 'two_shot' kernel based on topology. - Writes the result into 'out' and returns it. """ + if not self.should_torch_symm_mem_allreduce(inp): + return None if out is None: out = torch.empty_like(inp) self.buffer[: inp.numel()].copy_(inp.view(-1)) - if self.world_size in self._WORLD_SIZES_MULTIMEM[self.device_capability]: + if self.world_size in self._WORLD_SIZES_MULTIMEM.get( + self.device_capability, () + ): torch.ops.symm_mem.multimem_all_reduce_( self.buffer[: inp.numel()], "sum", self.group.group_name ) diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index 62dc01641..74c52a77c 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/python/sglang/srt/distributed/parallel_state.py @@ -765,8 +765,12 @@ class GroupCoordinator: torch_symm_mem_comm = self.torch_symm_mem_comm if pynccl_comm is not None and not pynccl_comm.disabled: pynccl_comm.all_reduce(input_) - elif torch_symm_mem_comm is not None and not torch_symm_mem_comm.disabled: - torch_symm_mem_comm.all_reduce(input_) + elif ( + torch_symm_mem_comm is not None + and not torch_symm_mem_comm.disabled + and torch_symm_mem_comm.should_torch_symm_mem_allreduce(input_) + ): + torch_symm_mem_comm.all_reduce(input_, out=input_) else: torch.distributed.all_reduce(input_, group=self.device_group)