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

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(
-                            inp_kernel, registered=False
-                        )
+                        result_kernel = custom_ar.custom_all_reduce(inp_kernel)
                         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
-        # 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)
+        result = custom_ar.custom_all_reduce(inp)
         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
-            # 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
-                )
+            result_flat = custom_ar.custom_all_reduce(batch_flat)
             torch.cuda.synchronize()
 
             # Reshape back to (bs, hidden_dim)