minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
DarkSharpness
2026-04-03 12:33:17 +08:00
committed by GitHub
parent 81efcc353a
commit d1b7c3907d
5 changed files with 50 additions and 142 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/distributed/device_communicators/custom_all_reduce.py
+39 -79
View File
@@ -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()
python/sglang/srt/distributed/parallel_state.py
-20
View File
@@ -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
+1 -1
View File
@@ -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)