chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Thomas Wang <1am9trash@gmail.com>
This commit is contained in:
sunxxuns
2025-12-19 02:36:03 -05:00
committed by GitHub
parent 0e869f0868
commit f2d64e6782
12 changed files with 1466 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

40
python/sglang/srt/distributed/device_communicators/custom_all_reduce.py
View File

@@ -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

10
python/sglang/srt/distributed/device_communicators/custom_all_reduce_ops.py
View File

@@ -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)

39
python/sglang/srt/distributed/parallel_state.py
View File

@@ -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)

5
python/sglang/srt/environ.py
View File

@@ -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)

18
python/sglang/srt/server_args.py
View File

@@ -2314,11 +2314,19 @@ class ServerArgs:
# Check TP size
if self.tp_size > 1:
os.environ["NCCL_ALGO"] = "allreduce:tree"
self.disable_custom_all_reduce = True
logger.warning(
"NCCL_ALGO is set to 'allreduce:tree' and custom all reduce is disabled for deterministic inference when 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(
"NCCL_ALGO is set to 'allreduce:tree' and custom all reduce is disabled for deterministic inference when TP size > 1."
)
def _handle_dllm_inference(self):
if self.dllm_algorithm is None: