[amd] Add deterministic all-reduce kernel for AMD (ROCm) (#15340)

Co-authored-by: Thomas Wang <1am9trash@gmail.com>

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

@@ -373,6 +373,23 @@ class CustomAllreduce:
             )
         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]:
         """The main allreduce API that provides support for cuda graph."""
         # When custom allreduce is disabled, this will be None.
@@ -411,20 +428,37 @@ class CustomAllreduce:
 
 
 def dispatch_custom_allreduce():
-    """Return the CustomAllreduce class to use (aiter on ROCm if enabled)."""
+    """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).
+    Otherwise use AiterCustomAllreduce if available.
+    """
+    # 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 is_hip() and use_1stage:
+        logger.info("[AR] Using sglang CustomAllreduce (1-stage kernel)")
+        return CustomAllreduce
+
     if is_hip() and get_bool_env_var("SGLANG_USE_AITER_AR", default="true"):
         try:
             from aiter.dist.device_communicators.custom_all_reduce import (
                 CustomAllreduce as AiterCustomAllreduce,
             )
 
-            logger.info("Using AiterCustomAllreduce for ROCm.")
+            logger.info("[AR] Using AiterCustomAllreduce (AMD default)")
             return AiterCustomAllreduce
         except ImportError as e:
             logger.warning(
-                "Aiter custom all-reduce not available (optional dependency missing); "
+                "[AR] Aiter custom all-reduce not available; "
                 "falling back to sglang CustomAllreduce. Details: %s",
                 e,
             )
             return CustomAllreduce
+    logger.info("[AR] Using sglang CustomAllreduce")
     return CustomAllreduce

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

@@ -90,6 +90,16 @@ elif _is_hip:
     ) -> None:
         _custom_ar.all_reduce_unreg(fa, inp, reg_buffer, out)
 
+    def deterministic_all_reduce_reg(
+        fa: int, inp: torch.Tensor, out: torch.Tensor
+    ) -> None:
+        _custom_ar.deterministic_all_reduce_reg(fa, inp, out)
+
+    def deterministic_all_reduce_unreg(
+        fa: int, inp: torch.Tensor, reg_buffer: torch.Tensor, out: torch.Tensor
+    ) -> None:
+        _custom_ar.deterministic_all_reduce_unreg(fa, inp, reg_buffer, out)
+
     def dispose(fa: int) -> None:
         _custom_ar.dispose(fa)
 

python/sglang/srt/distributed/parallel_state.py

@@ -376,6 +376,23 @@ class GroupCoordinator:
                     group=self.cpu_group,
                     device=self.device,
                 )
+                # Log which all-reduce mode will be used
+                if is_hip():
+                    if envs.SGLANG_USE_1STAGE_ALLREDUCE.is_set():
+                        if envs.SGLANG_USE_1STAGE_ALLREDUCE.get():
+                            logger.info(
+                                "[AR] All-reduce: 1-stage kernel (SGLANG_USE_1STAGE_ALLREDUCE=1)"
+                            )
+                        else:
+                            logger.info(
+                                "[AR] All-reduce: default (SGLANG_USE_1STAGE_ALLREDUCE=0)"
+                            )
+                    elif envs.SGLANG_ENABLE_DETERMINISTIC_INFERENCE.get():
+                        logger.info(
+                            "[AR] All-reduce: 1-stage kernel (deterministic inference enabled)"
+                        )
+                    else:
+                        logger.info("[AR] All-reduce: default")
             except Exception as e:
                 logger.warning(
                     f"Setup Custom allreduce failed with {e}. To silence this "
@@ -393,6 +410,8 @@ class GroupCoordinator:
                         )
                 except Exception as e:
                     logger.warning(f"Failed to initialize QuickAllReduce: {e}")
+        elif self.world_size > 1 and is_hip():
+            logger.info("[AR] All-reduce call path: NCCL (custom AR disabled)")
 
         self.torch_symm_mem_comm: Optional[TorchSymmMemCommunicator] = None
         if self.use_torch_symm_mem_all_reduce and self.world_size > 1:
@@ -560,6 +579,26 @@ 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/environ.py

@@ -308,6 +308,11 @@ class Envs:
 
     # Deterministic inference
     SGLANG_ENABLE_DETERMINISTIC_INFERENCE = EnvBool(False)
+    # Use 1-stage all-reduce kernel on AMD (deterministic, fixed accumulation order)
+    # If not set: auto (enabled when --enable-deterministic-inference is on)
+    # Set to 1: force enable (even without --enable-deterministic-inference)
+    # Set to 0: force disable (use default Aiter AR even with --enable-deterministic-inference)
+    SGLANG_USE_1STAGE_ALLREDUCE = EnvBool(False)
     SGLANG_FLASHINFER_PREFILL_SPLIT_TILE_SIZE = EnvInt(4096)
     SGLANG_FLASHINFER_DECODE_SPLIT_TILE_SIZE = EnvInt(2048)
     SGLANG_TRITON_PREFILL_TRUNCATION_ALIGN_SIZE = EnvInt(4096)

python/sglang/srt/server_args.py

@@ -2314,6 +2314,14 @@ class ServerArgs:
 
             # Check TP size
             if self.tp_size > 1:
