diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py index d7e0a3146..c9e055bb2 100644 --- a/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py +++ b/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, - # or, in the context of cuda graphs, register_graph_buffers - def all_reduce_reg(self, inp: torch.Tensor, out: torch.Tensor = None): - if out is None: - out = torch.empty_like(inp) - ops.all_reduce_reg(self._ptr, inp, out) - return out - - # all reduce, assuming inp tensor is NOT IPC registered - def all_reduce_unreg(self, inp: torch.Tensor, out: torch.Tensor = None): - if out is None: - out = torch.empty_like(inp) - ops.all_reduce_unreg(self._ptr, inp, self.buffer, out) - return out - - def all_reduce( - self, - inp: torch.Tensor, - *, - out: torch.Tensor = None, - registered: bool = False, - ): - """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) + def _all_reduce_impl(self, inp: torch.Tensor, registered: bool): + out = torch.empty_like(inp) + if not _is_hip: # CUDA-like + 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 + ) + elif self.use_amd_deterministic_impl: + inp_size = inp.numel() * inp.element_size() + if inp_size < self.max_size: + reg_buffer = self.buffer.view(inp.dtype)[: inp.numel()] + ops.deterministic_all_reduce_unreg(self._ptr, inp, reg_buffer, out) + else: + self.register_buffer(inp) + ops.deterministic_all_reduce_reg(self._ptr, inp, out) + else: # normal AMD ROCm path + if registered: + ops.all_reduce_reg(self._ptr, inp, out) + else: + ops.all_reduce_unreg(self._ptr, inp, self.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: - 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) + return self._all_reduce_impl(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_unreg(input) - else: - return self.all_reduce(input, registered=False) + return self._all_reduce_impl(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: - # 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) + return self._all_reduce_impl(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() diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index abe3c218a..b6f3987b7 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/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) diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 74e3da356..7d8721036 100644 --- a/python/sglang/srt/utils/common.py +++ b/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) diff --git a/sgl-kernel/benchmark/bench_amd_deterministic_allreduce.py b/sgl-kernel/benchmark/bench_amd_deterministic_allreduce.py index 647543089..234d09342 100644 --- a/sgl-kernel/benchmark/bench_amd_deterministic_allreduce.py +++ b/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) diff --git a/sgl-kernel/tests/test_amd_deterministic_custom_allreduce.py b/sgl-kernel/tests/test_amd_deterministic_custom_allreduce.py index 7e3e82c9b..aa7125925 100644 --- a/sgl-kernel/tests/test_amd_deterministic_custom_allreduce.py +++ b/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)