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

Refactor custom allreduce logics (#13710)

Browse Source 
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Kangyan-Zhou <zky314343421@gmail.com>
Co-authored-by: Lianmin Zheng <lianminzheng@gmail.com>
This commit is contained in:
Even Zhou
2025-12-03 09:20:05 +08:00
committed by GitHub
parent 7ae368efde
commit 7d1a130cde
5 changed files with 99 additions and 97 deletions
Download Patch File Download Diff File Expand all files Collapse all files

156
python/sglang/srt/_custom_ops.py
View File

@@ -4,29 +4,42 @@ from typing import List, Optional, Tuple
import torch
from sglang.srt.utils import is_hip, is_hpu, is_npu
from sglang.srt.utils import is_cuda, is_hip
logger = logging.getLogger(__name__)
_is_cuda = is_cuda()
_is_hip = is_hip()
if not is_hpu():
try:
import sgl_kernel
except ImportError as e:
IS_CUSTOM_AR_AVAILABLE = _is_cuda or _is_hip
IS_QUICK_AR_AVAILABLE = _is_hip
# TODO(zyksir): mscclpp is untested on AMD and therefore disabled.
IS_MSCCLPP_AR_AVAILABLE = _is_cuda
try:
import sgl_kernel.allreduce as _custom_ar
except ImportError as e:
if _is_cuda or _is_hip:
logger.warning("Failed to import from custom_ar with %r", e)
IS_CUSTOM_AR_AVAILABLE = False
IS_QUICK_AR_AVAILABLE = False
IS_MSCCLPP_AR_AVAILABLE = False
# region IS_CUSTOM_AR_AVAILABLE
if not is_hip() and not is_npu():
custom_op = sgl_kernel.allreduce
if not IS_CUSTOM_AR_AVAILABLE:
pass
elif _is_cuda:
# CUDA custom allreduce
# custom allreduce
def init_custom_ar(
ipc_tensors: List[torch.Tensor],
rank_data: torch.Tensor,
rank: int,
full_nvlink: bool,
) -> int:
return custom_op.init_custom_ar(ipc_tensors, rank_data, rank, full_nvlink)
return _custom_ar.init_custom_ar(ipc_tensors, rank_data, rank, full_nvlink)
def all_reduce(
fa: int,
@@ -35,26 +48,26 @@ if not is_hip() and not is_npu():
reg_buffer: int,
reg_buffer_sz_bytes: int,
) -> None:
custom_op.all_reduce(fa, inp, out, reg_buffer, reg_buffer_sz_bytes)
_custom_ar.all_reduce(fa, inp, out, reg_buffer, reg_buffer_sz_bytes)
def dispose(fa: int) -> None:
custom_op.dispose(fa)
_custom_ar.dispose(fa)
def meta_size() -> int:
return custom_op.meta_size()
return _custom_ar.meta_size()
def register_buffer(fa: int, ipc_tensors: List[int]) -> None:
return custom_op.register_buffer(fa, ipc_tensors)
return _custom_ar.register_buffer(fa, ipc_tensors)
def get_graph_buffer_ipc_meta(fa: int) -> Tuple[List[int], List[int]]:
return custom_op.get_graph_buffer_ipc_meta(fa)
return _custom_ar.get_graph_buffer_ipc_meta(fa)
def register_graph_buffers(
fa: int, handles: List[List[int]], offsets: List[List[int]]
) -> None:
custom_op.register_graph_buffers(fa, handles, offsets)
_custom_ar.register_graph_buffers(fa, handles, offsets)
else:
elif _is_hip:
# ROCM custom allreduce
def init_custom_ar(
@@ -65,55 +78,64 @@ else:
rank: int,
full_nvlink: bool,
) -> int:
return sgl_kernel.allreduce.init_custom_ar(
return _custom_ar.init_custom_ar(
meta, rank_data, handles, offsets, rank, full_nvlink
)
def all_reduce_reg(fa: int, inp: torch.Tensor, out: torch.Tensor) -> None:
sgl_kernel.allreduce.all_reduce_reg(fa, inp, out)
_custom_ar.all_reduce_reg(fa, inp, out)
def all_reduce_unreg(
fa: int, inp: torch.Tensor, reg_buffer: torch.Tensor, out: torch.Tensor
) -> None:
sgl_kernel.allreduce.all_reduce_unreg(fa, inp, reg_buffer, out)
_custom_ar.all_reduce_unreg(fa, inp, reg_buffer, out)
def dispose(fa: int) -> None:
sgl_kernel.allreduce.dispose(fa)
_custom_ar.dispose(fa)
def meta_size() -> int:
return sgl_kernel.allreduce.meta_size()
return _custom_ar.meta_size()
def register_buffer(
fa: int, t: torch.Tensor, handles: List[str], offsets: List[int]
) -> None:
return sgl_kernel.allreduce.register_buffer(fa, t, handles, offsets)
return _custom_ar.register_buffer(fa, t, handles, offsets)
def get_graph_buffer_ipc_meta(fa: int) -> Tuple[torch.Tensor, List[int]]:
return sgl_kernel.allreduce.get_graph_buffer_ipc_meta(fa)
return _custom_ar.get_graph_buffer_ipc_meta(fa)
def register_graph_buffers(
fa: int, handles: List[str], offsets: List[List[int]]
) -> None:
sgl_kernel.allreduce.register_graph_buffers(fa, handles, offsets)
_custom_ar.register_graph_buffers(fa, handles, offsets)
def allocate_meta_buffer(size: int) -> torch.Tensor:
return sgl_kernel.allreduce.allocate_meta_buffer(size)
return _custom_ar.allocate_meta_buffer(size)
def get_meta_buffer_ipc_handle(inp: torch.Tensor) -> torch.Tensor:
return sgl_kernel.allreduce.get_meta_buffer_ipc_handle(inp)
return _custom_ar.get_meta_buffer_ipc_handle(inp)
# endregion
# region IS_QUICK_AR_AVAILABLE
if not IS_QUICK_AR_AVAILABLE:
pass
elif _is_hip:
# ROCM custom quick allreduce
def init_custom_qr(
rank: int, world_size: int, qr_max_size: Optional[int] = None
) -> int:
return sgl_kernel.allreduce.init_custom_qr(world_size, rank, qr_max_size)
return _custom_ar.init_custom_qr(world_size, rank, qr_max_size)
def qr_get_handle(fa: int) -> torch.Tensor:
return sgl_kernel.allreduce.qr_get_handle(fa)
return _custom_ar.qr_get_handle(fa)
def qr_open_handles(fa: int, handles: list[torch.Tensor]) -> None:
sgl_kernel.allreduce.qr_open_handles(fa, handles)
_custom_ar.qr_open_handles(fa, handles)
def qr_all_reduce(
fa: int,
@@ -122,44 +144,54 @@ else:
quant_level: int,
cast_bf2half: bool,
) -> None:
sgl_kernel.allreduce.qr_all_reduce(fa, inp, out, quant_level, cast_bf2half)
_custom_ar.qr_all_reduce(fa, inp, out, quant_level, cast_bf2half)
def qr_destroy(fa: int) -> None:
sgl_kernel.allreduce.qr_destroy(fa)
_custom_ar.qr_destroy(fa)
def qr_max_size() -> int:
return sgl_kernel.allreduce.qr_max_size()
return _custom_ar.qr_max_size()
def mscclpp_generate_unique_id() -> bytes:
return sgl_kernel.allreduce.mscclpp_generate_unique_id()
# endregion
# region IS_MSCCLPP_AR_AVAILABLE
if not IS_MSCCLPP_AR_AVAILABLE:
pass
elif _is_cuda:
def mscclpp_generate_unique_id() -> bytes:
return _custom_ar.mscclpp_generate_unique_id()
def mscclpp_init_context(
unique_id: bytes,
rank: int,
world_size: int,
scratch: torch.Tensor,
put_buffer: torch.Tensor,
nranks_per_node: int,
rank_to_node: List[int],
rank_to_ib: List[int],
context_selection: int,
) -> int:
return _custom_ar.mscclpp_init_context(
unique_id,
rank,
world_size,
scratch,
put_buffer,
nranks_per_node,
rank_to_node,
rank_to_ib,
context_selection,
)
def mscclpp_allreduce(
context: int, inp: torch.Tensor, out: torch.Tensor, nthreads: int, nblocks: int
) -> None:
return _custom_ar.mscclpp_allreduce(context, inp, out, nthreads, nblocks)
def mscclpp_init_context(
unique_id: bytes,
rank: int,
world_size: int,
scratch: torch.Tensor,
put_buffer: torch.Tensor,
nranks_per_node: int,
rank_to_node: List[int],
rank_to_ib: List[int],
context_selection: int,
) -> int:
return sgl_kernel.allreduce.mscclpp_init_context(
unique_id,
rank,
world_size,
scratch,
put_buffer,
nranks_per_node,
rank_to_node,
rank_to_ib,
context_selection,
)
def mscclpp_allreduce(
context: int, inp: torch.Tensor, out: torch.Tensor, nthreads: int, nblocks: int
) -> None:
return sgl_kernel.allreduce.mscclpp_allreduce(context, inp, out, nthreads, nblocks)
# endregion

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

@@ -21,16 +21,6 @@ from sglang.srt.distributed.parallel_state import in_the_same_node_as
from sglang.srt.environ import envs
from sglang.srt.utils import get_bool_env_var, is_cuda, is_hip, log_info_on_rank0
try:
# Use custom allreduce from sgl kernel (ROCM and TRT-LLM)
import sgl_kernel # noqa: F401
custom_ar = True
except ImportError:
# For CPUs
custom_ar = False
_is_cuda = is_cuda()
_is_hip = is_hip()
@@ -79,7 +69,7 @@ class CustomAllreduce:
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
if not custom_ar:
if not ops.IS_CUSTOM_AR_AVAILABLE:
# disable because of missing custom allreduce library
# e.g. in a non-cuda environment
return

17
python/sglang/srt/distributed/device_communicators/pymscclpp.py
View File

@@ -11,25 +11,12 @@ import torch.distributed as dist
from torch.distributed import ProcessGroup, ReduceOp
from sglang.srt import _custom_ops as ops
from sglang.srt.utils import is_cuda, is_hip
from sglang.srt.utils import is_hip
logger = logging.getLogger(__name__)
_is_cuda = is_cuda()
_is_hip = is_hip()
mscclpp_is_available = False
if _is_hip:
# TODO(zyksir): mscclpp is untested on AMD and therefore disabled.
mscclpp_is_available = False
if _is_cuda:
try:
import sgl_kernel # noqa: F401
mscclpp_is_available = True
except:
mscclpp_is_available = False
class MscclContextSelection(IntEnum):
MSCCL1SHOT1NODELL = 1
@@ -127,7 +114,7 @@ class PyMscclppCommunicator:
self._IS_CAPTURING = False
self.disabled = True
if not mscclpp_is_available:
if not ops.IS_MSCCLPP_AR_AVAILABLE:
# disable because of missing mscclpp library
# e.g. in a non-cuda environment
return

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

@@ -24,14 +24,6 @@ _is_cuda = is_cuda()
_is_hip = is_hip()
try:
ops.qr_max_size()
quick_ar = True
except Exception:
# For CPUs and CUDA
quick_ar = False
@cache
def qr_rocm_arch_available():
if not _is_hip:
@@ -101,7 +93,7 @@ class QuickAllReduce:
)
return
if not quick_ar:
if not ops.IS_QUICK_AR_AVAILABLE:
# disable because of missing quick reduce library
# e.g. in a cuda environment
logger.info(

1
python/sglang/srt/utils/common.py
View File

@@ -122,6 +122,7 @@ HIP_FP8_E4M3_FNUZ_MAX = 224.0
# https://pytorch.org/docs/stable/notes/hip.html#checking-for-hip
@lru_cache(maxsize=1)
def is_hip() -> bool:
return torch.version.hip is not None