+                if is_hip():
+                    # AMD: use 1-stage all-reduce kernel which is inherently deterministic
+                    # (each GPU reads all data from all GPUs, reduces locally in fixed order)
+                    logger.info(
+                        "AMD/ROCm: Using 1-stage all-reduce kernel (deterministic)"
+                    )
+                else:
+                    # CUDA: use NCCL tree algorithm
                     os.environ["NCCL_ALGO"] = "allreduce:tree"
                     self.disable_custom_all_reduce = True
                     logger.warning(

sgl-kernel/benchmark/bench_amd_deterministic_allreduce.py

@@ -0,0 +1,692 @@
+"""
+Benchmark latency comparison between different all-reduce implementations.
+
+Compares:
+- NCCL all-reduce (may be non-deterministic)
+- Reduce-scatter + all-gather (RS+AG, deterministic but slower)
+- Deterministic 1-stage kernel (forces fixed accumulation order, deterministic)
+
+Note: The "deterministic kernel" is NOT RS+AG. It uses the 1-stage kernel where
+each GPU reads all data from all GPUs and reduces locally in a fixed order.
+
+Usage:
+    python bench_amd_deterministic_allreduce.py
+"""
+
+import multiprocessing as mp
+import os
+import socket
+import statistics
+import sys
+import time
+
+import torch
+import torch.distributed as dist
+
+# Add python directory to path to import sglang modules
+script_dir = os.path.dirname(os.path.abspath(__file__))
+python_dir = os.path.join(script_dir, "python")
+sys.path.insert(0, python_dir)
+
+# Try to import custom all-reduce if available
+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
+SGLANG_AVAILABLE = False
+
+
+def get_open_port():
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("127.0.0.1", 0))
+        return s.getsockname()[1]
+
+
+def init_custom_ar_if_available(rank, world_size, device):
+    """Check if custom all-reduce is available and applicable."""
+    if not CUSTOM_AR_AVAILABLE or CustomAllreduce is None:
+        return False
+
+    # Custom AR works best for single-node, even number of GPUs, world_size <= 8
+    if world_size <= 8 and world_size % 2 == 0:
+        return True
+
+    return False
+
+
+def reduce_scatter_then_all_gather(tensor, rank, world_size, custom_ar=None):
+    """
+    Deterministic all-reduce using reduce-scatter + all-gather.
+    This is deterministic because it uses fixed ordering (no atomics).
+    """
+    total_size = tensor.numel()
+    if total_size % world_size != 0:
+        # Fallback to all-gather + local reduce if not divisible
+        gather_list = [torch.empty_like(tensor) for _ in range(world_size)]
+        dist.all_gather(gather_list, tensor)
+        stacked = torch.stack(gather_list, dim=0)
+        tensor.copy_(stacked.sum(dim=0))
+        return
+
+    chunk_size = total_size // world_size
+
+    # Flatten to 1D
+    tensor_flat = tensor.view(-1)
+
+    # Reduce-scatter: each rank gets its chunk of the reduced result
+    output_chunk = torch.empty(chunk_size, dtype=tensor.dtype, device=tensor.device)
+
+    # Split input into chunks for reduce-scatter
+    input_chunks = [
+        tensor_flat[i * chunk_size : (i + 1) * chunk_size].clone()
+        for i in range(world_size)
+    ]
+
+    dist.reduce_scatter(output_chunk, input_chunks)
+
+    # All-gather: broadcast each rank's chunk to all ranks
+    output_chunks = [
+        torch.empty(chunk_size, dtype=tensor.dtype, device=tensor.device)
+        for _ in range(world_size)
+    ]
+    dist.all_gather(output_chunks, output_chunk)
+
+    # Concatenate results back
+    result_flat = torch.cat(output_chunks, dim=0)
+    tensor.copy_(result_flat.view(tensor.shape))
+
+
+def worker(world_size, rank, port, results_queue):
+    device = torch.device(f"cuda:{rank}")
+    torch.cuda.set_device(device)
+
+    dist.init_process_group(
+        backend="nccl",
+        init_method=f"tcp://localhost:{port}",
+        rank=rank,
+        world_size=world_size,
+    )
+
+    # Try to initialize custom all-reduce if available
+    custom_ar = None
+    use_custom_ar = init_custom_ar_if_available(rank, world_size, device)
+    if use_custom_ar and CUSTOM_AR_AVAILABLE:
+        try:
+            # Create a gloo group for custom AR (it requires non-NCCL backend)
+            # All ranks must call new_group with the same parameters
+            from torch.distributed import new_group
+
+            dist.barrier()  # Ensure all ranks are ready
+            ar_group = new_group(backend="gloo")
+            dist.barrier()  # Ensure group creation is complete
+            custom_ar = CustomAllreduce(group=ar_group, device=device)
+            if rank == 0:
+                print("  Using custom all-reduce (deterministic)")
+        except Exception as e:
+            if rank == 0:
+                print(f"  Custom AR init failed: {e}, using NCCL fallback")
+            custom_ar = None
+            dist.barrier()  # Ensure all ranks continue even if one fails
+
+    # Test different batch sizes - similar to test_ar.py
+    batch_sizes = [1, 4, 8, 16, 32, 64, 128, 256, 512]
+    hidden_dim = 16384  # Fixed hidden dimension
+
+    num_trials = 10  # Same as test_ar.py
+
+    # Different seed per rank - each GPU has DIFFERENT input (like test_ar.py)
+    torch.manual_seed(42 + rank)
+
+    results = {}
+
+    for bs in batch_sizes:
+        # Create fixed input for all trials (like test_ar.py)
+        base_input = torch.randn(bs, hidden_dim, dtype=torch.bfloat16, device=device)
+
+        dist.barrier()
+
+        if rank == 0:
+            print(f"\nBatch size {bs:4d}:")
+            print(f"  Testing determinism across {num_trials} trials...")
+
+        # Test all-reduce determinism
+        results_ar = []
+        latencies_ar = []
+        for trial in range(num_trials):
+            # Clone the same input for each trial
+            inp_ar = base_input.clone()
+            inp_flat_ar = inp_ar.view(-1)
+
+            # Measure latency
+            torch.cuda.synchronize()
+            start = time.perf_counter()
+            dist.all_reduce(inp_flat_ar, op=dist.ReduceOp.SUM)
+            torch.cuda.synchronize()
+            end = time.perf_counter()
+            latencies_ar.append(end - start)
+
+            # Store checksum and first values (like test_ar.py)
+            checksum = inp_flat_ar.sum().item()
+            first_vals = inp_flat_ar[:5].clone()
+            results_ar.append((checksum, first_vals))
+
+        # Test reduce-scatter + all-gather determinism
+        results_rs_ag = []
+        latencies_rs_ag = []
+        for trial in range(num_trials):
+            # Clone the same input for each trial
+            inp_rs_ag = base_input.clone()
+            inp_flat_rs_ag = inp_rs_ag.view(-1)
+
+            # Measure latency
+            torch.cuda.synchronize()
+            start = time.perf_counter()
+            reduce_scatter_then_all_gather(
+                inp_flat_rs_ag, rank, world_size, custom_ar=None
+            )
+            torch.cuda.synchronize()
+            end = time.perf_counter()
+            latencies_rs_ag.append(end - start)
+
+            # Store checksum and first values (like test_ar.py)
+            checksum = inp_flat_rs_ag.sum().item()
+            first_vals = inp_flat_rs_ag[:5].clone()
+            results_rs_ag.append((checksum, first_vals))
+
+        # Note: sglang's optimized all-reduce requires full runtime initialization
+        # and is not tested in this standalone benchmark
+        use_sglang_optimized = False
+        results_optimized_rs_ag = []
+        latencies_optimized_rs_ag = []
+
+        # Test custom all-reduce determinism (if available)
+        results_custom_ar = []
+        latencies_custom_ar = []
+        if custom_ar is not None:
+            for trial in range(num_trials):
+                # Clone the same input for each trial
+                inp_custom = base_input.clone()
+                inp_flat_custom = inp_custom.view(-1)
+
+                # Measure latency
+                torch.cuda.synchronize()
+                start = time.perf_counter()
+                reduce_scatter_then_all_gather(
+                    inp_flat_custom, rank, world_size, custom_ar=custom_ar
+                )
+                torch.cuda.synchronize()
+                end = time.perf_counter()
+                latencies_custom_ar.append(end - start)
+
+                # Store checksum and first values (like test_ar.py)
+                checksum = inp_flat_custom.sum().item()
+                first_vals = inp_flat_custom[:5].clone()
+                results_custom_ar.append((checksum, first_vals))
+
+        # Test deterministic kernel (if available)
+        results_deterministic_kernel = []
+        latencies_deterministic_kernel = []
+        deterministic_kernel_available = False
+        if custom_ar is not None and hasattr(custom_ar, "deterministic_all_reduce"):
+            # 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:
+                if rank == 0:
+                    print(
+                        f"  Deterministic kernel skipped: input size ({input_size_bytes/(1024*1024):.1f} MB) > buffer size ({custom_ar.max_size/(1024*1024):.1f} MB)"
+                    )
+                deterministic_kernel_available = False
+            else:
+                try:
+                    deterministic_kernel_available = True
+                    for trial in range(num_trials):
+                        # Clone the same input for each trial
+                        inp_kernel = base_input.clone()
+
+                        # Measure latency
+                        torch.cuda.synchronize()
+                        start = time.perf_counter()
+                        result_kernel = custom_ar.deterministic_all_reduce(
+                            inp_kernel, registered=False
+                        )
+                        torch.cuda.synchronize()
+                        end = time.perf_counter()
+                        latencies_deterministic_kernel.append(end - start)
+
+                        # Store checksum and first values
+                        result_flat_kernel = result_kernel.view(-1)
+                        checksum = result_flat_kernel.sum().item()
+                        first_vals = result_flat_kernel[:5].clone()
+                        results_deterministic_kernel.append((checksum, first_vals))
+                except Exception as e:
+                    if rank == 0:
+                        print(
+                            f"  Deterministic kernel test failed for batch size {bs}: {e}"
+                        )
+                    deterministic_kernel_available = False
+
+        dist.barrier()
+
+        if rank == 0:
+            # Check determinism for all-reduce
+            ar_deterministic = True
+            ar_ref_sum, ar_ref_vals = results_ar[0]
+            ar_variance = []
+            for i, (s, vals) in enumerate(results_ar[1:], 1):
+                if abs(ar_ref_sum - s) > 1e-3 or not torch.allclose(
+                    ar_ref_vals, vals, rtol=1e-3
+                ):
+                    ar_deterministic = False
+                ar_variance.append(abs(ar_ref_sum - s))
+
+            # Check determinism for reduce-scatter + all-gather
+            rs_ag_deterministic = True
+            rs_ag_ref_sum, rs_ag_ref_vals = results_rs_ag[0]
+            rs_ag_variance = []
+            for i, (s, vals) in enumerate(results_rs_ag[1:], 1):
+                if abs(rs_ag_ref_sum - s) > 1e-3 or not torch.allclose(
+                    rs_ag_ref_vals, vals, rtol=1e-3
+                ):
+                    rs_ag_deterministic = False
+                rs_ag_variance.append(abs(rs_ag_ref_sum - s))
+
+            # Check determinism for optimized RS+AG (if available)
+            optimized_rs_ag_deterministic = None
+            optimized_rs_ag_max_variance = None
+            lat_optimized_rs_ag_median = None
+            if use_sglang_optimized and results_optimized_rs_ag:
+                optimized_rs_ag_deterministic = True
+                opt_rs_ag_ref_sum, opt_rs_ag_ref_vals = results_optimized_rs_ag[0]
+                opt_rs_ag_variance = []
+                for i, (s, vals) in enumerate(results_optimized_rs_ag[1:], 1):
+                    if abs(opt_rs_ag_ref_sum - s) > 1e-3 or not torch.allclose(
+                        opt_rs_ag_ref_vals, vals, rtol=1e-3
+                    ):
+                        optimized_rs_ag_deterministic = False
+                    opt_rs_ag_variance.append(abs(opt_rs_ag_ref_sum - s))
+                optimized_rs_ag_max_variance = (
+                    max(opt_rs_ag_variance) if opt_rs_ag_variance else 0.0
+                )
+                lat_optimized_rs_ag_median = statistics.median(
+                    latencies_optimized_rs_ag
+                )
+
+            # Check determinism for custom all-reduce (if available)
+            custom_ar_deterministic = None
+            custom_ar_max_variance = None
+            lat_custom_ar_median = None
+            if custom_ar is not None and results_custom_ar:
+                custom_ar_deterministic = True
+                custom_ar_ref_sum, custom_ar_ref_vals = results_custom_ar[0]
+                custom_ar_variance = []
+                for i, (s, vals) in enumerate(results_custom_ar[1:], 1):
+                    if abs(custom_ar_ref_sum - s) > 1e-3 or not torch.allclose(
+                        custom_ar_ref_vals, vals, rtol=1e-3
+                    ):
+                        custom_ar_deterministic = False
+                    custom_ar_variance.append(abs(custom_ar_ref_sum - s))
+                custom_ar_max_variance = (
+                    max(custom_ar_variance) if custom_ar_variance else 0.0
+                )
+                lat_custom_ar_median = statistics.median(latencies_custom_ar)
+
+            # Check determinism for deterministic kernel (if available)
+            deterministic_kernel_deterministic = None
+            deterministic_kernel_max_variance = None
+            lat_deterministic_kernel_median = None
+            if deterministic_kernel_available and results_deterministic_kernel:
+                deterministic_kernel_deterministic = True
+                kernel_ref_sum, kernel_ref_vals = results_deterministic_kernel[0]
+                kernel_variance = []
+                for i, (s, vals) in enumerate(results_deterministic_kernel[1:], 1):
+                    if abs(kernel_ref_sum - s) > 1e-3 or not torch.allclose(
+                        kernel_ref_vals, vals, rtol=1e-3
+                    ):
+                        deterministic_kernel_deterministic = False
+                    kernel_variance.append(abs(kernel_ref_sum - s))
+                deterministic_kernel_max_variance = (
+                    max(kernel_variance) if kernel_variance else 0.0
+                )
+                lat_deterministic_kernel_median = statistics.median(
+                    latencies_deterministic_kernel
+                )
+
+            # Calculate latency statistics
+            lat_ar_median = statistics.median(latencies_ar)
+            lat_rs_ag_median = statistics.median(latencies_rs_ag)
+            overhead_rs_ag = ((lat_rs_ag_median - lat_ar_median) / lat_ar_median) * 100
+
+            # Calculate variance statistics
+            ar_max_variance = max(ar_variance) if ar_variance else 0.0
+            rs_ag_max_variance = max(rs_ag_variance) if rs_ag_variance else 0.0
+
+            results[bs] = {
+                "all_reduce": {
+                    "latency_median": lat_ar_median,
+                    "deterministic": ar_deterministic,
+                    "max_variance": ar_max_variance,
+                },
+                "rs_ag": {
+                    "latency_median": lat_rs_ag_median,
+                    "deterministic": rs_ag_deterministic,
+                    "max_variance": rs_ag_max_variance,
+                },
+                "custom_ar": (
+                    {
+                        "latency_median": lat_custom_ar_median,
+                        "deterministic": custom_ar_deterministic,
+                        "max_variance": custom_ar_max_variance,
+                    }
+                    if custom_ar is not None
+                    else None
+                ),
+                "deterministic_kernel": (
+                    {
+                        "latency_median": lat_deterministic_kernel_median,
+                        "deterministic": deterministic_kernel_deterministic,
+                        "max_variance": deterministic_kernel_max_variance,
+                    }
+                    if lat_deterministic_kernel_median is not None
+                    else None
+                ),
+                "optimized_rs_ag": (
+                    {
+                        "latency_median": lat_optimized_rs_ag_median,
+                        "deterministic": optimized_rs_ag_deterministic,
+                        "max_variance": optimized_rs_ag_max_variance,
+                    }
+                    if lat_optimized_rs_ag_median is not None
+                    else None
+                ),
+                "overhead_rs_ag_pct": overhead_rs_ag,
+            }
+
+            print(
+                f"    All-Reduce:     {lat_ar_median*1000:.3f}ms, Deterministic: {ar_deterministic}, Max variance: {ar_max_variance:.6f}"
+            )
+            print(
+                f"    RS+All-Gather:   {lat_rs_ag_median*1000:.3f}ms, Deterministic: {rs_ag_deterministic}, Max variance: {rs_ag_max_variance:.6f}"
+            )
+            if custom_ar is not None and lat_custom_ar_median is not None:
+                overhead_custom = (
+                    (lat_custom_ar_median - lat_ar_median) / lat_ar_median
+                ) * 100
+                print(
+                    f"    Custom AR:       {lat_custom_ar_median*1000:.3f}ms, Deterministic: {custom_ar_deterministic}, Max variance: {custom_ar_max_variance:.6f}, Overhead: {overhead_custom:+.1f}%"
+                )
+            if lat_deterministic_kernel_median is not None:
+                overhead_kernel = (
+                    (lat_deterministic_kernel_median - lat_ar_median) / lat_ar_median
+                ) * 100
+                speedup_kernel_vs_rs_ag = (
+                    (lat_rs_ag_median - lat_deterministic_kernel_median)
+                    / lat_rs_ag_median
+                ) * 100
+                print(
+                    f"    Deterministic Kernel: {lat_deterministic_kernel_median*1000:.3f}ms, Deterministic: {deterministic_kernel_deterministic}, Max variance: {deterministic_kernel_max_variance:.6f}, Overhead: {overhead_kernel:+.1f}%, Speedup vs RS+AG: {speedup_kernel_vs_rs_ag:+.1f}%"
+                )
+            if lat_optimized_rs_ag_median is not None:
+                overhead_opt = (
+                    (lat_optimized_rs_ag_median - lat_ar_median) / lat_ar_median
+                ) * 100
+                speedup_vs_rs_ag = (
+                    (lat_rs_ag_median - lat_optimized_rs_ag_median) / lat_rs_ag_median
+                ) * 100
+                print(
+                    f"    Optimized RS+AG: {lat_optimized_rs_ag_median*1000:.3f}ms, Deterministic: {optimized_rs_ag_deterministic}, Max variance: {optimized_rs_ag_max_variance:.6f}, Overhead: {overhead_opt:+.1f}%, Speedup vs RS+AG: {speedup_vs_rs_ag:+.1f}%"
+                )
+            print(f"    RS+AG Overhead: {overhead_rs_ag:+.1f}%")
+
+    if rank == 0:
+        results_queue.put(results)
+
+    dist.destroy_process_group()
+
+
+def main():
+    world_size = 8
+    available_gpus = torch.cuda.device_count()
+
+    print("=" * 80)
+    print("All-Reduce vs Reduce-Scatter + All-Gather Determinism & Latency Benchmark")
+    print("=" * 80)
+    print(f"Available GPUs: {available_gpus}")
+    print(f"Using world_size: {world_size}")
+    print(f"Hidden dimension: 16384")
+    print(f"Tensor dtype: bfloat16")
+    print(f"Trials per batch size: 10 (testing determinism)")
+    print(f"Testing batch sizes: [1, 4, 8, 16, 32, 64, 128, 256, 512]")
+    print("=" * 80)
+
+    if available_gpus < world_size:
+        print(
+            f"WARNING: Only {available_gpus} GPUs available, using {available_gpus} instead"
+        )
+        world_size = available_gpus
+
+    if world_size < 2:
+        print("ERROR: Need at least 2 GPUs for this benchmark")
+        return
+
+    mp.set_start_method("spawn", force=True)
+    port = get_open_port()
+
+    results_queue = mp.Queue()
+    procs = []
+    for rank in range(world_size):
+        p = mp.Process(target=worker, args=(world_size, rank, port, results_queue))
+        p.start()
+        procs.append(p)
+
+    for p in procs:
+        p.join()
+
+    # Collect results
+    if not results_queue.empty():
+        results = results_queue.get()
+
+        print("\n" + "=" * 80)
+        print("SUMMARY")
+        print("=" * 80)
+        header = f"{'Batch':<8} {'AR (ms)':<12} {'AR Det':<8} {'RS+AG (ms)':<15} {'RS+AG Det':<10} {'RS+AG Ovh':<12}"
+        if any(r.get("custom_ar") is not None for r in results.values()):
+            header += (
+                f" {'Custom AR (ms)':<18} {'Custom AR Det':<15} {'Custom AR Ovh':<15}"
+            )
+        if any(r.get("deterministic_kernel") is not None for r in results.values()):
+            header += f" {'Det Kernel (ms)':<18} {'Det Kernel Det':<15} {'Det Kernel Ovh':<15} {'Speedup':<10}"
+        if any(r.get("optimized_rs_ag") is not None for r in results.values()):
+            header += f" {'Opt RS+AG (ms)':<18} {'Opt RS+AG Det':<15} {'Opt RS+AG Ovh':<15} {'Speedup':<10}"
+        print(header)
+        print("-" * 150)
+
+        for bs in sorted(results.keys()):
+            r = results[bs]
+            ar_det_str = "✓" if r["all_reduce"]["deterministic"] else "✗"
+            rs_ag_det_str = "✓" if r["rs_ag"]["deterministic"] else "✗"
+            line = (
+                f"{bs:<8} {r['all_reduce']['latency_median']*1000:<12.3f} {ar_det_str:<8} "
+                f"{r['rs_ag']['latency_median']*1000:<15.3f} {rs_ag_det_str:<10} "
+                f"{r['overhead_rs_ag_pct']:<12.1f}"
+            )
+            if r.get("custom_ar") is not None:
+                custom_ar = r["custom_ar"]
+                custom_ar_det_str = "✓" if custom_ar["deterministic"] else "✗"
+                custom_ar_overhead = (
+                    (custom_ar["latency_median"] - r["all_reduce"]["latency_median"])
+                    / r["all_reduce"]["latency_median"]
+                ) * 100
+                line += f" {custom_ar['latency_median']*1000:<18.3f} {custom_ar_det_str:<15} {custom_ar_overhead:<15.1f}"
+            if r.get("deterministic_kernel") is not None:
+                det_kernel = r["deterministic_kernel"]
+                det_kernel_det_str = "✓" if det_kernel["deterministic"] else "✗"
+                det_kernel_overhead = (
+                    (det_kernel["latency_median"] - r["all_reduce"]["latency_median"])
+                    / r["all_reduce"]["latency_median"]
+                ) * 100
+                speedup_kernel = (
+                    (r["rs_ag"]["latency_median"] - det_kernel["latency_median"])
+                    / r["rs_ag"]["latency_median"]
+                ) * 100
+                line += f" {det_kernel['latency_median']*1000:<18.3f} {det_kernel_det_str:<15} {det_kernel_overhead:<15.1f} {speedup_kernel:<10.1f}"
+            if r.get("optimized_rs_ag") is not None:
+                opt_rs_ag = r["optimized_rs_ag"]
+                opt_rs_ag_det_str = "✓" if opt_rs_ag["deterministic"] else "✗"
+                opt_rs_ag_overhead = (
+                    (opt_rs_ag["latency_median"] - r["all_reduce"]["latency_median"])
+                    / r["all_reduce"]["latency_median"]
+                ) * 100
+                speedup = (
+                    (r["rs_ag"]["latency_median"] - opt_rs_ag["latency_median"])
+                    / r["rs_ag"]["latency_median"]
+                ) * 100
+                line += f" {opt_rs_ag['latency_median']*1000:<18.3f} {opt_rs_ag_det_str:<15} {opt_rs_ag_overhead:<15.1f} {speedup:<10.1f}"
+            print(line)
+
+        print("=" * 80)
+
+        # Calculate statistics
+        overheads_rs_ag = [r["overhead_rs_ag_pct"] for r in results.values()]
+        ar_deterministic_count = sum(
+            1 for r in results.values() if r["all_reduce"]["deterministic"]
+        )
+        rs_ag_deterministic_count = sum(
+            1 for r in results.values() if r["rs_ag"]["deterministic"]
+        )
+        custom_ar_deterministic_count = sum(
+            1
+            for r in results.values()
+            if r.get("custom_ar") and r["custom_ar"]["deterministic"]
+        )
+        custom_ar_total_count = sum(
+            1 for r in results.values() if r.get("custom_ar") is not None
+        )
+
+        deterministic_kernel_deterministic_count = sum(
+            1
+            for r in results.values()
+            if r.get("deterministic_kernel")
+            and r["deterministic_kernel"]["deterministic"]
+        )
+        deterministic_kernel_total_count = sum(
+            1 for r in results.values() if r.get("deterministic_kernel") is not None
+        )
+
+        print(f"\nDeterminism Summary:")
+        print(
+            f"  All-Reduce deterministic: {ar_deterministic_count}/{len(results)} batch sizes"
+        )
+        print(
+            f"  RS+All-Gather deterministic: {rs_ag_deterministic_count}/{len(results)} batch sizes"
+        )
+        if custom_ar_total_count > 0:
+            print(
+                f"  Custom AR deterministic: {custom_ar_deterministic_count}/{custom_ar_total_count} batch sizes"
+            )
+        if deterministic_kernel_total_count > 0:
+            print(
+                f"  Deterministic Kernel deterministic: {deterministic_kernel_deterministic_count}/{deterministic_kernel_total_count} batch sizes"
+            )
+
+        print(f"\nLatency Overhead Statistics (RS+AG vs All-Reduce):")
+        avg_overhead = statistics.mean(overheads_rs_ag)
+        median_overhead = statistics.median(overheads_rs_ag)
+        min_overhead = min(overheads_rs_ag)
+        max_overhead = max(overheads_rs_ag)
+        print(f"  Average: {avg_overhead:.1f}%")
+        print(f"  Median:  {median_overhead:.1f}%")
+        print(f"  Min:     {min_overhead:.1f}%")
+        print(f"  Max:     {max_overhead:.1f}%")
+
+        if custom_ar_total_count > 0:
+            overheads_custom = []
+            for r in results.values():
+                if r.get("custom_ar") is not None:
+                    overhead = (
+                        (
+                            r["custom_ar"]["latency_median"]
+                            - r["all_reduce"]["latency_median"]
+                        )
+                        / r["all_reduce"]["latency_median"]
+                    ) * 100
+                    overheads_custom.append(overhead)
+            print(f"\nLatency Overhead Statistics (Custom AR vs All-Reduce):")
+            print(f"  Average: {statistics.mean(overheads_custom):.1f}%")
+            print(f"  Median:  {statistics.median(overheads_custom):.1f}%")
+            print(f"  Min:     {min(overheads_custom):.1f}%")
+            print(f"  Max:     {max(overheads_custom):.1f}%")
+
+        if deterministic_kernel_total_count > 0:
+            overheads_kernel = []
+            speedups_kernel = []
+            for r in results.values():
+                if r.get("deterministic_kernel") is not None:
+                    overhead = (
+                        (
+                            r["deterministic_kernel"]["latency_median"]
+                            - r["all_reduce"]["latency_median"]
+                        )
+                        / r["all_reduce"]["latency_median"]
+                    ) * 100
+                    overheads_kernel.append(overhead)
+                    speedup = (
+                        (
+                            r["rs_ag"]["latency_median"]
+                            - r["deterministic_kernel"]["latency_median"]
+                        )
+                        / r["rs_ag"]["latency_median"]
+                    ) * 100
+                    speedups_kernel.append(speedup)
+            print(
+                f"\nLatency Overhead Statistics (Deterministic Kernel vs All-Reduce):"
+            )
+            print(f"  Average: {statistics.mean(overheads_kernel):.1f}%")
+            print(f"  Median:  {statistics.median(overheads_kernel):.1f}%")
+            print(f"  Min:     {min(overheads_kernel):.1f}%")
+            print(f"  Max:     {max(overheads_kernel):.1f}%")
+            print(f"\nSpeedup Statistics (Deterministic Kernel vs RS+AG):")
+            print(f"  Average: {statistics.mean(speedups_kernel):.1f}%")
+            print(f"  Median:  {statistics.median(speedups_kernel):.1f}%")
+            print(f"  Min:     {min(speedups_kernel):.1f}%")
+            print(f"  Max:     {max(speedups_kernel):.1f}%")
+
+        # Show variance for non-deterministic cases
+        print(f"\nVariance Analysis (non-deterministic cases):")
+        for bs in sorted(results.keys()):
+            r = results[bs]
+            if not r["all_reduce"]["deterministic"]:
+                print(
+                    f"  Batch {bs}: All-Reduce max variance: {r['all_reduce']['max_variance']:.6f}"
+                )
+            if not r["rs_ag"]["deterministic"]:
+                print(
+                    f"  Batch {bs}: RS+All-Gather max variance: {r['rs_ag']['max_variance']:.6f}"
+                )
+            if r.get("custom_ar") is not None and not r["custom_ar"]["deterministic"]:
+                print(
+                    f"  Batch {bs}: Custom AR max variance: {r['custom_ar']['max_variance']:.6f}"
+                )
+            if (
+                r.get("deterministic_kernel") is not None
+                and not r["deterministic_kernel"]["deterministic"]
+            ):
+                print(
+                    f"  Batch {bs}: Deterministic Kernel max variance: {r['deterministic_kernel']['max_variance']:.6f}"
+                )
+
+
+if __name__ == "__main__":
+    main()

