[Parallel State Refactor 2/n] Unify code path of AMD deterministic all reduce (#20871)

python/sglang/srt/distributed/device_communicators/custom_all_reduce.py

@@ -63,6 +63,7 @@ class CustomAllreduce:
         self._IS_CAPTURING = False
         self.disabled = True  # This can be modified in-place by context manager in piecewise cuda graph runner
         self.original_disabled = True  # To store the original state
+        self.use_amd_deterministic_impl = _use_amd_deterministic_impl()
 
         if not ops.IS_CUSTOM_AR_AVAILABLE:
             # disable because of missing custom allreduce library
@@ -269,65 +270,36 @@ class CustomAllreduce:
             return False
 
         if _is_hip:
+            if self.use_amd_deterministic_impl:
+                return True
             if self.full_nvlink:
                 return inp_size <= self.max_size
             return False
 
         return False
 
-    # all reduce, assuming inp tensor is IPC registered with register_buffer,
+    def _all_reduce_impl(self, inp: torch.Tensor, registered: bool):
-    # or, in the context of cuda graphs, register_graph_buffers
+        out = torch.empty_like(inp)
-    def all_reduce_reg(self, inp: torch.Tensor, out: torch.Tensor = None):
+        if not _is_hip:  # CUDA-like
-        if out is None:
+            if registered:
-            out = torch.empty_like(inp)
+                ops.all_reduce(self._ptr, inp, out, 0, 0)
-        ops.all_reduce_reg(self._ptr, inp, out)
+            else:
-        return out
+                ops.all_reduce(
-
+                    self._ptr, inp, out, self.buffer_ptrs[self.rank], self.max_size
-    # all reduce, assuming inp tensor is NOT IPC registered
+                )
-    def all_reduce_unreg(self, inp: torch.Tensor, out: torch.Tensor = None):
+        elif self.use_amd_deterministic_impl:
-        if out is None:
+            inp_size = inp.numel() * inp.element_size()
-            out = torch.empty_like(inp)
+            if inp_size < self.max_size:
-        ops.all_reduce_unreg(self._ptr, inp, self.buffer, out)
+                reg_buffer = self.buffer.view(inp.dtype)[: inp.numel()]
-        return out
+                ops.deterministic_all_reduce_unreg(self._ptr, inp, reg_buffer, out)
-
+            else:
-    def all_reduce(
+                self.register_buffer(inp)
-        self,
+                ops.deterministic_all_reduce_reg(self._ptr, inp, out)
-        inp: torch.Tensor,
+        else:  # normal AMD ROCm path
-        *,
+            if registered:
-        out: torch.Tensor = None,
+                ops.all_reduce_reg(self._ptr, inp, out)
-        registered: bool = False,
+            else:
-    ):
+                ops.all_reduce_unreg(self._ptr, inp, self.buffer, out)
-        """Performs an out-of-place all reduce.
-
-        If registered is True, this assumes inp's pointer is already
-        IPC-registered. Otherwise, inp is first copied into a pre-registered
-        buffer.
-        """
-        if out is None:
-            out = torch.empty_like(inp)
-        if registered:
-            ops.all_reduce(self._ptr, inp, out, 0, 0)
-        else:
-            ops.all_reduce(
-                self._ptr, inp, out, self.buffer_ptrs[self.rank], self.max_size
-            )
-        return out
-
-    def deterministic_all_reduce(
-        self,
-        inp: torch.Tensor,
-        *,
-        out: torch.Tensor = None,
-        registered: bool = False,
-    ):
-        """Deterministic all-reduce using 1-stage kernel with fixed ordering (AMD only)."""
-        if out is None:
-            out = torch.empty_like(inp)
-        if registered:
-            ops.deterministic_all_reduce_reg(self._ptr, inp, out)
-        else:
-            reg_buffer = self.buffer.view(inp.dtype)[: inp.numel()]
-            ops.deterministic_all_reduce_unreg(self._ptr, inp, reg_buffer, out)
         return out
 
     def custom_all_reduce(self, input: torch.Tensor) -> Optional[torch.Tensor]:
@@ -337,35 +309,20 @@ class CustomAllreduce:
             return None
         if self._IS_CAPTURING:
             if torch.cuda.is_current_stream_capturing():
-                if _is_hip:
+                return self._all_reduce_impl(input, registered=not self.tms_cudagraph)
-                    if self.tms_cudagraph:
-                        return self.all_reduce_unreg(input)
-                    return self.all_reduce_reg(input)
-                else:
-                    return self.all_reduce(input, registered=not self.tms_cudagraph)
             else:
                 # Could be warmup OR piecewise cuda graph split op execution.
                 # In piecewise cuda graph, split ops run eagerly outside the graph
                 # but _IS_CAPTURING is still True. We need to do real all-reduce.
                 if is_in_piecewise_cuda_graph():
                     # Split op execution - do real all-reduce
-                    if _is_hip:
+                    return self._all_reduce_impl(input, registered=False)
-                        return self.all_reduce_unreg(input)
-                    else:
-                        return self.all_reduce(input, registered=False)
                 else:
                     # True warmup - mimic the allocation pattern since custom
                     # allreduce is out-of-place.
                     return torch.zeros_like(input)
         else:
-            if _is_hip:
+            return self._all_reduce_impl(input, registered=False)
-                # note: outside of cuda graph context,
-                # custom allreduce incurs a cost of cudaMemcpy, which should
-                # be small(<=1% of overall latency) compared to the performance
-                # gains of using custom kernels
-                return self.all_reduce_unreg(input)
-            else:
-                return self.all_reduce(input, registered=False)
 
     def close(self):
         if not self.disabled and self._ptr:
@@ -382,7 +339,7 @@ class CustomAllreduce:
 def dispatch_custom_allreduce():
     """Return the CustomAllreduce class to use (aiter on ROCm if enabled).
 
-    On AMD with 1-stage AR enabled, use sglang's CustomAllreduce (has deterministic_all_reduce method).
+    On AMD with 1-stage AR enabled, use sglang's CustomAllreduce.
     Otherwise use AiterCustomAllreduce if available.
 
     Set SGLANG_USE_JIT_ALL_REDUCE=1 to use the JIT-compiled v2 implementation.
@@ -414,15 +371,9 @@ def dispatch_custom_allreduce():
     else:
         logger.debug("[AR] All-reduce: default")
 
-    # Check if 1-stage AR should be used
-    if envs.SGLANG_USE_1STAGE_ALLREDUCE.is_set():
-        use_1stage = envs.SGLANG_USE_1STAGE_ALLREDUCE.get()
-    else:
-        use_1stage = envs.SGLANG_ENABLE_DETERMINISTIC_INFERENCE.get()
-
     # On AMD with 1-stage AR, use sglang's CustomAllreduce
     # (AiterCustomAllreduce doesn't have deterministic_all_reduce method)
-    if use_1stage:
+    if _use_amd_deterministic_impl():
         return CustomAllreduce
 
     if get_bool_env_var("SGLANG_USE_AITER_AR", default="true"):
@@ -446,3 +397,12 @@ def dispatch_custom_allreduce():
             return CustomAllreduce
 
     return CustomAllreduce
+
+
+def _use_amd_deterministic_impl() -> bool:
+    if not _is_hip:  # CUDA is always deterministic
+        return False
+    if envs.SGLANG_USE_1STAGE_ALLREDUCE.is_set():
+        return envs.SGLANG_USE_1STAGE_ALLREDUCE.get()
+    else:
+        return envs.SGLANG_ENABLE_DETERMINISTIC_INFERENCE.get()

python/sglang/srt/distributed/parallel_state.py

@@ -560,26 +560,6 @@ class GroupCoordinator:
         if self.world_size == 1:
             return input_
 
-        # On AMD, use the deterministic 1-stage kernel when:
-        # - SGLANG_USE_1STAGE_ALLREDUCE=1 (explicitly enabled), OR
-        # - SGLANG_USE_1STAGE_ALLREDUCE not set AND --enable-deterministic-inference is on
-        if envs.SGLANG_USE_1STAGE_ALLREDUCE.is_set():
-            use_1stage_ar = envs.SGLANG_USE_1STAGE_ALLREDUCE.get()
-        else:
-            use_1stage_ar = envs.SGLANG_ENABLE_DETERMINISTIC_INFERENCE.get()
-        use_deterministic_ar = is_hip() and use_1stage_ar
-        if use_deterministic_ar:
-            if not input_.is_cpu and self.ca_comm is not None:
-                inp_size = input_.numel() * input_.element_size()
-                # Try unregistered mode first (faster for smaller tensors)
-                if inp_size < self.ca_comm.max_size:
-                    return self.ca_comm.deterministic_all_reduce(
-                        input_, registered=False
-                    )
-                # Use registered mode for larger tensors
-                self.ca_comm.register_buffer(input_)
-                return self.ca_comm.deterministic_all_reduce(input_, registered=True)
-
         if input_.is_cpu:
             if is_shm_available(input_.dtype, self.world_size, self.local_size):
                 torch.ops.sgl_kernel.shm_allreduce(input_, REDUCE_OP_SUM)

python/sglang/srt/utils/common.py

@@ -127,7 +127,7 @@ BAR_FORMAT = "{desc}: {percentage:3.0f}% Completed | {n_fmt}/{total_fmt} [{elaps
 
 @lru_cache(maxsize=1)
 def is_cuda():
-    return torch.cuda.is_available() and torch.version.cuda
+    return torch.cuda.is_available() and torch.version.cuda is not None
 
 
 @lru_cache(maxsize=1)

sgl-kernel/benchmark/bench_amd_deterministic_allreduce.py

@@ -29,20 +29,18 @@ python_dir = os.path.join(script_dir, "python")
 sys.path.insert(0, python_dir)
 
 # Try to import custom all-reduce if available
+from sglang.srt.environ import envs
+
 try:
     import sglang.srt.distributed.device_communicators.custom_all_reduce_ops as custom_ar_ops
     from sglang.srt.distributed.device_communicators.custom_all_reduce import (
         CustomAllreduce,
     )
-    from sglang.srt.distributed.device_communicators.custom_all_reduce_utils import (
-        is_weak_contiguous,
-    )
 
     CUSTOM_AR_AVAILABLE = custom_ar_ops.IS_CUSTOM_AR_AVAILABLE
 except (ImportError, AttributeError):
     CUSTOM_AR_AVAILABLE = False
     CustomAllreduce = None
-    is_weak_contiguous = None
 
 # Note: sglang's optimized all-reduce requires full runtime initialization
 # and won't work in standalone benchmarks, so we skip it
@@ -110,6 +108,7 @@ def reduce_scatter_then_all_gather(tensor, rank, world_size, custom_ar=None):
 
 
 def worker(world_size, rank, port, results_queue):
+    envs.SGLANG_USE_1STAGE_ALLREDUCE.set("1")
     device = torch.device(f"cuda:{rank}")
     torch.cuda.set_device(device)
 
@@ -240,7 +239,7 @@ def worker(world_size, rank, port, results_queue):
         results_deterministic_kernel = []
         latencies_deterministic_kernel = []
         deterministic_kernel_available = False
-        if custom_ar is not None and hasattr(custom_ar, "deterministic_all_reduce"):
+        if custom_ar is not None:
             # Check if input size fits in buffer
             input_size_bytes = base_input.numel() * base_input.element_size()
             if input_size_bytes > custom_ar.max_size:
@@ -259,9 +258,7 @@ def worker(world_size, rank, port, results_queue):
                         # Measure latency
                         torch.cuda.synchronize()
                         start = time.perf_counter()
-                        result_kernel = custom_ar.deterministic_all_reduce(
+                        result_kernel = custom_ar.custom_all_reduce(inp_kernel)
-                            inp_kernel, registered=False
-                        )
                         torch.cuda.synchronize()
                         end = time.perf_counter()
                         latencies_deterministic_kernel.append(end - start)

sgl-kernel/tests/test_amd_deterministic_custom_allreduce.py

@@ -22,6 +22,8 @@ import pytest
 import torch
 import torch.distributed as dist
 
+from sglang.srt.environ import envs
+
 
 def get_open_port():
     with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
@@ -30,6 +32,7 @@ def get_open_port():
 
 
 def worker(world_size, rank, port):
+    envs.SGLANG_USE_1STAGE_ALLREDUCE.set("1")
     device = torch.device(f"cuda:{rank}")
     torch.cuda.set_device(device)
 
@@ -60,12 +63,6 @@ def worker(world_size, rank, port):
                 print("✗ Custom AR not available or disabled")
             dist.destroy_process_group()
             return
-
-        if not hasattr(custom_ar, "deterministic_all_reduce"):
-            if rank == 0:
-                print("✗ Deterministic kernel not available")
-            dist.destroy_process_group()
-            return
     except Exception as e:
         if rank == 0:
             print(f"✗ Failed to initialize deterministic kernel: {e}")
@@ -115,18 +112,7 @@ def worker(world_size, rank, port):
         # Clone the same input
         inp = base_input.clone()
 
-        # Use deterministic kernel
+        result = custom_ar.custom_all_reduce(inp)
-        # Check if input fits in buffer, use registered mode if too large
-        input_size_bytes = inp.numel() * inp.element_size()
-        use_registered = input_size_bytes > custom_ar.max_size
-
-        if use_registered:
-            # For large inputs, register buffer first
-            custom_ar.register_buffer(inp)
-            result = custom_ar.deterministic_all_reduce(inp, registered=True)
-        else:
-            # For smaller inputs, use unregistered mode (copies to internal buffer)
-            result = custom_ar.deterministic_all_reduce(inp, registered=False)
         torch.cuda.synchronize()
 
         # Store checksum
@@ -179,22 +165,7 @@ def worker(world_size, rank, port):
             # Flatten for all-reduce: (bs * hidden_dim,)
             batch_flat = batch.view(-1)
 
-            # Use deterministic kernel
+            result_flat = custom_ar.custom_all_reduce(batch_flat)
-            # Check if input fits in buffer, use registered mode if too large
-            input_size_bytes = batch_flat.numel() * batch_flat.element_size()
-            use_registered = input_size_bytes > custom_ar.max_size
-
-            if use_registered:
-                # For large inputs, register buffer first
-                custom_ar.register_buffer(batch_flat)
-                result_flat = custom_ar.deterministic_all_reduce(
-                    batch_flat, registered=True
-                )
-            else:
-                # For smaller inputs, use unregistered mode
-                result_flat = custom_ar.deterministic_all_reduce(
-                    batch_flat, registered=False
-                )
             torch.cuda.synchronize()
 
             # Reshape back to (bs, hidden_dim)