[Distributed] Guard torch symm mem all-reduce sizes (#24548)

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 (
+        if self.world_size not in supported_max_sizes:
-            self.world_size
-            not in TORCH_SYMM_MEM_ALL_REDUCE_MAX_SIZES[self.device_capability]
-        ):
             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.max_size = supported_max_sizes[self.world_size]
-            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
             )

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:
+        elif (
-            torch_symm_mem_comm.all_reduce(input_)
+            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)