sgl-kernel/csrc/allreduce/deterministic_all_reduce.hip

@@ -0,0 +1,178 @@
+// Deterministic All-Reduce for ROCm/HIP
+//
+// This is a wrapper that forces the use of the existing 1-stage all-reduce kernel
+// (cross_device_reduce_1stage) which is inherently deterministic due to fixed
+// accumulation ordering (no atomics, no race conditions).
+//
+// How the 1-stage kernel works:
+//   - Each GPU reads ALL data from ALL other GPUs via direct memory access
+//   - Each GPU reduces the data locally in a fixed order
+//   - Result: every GPU has the complete reduced output
+//
+// This is NOT a reduce-scatter + all-gather (RS+AG) approach.
+// The 2-stage kernel (cross_device_reduce_2stage) implements RS+AG but may have
+// non-deterministic behavior, so we explicitly avoid it here.
+
+#include <ATen/hip/Exceptions.h>
+#include <ATen/hip/impl/HIPGuardImplMasqueradingAsCUDA.h>
+#include <ATen/hip/impl/HIPStreamMasqueradingAsCUDA.h>
+#include <torch/all.h>
+
+#include "custom_all_reduce_hip.cuh"
+
+using fptr_t = int64_t;
+static_assert(sizeof(void*) == sizeof(fptr_t));
+
+// Helper function for weak contiguity check
+bool _is_weak_contiguous_det(torch::Tensor& t) {
+  return t.is_contiguous() ||
+         (t.storage().nbytes() - t.storage_offset() * t.element_size() == t.numel() * t.element_size());
+}
+
+// Deterministic all-reduce for registered buffers (ROCm)
+// Uses the 1-stage kernel which is deterministic (fixed ordering)
+void deterministic_all_reduce_reg(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out) {
+  const at::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard(device_of(inp));
+  auto stream = c10::hip::getCurrentHIPStreamMasqueradingAsCUDA().stream();
+  TORCH_CHECK_EQ(inp.scalar_type(), out.scalar_type());
+  TORCH_CHECK_EQ(inp.numel(), out.numel());
+  TORCH_CHECK(_is_weak_contiguous_det(out));
+  TORCH_CHECK(_is_weak_contiguous_det(inp));
+
+  auto fa = reinterpret_cast<sglang::CustomAllreduce*>(_fa);
+
+  // For ROCm, manually call the 1-stage kernel to ensure deterministic ordering
+  // Get rank data pointer
+  sglang::RankData* ptrs;
+  hipStreamCaptureStatus status;
+  AT_CUDA_CHECK(hipStreamIsCapturing(stream, &status));
+  if (status == hipStreamCaptureStatusActive) {
+    ptrs = fa->d_rank_data_base_ + fa->graph_unreg_buffers_.size();
+    fa->graph_unreg_buffers_.push_back(inp.data_ptr());
+  } else {
+    auto it = fa->buffers_.find(inp.data_ptr());
+    if (it == fa->buffers_.end()) {
+      throw std::runtime_error("buffer not registered!");
+    }
+    ptrs = it->second;
+  }
+
+  int size = out.numel();
+  int threads = 512;
+
+  switch (out.scalar_type()) {
+    case at::ScalarType::Float: {
+      using T = float;
+      using P = typename sglang::packed_t<T>::P;
+      auto d = P::size;
+      if (size % d != 0) {
+        throw std::runtime_error("size must be multiple of " + std::to_string(d));
+      }
+      size /= d;
+      int blocks = std::min(16, (size + threads - 1) / threads);
+      // Always use 1-stage kernel for determinism
+      switch (fa->world_size_) {
+        case 2:
+          hipLaunchKernelGGL((sglang::cross_device_reduce_1stage<T, 2>), dim3(blocks), dim3(threads), 0, stream,
+                              ptrs, fa->sg_, fa->self_sg_, reinterpret_cast<T*>(out.data_ptr()), fa->rank_, size);
+          break;
+        case 4:
+          hipLaunchKernelGGL((sglang::cross_device_reduce_1stage<T, 4>), dim3(blocks), dim3(threads), 0, stream,
+                              ptrs, fa->sg_, fa->self_sg_, reinterpret_cast<T*>(out.data_ptr()), fa->rank_, size);
+          break;
+        case 6:
+          hipLaunchKernelGGL((sglang::cross_device_reduce_1stage<T, 6>), dim3(blocks), dim3(threads), 0, stream,
+                              ptrs, fa->sg_, fa->self_sg_, reinterpret_cast<T*>(out.data_ptr()), fa->rank_, size);
+          break;
+        case 8:
+          hipLaunchKernelGGL((sglang::cross_device_reduce_1stage<T, 8>), dim3(blocks), dim3(threads), 0, stream,
+                              ptrs, fa->sg_, fa->self_sg_, reinterpret_cast<T*>(out.data_ptr()), fa->rank_, size);
+          break;
+        default:
+          throw std::runtime_error("world_size must be in (2,4,6,8)");
+      }
+      break;
+    }
+    case at::ScalarType::Half: {
+      using T = half;
+      using P = typename sglang::packed_t<T>::P;
+      auto d = P::size;
+      if (size % d != 0) {
+        throw std::runtime_error("size must be multiple of " + std::to_string(d));
+      }
+      size /= d;
+      int blocks = std::min(16, (size + threads - 1) / threads);
+      switch (fa->world_size_) {
+        case 2:
+          hipLaunchKernelGGL((sglang::cross_device_reduce_1stage<T, 2>), dim3(blocks), dim3(threads), 0, stream,
+                              ptrs, fa->sg_, fa->self_sg_, reinterpret_cast<T*>(out.data_ptr()), fa->rank_, size);
+          break;
+        case 4:
+          hipLaunchKernelGGL((sglang::cross_device_reduce_1stage<T, 4>), dim3(blocks), dim3(threads), 0, stream,
+                              ptrs, fa->sg_, fa->self_sg_, reinterpret_cast<T*>(out.data_ptr()), fa->rank_, size);
+          break;
+        case 6:
+          hipLaunchKernelGGL((sglang::cross_device_reduce_1stage<T, 6>), dim3(blocks), dim3(threads), 0, stream,
+                              ptrs, fa->sg_, fa->self_sg_, reinterpret_cast<T*>(out.data_ptr()), fa->rank_, size);
+          break;
+        case 8:
+          hipLaunchKernelGGL((sglang::cross_device_reduce_1stage<T, 8>), dim3(blocks), dim3(threads), 0, stream,
+                              ptrs, fa->sg_, fa->self_sg_, reinterpret_cast<T*>(out.data_ptr()), fa->rank_, size);
+          break;
+        default:
+          throw std::runtime_error("world_size must be in (2,4,6,8)");
+      }
+      break;
+    }
+#if (__HIP_ARCH__ >= 800 || !defined(__HIP_ARCH__))
+    case at::ScalarType::BFloat16: {
+      using T = nv_bfloat16;
+      using P = typename sglang::packed_t<T>::P;
+      auto d = P::size;
+      if (size % d != 0) {
+        throw std::runtime_error("size must be multiple of " + std::to_string(d));
+      }
+      size /= d;
+      int blocks = std::min(16, (size + threads - 1) / threads);
+      switch (fa->world_size_) {
+        case 2:
+          hipLaunchKernelGGL((sglang::cross_device_reduce_1stage<T, 2>), dim3(blocks), dim3(threads), 0, stream,
+                              ptrs, fa->sg_, fa->self_sg_, reinterpret_cast<T*>(out.data_ptr()), fa->rank_, size);
+          break;
+        case 4:
+          hipLaunchKernelGGL((sglang::cross_device_reduce_1stage<T, 4>), dim3(blocks), dim3(threads), 0, stream,
+                              ptrs, fa->sg_, fa->self_sg_, reinterpret_cast<T*>(out.data_ptr()), fa->rank_, size);
+          break;
+        case 6:
+          hipLaunchKernelGGL((sglang::cross_device_reduce_1stage<T, 6>), dim3(blocks), dim3(threads), 0, stream,
+                              ptrs, fa->sg_, fa->self_sg_, reinterpret_cast<T*>(out.data_ptr()), fa->rank_, size);
+          break;
+        case 8:
+          hipLaunchKernelGGL((sglang::cross_device_reduce_1stage<T, 8>), dim3(blocks), dim3(threads), 0, stream,
+                              ptrs, fa->sg_, fa->self_sg_, reinterpret_cast<T*>(out.data_ptr()), fa->rank_, size);
+          break;
+        default:
+          throw std::runtime_error("world_size must be in (2,4,6,8)");
+      }
+      break;
+    }
+#endif
+    default:
+      throw std::runtime_error("deterministic allreduce only supports float32, float16 and bfloat16");
+  }
+}
+
+// Deterministic all-reduce for unregistered buffers (ROCm)
+void deterministic_all_reduce_unreg(fptr_t _fa, torch::Tensor& inp, torch::Tensor& reg_buffer, torch::Tensor& out) {
+  const at::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard(device_of(inp));
+  auto stream = c10::hip::getCurrentHIPStreamMasqueradingAsCUDA().stream();
+
+  auto input_size = inp.numel() * inp.element_size();
+  TORCH_CHECK_EQ(inp.scalar_type(), out.scalar_type());
+  TORCH_CHECK_EQ(inp.numel(), out.numel());
+  TORCH_CHECK(input_size <= reg_buffer.numel() * reg_buffer.element_size(),
+              "registered buffer is too small to contain the input");
+  AT_CUDA_CHECK(hipMemcpyAsync(reg_buffer.data_ptr(), inp.data_ptr(),
+                                input_size, hipMemcpyDeviceToDevice, stream));
+  deterministic_all_reduce_reg(_fa, reg_buffer, out);
+}

