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)
sgl-kernel/benchmark/bench_amd_deterministic_allreduce.py
+5 -8
View File
@@ -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
+5 -34
View File
@@ -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)