minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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-02 17:20:05 -08:00
committed by GitHub
co-authored by gemini-code-assist[bot] Kangyan-Zhou Lianmin Zheng
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
python/sglang/srt/_custom_ops.py
+73 -41
View File
@@ -4,29 +4,42 @@ from typing import List, Optional, Tuple
import torch 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__) logger = logging.getLogger(__name__)
_is_cuda = is_cuda()
_is_hip = is_hip()
if not is_hpu(): IS_CUSTOM_AR_AVAILABLE = _is_cuda or _is_hip
try: IS_QUICK_AR_AVAILABLE = _is_hip
import sgl_kernel # TODO(zyksir): mscclpp is untested on AMD and therefore disabled.
except ImportError as e: 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) 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(): if not IS_CUSTOM_AR_AVAILABLE:
custom_op = sgl_kernel.allreduce pass
elif _is_cuda:
# CUDA custom allreduce
# custom allreduce
def init_custom_ar( def init_custom_ar(
ipc_tensors: List[torch.Tensor], ipc_tensors: List[torch.Tensor],
rank_data: torch.Tensor, rank_data: torch.Tensor,
rank: int, rank: int,
full_nvlink: bool, full_nvlink: bool,
) -> int: ) -> 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( def all_reduce(
fa: int, fa: int,
@@ -35,26 +48,26 @@ if not is_hip() and not is_npu():
reg_buffer: int, reg_buffer: int,
reg_buffer_sz_bytes: int, reg_buffer_sz_bytes: int,
) -> None: ) -> 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: def dispose(fa: int) -> None:
custom_op.dispose(fa) _custom_ar.dispose(fa)
def meta_size() -> int: def meta_size() -> int:
return custom_op.meta_size() return _custom_ar.meta_size()
def register_buffer(fa: int, ipc_tensors: List[int]) -> None: 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]]: 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( def register_graph_buffers(
fa: int, handles: List[List[int]], offsets: List[List[int]] fa: int, handles: List[List[int]], offsets: List[List[int]]
) -> None: ) -> None:
custom_op.register_graph_buffers(fa, handles, offsets) _custom_ar.register_graph_buffers(fa, handles, offsets)
else: elif _is_hip:
# ROCM custom allreduce # ROCM custom allreduce
def init_custom_ar( def init_custom_ar(
@@ -65,55 +78,64 @@ else:
rank: int, rank: int,
full_nvlink: bool, full_nvlink: bool,
) -> int: ) -> int:
return sgl_kernel.allreduce.init_custom_ar( return _custom_ar.init_custom_ar(
meta, rank_data, handles, offsets, rank, full_nvlink meta, rank_data, handles, offsets, rank, full_nvlink
) )
def all_reduce_reg(fa: int, inp: torch.Tensor, out: torch.Tensor) -> None: 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( def all_reduce_unreg(
fa: int, inp: torch.Tensor, reg_buffer: torch.Tensor, out: torch.Tensor fa: int, inp: torch.Tensor, reg_buffer: torch.Tensor, out: torch.Tensor
) -> None: ) -> 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: def dispose(fa: int) -> None:
sgl_kernel.allreduce.dispose(fa) _custom_ar.dispose(fa)
def meta_size() -> int: def meta_size() -> int:
return sgl_kernel.allreduce.meta_size() return _custom_ar.meta_size()
def register_buffer( def register_buffer(
fa: int, t: torch.Tensor, handles: List[str], offsets: List[int] fa: int, t: torch.Tensor, handles: List[str], offsets: List[int]
) -> None: ) -> 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]]: 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( def register_graph_buffers(
fa: int, handles: List[str], offsets: List[List[int]] fa: int, handles: List[str], offsets: List[List[int]]
) -> None: ) -> 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: 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: 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 # ROCM custom quick allreduce
def init_custom_qr( def init_custom_qr(
rank: int, world_size: int, qr_max_size: Optional[int] = None rank: int, world_size: int, qr_max_size: Optional[int] = None
) -> int: ) -> 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: 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: 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( def qr_all_reduce(
fa: int, fa: int,
@@ -122,20 +144,28 @@ else:
quant_level: int, quant_level: int,
cast_bf2half: bool, cast_bf2half: bool,
) -> None: ) -> 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: def qr_destroy(fa: int) -> None:
sgl_kernel.allreduce.qr_destroy(fa) _custom_ar.qr_destroy(fa)
def qr_max_size() -> int: def qr_max_size() -> int:
return sgl_kernel.allreduce.qr_max_size() return _custom_ar.qr_max_size()
def mscclpp_generate_unique_id() -> bytes: # endregion
return sgl_kernel.allreduce.mscclpp_generate_unique_id()
# region IS_MSCCLPP_AR_AVAILABLE
def mscclpp_init_context( 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, unique_id: bytes,
rank: int, rank: int,
world_size: int, world_size: int,
@@ -145,8 +175,8 @@ def mscclpp_init_context(
rank_to_node: List[int], rank_to_node: List[int],
rank_to_ib: List[int], rank_to_ib: List[int],
context_selection: int, context_selection: int,
) -> int: ) -> int:
return sgl_kernel.allreduce.mscclpp_init_context( return _custom_ar.mscclpp_init_context(
unique_id, unique_id,
rank, rank,
world_size, world_size,
@@ -158,8 +188,10 @@ def mscclpp_init_context(
context_selection, context_selection,
) )
def mscclpp_allreduce(
def mscclpp_allreduce(
context: int, inp: torch.Tensor, out: torch.Tensor, nthreads: int, nblocks: int context: int, inp: torch.Tensor, out: torch.Tensor, nthreads: int, nblocks: int
) -> None: ) -> None:
return sgl_kernel.allreduce.mscclpp_allreduce(context, inp, out, nthreads, nblocks) return _custom_ar.mscclpp_allreduce(context, inp, out, nthreads, nblocks)
# endregion
python/sglang/srt/distributed/device_communicators/custom_all_reduce.py
+1 -11
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.environ import envs
from sglang.srt.utils import get_bool_env_var, is_cuda, is_hip, log_info_on_rank0 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_cuda = is_cuda()
_is_hip = is_hip() _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.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.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 # disable because of missing custom allreduce library
# e.g. in a non-cuda environment # e.g. in a non-cuda environment
return return
python/sglang/srt/distributed/device_communicators/pymscclpp.py
+2 -15
View File
@@ -11,25 +11,12 @@ import torch.distributed as dist
from torch.distributed import ProcessGroup, ReduceOp from torch.distributed import ProcessGroup, ReduceOp
from sglang.srt import _custom_ops as ops 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__) logger = logging.getLogger(__name__)
_is_cuda = is_cuda()
_is_hip = is_hip() _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): class MscclContextSelection(IntEnum):
MSCCL1SHOT1NODELL = 1 MSCCL1SHOT1NODELL = 1
@@ -127,7 +114,7 @@ class PyMscclppCommunicator:
self._IS_CAPTURING = False self._IS_CAPTURING = False
self.disabled = True self.disabled = True
if not mscclpp_is_available: if not ops.IS_MSCCLPP_AR_AVAILABLE:
# disable because of missing mscclpp library # disable because of missing mscclpp library
# e.g. in a non-cuda environment # e.g. in a non-cuda environment
return return
python/sglang/srt/distributed/device_communicators/quick_all_reduce.py
+1 -9
View File
@@ -24,14 +24,6 @@ _is_cuda = is_cuda()
_is_hip = is_hip() _is_hip = is_hip()
try:
ops.qr_max_size()
quick_ar = True
except Exception:
# For CPUs and CUDA
quick_ar = False
@cache @cache
def qr_rocm_arch_available(): def qr_rocm_arch_available():
if not _is_hip: if not _is_hip:
@@ -101,7 +93,7 @@ class QuickAllReduce:
) )
return return
if not quick_ar: if not ops.IS_QUICK_AR_AVAILABLE:
# disable because of missing quick reduce library # disable because of missing quick reduce library
# e.g. in a cuda environment # e.g. in a cuda environment
logger.info( logger.info(
python/sglang/srt/utils/common.py
+1
View File
@@ -122,6 +122,7 @@ HIP_FP8_E4M3_FNUZ_MAX = 224.0
# https://pytorch.org/docs/stable/notes/hip.html#checking-for-hip # https://pytorch.org/docs/stable/notes/hip.html#checking-for-hip
@lru_cache(maxsize=1)
def is_hip() -> bool: def is_hip() -> bool:
return torch.version.hip is not None return torch.version.hip is not None