sgl-kernel/csrc/common_extension_rocm.cc

@@ -51,6 +51,17 @@ TORCH_LIBRARY_EXPAND(sgl_kernel, m) {
       "()");
   m.impl("all_reduce_unreg", torch::kCUDA, &all_reduce_unreg);
 
+  // Deterministic all-reduce for ROCm
+  extern void deterministic_all_reduce_reg(int64_t _fa, torch::Tensor & inp, torch::Tensor & out);
+  extern void deterministic_all_reduce_unreg(
+      int64_t _fa, torch::Tensor & inp, torch::Tensor & reg_buffer, torch::Tensor & out);
+
+  m.def("deterministic_all_reduce_reg(int fa, Tensor inp, Tensor! out) -> ()");
+  m.impl("deterministic_all_reduce_reg", torch::kCUDA, &deterministic_all_reduce_reg);
+
+  m.def("deterministic_all_reduce_unreg(int fa, Tensor inp, Tensor reg_buffer, Tensor! out) -> ()");
+  m.impl("deterministic_all_reduce_unreg", torch::kCUDA, &deterministic_all_reduce_unreg);
+
   m.def("dispose", &dispose);
 
   m.def("meta_size", &meta_size);

sgl-kernel/python/sgl_kernel/allreduce.py

@@ -24,6 +24,18 @@ if torch.version.hip is not None:
     ) -> None:
         torch.ops.sgl_kernel.all_reduce_unreg.default(fa, inp, reg_buffer, out)
 
+    def deterministic_all_reduce_reg(
+        fa: int, inp: torch.Tensor, out: torch.Tensor
+    ) -> None:
+        torch.ops.sgl_kernel.deterministic_all_reduce_reg.default(fa, inp, out)
+
+    def deterministic_all_reduce_unreg(
+        fa: int, inp: torch.Tensor, reg_buffer: torch.Tensor, out: torch.Tensor
+    ) -> None:
+        torch.ops.sgl_kernel.deterministic_all_reduce_unreg.default(
+            fa, inp, reg_buffer, out
+        )
+
     def dispose(fa: int) -> None:
         torch.ops.sgl_kernel.dispose.default(fa)
 

sgl-kernel/setup_rocm.py

@@ -42,6 +42,7 @@ include_dirs = [
 
 sources = [
     "csrc/allreduce/custom_all_reduce.hip",
+    "csrc/allreduce/deterministic_all_reduce.hip",
     "csrc/allreduce/quick_all_reduce.cu",
     "csrc/common_extension_rocm.cc",
     "csrc/elementwise/activation.cu",

sgl-kernel/tests/test_amd_deterministic_custom_allreduce.py

@@ -0,0 +1,270 @@
+"""
+Test deterministic custom all-reduce kernel behavior with batch size invariance.
+
+This test uses the 1-stage all-reduce kernel which is inherently deterministic
+due to fixed accumulation ordering (each GPU reads all data from all GPUs and
+reduces locally in a fixed order - no atomics, no race conditions).
+
+Note: This is NOT a reduce-scatter + all-gather (RS+AG) approach.
+
+This test compares:
+1. Deterministic kernel (same batch size)
+2. Deterministic kernel (different batch size)
+
+Usage:
+    python test_amd_deterministic_custom_allreduce.py
+"""
+
+import multiprocessing as mp
+import socket
+
+import torch
+import torch.distributed as dist
+
+
+def get_open_port():
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("127.0.0.1", 0))
+        return s.getsockname()[1]
+
+
+def worker(world_size, rank, port):
+    device = torch.device(f"cuda:{rank}")
+    torch.cuda.set_device(device)
+
+    dist.init_process_group(
+        backend="nccl",
+        init_method=f"tcp://localhost:{port}",
+        rank=rank,
+        world_size=world_size,
+    )
+
+    # Try to import and use deterministic kernel
+    try:
+        from torch.distributed import new_group
+
+        from sglang.srt.distributed.device_communicators.custom_all_reduce import (
+            CustomAllreduce,
+        )
+
+        # Create gloo group for custom AR
+        dist.barrier()
+        ar_group = new_group(backend="gloo")
+        dist.barrier()
+
+        custom_ar = CustomAllreduce(group=ar_group, device=device)
+
+        if custom_ar is None or custom_ar.disabled:
+            if rank == 0:
+                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}")
+            import traceback
+
+            traceback.print_exc()
+        dist.destroy_process_group()
+        return
+
+    num_trials = 10
+
+    # Matrix sizes similar to real model layers
+    # Format: (batch_size, hidden_dim) - typical tensor shape for all-reduce
+    BS = 50  # max batch_size (1..BS)
+    hidden_dim = 16384  # hidden dimension / intermediate dimension
+
+    # Different seed per rank - each GPU has DIFFERENT input
+    torch.manual_seed(42 + rank)
+
+    # Create fixed inputs for all trials
+    # Single request: (hidden_dim,)
+    base_input = torch.randn(hidden_dim, dtype=torch.bfloat16, device=device)
+    base_input_rand = torch.randn(hidden_dim, dtype=torch.bfloat16, device=device)
+
+    # Check if inputs fit in buffer
+    # Buffer size is max_size bytes, input size is numel * element_size bytes
+    input_size_bytes = base_input.numel() * base_input.element_size()
+    if input_size_bytes > custom_ar.max_size and rank == 0:
+        print(
+            f"Warning: Input size ({input_size_bytes/(1024*1024):.1f} MB) exceeds buffer size ({custom_ar.max_size/(1024*1024):.1f} MB)"
+        )
+        print("  Using unregistered mode (will copy to buffer)")
+
+    dist.barrier()
+
+    # =========================================================================
+    # TEST 1: Deterministic kernel (same batch size) - should be DETERMINISTIC
+    # =========================================================================
+    if rank == 0:
+        print(f"\n{'='*70}")
+        print("TEST 1: Deterministic kernel (same batch size)")
+        print(f"{'='*70}")
+    dist.barrier()
+
+    results_allreduce_only = []
+    for trial in range(num_trials):
+        # Clone the same input
+        inp = base_input.clone()
+
+        # Use deterministic kernel
+        # 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
+        checksum = result.view(-1).sum().item()
+        first_vals = result.view(-1)[:5].clone()
+        results_allreduce_only.append((checksum, first_vals))
+
+        if rank == 0:
+            print(
+                f"  Trial {trial+1:2d}: sum={checksum:.6f}, first5={first_vals.tolist()}"
+            )
+
+    # Check determinism
+    if rank == 0:
+        ref_sum, ref_vals = results_allreduce_only[0]
+        all_match = True
+        for i, (s, vals) in enumerate(results_allreduce_only[1:], 1):
+            if abs(ref_sum - s) > 1e-3 or not torch.allclose(ref_vals, vals, rtol=1e-3):
+                all_match = False
+                print(f"  Trial {i+1} DIFFERS! ref_sum={ref_sum:.6f}, got={s:.6f}")
+
+        if all_match:
+            print("  ✓ DETERMINISTIC KERNEL (fixed BS): DETERMINISTIC (as expected)")
+        else:
+            print(
+                "  ✗ DETERMINISTIC KERNEL (fixed BS): NON-DETERMINISTIC (unexpected!)"
+            )
+
+    dist.barrier()
+
+    # =========================================================================
+    # TEST 2: Deterministic kernel (different batch size) - should be DETERMINISTIC
+    # [a], [a, x], [a, x, x], ...
+    # =========================================================================
+    if rank == 0:
+        print(f"\n{'='*70}")
+        print("TEST 2: Deterministic kernel (different batch size)")
+        print("Batches: [a], [a,x], [a,x,x], ...")
+        print(f"{'='*70}")
+    dist.barrier()
+
+    results_allreduce_only = {trial: [] for trial in range(num_trials)}
+    for trial in range(num_trials):
+        for bs in range(1, BS + 1):
+            # Construct batch: (batch_size, hidden_dim)
+            # First element is base_input, rest are base_input_rand
+            batch = torch.stack([base_input] + [base_input_rand] * (bs - 1), dim=0)
+            # Shape: (bs, hidden_dim)
+
+            # Flatten for all-reduce: (bs * hidden_dim,)
+            batch_flat = batch.view(-1)
+
+            # Use deterministic kernel
+            # 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)
+            batch_out = result_flat.view(bs, hidden_dim)
+
+            # Only compare output corresponding to first request
+            out_first_req = batch_out[0].clone()
+            checksum = out_first_req.sum().item()
+            first_vals = out_first_req[:5].clone()
+            results_allreduce_only[trial].append((bs, checksum, first_vals))
+
+            if rank == 0:
+                print(
+                    f"  Batch size {bs:2d}: sum={checksum:.6f}, first5={first_vals.tolist()}"
+                )
+
+    # Check determinism
+    if rank == 0:
+        for trial in range(num_trials):
+            results = results_allreduce_only[trial]
+
+            _, ref_sum, ref_vals = results[0]
+            all_match = True
+            for _, s, vals in results[1:]:
+                if abs(ref_sum - s) > 1e-3 or not torch.allclose(
+                    ref_vals, vals, rtol=1e-3
+                ):
+                    all_match = False
+
+        if all_match:
+            print("  ✓ DETERMINISTIC KERNEL (variant BS): DETERMINISTIC")
+        else:
+            print("  ✗ DETERMINISTIC KERNEL (variant BS): NON-DETERMINISTIC")
+
+    dist.barrier()
+
+    dist.destroy_process_group()
+
+
+def main():
+    world_size = 8
+    available_gpus = torch.cuda.device_count()
+
+    print("=" * 70)
+    print("Deterministic Kernel All-Reduce Determinism Test")
+    print("=" * 70)
+    print(f"Available GPUs: {available_gpus}")
+    print(f"Using world_size: {world_size}")
+
+    if available_gpus < world_size:
+        print(
+            f"WARNING: Only {available_gpus} GPUs available, using {available_gpus} instead"
+        )
+        world_size = available_gpus
+
+    if world_size < 2:
+        print("ERROR: Need at least 2 GPUs for this test")
+        return
+
+    mp.set_start_method("spawn", force=True)
+    port = get_open_port()
+
+    procs = []
+    for rank in range(world_size):
+        p = mp.Process(target=worker, args=(world_size, rank, port))
+        p.start()
+        procs.append(p)
+
+    for p in procs:
+        p.join()
+
+
+if __name__ == "__main__":
+    main()

sgl-kernel/tests/test_amd_nccl_allreduce_determinism.py

@@ -0,0 +1,198 @@
+"""
+Test to confirm non-determinism of default NCCL all-reduce with batch size invariance.
+
+This test uses the default torch.distributed.all_reduce (NCCL) which can be
+NON-DETERMINISTIC due to tree-based reduction algorithms that don't guarantee
+fixed accumulation order for bfloat16/float16.
+
+This test compares:
+1. Default all-reduce (same batch size) - should be DETERMINISTIC
+2. Default all-reduce (different batch size) - typically NON-DETERMINISTIC for bfloat16
+
+Usage:
+    python test_ar.py
+"""
+
+import multiprocessing as mp
+import socket
+
+import torch
+import torch.distributed as dist
+
+
+def get_open_port():
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("127.0.0.1", 0))
+        return s.getsockname()[1]
+
+
+def worker(world_size, rank, port):
+    device = torch.device(f"cuda:{rank}")
+    torch.cuda.set_device(device)
+
+    dist.init_process_group(
+        backend="nccl",
+        init_method=f"tcp://localhost:{port}",
+        rank=rank,
+        world_size=world_size,
+    )
+
+    num_trials = 10
+
+    # Matrix sizes similar to real model layers
+    # Format: (batch_size, hidden_dim) - typical tensor shape for all-reduce
+    BS = 50  # max batch_size (1..BS)
+    hidden_dim = 16384  # hidden dimension / intermediate dimension
+
+    # Different seed per rank - each GPU has DIFFERENT input
+    torch.manual_seed(42 + rank)
+
+    # Create fixed inputs for all trials
+    # Single request: (hidden_dim,)
+    base_input = torch.randn(hidden_dim, dtype=torch.bfloat16, device=device)
+    base_input_rand = torch.randn(hidden_dim, dtype=torch.bfloat16, device=device)
+
+    dist.barrier()
+
+    # =========================================================================
+    # TEST 1: Default all-reduce (same batch size) - should be DETERMINISTIC
+    # =========================================================================
+    if rank == 0:
+        print(f"\n{'='*70}")
+        print("TEST 1: Default NCCL all_reduce (same batch size)")
+        print(f"{'='*70}")
+    dist.barrier()
+
+    results_allreduce_only = []
+    for trial in range(num_trials):
+        # Clone the same input
+        inp = base_input.clone()
+
+        # Use default NCCL all-reduce
+        dist.all_reduce(inp)
+        torch.cuda.synchronize()
+
+        # Store checksum
+        checksum = inp.view(-1).sum().item()
+        first_vals = inp.view(-1)[:5].clone()
+        results_allreduce_only.append((checksum, first_vals))
+
+        if rank == 0:
+            print(
+                f"  Trial {trial+1:2d}: sum={checksum:.6f}, first5={first_vals.tolist()}"
+            )
+
+    # Check determinism
+    if rank == 0:
+        ref_sum, ref_vals = results_allreduce_only[0]
+        all_match = True
+        for i, (s, vals) in enumerate(results_allreduce_only[1:], 1):
+            if abs(ref_sum - s) > 1e-3 or not torch.allclose(ref_vals, vals, rtol=1e-3):
+                all_match = False
+                print(f"  Trial {i+1} DIFFERS! ref_sum={ref_sum:.6f}, got={s:.6f}")
+
+        if all_match:
+            print("  ✓ DEFAULT ALL_REDUCE (fixed BS): DETERMINISTIC (as expected)")
+        else:
+            print("  ✗ DEFAULT ALL_REDUCE (fixed BS): NON-DETERMINISTIC (unexpected!)")
+
+    dist.barrier()
+
+    # =========================================================================
+    # TEST 2: Default all-reduce (different batch size) - typically NON-DETERMINISTIC
+    # [a], [a, x], [a, x, x], ...
+    # =========================================================================
+    if rank == 0:
+        print(f"\n{'='*70}")
+        print("TEST 2: Default NCCL all_reduce (different batch size)")
+        print("Batches: [a], [a,x], [a,x,x], ...")
+        print(f"{'='*70}")
+    dist.barrier()
+
+    results_allreduce_only = {trial: [] for trial in range(num_trials)}
+    for trial in range(num_trials):
+        for bs in range(1, BS + 1):
+            # Construct batch: (batch_size, hidden_dim)
+            # First element is base_input, rest are base_input_rand
+            batch = torch.stack([base_input] + [base_input_rand] * (bs - 1), dim=0)
+            # Shape: (bs, hidden_dim)
+
+            # Flatten for all-reduce: (bs * hidden_dim,)
+            batch_flat = batch.view(-1)
+
+            # Use default NCCL all-reduce
+            dist.all_reduce(batch_flat)
+            torch.cuda.synchronize()
+
+            # Reshape back to (bs, hidden_dim)
+            batch_out = batch_flat.view(bs, hidden_dim)
+
+            # Only compare output corresponding to first request
+            out_first_req = batch_out[0].clone()
+            checksum = out_first_req.sum().item()
+            first_vals = out_first_req[:5].clone()
+            results_allreduce_only[trial].append((bs, checksum, first_vals))
+
+            if rank == 0:
+                print(
+                    f"  Batch size {bs:2d}: sum={checksum:.6f}, first5={first_vals.tolist()}"
+                )
+
+    # Check determinism
+    if rank == 0:
+        for trial in range(num_trials):
+            results = results_allreduce_only[trial]
+
+            _, ref_sum, ref_vals = results[0]
+            all_match = True
+            for _, s, vals in results[1:]:
+                if abs(ref_sum - s) > 1e-3 or not torch.allclose(
+                    ref_vals, vals, rtol=1e-3
+                ):
+                    all_match = False
+
+        if all_match:
+            print("  ✓ DEFAULT ALL_REDUCE (variant BS): DETERMINISTIC")
+        else:
+            print("  ✗ DEFAULT ALL_REDUCE (variant BS): NON-DETERMINISTIC")
+
+    dist.barrier()
+
+    dist.destroy_process_group()
+
+
+def main():
+    world_size = 8
+    available_gpus = torch.cuda.device_count()
+
+    print("=" * 70)
+    print("Default NCCL All-Reduce Determinism Test")
+    print("=" * 70)
+    print(f"Available GPUs: {available_gpus}")
+    print(f"Using world_size: {world_size}")
+
+    if available_gpus < world_size:
+        print(
+            f"WARNING: Only {available_gpus} GPUs available, using {available_gpus} instead"
+        )
+        world_size = available_gpus
+
+    if world_size < 2:
+        print("ERROR: Need at least 2 GPUs for this test")
+        return
+
+    mp.set_start_method("spawn", force=True)
+    port = get_open_port()
+
+    procs = []
+    for rank in range(world_size):
+        p = mp.Process(target=worker, args=(world_size, rank, port))
+        p.start()
+        procs.append(p)
+
+    for p in procs:
+        p.join()
+
+
+if __name__ == "__main__":
+    main()