chore: import upstream snapshot with attribution

python/sglang/srt/distributed/__init__.py

@@ -0,0 +1,3 @@
+from sglang.srt.distributed.communication_op import *
+from sglang.srt.distributed.parallel_state import *
+from sglang.srt.distributed.utils import *

python/sglang/srt/distributed/communication_op.py

@@ -0,0 +1,84 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/distributed/communication_op.py
+
+from typing import Any, Dict, Optional, Tuple, Union
+
+import torch
+import torch.distributed
+
+from .parallel_state import (
+    get_attn_tp_group,
+    get_moe_ep_group,
+    get_moe_tp_group,
+    get_tp_group,
+)
+
+
+def tensor_model_parallel_all_reduce(input_: torch.Tensor) -> torch.Tensor:
+    """All-reduce the input tensor across model parallel group."""
+    return get_tp_group().all_reduce(input_)
+
+
+def tensor_model_parallel_quant_all_reduce(input_: torch.Tensor) -> torch.Tensor:
+    """All-reduce the input tensor across model parallel group."""
+    return get_tp_group().quant_all_reduce(input_)
+
+
+def tensor_model_parallel_fused_allreduce_rmsnorm(
+    input_: torch.Tensor,
+    residual_inp_: torch.Tensor,
+    weight_: torch.Tensor,
+    eps: float,
+) -> Optional[Tuple[torch.Tensor, torch.Tensor]]:
+    """Fused TP all-reduce + RMSNorm.
+
+    Policy and backend selection are owned by GroupCoordinator:
+    it may dispatch to communicator-native fused APIs, custom fused kernels,
+    or return None so callers can run generic fallback paths.
+    """
+    return get_tp_group().fused_allreduce_rmsnorm(input_, residual_inp_, weight_, eps)
+
+
+def tensor_model_parallel_all_gather(
+    input_: torch.Tensor, dim: int = -1
+) -> torch.Tensor:
+    """All-gather the input tensor across model parallel group."""
+    return get_tp_group().all_gather(input_, dim)
+
+
+def tensor_model_parallel_gather(
+    input_: torch.Tensor, dst: int = 0, dim: int = -1
+) -> Optional[torch.Tensor]:
+    """Gather the input tensor across model parallel group."""
+    return get_tp_group().gather(input_, dst, dim)
+
+
+def broadcast_tensor_dict(
+    tensor_dict: Optional[Dict[Any, Union[torch.Tensor, Any]]] = None, src: int = 0
+):
+    if not torch.distributed.is_initialized():
+        return tensor_dict
+    return get_tp_group().broadcast_tensor_dict(tensor_dict, src)
+
+
+def attention_tensor_model_parallel_all_reduce(input_: torch.Tensor) -> torch.Tensor:
+    """All-reduce the input tensor across attention parallel group."""
+    return get_attn_tp_group().all_reduce(input_)
+
+
+def attention_tensor_model_parallel_quant_all_reduce(
+    input_: torch.Tensor,
+) -> torch.Tensor:
+    """All-reduce the input tensor across attention parallel group."""
+    return get_attn_tp_group().quant_all_reduce(input_)
+
+
+def moe_tensor_model_parallel_all_reduce(input_: torch.Tensor) -> torch.Tensor:
+    """All-reduce the input tensor across moe parallel group."""
+    return get_moe_tp_group().all_reduce(input_)
+
+
+def moe_expert_parallel_all_reduce(input_: torch.Tensor) -> torch.Tensor:
+    """All-reduce the input tensor across moe expert parallel group."""
+    return get_moe_ep_group().all_reduce(input_)

python/sglang/srt/distributed/communication_tags.py

@@ -0,0 +1,15 @@
+from enum import IntEnum, unique
+
+
+@unique
+class P2PTag(IntEnum):
+    """
+    Tags reserved for point-to-point communication protocols.
+
+    Communications introduced outside existing scheduler loops need explicit
+    tags to avoid being consumed by unrelated send/recv paths.
+    """
+
+    DEFAULT = 0
+    HIRADIX_PP_SYNC = int.from_bytes(b"PpHi", byteorder="big")
+    GRAMMAR_PP_SYNC = int.from_bytes(b"PpGr", byteorder="big")

python/sglang/srt/distributed/device_communicators/all_reduce_utils.py

@@ -0,0 +1,16 @@
+MiB = 1024 * 1024
+
+TORCH_SYMM_MEM_ALL_REDUCE_MAX_SIZES = {
+    9: {
+        2: 64 * MiB,  # 64 MB
+        4: 64 * MiB,  # 64 MB
+        6: 128 * MiB,  # 128 MB
+        8: 128 * MiB,  # 128 MB
+    },
+    10: {
+        2: 64 * MiB,  # 64 MB
+        4: 64 * MiB,  # 64 MB
+        6: 128 * MiB,  # 128 MB
+        8: 128 * MiB,  # 128 MB
+    },
+}

python/sglang/srt/distributed/device_communicators/cuda_wrapper.py

@@ -0,0 +1,189 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/distributed/device_communicators/cuda_wrapper.py
+
+"""This file is a pure Python wrapper for the cudart library.
+It avoids the need to compile a separate shared library, and is
+convenient for use when we just need to call a few functions.
+"""
+
+import ctypes
+import logging
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional
+
+# this line makes it possible to directly load `libcudart.so` using `ctypes`
+import torch  # noqa
+
+from sglang.srt.utils import is_musa
+
+_is_musa = is_musa()
+
+logger = logging.getLogger(__name__)
+
+# === export types and functions from cudart to Python ===
+# for the original cudart definition, please check
+# https://docs.nvidia.com/cuda/cuda-runtime-api/index.html
+
+cudaError_t = ctypes.c_int
+cudaMemcpyKind = ctypes.c_int
+
+
+class cudaIpcMemHandle_t(ctypes.Structure):
+    _fields_ = [("internal", ctypes.c_byte * 128)]
+
+
+@dataclass
+class Function:
+    name: str
+    restype: Any
+    argtypes: List[Any]
+
+
+def find_loaded_library(lib_name) -> Optional[str]:
+    """
+    According to according to https://man7.org/linux/man-pages/man5/proc_pid_maps.5.html,
+    the file `/proc/self/maps` contains the memory maps of the process, which includes the
+    shared libraries loaded by the process. We can use this file to find the path of the
+    a loaded library.
+    """  # noqa
+    found = False
+    with open("/proc/self/maps") as f:
+        for line in f:
+            if lib_name in line:
+                found = True
+                break
+    if not found:
+        # the library is not loaded in the current process
+        return None
+    # if lib_name is libcudart, we need to match a line with:
+    # address /path/to/libcudart-hash.so.11.0
+    start = line.index("/")
+    path = line[start:].strip()
+    filename = path.split("/")[-1]
+    assert filename.rpartition(".so")[0].startswith(
+        lib_name
+    ), f"Unexpected filename: {filename} for library {lib_name}"
+    return path
+
+
+class CudaRTLibrary:
+    exported_functions = [
+        # ​cudaError_t cudaSetDevice ( int  device )
+        Function("cudaSetDevice", cudaError_t, [ctypes.c_int]),
+        # cudaError_t 	cudaDeviceSynchronize ( void )
+        Function("cudaDeviceSynchronize", cudaError_t, []),
+        # ​cudaError_t cudaDeviceReset ( void )
+        Function("cudaDeviceReset", cudaError_t, []),
+        # const char* 	cudaGetErrorString ( cudaError_t error )
+        Function("cudaGetErrorString", ctypes.c_char_p, [cudaError_t]),
+        # ​cudaError_t 	cudaMalloc ( void** devPtr, size_t size )
+        Function(
+            "cudaMalloc",
+            cudaError_t,
+            [ctypes.POINTER(ctypes.c_void_p), ctypes.c_size_t],
+        ),
+        # ​cudaError_t 	cudaFree ( void* devPtr )
+        Function("cudaFree", cudaError_t, [ctypes.c_void_p]),
+        # ​cudaError_t cudaMemset ( void* devPtr, int  value, size_t count )
+        Function(
+            "cudaMemset", cudaError_t, [ctypes.c_void_p, ctypes.c_int, ctypes.c_size_t]
+        ),
+        # ​cudaError_t cudaMemcpy ( void* dst, const void* src, size_t count, cudaMemcpyKind kind ) # noqa
+        Function(
+            "cudaMemcpy",
+            cudaError_t,
+            [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_size_t, cudaMemcpyKind],
+        ),
+        # cudaError_t cudaIpcGetMemHandle ( cudaIpcMemHandle_t* handle, void* devPtr ) # noqa
+        Function(
+            "cudaIpcGetMemHandle",
+            cudaError_t,
+            [ctypes.POINTER(cudaIpcMemHandle_t), ctypes.c_void_p],
+        ),
+        # ​cudaError_t cudaIpcOpenMemHandle ( void** devPtr, cudaIpcMemHandle_t handle, unsigned int  flags ) # noqa
+        Function(
+            "cudaIpcOpenMemHandle",
+            cudaError_t,
+            [ctypes.POINTER(ctypes.c_void_p), cudaIpcMemHandle_t, ctypes.c_uint],
+        ),
+    ]
+
+    # class attribute to store the mapping from the path to the library
+    # to avoid loading the same library multiple times
+    path_to_library_cache: Dict[str, Any] = {}
+
+    # class attribute to store the mapping from library path
+    #  to the corresponding dictionary
+    path_to_dict_mapping: Dict[str, Dict[str, Any]] = {}
+
+    def __init__(self, so_file: Optional[str] = None):
+        if so_file is None:
+            so_file = find_loaded_library("libcudart" if not _is_musa else "libmusart")
+            assert so_file is not None, "libcudart is not loaded in the current process"
+        if so_file not in CudaRTLibrary.path_to_library_cache:
+            lib = ctypes.CDLL(so_file)
+            CudaRTLibrary.path_to_library_cache[so_file] = lib
+        self.lib = CudaRTLibrary.path_to_library_cache[so_file]
+
+        if so_file not in CudaRTLibrary.path_to_dict_mapping:
+            _funcs = {}
+            for func in CudaRTLibrary.exported_functions:
+                f = getattr(self.lib, func.name)
+                f.restype = func.restype
+                f.argtypes = func.argtypes
+                _funcs[func.name] = f
+            CudaRTLibrary.path_to_dict_mapping[so_file] = _funcs
+        self.funcs = CudaRTLibrary.path_to_dict_mapping[so_file]
+
+    def CUDART_CHECK(self, result: cudaError_t) -> None:
+        if result != 0:
+            error_str = self.cudaGetErrorString(result)
+            raise RuntimeError(f"CUDART error: {error_str}")
+
+    def cudaGetErrorString(self, error: cudaError_t) -> str:
+        return self.funcs["cudaGetErrorString"](error).decode("utf-8")
+
+    def cudaSetDevice(self, device: int) -> None:
+        self.CUDART_CHECK(self.funcs["cudaSetDevice"](device))
+
+    def cudaDeviceSynchronize(self) -> None:
+        self.CUDART_CHECK(self.funcs["cudaDeviceSynchronize"]())
+
+    def cudaDeviceReset(self) -> None:
+        self.CUDART_CHECK(self.funcs["cudaDeviceReset"]())
+
+    def cudaMalloc(self, size: int) -> ctypes.c_void_p:
+        devPtr = ctypes.c_void_p()
+        self.CUDART_CHECK(self.funcs["cudaMalloc"](ctypes.byref(devPtr), size))
+        return devPtr
+
+    def cudaFree(self, devPtr: ctypes.c_void_p) -> None:
+        self.CUDART_CHECK(self.funcs["cudaFree"](devPtr))
+
+    def cudaMemset(self, devPtr: ctypes.c_void_p, value: int, count: int) -> None:
+        self.CUDART_CHECK(self.funcs["cudaMemset"](devPtr, value, count))
+
+    def cudaMemcpy(
+        self, dst: ctypes.c_void_p, src: ctypes.c_void_p, count: int
+    ) -> None:
+        cudaMemcpyDefault = 4
+        kind = cudaMemcpyDefault
+        self.CUDART_CHECK(self.funcs["cudaMemcpy"](dst, src, count, kind))
+
+    def cudaIpcGetMemHandle(self, devPtr: ctypes.c_void_p) -> cudaIpcMemHandle_t:
+        handle = cudaIpcMemHandle_t()
+        self.CUDART_CHECK(
+            self.funcs["cudaIpcGetMemHandle"](ctypes.byref(handle), devPtr)
+        )
+        return handle
+
+    def cudaIpcOpenMemHandle(self, handle: cudaIpcMemHandle_t) -> ctypes.c_void_p:
+        cudaIpcMemLazyEnablePeerAccess = 1
+        devPtr = ctypes.c_void_p()
+        self.CUDART_CHECK(
+            self.funcs["cudaIpcOpenMemHandle"](
+                ctypes.byref(devPtr), handle, cudaIpcMemLazyEnablePeerAccess
+            )
+        )
+        return devPtr

python/sglang/srt/distributed/device_communicators/custom_all_reduce.py

@@ -0,0 +1,421 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/distributed/device_communicators/custom_all_reduce.py
+
+import ctypes
+import logging
+from contextlib import contextmanager
+from functools import partial
+from typing import Any, List, Optional, Union
+
+import torch
+import torch.distributed as dist
+from torch.distributed import ProcessGroup
+
+import sglang.srt.distributed.device_communicators.custom_all_reduce_ops as ops
+from sglang.srt.distributed.device_communicators.cuda_wrapper import CudaRTLibrary
+from sglang.srt.distributed.device_communicators.custom_all_reduce_utils import (
+    can_use_custom_all_reduce_with_nvlink,
+    is_weak_contiguous,
+)
+from sglang.srt.environ import envs
+from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
+    is_in_tc_piecewise_cuda_graph,
+)
+from sglang.srt.utils import (
+    get_bool_env_var,
+    is_cuda,
+    is_hip,
+    is_musa,
+    log_info_on_rank0,
+)
+
+_is_cuda = is_cuda()
+_is_hip = is_hip()
+_is_musa = is_musa()
+
+logger = logging.getLogger(__name__)
+
+
+class CustomAllreduce:
+    _SUPPORTED_WORLD_SIZES = [2, 4, 6, 8]
+    _MAX_CAR_SIZE = 8192 * 1024
+    if _is_hip:
+        # crossover is at 16MB buffer size for ROCm
+        _MAX_CAR_SIZE = 2 * 8192 * 1024
+    if _is_musa:
+        # crossover is at 128MB buffer size for MUSA
+        _MAX_CAR_SIZE = 16 * 8196 * 1024
+
+    # max_size: max supported allreduce size
+    def __init__(
+        self,
+        group: ProcessGroup,
+        device: Union[int, str, torch.device],
+        max_size=_MAX_CAR_SIZE,
+    ) -> None:
+        """
+        Args:
+            group: the process group to work on. If None, it will use the
+                default process group.
+            device: the device to bind the CustomAllreduce to. If None,
+                it will be bind to f"cuda:{local_rank}".
+        It is the caller's responsibility to make sure each communicator
+        is bind to a unique device, and all communicators in this group
+        are in the same node.
+        """
+        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
+            # e.g. in a non-cuda environment
+            return
+
+        rank = dist.get_rank(group=group)
+        world_size = dist.get_world_size(group=group)
+
+        if isinstance(device, int):
+            device = torch.device(f"cuda:{device}")
+        elif isinstance(device, str):
+            device = torch.device(device)
+        # now `device` is a `torch.device` object
+        assert isinstance(device, torch.device)
+        self.device = device
+        full_nvlink = can_use_custom_all_reduce_with_nvlink(
+            group=group,
+            device=device,
+            supported_world_size=self._SUPPORTED_WORLD_SIZES,
+            cls_name="CustomAllreduce",
+        )
+        if full_nvlink is None:
+            return  # fail to get nvlink status
+
+        self.group = group
+        self.max_size = max_size
+        self.rank = rank
+        self.world_size = world_size
+        self.full_nvlink = full_nvlink
+
+        if not _is_hip:
+            # Buffers memory are owned by this Python class and passed to C++.
+            # Meta data composes of two parts: meta data for synchronization and a
+            # temporary buffer for storing intermediate allreduce results.
+            self.meta_ptrs = self.create_shared_buffer(
+                ops.meta_size() + max_size, group=group
+            )
+            # This is a pre-registered IPC buffer. In eager mode, input tensors
+            # are first copied into this buffer before allreduce is performed
+            self.buffer_ptrs = self.create_shared_buffer(max_size, group=group)
+            # This is a buffer for storing the tuples of pointers pointing to
+            # IPC buffers from all ranks. Each registered tuple has size of
+            # 8*world_size bytes where world_size is at most 8. Allocating 8MB
+            # is enough for 131072 such tuples. The largest model I've seen only
+            # needs less than 10000 of registered tuples.
+            self.rank_data = torch.empty(
+                max_size, dtype=torch.uint8, device=self.device
+            )
+            self._ptr = ops.init_custom_ar(
+                self.meta_ptrs, self.rank_data, rank, self.full_nvlink
+            )
+            ops.register_buffer(self._ptr, self.buffer_ptrs)
+        else:
+            # meta data buffers need to be "uncached" for signal on MI200
+            self.meta = ops.allocate_meta_buffer(ops.meta_size() + max_size)
+            self.buffer = torch.empty(max_size, dtype=torch.uint8, device=self.device)
+            handle = ops.get_meta_buffer_ipc_handle(self.meta)
+            shard_data = (
+                bytes(handle),  # ipc handle to base ptr
+                0,  # offset of base ptr
+            )
+            handles, offsets = self._gather_ipc_meta(shard_data)
+            self.rank_data = torch.empty(
+                max_size, dtype=torch.uint8, device=self.device
+            )
+            self._ptr = ops.init_custom_ar(
+                self.meta, self.rank_data, handles, offsets, rank, self.full_nvlink
+            )
+            self.register_buffer(self.buffer)
+
+        self.disabled = False
+        self.original_disabled = False  # Ensure original_disabled == disabled
+        self.tms_cudagraph = envs.SGLANG_MEMORY_SAVER_CUDA_GRAPH.get()
+
+    @staticmethod
+    def create_shared_buffer(
+        size_in_bytes: int, group: Optional[ProcessGroup] = None
+    ) -> List[int]:
+        """
+        Creates a shared buffer and returns a list of pointers
+        representing the buffer on all processes in the group.
+        """
+        lib = CudaRTLibrary()
+        pointer = lib.cudaMalloc(size_in_bytes)
+        if _is_musa:
+            lib.cudaMemset(pointer, 0, size_in_bytes)
+        handle = lib.cudaIpcGetMemHandle(pointer)
+        world_size = dist.get_world_size(group=group)
+        rank = dist.get_rank(group=group)
+        handles = [None] * world_size
+        dist.all_gather_object(handles, handle, group=group)
+
+        pointers: List[int] = []
+        for i, h in enumerate(handles):
+            if i == rank:
+                pointers.append(pointer.value)  # type: ignore
+            else:
+                pointers.append(lib.cudaIpcOpenMemHandle(h).value)  # type: ignore
+
+        return pointers
+
+    @staticmethod
+    def free_shared_buffer(
+        pointers: List[int], group: Optional[ProcessGroup] = None
+    ) -> None:
+        rank = dist.get_rank(group=group)
+        lib = CudaRTLibrary()
+        lib.cudaFree(ctypes.c_void_p(pointers[rank]))
+
+    @contextmanager
+    def capture(self):
+        """
+        The main responsibility of this context manager is the
+        `register_graph_buffers` call at the end of the context.
+        It records all the buffer addresses used in the CUDA graph.
+        """
+        try:
+            self._IS_CAPTURING = True
+            yield
+        finally:
+            self._IS_CAPTURING = False
+            if not self.disabled:
+                self.register_graph_buffers()
+
+    def _get_ipc_meta(self, inp: torch.Tensor):
+        # _share_cuda_() doesn't accept meta buffer not allocated from
+        # PyTorch cache allocator, use direct HIP call to get IPC handle
+        handle = ops.get_meta_buffer_ipc_handle(inp)
+        shard_data = (
+            bytes(handle),  # ipc handle to base ptr
+            0,  # offset of base ptr
+        )
+        return self._gather_ipc_meta(shard_data)
+
+    def _gather_ipc_meta(self, shard_data):
+        # Note: don't use `[[None]] * self.world_size` here
+        # because it will create a list of the same reference
+        all_data: List[Optional[Any]] = [[None] for i in range(self.world_size)]
+        all_data[self.rank][0] = shard_data
+
+        ranks = dist.get_process_group_ranks(group=self.group)
+        ranks.sort()
+        for i, rank in enumerate(ranks):
+            dist.broadcast_object_list(
+                all_data[i], src=rank, group=self.group, device="cpu"
+            )
+
+        # we cannot directly use `dist.all_gather_object` here
+        # because it is incompatible with `gloo` backend under inference mode.
+        # see https://github.com/pytorch/pytorch/issues/126032 for details.
+
+        handles = []
+        offsets = []
+        for i in range(len(all_data)):
+            handles.append(all_data[i][0][0])  # type: ignore
+            offsets.append(all_data[i][0][1])  # type: ignore
+        return handles, offsets
+
+    def register_buffer(self, inp: torch.Tensor):
+        handles, offsets = self._get_ipc_meta(inp)
+        ops.register_buffer(self._ptr, inp, handles, offsets)
+
+    def register_graph_buffers(self):
+        if _is_hip:
+            handle, offset = ops.get_graph_buffer_ipc_meta(self._ptr)
+            handles, offsets = self._gather_ipc_meta((bytes(handle), offset))
+            log_info_on_rank0(logger, f"Registering {len(offset)} cuda graph addresses")
+            ops.register_graph_buffers(self._ptr, handles, offsets)
+        else:
+            handle, offset = ops.get_graph_buffer_ipc_meta(self._ptr)
+            log_info_on_rank0(logger, f"Registering {len(offset)} cuda graph addresses")
+            # We cannot directly use `dist.all_gather_object` here
+            # because it is incompatible with `gloo` backend under inference mode.
+            # see https://github.com/pytorch/pytorch/issues/126032 for details.
+            all_data = [
+                [None, None] for _ in range(dist.get_world_size(group=self.group))
+            ]
+            all_data[self.rank] = [handle, offset]
+            ranks = sorted(dist.get_process_group_ranks(group=self.group))
+            for i, rank in enumerate(ranks):
+                dist.broadcast_object_list(
+                    all_data[i], src=rank, group=self.group, device="cpu"
+                )
+            # Unpack list of tuples to tuple of lists.
+            handles = [d[0] for d in all_data]  # type: ignore
+            offsets = [d[1] for d in all_data]  # type: ignore
+            ops.register_graph_buffers(self._ptr, handles, offsets)
+
+    def should_custom_ar(self, inp: torch.Tensor):
+        if self.disabled:
+            return False
+        inp_size = inp.numel() * inp.element_size()
+        # custom allreduce requires input byte size to be multiples of 16
+        if inp_size % 16 != 0:
+            return False
+        if not is_weak_contiguous(inp):
+            return False
+        # for 4 or more non NVLink-capable GPUs, custom allreduce provides
+        # little performance improvement over NCCL.
+        if not _is_hip:
+            if self.world_size == 2 or self.full_nvlink:
+                return inp_size <= self.max_size
+            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
+
+    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]:
+        """The main allreduce API that provides support for cuda graph."""
+        # When custom allreduce is disabled, this will be None.
+        if self.disabled or not self.should_custom_ar(input):
+            return None
+        if self._IS_CAPTURING:
+            if torch.cuda.is_current_stream_capturing():
+                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_tc_piecewise_cuda_graph():
+                    # Split op execution - do real all-reduce
+                    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:
+            return self._all_reduce_impl(input, registered=False)
+
+    def close(self):
+        if not self.disabled and self._ptr:
+            if ops is not None:
+                ops.dispose(self._ptr)
+            if _is_cuda:
+                self.free_shared_buffer(self.meta_ptrs)
+                self.free_shared_buffer(self.buffer_ptrs)
+            self._ptr = 0
+
+    def __del__(self):
+        self.close()
+
+
+def dispatch_custom_allreduce(
+    group: ProcessGroup,
+    device: torch.device,
+):
+    """Return the CustomAllreduce class to use (aiter on ROCm if enabled).
+
+    On AMD with 1-stage AR enabled, use sglang's CustomAllreduce.
+    Otherwise use AiterCustomAllreduce if available.
+
+    On CUDA, the JIT-compiled v2 implementation is used by default.
+    Set SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2=0 to fall back to the legacy CustomAllreduce.
+    Note: ServerArgs._handle_environment_variables forces this env to "0" when
+    nnodes > 1 since custom AR is intra-node only.
+    """
+    if _is_cuda and envs.SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2.get():
+        from .custom_all_reduce_v2 import (
+            CustomAllReduceV2,
+            can_use_custom_all_reduce_v2,
+        )
+
+        if can_use_custom_all_reduce_v2(group=group, device=device):
+            logger.debug("[AR] Using CustomAllReduceV2 (JIT-compiled)")
+            return CustomAllReduceV2
+
+    if _is_cuda or _is_musa:
+        return CustomAllreduce
+
+    assert _is_hip
+
+    if envs.SGLANG_USE_1STAGE_ALLREDUCE.is_set():
+        if envs.SGLANG_USE_1STAGE_ALLREDUCE.get():
+            logger.debug(
+                "[AR] All-reduce: 1-stage kernel (SGLANG_USE_1STAGE_ALLREDUCE=1)"
+            )
+        else:
+            logger.debug("[AR] All-reduce: default (SGLANG_USE_1STAGE_ALLREDUCE=0)")
+    elif envs.SGLANG_ENABLE_DETERMINISTIC_INFERENCE.get():
+        logger.debug(
+            "[AR] All-reduce: 1-stage kernel (deterministic inference enabled)"
+        )
+    else:
+        logger.debug("[AR] All-reduce: default")
+
+    # On AMD with 1-stage AR, use sglang's CustomAllreduce
+    # (AiterCustomAllreduce doesn't have deterministic_all_reduce method)
+    if _use_amd_deterministic_impl():
+        return CustomAllreduce
+
+    if 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("[AR] Using AiterCustomAllreduce (AMD default)")
+            tms_cudagraph = envs.SGLANG_MEMORY_SAVER_CUDA_GRAPH.get()
+            return partial(
+                AiterCustomAllreduce,
+                enable_register_for_capturing=not tms_cudagraph,
+            )
+        except ImportError as e:
+            logger.warning(
+                "[AR] Aiter custom all-reduce not available; "
+                "falling back to sglang CustomAllreduce. Details: %s",
+                e,
+            )
+            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/device_communicators/custom_all_reduce_ops.py

@@ -0,0 +1,166 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/_custom_ops.py
+import logging
+from typing import List, Optional, Tuple
+
+import torch
+
+from sglang.srt.utils import is_cuda, is_hip, is_musa
+
+logger = logging.getLogger(__name__)
+
+_is_cuda = is_cuda()
+_is_hip = is_hip()
+_is_musa = is_musa()
+
+IS_CUSTOM_AR_AVAILABLE = _is_cuda or _is_hip or _is_musa
+IS_QUICK_AR_AVAILABLE = _is_hip
+
+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
+
+# region IS_CUSTOM_AR_AVAILABLE
+
+if not IS_CUSTOM_AR_AVAILABLE:
+    pass
+
+elif _is_cuda or _is_musa:
+    # CUDA custom allreduce
+
+    def init_custom_ar(
+        ipc_tensors: List[torch.Tensor],
+        rank_data: torch.Tensor,
+        rank: int,
+        full_nvlink: bool,
+    ) -> int:
+        return _custom_ar.init_custom_ar(ipc_tensors, rank_data, rank, full_nvlink)
+
+    def all_reduce(
+        fa: int,
+        inp: torch.Tensor,
+        out: torch.Tensor,
+        reg_buffer: int,
+        reg_buffer_sz_bytes: int,
+    ) -> None:
+        _custom_ar.all_reduce(fa, inp, out, reg_buffer, reg_buffer_sz_bytes)
+
+    def dispose(fa: int) -> None:
+        _custom_ar.dispose(fa)
+
+    def meta_size() -> int:
+        return _custom_ar.meta_size()
+
+    def register_buffer(fa: int, ipc_tensors: List[int]) -> None:
+        return _custom_ar.register_buffer(fa, ipc_tensors)
+
+    def get_graph_buffer_ipc_meta(fa: int) -> Tuple[List[int], List[int]]:
+        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_ar.register_graph_buffers(fa, handles, offsets)
+
+elif _is_hip:
+    # ROCM custom allreduce
+
+    def init_custom_ar(
+        meta: torch.Tensor,
+        rank_data: torch.Tensor,
+        handles: List[str],
+        offsets: List[int],
+        rank: int,
+        full_nvlink: bool,
+    ) -> int:
+        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:
+        _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:
+        _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)
+
+    def meta_size() -> int:
+        return _custom_ar.meta_size()
+
+    def register_buffer(
+        fa: int, t: torch.Tensor, handles: List[str], offsets: List[int]
+    ) -> None:
+        return _custom_ar.register_buffer(fa, t, handles, offsets)
+
+    def get_graph_buffer_ipc_meta(fa: int) -> Tuple[torch.Tensor, List[int]]:
+        return _custom_ar.get_graph_buffer_ipc_meta(fa)
+
+    def register_graph_buffers(
+        fa: int, handles: List[str], offsets: List[List[int]]
+    ) -> None:
+        _custom_ar.register_graph_buffers(fa, handles, offsets)
+
+    def allocate_meta_buffer(size: int) -> torch.Tensor:
+        return _custom_ar.allocate_meta_buffer(size)
+
+    def get_meta_buffer_ipc_handle(inp: torch.Tensor) -> torch.Tensor:
+        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 _custom_ar.init_custom_qr(world_size, rank, qr_max_size)
+
+    def qr_get_handle(fa: int) -> torch.Tensor:
+        return _custom_ar.qr_get_handle(fa)
+
+    def qr_open_handles(fa: int, handles: list[torch.Tensor]) -> None:
+        _custom_ar.qr_open_handles(fa, handles)
+
+    def qr_all_reduce(
+        fa: int,
+        inp: torch.Tensor,
+        out: torch.Tensor,
+        quant_level: int,
+        cast_bf2half: bool,
+    ) -> None:
+        _custom_ar.qr_all_reduce(fa, inp, out, quant_level, cast_bf2half)
+
+    def qr_destroy(fa: int) -> None:
+        _custom_ar.qr_destroy(fa)
+
+    def qr_max_size() -> int:
+        return _custom_ar.qr_max_size()
+
+
+# endregion

python/sglang/srt/distributed/device_communicators/custom_all_reduce_utils.py

@@ -0,0 +1,492 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/distributed/device_communicators/custom_all_reduce_utils.py
+
+import ctypes
+import json
+import logging
+import os
+import pickle
+import subprocess
+import sys
+import tempfile
+from functools import wraps
+from itertools import product
+from typing import Callable, Dict, List, Optional, Sequence, TypeVar
+
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+from typing_extensions import ParamSpec
+
+from sglang.srt.distributed.device_communicators.cuda_wrapper import CudaRTLibrary
+from sglang.srt.distributed.parallel_state import in_the_same_node_as
+from sglang.srt.environ import envs as sglang_envs
+from sglang.srt.utils import is_cuda, is_hip, is_musa
+
+logger = logging.getLogger(__name__)
+
+_is_cuda = is_cuda()
+_is_hip = is_hip()
+_is_musa = is_musa()
+
+if _is_cuda:
+    try:
+        import pynvml
+    except ImportError as e:
+        logger.warning("Failed to import pynvml with %r", e)
+
+if _is_musa:
+    try:
+        import pymtml as pynvml
+    except ImportError as e:
+        logger.warning("Failed to import pymtml with %r", e)
+
+if _is_hip:
+    try:
+        from amdsmi import (
+            AmdSmiException,
+            amdsmi_get_processor_handles,
+            amdsmi_init,
+            amdsmi_shut_down,
+            amdsmi_topo_get_link_type,
+        )
+    except ImportError as e:
+        logger.warning("Failed to import amdsmi with %r", e)
+
+_P = ParamSpec("_P")
+_R = TypeVar("_R")
+
+
+def update_environment_variables(envs: Dict[str, str]):
+    for k, v in envs.items():
+        if k in os.environ and os.environ[k] != v:
+            logger.warning(
+                "Overwriting environment variable %s " "from '%s' to '%s'",
+                k,
+                os.environ[k],
+                v,
+            )
+        os.environ[k] = v
+
+
+def producer(
+    batch_src: Sequence[int],
+    producer_queue,
+    consumer_queue,
+    result_queue,
+    cuda_visible_devices: Optional[str] = None,
+):
+    if cuda_visible_devices is not None:
+        update_environment_variables({"CUDA_VISIBLE_DEVICES": cuda_visible_devices})
+
+    lib = CudaRTLibrary()
+    for i in batch_src:
+        lib.cudaSetDevice(i)
+        pointer = lib.cudaMalloc(1024)
+        lib.cudaMemset(pointer, 1, 1024)
+        lib.cudaDeviceSynchronize()
+        handle = lib.cudaIpcGetMemHandle(pointer)
+        producer_queue.put(handle)
+        open_success = consumer_queue.get()
+        if open_success:
+            # use two queues to simulate barrier
+            producer_queue.put(0)
+            consumer_queue.get()
+            # check if the memory is modified
+            host_data = (ctypes.c_char * 1024)()
+            lib.cudaMemcpy(host_data, pointer, 1024)  # type: ignore
+            for i in range(1024):
+                if ord(host_data[i]) != 2:
+                    open_success = False
+                    break
+        result_queue.put(open_success)
+        lib.cudaDeviceReset()
+
+
+def consumer(
+    batch_tgt: Sequence[int],
+    producer_queue,
+    consumer_queue,
+    result_queue,
+    cuda_visible_devices: Optional[str] = None,
+):
+    if cuda_visible_devices is not None:
+        update_environment_variables({"CUDA_VISIBLE_DEVICES": cuda_visible_devices})
+
+    lib = CudaRTLibrary()
+    for j in batch_tgt:
+        lib.cudaSetDevice(j)
+        handle = producer_queue.get()
+        open_success = False
+        try:
+            pointer = lib.cudaIpcOpenMemHandle(handle)  # type: ignore
+            open_success = True
+        except RuntimeError:
+            # cannot error out here, because the producer process
+            # is still waiting for the response.
+            pass
+        consumer_queue.put(open_success)
+        if open_success:
+            # modify the memory
+            lib.cudaMemset(pointer, 2, 1024)
+            lib.cudaDeviceSynchronize()
+            # use two queues to simulate barrier
+            producer_queue.get()
+            consumer_queue.put(0)
+            # check if the memory is modified
+            host_data = (ctypes.c_char * 1024)()
+            lib.cudaMemcpy(host_data, pointer, 1024)  # type: ignore
+            for i in range(1024):
+                if ord(host_data[i]) != 2:
+                    open_success = False
+                    break
+        result_queue.put(open_success)
+        lib.cudaDeviceReset()
+
+
+def can_actually_p2p(
+    batch_src: Sequence[int],
+    batch_tgt: Sequence[int],
+) -> Sequence[bool]:
+    """
+    Usually, checking if P2P access is enabled can be done by
+    `torch.cuda.can_device_access_peer(src, tgt)`. However, sometimes
+    the driver might be broken, and `torch.cuda.can_device_access_peer(src, tgt)`
+    returns `True` even if P2P access is not actually possible.
+    See https://github.com/vllm-project/vllm/issues/2728 and
+    https://forums.developer.nvidia.com/t/direct-gpu-gpu-communication-does-not-seem-to-work-properly/283264/10
+    Therefore, we have to perform a real P2P access to check if it is actually
+    possible.
+
+    Note on p2p and cuda IPC:
+    Usually, one process uses one GPU:
+    GPU src --> cuda context src --> tensor src --> process src
+
+    We need to combine p2p and cuda IPC, so that:
+    GPU src --> cuda context src --> tensor src --> process src
+                                      |shared|
+    GPU tgt --> cuda context tgt --> tensor tgt --> process tgt
+    That is to say, process src creates a tensor in GPU src, passes IPC handle to
+    process tgt, and process tgt accesses the tensor in GPU tgt. Any operation on the
+    tensor in process tgt will be reflected in the tensor in process src, because
+    they are the same memory segment.
+    It is important to note that process tgt accesses the tensor in GPU tgt, not
+    GPU src. That's why we need p2p access.
+
+    The most time-consuming part is the process creation. To avoid creating
+    processes for every pair of GPUs, we use batched testing. We create two
+    processes for testing all pairs of GPUs in batch. The trick is to reset
+    the device after each test (which is not available in PyTorch).
+    """  # noqa
+    cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES", None)
+    # pass the CUDA_VISIBLE_DEVICES to the child process
+    # to make sure they see the same set of GPUs
+
+    # make sure the processes are spawned
+    smp = mp.get_context("spawn")
+    producer_queue = smp.Queue()
+    consumer_queue = smp.Queue()
+    result_queue = smp.Queue()
+    p_src = smp.Process(
+        target=producer,
+        args=(
+            batch_src,
+            producer_queue,
+            consumer_queue,
+            result_queue,
+            cuda_visible_devices,
+        ),
+    )
+    p_tgt = smp.Process(
+        target=consumer,
+        args=(
+            batch_tgt,
+            producer_queue,
+            consumer_queue,
+            result_queue,
+            cuda_visible_devices,
+        ),
+    )
+    p_src.start()
+    p_tgt.start()
+    p_src.join()
+    p_tgt.join()
+    assert p_src.exitcode == 0 and p_tgt.exitcode == 0
+    result: List[bool] = []
+    for src, tgt in zip(batch_src, batch_tgt):
+        a = result_queue.get()
+        b = result_queue.get()
+        if a != b:
+            logger.warning(
+                "Two processes do not agree on the P2P access"
+                " status on %d -> %d, treat as disabled.",
+                src,
+                tgt,
+            )
+            result.append(False)
+        else:
+            result.append(a)
+    return result
+
+
+# why do we need this cache?
+# we are testing peer-to-peer (p2p) access between GPUs,across processes.
+# if we test it every time, it will be very slow, because we need to create
+#  N * N * 2 processes, where N is the world size. This is very slow.
+# to reduce the time, we use a cache file to store the p2p access status.
+# the cache file is generated by the master process if it does not exist.
+# then all the processes can read the cache file to check the p2p access status.
+# Note that the cache file is suffixed by the CUDA_VISIBLE_DEVICES, so that we
+#  can have different cache files for different CUDA_VISIBLE_DEVICES settings,
+#  e.g. used by different vllm engines. The device id in the cache file is a
+#  **local** device id, i.e. from 0 to num_dev-1, where num_dev is the number
+#  of visible devices in the vllm engine.
+_gpu_p2p_access_cache: Optional[Dict[str, bool]] = None
+
+
+def gpu_p2p_access_check(src: int, tgt: int) -> bool:
+    """Check if GPU src can access GPU tgt."""
+
+    # if the cache variable is already calculated,
+    # read from the cache instead of checking it again
+    global _gpu_p2p_access_cache
+    if _gpu_p2p_access_cache is not None:
+        return _gpu_p2p_access_cache[f"{src}->{tgt}"]
+
+    is_distributed = dist.is_initialized()
+
+    num_dev = torch.cuda.device_count()
+    cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES", None)
+    if cuda_visible_devices is None:
+        cuda_visible_devices = ",".join(str(i) for i in range(num_dev))
+
+    # VLLM_CACHE_ROOT -> SGLANG_CACHE_ROOT
+    # "~/.cache/vllm" -> envs.SGLANG_CACHE_DIR
+    SGLANG_CACHE_ROOT = os.path.expanduser(sglang_envs.SGLANG_CACHE_DIR.get())
+    path = os.path.join(
+        SGLANG_CACHE_ROOT, f"gpu_p2p_access_cache_for_{cuda_visible_devices}.json"
+    )
+    cache_dir = os.path.dirname(path)
+    try:
+        os.makedirs(cache_dir, exist_ok=True)
+    except (FileExistsError, NotADirectoryError):
+        if not os.path.isdir(cache_dir):
+            # Path exists as a file (stale cache/lock). Remove and retry.
+            try:
+                os.remove(cache_dir)
+            except OSError:
+                pass
+            os.makedirs(cache_dir, exist_ok=True)
+    from sglang.srt.distributed.parallel_state import get_world_group
+
+    if (not is_distributed or get_world_group().local_rank == 0) and (
+        not os.path.exists(path)
+    ):
+        # only the local master process (with local_rank == 0) can
+        #  enter this block to calculate the cache
+        logger.info("generating GPU P2P access cache in %s", path)
+        cache: Dict[str, bool] = {}
+        ids = list(range(num_dev))
+        # batch of all pairs of GPUs
+        batch_src, batch_tgt = zip(*list(product(ids, ids)))
+        # NOTE: we use `subprocess` rather than `multiprocessing` here
+        # because the caller might not have `if __name__ == "__main__":`,
+        # in that case we cannot use spawn method in multiprocessing.
+        # However, `can_actually_p2p` requires spawn method.
+        # The fix is, we use `subprocess` to call the function,
+        # where we have `if __name__ == "__main__":` in this file.
+
+        # use a temporary file to store the result
+        # we don't use the output of the subprocess directly,
+        # because the subprocess might produce logging output
+        with tempfile.NamedTemporaryFile() as output_file:
+            input_bytes = pickle.dumps((batch_src, batch_tgt, output_file.name))
+            returned = subprocess.run(
+                [sys.executable, __file__], input=input_bytes, capture_output=True
+            )
+            # check if the subprocess is successful
+            try:
+                returned.check_returncode()
+            except Exception as e:
+                # wrap raised exception to provide more information
+                raise RuntimeError(
+                    f"Error happened when batch testing "
+                    f"peer-to-peer access from {batch_src} to {batch_tgt}:\n"
+                    f"{returned.stderr.decode()}"
+                ) from e
+            with open(output_file.name, "rb") as f:
+                result = pickle.load(f)
+        for _i, _j, r in zip(batch_src, batch_tgt, result):
+            cache[f"{_i}->{_j}"] = r
+        with open(path, "w") as f:
+            json.dump(cache, f, indent=4)
+    if is_distributed:
+        get_world_group().barrier()
+    logger.info("reading GPU P2P access cache from %s", path)
+    with open(path) as f:
+        cache = json.load(f)
+    _gpu_p2p_access_cache = cache
+    return _gpu_p2p_access_cache[f"{src}->{tgt}"]
+
+
+def with_nvml_context(fn: Callable[_P, _R]) -> Callable[_P, _R]:
+    @wraps(fn)
+    def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
+        if _is_hip:
+            try:
+                amdsmi_init()
+                return fn(*args, **kwargs)
+            finally:
+                amdsmi_shut_down()
+        else:
+            pynvml.nvmlInit()
+            try:
+                return fn(*args, **kwargs)
+            finally:
+                pynvml.nvmlShutdown()
+
+    return wrapper
+
+
+@with_nvml_context
+def is_full_nvlink(physical_device_ids: List[int], world_size: int) -> bool:
+    if _is_hip:
+        """
+        query if the set of gpus are fully connected by xgmi (1 hop)
+        """
+        handles = [amdsmi_get_processor_handles()[i] for i in physical_device_ids]
+        for i, handle in enumerate(handles):
+            for j, peer_handle in enumerate(handles):
+                if i < j:
+                    try:
+                        link_type = amdsmi_topo_get_link_type(handle, peer_handle)
+                        # type is 2 for XGMI
+                        if link_type["hops"] != 1 or link_type["type"] != 2:
+                            return False
+                    except AmdSmiException as error:
+                        logger.error("AMD 1 hop XGMI detection failed.", exc_info=error)
+                        return False
+        return True
+    else:
+        """
+        query if the set of gpus are fully connected by nvlink (1 hop)
+        """
+        handles = [pynvml.nvmlDeviceGetHandleByIndex(i) for i in physical_device_ids]
+        for i, handle in enumerate(handles):
+            for j, peer_handle in enumerate(handles):
+                if i < j:
+                    try:
+                        p2p_status = pynvml.nvmlDeviceGetP2PStatus(
+                            handle, peer_handle, pynvml.NVML_P2P_CAPS_INDEX_NVLINK
+                        )
+                        if p2p_status != pynvml.NVML_P2P_STATUS_OK:
+                            return False
+                    except pynvml.NVMLError:
+                        logger.exception(
+                            "NVLink detection failed. This is normal if your"
+                            " machine has no NVLink equipped."
+                        )
+                        return False
+        return True
+
+
+def is_weak_contiguous(inp: torch.Tensor):
+    return inp.is_contiguous() or (
+        inp.storage().nbytes() - inp.storage_offset() * inp.element_size()
+        == inp.numel() * inp.element_size()
+    )
+
+
+def can_p2p(rank: int, world_size: int) -> bool:
+    # SGLANG_SKIP_P2P_CHECK can be set to False in sglang
+    SGLANG_SKIP_P2P_CHECK = os.getenv("SGLANG_SKIP_P2P_CHECK", "0") == "1"
+    for i in range(world_size):
+        if i == rank:
+            continue
+        if SGLANG_SKIP_P2P_CHECK:
+            logger.info("Skipping P2P check and trusting the driver's P2P report.")
+            return torch.cuda.can_device_access_peer(rank, i)
+        if not gpu_p2p_access_check(rank, i):
+            return False
+    return True
+
+
+def can_use_custom_all_reduce_with_nvlink(
+    group: torch.distributed.ProcessGroup,
+    device: torch.device,
+    supported_world_size: List[int],
+    cls_name: str,
+) -> Optional[bool]:  # None if fail; otherwise return whether NVLink is available
+    assert (
+        dist.get_backend(group) != dist.Backend.NCCL
+    ), f"{cls_name} should be attached to a non-NCCL group."
+
+    rank = dist.get_rank(group=group)
+    world_size = dist.get_world_size(group=group)
+
+    # No need to initialize custom allreduce for single GPU case.
+    if world_size == 1:
+        return
+
+    # No need to initialize custom allreduce for multi-node case.
+    if not all(in_the_same_node_as(group, source_rank=0)):
+        logger.warning(
+            f"{cls_name} is disabled because this process group" " spans across nodes."
+        )
+        return
+
+    # For not supported world size, we disable custom allreduce.
+    if world_size not in supported_world_size:
+        logger.warning(
+            f"{cls_name} is disabled due to an unsupported world"
+            f" size: {world_size}. Supported world sizes: {supported_world_size}. "
+            "To silence this warning, specify disable_custom_all_reduce=True explicitly.",
+        )
+        return
+
+    cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES", None)
+    if cuda_visible_devices:
+        device_ids = list(map(int, cuda_visible_devices.split(",")))
+    else:
+        device_ids = list(range(torch.cuda.device_count()))
+    physical_device_id = device_ids[device.index]
+    tensor = torch.tensor([physical_device_id], dtype=torch.int, device="cpu")
+    gather_list = [
+        torch.tensor([0], dtype=torch.int, device="cpu") for _ in range(world_size)
+    ]
+    dist.all_gather(gather_list, tensor, group=group)
+    physical_device_ids = [int(t) for t in gather_list]
+    full_nvlink = is_full_nvlink(physical_device_ids, world_size)
+
+    # test nvlink first, this will filter out most of the cases
+    # where custom allreduce is not supported
+    # this checks hardware and driver support for NVLink
+    if world_size > 2 and not full_nvlink:
+        logger.warning(
+            f"{cls_name} is disabled because it's not supported on"
+            " more than two PCIe-only GPUs. To silence this warning, "
+            "specify disable_custom_all_reduce=True explicitly."
+        )
+        return
+
+    # test P2P capability, this checks software/cudaruntime support
+    # this is expensive to compute at the first time
+    # then we cache the result
+    # On AMD GPU, p2p is always enabled between XGMI connected GPUs
+    if not _is_hip and not can_p2p(rank, world_size):
+        logger.warning(
+            f"{cls_name} is disabled because your platform lacks "
+            "GPU P2P capability or P2P test failed. To silence this "
+            "warning, specify disable_custom_all_reduce=True explicitly."
+        )
+        return
+
+    return full_nvlink
+
+
+if __name__ == "__main__":
+    batch_src, batch_tgt, output_file = pickle.loads(sys.stdin.buffer.read())
+    result = can_actually_p2p(batch_src, batch_tgt)
+    with open(output_file, "wb") as f:
+        f.write(pickle.dumps(result))

python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py

@@ -0,0 +1,240 @@
+import logging
+from contextlib import contextmanager
+from dataclasses import dataclass, replace
+from typing import Dict, List, Optional, TypeVar
+
+import torch
+import torch.distributed as dist
+from torch.distributed import ProcessGroup
+
+from sglang.jit_kernel.all_reduce import AllReduceAlgo, get_custom_all_reduce_cls
+from sglang.srt.distributed.device_communicators.custom_all_reduce_utils import (
+    can_use_custom_all_reduce_with_nvlink,
+    is_weak_contiguous,
+)
+from sglang.srt.distributed.device_communicators.vmm_utils import (
+    VmmGraphInputManager,
+    is_vmm_pointer,
+)
+from sglang.srt.environ import envs
+from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
+    is_in_tc_piecewise_cuda_graph,
+)
+from sglang.srt.utils import is_sm100_supported
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+INF = 1 << 60
+
+
+@dataclass(frozen=True)
+class ModeConfig:
+    one_shot_push_threshold: int  # below this, use one-shot push
+    one_shot_pull_threshold: int  # below this, use one-shot pull
+
+
+class CustomAllReduceV2:
+    def __init__(
+        self,
+        group: ProcessGroup,
+        device: torch.device,
+        max_pull_size: Optional[int] = None,
+        max_push_size: Optional[int] = None,
+        max_pull_blocks: Optional[int] = None,
+        max_push_blocks: Optional[int] = None,
+    ) -> None:
+        _maybe_init_config()
+        self.disabled = True
+        if not can_use_custom_all_reduce_v2(group=group, device=device):
+            return
+
+        self.group = group
+        self.rank = dist.get_rank(group=self.group)
+        self.world_size = dist.get_world_size(group=self.group)
+        if max_pull_size is None:  # default to 16MB
+            max_pull_size = 16 * 1024 * 1024
+        if max_push_size is None:  # default to recommended size
+            config = THRESHOLD_2_SHOT_MAP[self.world_size]
+            max_push_size = config.one_shot_push_threshold
+        self.max_pull_size = max_pull_size
+        self.max_push_size = max_push_size
+        self.max_size = max(max_pull_size, max_push_size)
+        self.override_shot(None)  # set default config based on world size
+        self.override_algo: Optional[AllReduceAlgo] = None
+        self.tms_cudagraph = envs.SGLANG_MEMORY_SAVER_CUDA_GRAPH.get()
+        self.obj = get_custom_all_reduce_cls()(
+            rank=self.rank,
+            world_size=self.world_size,
+            pull_buffer_bytes=self.max_pull_size,
+            push_buffer_bytes=self.max_push_size,
+            graph_input_count=131072,
+            max_pull_blocks=max_pull_blocks,
+            max_push_blocks=max_push_blocks,
+        )
+        self._vmm_graph_input_manager = VmmGraphInputManager(
+            obj=self.obj,
+            group=self.group,
+            rank=self.rank,
+            world_size=self.world_size,
+        )
+        self._post_init_obj()
+        self.disabled = False
+
+    def override_shot(self, shot: int | None):
+        if shot is None:
+            config = THRESHOLD_2_SHOT_MAP[self.world_size]
+        else:
+            assert shot in (1, 2)
+            threshold = INF if shot == 1 else 0
+            config = replace(self.config, one_shot_pull_threshold=threshold)
+        # need to clip the config thresholds to max sizes to avoid invalid config
+        push_threshold = min(config.one_shot_push_threshold, self.max_push_size)
+        pull_threshold = min(config.one_shot_pull_threshold, self.max_pull_size)
+        self.config: ModeConfig = replace(
+            config,
+            one_shot_push_threshold=push_threshold,
+            one_shot_pull_threshold=pull_threshold,
+        )
+
+    @contextmanager
+    def capture(self):
+        if self.disabled:
+            yield
+            return
+        try:
+            self.obj.set_cuda_graph_capture(not self.tms_cudagraph)
+            yield
+        finally:
+            self.obj.set_cuda_graph_capture(False)
+        assert (
+            not torch.cuda.is_current_stream_capturing()
+        ), "Cannot register graph inputs while capturing CUDA graph"
+        raw_ptrs = self.obj.get_graph_capture_ptrs()
+        if raw_ptrs and is_vmm_pointer(raw_ptrs[0]):
+            self._vmm_graph_input_manager.register_graph_inputs()
+        else:
+            self._register_graph_inputs_ipc()
+
+    def _register_graph_inputs_ipc(self):
+        """Register graph capture inputs via cudaIpcGetMemHandle.
+
+        This is the fast path for cudaMalloc-backed allocations. Fails
+        on VMM pointers (expandable_segments).
+        """
+        pairs = self.obj.share_graph_inputs()
+        handles = [handle for _, handle in pairs]
+        offsets = [offset for offset, _ in pairs]
+        handles_all = self._share_list(handles)
+        offsets_all = self._share_list(offsets)
+        result = [list(zip(o, h)) for o, h in zip(offsets_all, handles_all)]
+        self.obj.register_inputs(result)
+
+    def should_custom_ar(self, inp: torch.Tensor) -> bool:
+        """Check if the input tensor is suitable for custom all-reduce."""
+        if self.disabled:
+            return False
+        inp_size = inp.numel() * inp.element_size()
+        # custom allreduce requires input byte size to be multiples of 16
+        if inp_size % 16 != 0:
+            return False
+        if not is_weak_contiguous(inp):
+            return False
+        return inp_size <= self.max_size
+
+    def custom_all_reduce(self, input: torch.Tensor) -> torch.Tensor:
+        if is_in_tc_piecewise_cuda_graph():  # disable inplace optimization
+            try:
+                self.obj.set_cuda_graph_capture(False)
+                return self._all_reduce(input)
+            finally:
+                self.obj.set_cuda_graph_capture(not self.tms_cudagraph)
+        return self._all_reduce(input)
+
+    def close(self):
+        if not self.disabled and hasattr(self, "obj"):
+            self.obj.free(self.group)
+        if hasattr(self, "_vmm_graph_input_manager"):
+            self._vmm_graph_input_manager.close()
+
+    def _all_reduce(self, input: torch.Tensor) -> torch.Tensor:
+        """Perform the actual all-reduce via JIT kernel."""
+        algo = self._determine_algo(input)
+        return torch.from_dlpack(self.obj.all_reduce(input, algo))
+
+    def _determine_algo(self, input: torch.Tensor) -> AllReduceAlgo:
+        if self.override_algo is not None:
+            return self.override_algo
+        input_bytes = input.numel() * input.element_size()
+        if input_bytes <= self.config.one_shot_push_threshold:
+            return AllReduceAlgo.ONE_SHOT_PUSH
+        if input_bytes <= self.config.one_shot_pull_threshold:
+            return AllReduceAlgo.ONE_SHOT_PULL
+        else:
+            return AllReduceAlgo.TWO_SHOT_PULL
+
+    def _post_init_obj(self):
+        handles = [self.obj.share_storage()]
+        result = self._share_list(handles)
+        assert all(len(r) == 1 for r in result)
+        result = [h[0] for h in result]
+        self.obj.post_init(result)
+
+    def _share_list(self, input: List[T]) -> List[List[T]]:
+        input_tensor = torch.tensor(input, dtype=torch.int64, device="cpu")
+        gather_list = [torch.empty_like(input_tensor) for _ in range(self.world_size)]
+        dist.all_gather(gather_list, input_tensor, group=self.group)
+        return [g.tolist() for g in gather_list]
+
+    def __del__(self):
+        self.close()
+
+
+def _maybe_init_config():
+    global THRESHOLD_2_SHOT_MAP
+    if THRESHOLD_2_SHOT_MAP:
+        return
+    KB, MB = 1024, 1024 * 1024
+
+    if is_sm100_supported():
+        # NOTE: This result is based on benchmarks on B200 GPUs
+        THRESHOLD_2_SHOT_MAP = {
+            2: ModeConfig(4 * MB, INF),
+            3: ModeConfig(4 * MB, 4 * MB),
+            4: ModeConfig(2 * MB, 2 * MB),
+            5: ModeConfig(2 * MB, 2 * MB),
+            6: ModeConfig(1 * MB, 1 * MB),
+            7: ModeConfig(896 * KB, 896 * KB),
+            8: ModeConfig(720 * KB, 720 * KB),
+        }
+    else:
+        # NOTE: This result is based on benchmarks on H200 GPUs
+        THRESHOLD_2_SHOT_MAP = {
+            2: ModeConfig(2 * MB, INF),
+            3: ModeConfig(512 * KB, 512 * KB),
+            4: ModeConfig(384 * KB, 256 * KB),
+            5: ModeConfig(256 * KB, 256 * KB),
+            6: ModeConfig(192 * KB, 192 * KB),
+            7: ModeConfig(192 * KB, 192 * KB),
+            8: ModeConfig(160 * KB, 160 * KB),
+        }
+    # TODO: tune on more GPUs, e.g A100
+
+
+def can_use_custom_all_reduce_v2(
+    group: ProcessGroup,
+    device: torch.device,
+) -> bool:
+    # call _maybe_init_config() to ensure THRESHOLD_2_SHOT_MAP is initialized, since can_use_custom_all_reduce_v2 can be called before CustomAllReduceV2 is initialized
+    _maybe_init_config()
+    full_nvlink = can_use_custom_all_reduce_with_nvlink(
+        group=group,
+        device=device,
+        supported_world_size=list(THRESHOLD_2_SHOT_MAP.keys()),
+        cls_name="CustomAllReduceV2",
+    )
+    return full_nvlink is True
+
+
+THRESHOLD_2_SHOT_MAP: Dict[int, ModeConfig] = {}

python/sglang/srt/distributed/device_communicators/hpu_communicator.py

@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/distributed/device_communicators/hpu_communicator.py
+
+import torch
+import torch.distributed as dist
+from torch.distributed import ProcessGroup
+
+from sglang.srt.utils import is_hpu
+
+if is_hpu():
+    import habana_frameworks.torch as htorch  # noqa: F401
+
+
+class HpuCommunicator:
+
+    def __init__(self, group: ProcessGroup):
+        if not is_hpu():
+            self.disabled = True
+            return
+        self.disabled = False
+        self.group = group
+        self.world_size = dist.get_world_size(self.group)
+
+    def all_reduce(self, x: torch.Tensor) -> torch.Tensor:
+        # FIXME(kzawora): this is a workaround for a bug in Habana PT bridge
+        # occurring when PT_HPU_ENABLE_LAZY_COLLECTIVES=true env var is used
+        # (which is required for tensor parallel HPUGraph inference)
+        htorch.core.mark_step()
+        dist.all_reduce(x, group=self.group)
+        return x
+
+    def all_gather(self, x: torch.Tensor, dim: int = -1) -> torch.Tensor:
+        world_size = self.world_size
+        if dim < 0:
+            # Convert negative dim to positive.
+            dim += x.dim()
+        input_size = x.size()
+        # Allocate output tensor.
+        output_tensor = torch.empty(
+            (world_size,) + input_size, dtype=x.dtype, device=x.device
+        )
+        # All-gather.
+        htorch.core.mark_step()
+        dist.all_gather_into_tensor(output_tensor, x, group=self.group)
+        # Reshape
+        output_tensor = output_tensor.movedim(0, dim)
+        output_tensor = output_tensor.reshape(
+            input_size[:dim] + (world_size * input_size[dim],) + input_size[dim + 1 :]
+        )
+        return output_tensor

python/sglang/srt/distributed/device_communicators/mooncake_transfer_engine.py

@@ -0,0 +1,297 @@
+import json
+import logging
+import os
+from typing import Dict, List, Optional, Union
+
+from sglang.srt.environ import envs
+from sglang.srt.utils.network import NetworkAddress, get_free_port
+
+logger = logging.getLogger(__name__)
+
+# Module-level shared engine instance, set by init_mooncake_transfer_engine().
+_mooncake_transfer_engine: Optional["MooncakeTransferEngine"] = None
+
+
+def parse_ib_device_config(
+    ib_device_str: Optional[str],
+) -> Optional[Union[str, Dict[int, str]]]:
+    """Parse IB device config from a shared string, JSON mapping, or JSON file."""
+    if ib_device_str is None or not ib_device_str.strip():
+        return None
+
+    normalized_input = ib_device_str.strip()
+    if not normalized_input.endswith(".json") and not normalized_input.startswith("{"):
+        return normalized_input
+
+    if normalized_input.endswith(".json"):
+        if not os.path.isfile(normalized_input):
+            raise RuntimeError(f"File {normalized_input} does not exist.")
+        try:
+            with open(normalized_input, "r", encoding="utf-8") as file:
+                mapping = json.load(file)
+        except json.JSONDecodeError as exc:
+            raise RuntimeError(
+                f"Failed to parse JSON content from file {normalized_input}"
+            ) from exc
+        except (IOError, OSError) as exc:
+            raise RuntimeError(
+                f"Failed to read JSON file {normalized_input}: {exc}"
+            ) from exc
+    else:
+        try:
+            mapping = json.loads(normalized_input)
+        except json.JSONDecodeError as exc:
+            raise ValueError(f"Invalid JSON mapping: {normalized_input}") from exc
+
+    if not isinstance(mapping, dict):
+        raise ValueError(
+            "Invalid format: expected a mapping from GPU id to IB device string"
+        )
+
+    normalized_mapping: Dict[int, str] = {}
+    for gpu_key, ib_devices in mapping.items():
+        normalized_key = int(gpu_key) if str(gpu_key).isdigit() else None
+        if normalized_key is None or not isinstance(ib_devices, str):
+            raise ValueError(
+                "Invalid format: keys must be integers (or string "
+                "representations of integers) and values must be strings"
+            )
+        normalized_mapping[normalized_key] = ib_devices.strip()
+
+    if not normalized_mapping:
+        raise ValueError("No valid GPU mappings found in JSON")
+
+    return normalized_mapping
+
+
+def get_ib_devices_for_gpu(ib_device_str: Optional[str], gpu_id: int) -> Optional[str]:
+    """
+    Parse IB device string and get IB devices for a specific GPU ID.
+
+    Supports all the following formats:
+    1. Old format: "ib0, ib1, ib2"
+    2. New format: {0: "ib0, ib1", 1: "ib2, ib3", 2: "ib4"}
+    3. JSON file: path to a JSON file containing the mapping
+
+    Args:
+        ib_device_str: The original IB device string or path to JSON file
+        gpu_id: The GPU ID to get devices for
+
+    Returns:
+        IB devices string for the GPU, or None if not available
+    """
+    parsed_config = parse_ib_device_config(ib_device_str)
+    if parsed_config is None:
+        return None
+
+    if isinstance(parsed_config, str):
+        return parsed_config
+
+    if gpu_id in parsed_config:
+        return parsed_config[gpu_id]
+
+    raise ValueError(
+        f"No IB devices configured for GPU {gpu_id}. "
+        f"Available GPUs: {list(parsed_config.keys())}"
+    )
+
+
+class MooncakeTransferEngine:
+    """Shared Mooncake transfer engine for RDMA/transfer operations."""
+
+    def __init__(
+        self,
+        hostname: str,
+        gpu_id: Optional[int] = None,
+        ib_device: Optional[str] = None,
+    ):
+        try:
+            from mooncake.engine import TransferEngine
+        except ImportError as e:
+            raise ImportError(
+                "Please install mooncake by following the instructions at "
+                "https://kvcache-ai.github.io/Mooncake/getting_started/build.html "
+                "to run SGLang with MooncakeTransferEngine."
+            ) from e
+
+        self.engine = TransferEngine()
+        self.hostname = hostname
+        self.gpu_id = gpu_id if gpu_id is not None else 0
+        # MC_FORCE_TCP=1 makes mooncake install TcpTransport instead of RDMA,
+        # in which case RDMA HCA selection is irrelevant; pass empty device.
+        if os.environ.get("MC_FORCE_TCP") == "1":
+            self.ib_device = ""
+        else:
+            self.ib_device = get_ib_devices_for_gpu(ib_device, self.gpu_id)
+
+        self.initialize(
+            hostname=self.hostname,
+            device_name=self.ib_device,
+        )
+        self.session_id = NetworkAddress(
+            self.hostname, self.engine.get_rpc_port()
+        ).to_host_port_str()
+
+    def register(self, ptr, length):
+        try:
+            ret_value = self.engine.register_memory(ptr, length)
+        except Exception:
+            # Mark register as failed
+            ret_value = -1
+
+        if ret_value != 0:
+            logger.debug("Mooncake memory registration %s failed.", ptr)
+
+    def deregister(self, ptr):
+        try:
+            ret_value = self.engine.unregister_memory(ptr)
+        except Exception:
+            # Mark deregister as failed
+            ret_value = -1
+
+        if ret_value != 0:
+            logger.debug("Mooncake memory deregistration %s failed.", ptr)
+
+    def batch_register(self, ptrs: List[int], lengths: List[int]) -> int:
+        """Batch register multiple memory regions."""
+        try:
+            ret_value = self.engine.batch_register_memory(ptrs, lengths)
+        except Exception:
+            # Mark batch register as failed
+            ret_value = -1
+            if not hasattr(self.engine, "batch_register_memory"):
+                raise RuntimeError(
+                    "Mooncake's batch register requires a newer version of "
+                    "mooncake-transfer-engine. Please upgrade Mooncake."
+                )
+
+        if ret_value != 0:
+            logger.debug("Mooncake batch memory registration failed.")
+        return ret_value
+
+    def batch_deregister(self, ptrs: List[int]) -> int:
+        """Batch deregister multiple memory regions."""
+        try:
+            ret_value = self.engine.batch_unregister_memory(ptrs)
+        except Exception:
+            # Mark batch deregister as failed
+            ret_value = -1
+
+        if ret_value != 0:
+            logger.debug("Mooncake batch memory deregistration failed.")
+        return ret_value
+
+    def initialize(
+        self,
+        hostname: str,
+        device_name: Optional[str],
+    ) -> None:
+        """Initialize the mooncake instance."""
+        if envs.ENABLE_ASCEND_TRANSFER_WITH_MOONCAKE.get():
+            npu_phy_id = envs.ASCEND_NPU_PHY_ID.get()
+            suffix = self.gpu_id if npu_phy_id == -1 else npu_phy_id
+            hostname += f":{get_free_port()}:npu_{suffix}"
+            protocol = "ascend"
+        else:
+            # MOONCAKE_PROTOCOL selects the transport (rdma | efa | tcp | ...).
+            # Default is "rdma"; set MOONCAKE_PROTOCOL=efa on AWS EFA hardware.
+            protocol = envs.MOONCAKE_PROTOCOL.get()
+
+        ret_value = self.engine.initialize(
+            hostname,
+            "P2PHANDSHAKE",
+            protocol,
+            device_name if device_name is not None else "",
+        )
+        if ret_value != 0:
+            logger.error("Mooncake Transfer Engine initialization failed.")
+            raise RuntimeError("Mooncake Transfer Engine initialization failed.")
+
+    def transfer_sync(
+        self, session_id: str, buffer: int, peer_buffer_address: int, length: int
+    ) -> int:
+        """Synchronously transfer data to the specified address."""
+        try:
+            ret = self.engine.transfer_sync_write(
+                session_id, buffer, peer_buffer_address, length
+            )
+        except Exception:
+            ret = -1
+
+        if ret < 0:
+            logger.debug(
+                "Failed to transfer data from %s to %s - %s.",
+                buffer,
+                session_id,
+                peer_buffer_address,
+            )
+
+        return ret
+
+    def batch_transfer_sync(
+        self,
+        session_id: str,
+        buffers: List[int],
+        peer_buffer_addresses: List[int],
+        lengths: List[int],
+    ) -> int:
+        """Synchronously transfer data to the specified addresses in batches."""
+        try:
+            ret = self.engine.batch_transfer_sync_write(
+                session_id, buffers, peer_buffer_addresses, lengths
+            )
+        except Exception:
+            ret = -1
+            if not hasattr(self.engine, "batch_transfer_sync_write"):
+                raise RuntimeError(
+                    "Mooncake's batch transfer requires mooncake-transfer-engine "
+                    ">= 0.3.4.post2. Please upgrade Mooncake by "
+                    "'pip install mooncake-transfer-engine --upgrade'"
+                )
+
+        if ret < 0:
+            logger.debug(
+                "Failed to batch transfer data. Buffers: %s, Session: %s, "
+                "Peer addresses: %s",
+                buffers,
+                session_id,
+                peer_buffer_addresses,
+            )
+        return ret
+
+    def get_session_id(self):
+        return self.session_id
+
+    def send_probe(self, peer_session_id: str) -> int:
+        return self.engine.send_probe(peer_session_id)
+
+    def get_engine(self):
+        return self.engine.get_engine()
+
+    def get_ib_device(self):
+        return self.ib_device
+
+
+def init_mooncake_transfer_engine(
+    hostname: str,
+    gpu_id: Optional[int] = None,
+    ib_device: Optional[str] = None,
+) -> MooncakeTransferEngine:
+    """
+    Initialize the shared MooncakeTransferEngine. Note: if already
+    initialized with the same (hostname, gpu_id, ib_device), returns existing
+    instance. Call from parallel_state when model parallel is set up and
+    mooncake transfer is needed.
+    """
+    global _mooncake_transfer_engine
+    if _mooncake_transfer_engine is not None:
+        return _mooncake_transfer_engine
+    _mooncake_transfer_engine = MooncakeTransferEngine(
+        hostname=hostname, gpu_id=gpu_id, ib_device=ib_device
+    )
+    return _mooncake_transfer_engine
+
+
+def get_mooncake_transfer_engine() -> Optional[MooncakeTransferEngine]:
+    """Return the shared MooncakeTransferEngine if initialized, else None."""
+    return _mooncake_transfer_engine

python/sglang/srt/distributed/device_communicators/npu_communicator.py

@@ -0,0 +1,71 @@
+import torch
+import torch.distributed as dist
+from torch.distributed import ProcessGroup
+
+from sglang.srt.utils import is_npu
+
+_is_npu = is_npu()
+
+if _is_npu:
+    from torch_npu import npu_dynamic_quant
+
+
+class NpuCommunicator:
+
+    def __init__(self, group: ProcessGroup):
+        if not _is_npu:
+            self.disabled = True
+            return
+        self.disabled = False
+        self.group = group
+        self.world_size = dist.get_world_size(self.group)
+
+    def all_reduce(self, x: torch.Tensor) -> torch.Tensor:
+        dist.all_reduce(x, group=self.group)
+        return x
+
+    def quant_all_reduce(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Note:
+        All reduce is split into All gather + reduce.
+        All gather is performed in low precision, but reduce in full precision.
+        """
+        world_size = self.world_size
+        input_size = x.size()
+        output_size = (input_size[0] * world_size,) + input_size[1:]
+        x_q, scale = npu_dynamic_quant(x, dst_type=torch.int8)
+        # Allocate output tensor.
+        output_tensor = torch.empty(output_size, dtype=x_q.dtype, device=x.device)
+        output_scale = torch.empty(
+            output_size[:1], dtype=scale.dtype, device=scale.device
+        )
+        # All-gather.
+        dist.all_gather_into_tensor(output_tensor, x_q, group=self.group)
+        dist.all_gather_into_tensor(output_scale, scale, group=self.group)
+
+        output_tensor = output_tensor.to(x.dtype) * output_scale.unsqueeze(-1).to(
+            x.dtype
+        )
+        # Reshape
+        output_tensor = output_tensor.reshape((world_size,) + input_size)
+
+        return output_tensor.sum(dim=0)
+
+    def all_gather(self, x: torch.Tensor, dim: int = -1) -> torch.Tensor:
+        world_size = self.world_size
+        if dim < 0:
+            # Convert negative dim to positive.
+            dim += x.dim()
+        input_size = x.size()
+        output_size = (input_size[0] * world_size,) + input_size[1:]
+        # Allocate output tensor.
+        output_tensor = torch.empty(output_size, dtype=x.dtype, device=x.device)
+        # All-gather.
+        dist.all_gather_into_tensor(output_tensor, x, group=self.group)
+        # Reshape
+        output_tensor = output_tensor.reshape((world_size,) + input_size)
+        output_tensor = output_tensor.movedim(0, dim)
+        output_tensor = output_tensor.reshape(
+            input_size[:dim] + (world_size * input_size[dim],) + input_size[dim + 1 :]
+        )
+        return output_tensor

python/sglang/srt/distributed/device_communicators/pymscclpp.py

@@ -0,0 +1,394 @@
+import importlib
+import logging
+from contextlib import contextmanager
+from typing import Optional, Union
+
+import torch
+import torch.distributed as dist
+from torch.distributed import ProcessGroup, ReduceOp
+
+from sglang.srt.compilation.compile_phase import (
+    get_pcg_capture_stream,
+    is_in_torch_compile_warmup,
+)
+from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
+    is_in_tc_piecewise_cuda_graph,
+)
+from sglang.srt.runtime_context import get_server_args
+
+logger = logging.getLogger(__name__)
+
+
+class PyMscclppCommunicator:
+    _SUPPORTED_WORLD_SIZES = [8, 16, 32]
+    _SUPPORTED_DTYPE = [torch.float, torch.float16, torch.bfloat16]
+
+    def _is_symm_mem_enabled(self) -> bool:
+        try:
+            return get_server_args().enable_symm_mem
+        except ValueError:
+            return False
+
+    def _is_weak_contiguous(self, inp: torch.Tensor):
+        return inp.is_contiguous() or (
+            inp.storage().nbytes() - inp.storage_offset() * inp.element_size()
+            == inp.numel() * inp.element_size()
+        )
+
+    def _get_tuned_config(self, size):
+        if size <= 512:
+            target_size = 512
+        elif size > 256 * 1024 * 1024:
+            target_size = 256 * 1024 * 1024
+        else:
+            target_size = 1 << (size - 1).bit_length()
+        return self.best_configs.get(target_size)
+
+    def _create_dsl_algorithms(self):
+        dsl_algos_config = []
+        n_nodes = self.world_size // self.nranks_per_node
+        if n_nodes == 2 or n_nodes == 4:
+            for tbg in [1, 2, 4, 8]:
+                for num_threads_per_block in [256, 512, 768, 1024]:
+                    spec = self.mscclpp.language.AlgoSpec(
+                        name=f"allreduce_{n_nodes}node_{tbg}TBG_{num_threads_per_block}TPB",
+                        collective=self.mscclpp.language.collectives.AllReduce(
+                            self.world_size, 1, True
+                        ),
+                        nranks_per_node=self.nranks_per_node,
+                        world_size=self.world_size,
+                        in_place=True,
+                        instances=1,
+                        protocol="LL",
+                        auto_sync=False,
+                        num_threads_per_block=num_threads_per_block,
+                        reuse_resources=True,
+                        use_double_scratch_buffer=True,
+                        min_message_size=tbg * (1 << 10),
+                        max_message_size=8 << 20,
+                        tags={"default": 1},
+                    )
+                    algo = self.mscclpp.compile(
+                        self.def_algo.allreduce_multi_nodes,
+                        spec,
+                        self.rank,
+                        thread_block_group_size=tbg,
+                    )
+                    dsl_algos_config.append((algo, [0], [0]))
+        return dsl_algos_config
+
+    def _create_native_algorithms(self):
+        navitve_algorithms_config = []
+        dlpack = self.mscclpp.RawGpuBuffer(1 << 27).to_dlpack(
+            data_type=str(torch.float16)
+        )
+        self.scratch_buffer = torch.utils.dlpack.from_dlpack(dlpack)
+        self.flag_buffer = torch.ones(128, dtype=torch.uint32, device="cuda")
+        algos = self.mscclpp_ext.AlgorithmCollectionBuilder().build_default_algorithms(
+            scratch_buffer=self.scratch_buffer.data_ptr(),
+            scratch_buffer_size=self.scratch_buffer.nbytes,
+            rank=self.rank,
+        )
+
+        for algo in algos:
+            if algo.name == "default_allreduce_nvls_packet":
+                algo.set_message_size_range(0, 512 << 10)
+                navitve_algorithms_config.append(
+                    (algo, [4, 8, 12, 16], [256, 512, 768, 1024])
+                )
+            if algo.name == "default_allreduce_packet":
+                algo.set_message_size_range(0, 2 << 20)
+                navitve_algorithms_config.append(
+                    (algo, [14, 21, 28, 42, 56], [256, 512, 768, 1024])
+                )
+            if algo.name == "default_allreduce_rsag_zero_copy":
+                algo.set_message_size_range(512 << 10, 4 << 30)
+                navitve_algorithms_config.append(
+                    (algo, [32, 48, 64, 128], [256, 512, 768, 1024])
+                )
+            if (
+                self.symm_mem_enabled
+                and algo.name == "default_allreduce_nvls_zero_copy"
+            ):
+                algo.set_message_size_range(512 << 10, 4 << 30)
+                navitve_algorithms_config.append(
+                    (algo, [4, 8, 12, 16, 32], [256, 512, 768, 1024])
+                )
+
+        return navitve_algorithms_config
+
+    def _create_algorithms(self):
+        if self.world_size == 8:
+            self.algos_config = self._create_native_algorithms()
+            self._tune(5, 10, 20, self.algos_config)
+        elif self.world_size == 16 or self.world_size == 32:
+            self.dsl_algos_config = self._create_dsl_algorithms()
+            self._tune(5, 10, 20, self.dsl_algos_config)
+
+    def _get_time(
+        self,
+        algo,
+        tune_tensor,
+        size,
+        nb,
+        nt,
+        n_warmup,
+        n_graph_launches,
+        n_ops_per_graph,
+    ):
+        # Check if the algorithm can run with the given configuration
+        if self._run_algo(algo, tune_tensor, size, nb, nt, True) != 0:
+            return float("inf")
+
+        # Warmup iterations to stabilize performance
+        for _ in range(n_warmup):
+            self._run_algo(algo, tune_tensor, size, nb, nt, True)
+
+        # Warmup on capture stream
+        capture_stream = torch.cuda.Stream()
+        capture_stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(capture_stream):
+            self._run_algo(algo, tune_tensor, size, nb, nt, True)
+        capture_stream.synchronize()
+
+        # Capture the algorithm execution in a CUDA graph
+        g = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(g, stream=capture_stream):
+            for _ in range(n_ops_per_graph):
+                self._run_algo(algo, tune_tensor, size, nb, nt, True)
+
+        # Measure the execution time of the captured graph
+        start_event = torch.cuda.Event(enable_timing=True)
+        end_event = torch.cuda.Event(enable_timing=True)
+        start_event.record(capture_stream)
+        with torch.cuda.stream(capture_stream):
+            for _ in range(n_graph_launches):
+                g.replay()
+        end_event.record(capture_stream)
+        end_event.synchronize()
+        elapsed = start_event.elapsed_time(end_event)
+
+        # Synchronize timing results across all ranks to ensure consistent algorithm selection
+        # replicate n times such due to algo limitations
+        time_tensor = torch.full(
+            (self.world_size,), elapsed, dtype=torch.float64, device="cuda"
+        ).to(dtype=torch.float32)
+        torch.cuda.current_stream().wait_stream(capture_stream)
+        if self.rank == 0:
+            avg_time = time_tensor[self.rank].item() / self.world_size
+            tensor = torch.tensor([avg_time])
+        else:
+            tensor = torch.empty(1)
+        dist.broadcast(tensor, src=0, group=self.group)
+        avg_time = tensor.item()
+
+        return avg_time
+
+    def _tune(self, n_warmup, n_graph_launches, n_ops_per_graph, algos_config):
+        sizes = [1 << i for i in range(9, 24)]
+        dlpack = self.mscclpp.RawGpuBuffer(1 << 27).to_dlpack(
+            data_type=str(torch.float16)
+        )
+        tune_tensor = torch.utils.dlpack.from_dlpack(dlpack)
+
+        for size in sizes:
+            best_time = float("inf")
+            best_config = None
+            for i in range(len(algos_config)):
+                algo, candidates_nblocks, candidates_nthreads = algos_config[i]
+                if (
+                    size >= algo.message_size_range[0]
+                    and size <= algo.message_size_range[1]
+                ):
+                    for nb in candidates_nblocks:
+                        for nt in candidates_nthreads:
+                            avg_time = self._get_time(
+                                algo,
+                                tune_tensor,
+                                size,
+                                nb,
+                                nt,
+                                n_warmup,
+                                n_graph_launches,
+                                n_ops_per_graph,
+                            )
+                            if avg_time < best_time:
+                                best_time = avg_time
+                                best_config = (algo, nb, nt)
+            if best_config:
+                self.best_configs[size] = best_config
+
+        torch.cuda.synchronize()
+        for algo, _, _ in algos_config:
+            algo.reset()
+
+    def _run_algo(self, algo, tensor, size, nblocks, nthreads, sym_mem_enabled=False):
+        return algo.execute(
+            comm=self.comm.communicator,
+            executor=self.executor,
+            input_buffer=tensor.data_ptr(),
+            output_buffer=tensor.data_ptr(),
+            input_size=size,
+            output_size=size,
+            dtype=self.dtype_to_mscclpp_dtype(tensor.dtype),
+            op=self.mscclpp.ReduceOp.SUM,
+            stream=torch.cuda.current_stream().cuda_stream,
+            nblocks=nblocks,
+            nthreads_per_block=nthreads,
+            symmetric_memory=sym_mem_enabled,
+        )
+
+    def __init__(
+        self,
+        group: ProcessGroup,
+        device: Union[int, str, torch.device],
+    ) -> None:
+        """Args:
+            group: the process group to work on. If None, it will use the
+                default process group.
+            device: the device to bind the CustomAllreduce to. If None,
+                it will be bind to f"cuda:{local_rank}".
+        It is the caller's responsibility to make sure each communicator
+        is bind to a unique device, and all communicators in this group
+        are in the same node.
+        """
+        self._IS_CAPTURING = False
+        self.disabled = True
+
+        try:
+            self.mscclpp = importlib.import_module("mscclpp")
+            self.mscclpp_ext = importlib.import_module("mscclpp.ext")
+            self.def_algo = importlib.import_module("mscclpp.default_algos")
+        except ImportError:
+            self.available = False
+            self.mscclpp = None
+            return
+
+        self.available = True
+        self.group = group
+
+        assert (
+            dist.get_backend(group) != dist.Backend.NCCL
+        ), "CustomAllreduce should be attached to a non-NCCL group."
+
+        rank = dist.get_rank(group=self.group)
+        world_size = dist.get_world_size(group=self.group)
+        if world_size == 1:
+            # No need to initialize mscclpp for single GPU case.
+            return
+
+        if world_size not in PyMscclppCommunicator._SUPPORTED_WORLD_SIZES:
+            logger.warning(
+                "PyMscclpp is disabled due to an unsupported world"
+                " size: %d. Supported world sizes: %s. To silence this "
+                "warning, specify disable_mscclpp=True explicitly.",
+                world_size,
+                str(PyMscclppCommunicator._SUPPORTED_WORLD_SIZES),
+            )
+            return
+
+        self.ranks = torch.distributed.get_process_group_ranks(group)
+        self.nranks_per_node = torch.cuda.device_count()
+        # for now mscclpp with stride in the communicator is not tested
+        if not (abs(self.ranks[-1] - self.ranks[0]) == world_size - 1):
+            logger.warning(
+                "PyMscclpp is disabled due to an unsupported group %s."
+                "Please ensure all ranks in the group are consecutive."
+                "To silence this warning, specify disable_mscclpp=True explicitly.",
+                str(self.ranks),
+            )
+            return
+
+        if isinstance(device, int):
+            device = torch.device(f"cuda:{device}")
+        elif isinstance(device, str):
+            device = torch.device(device)
+        # now `device` is a `torch.device` object
+        assert isinstance(device, torch.device)
+        self.device = device
+
+        self.rank = rank
+        self.world_size = world_size
+        self.comm = self.mscclpp.CommGroup(
+            torch_group=self.group, rank=rank, size=world_size
+        )
+        self.executor = self.mscclpp.Executor(self.comm.communicator)
+        self.symm_mem_enabled = self._is_symm_mem_enabled()
+        self.best_configs = {}
+        self._create_algorithms()
+
+    def destroy(self):
+        self.algos_config = None
+        self.best_configs = None
+        self.executor = None
+        self.scratch_buffer = None
+        self.flag_buffer = None
+        self.comm = None
+
+    def should_mscclpp_allreduce(
+        self, inp: torch.Tensor, op: ReduceOp = ReduceOp.SUM
+    ) -> bool:
+        if (
+            self.disabled
+            or self.world_size not in PyMscclppCommunicator._SUPPORTED_WORLD_SIZES
+        ):
+            return False
+        if inp.dtype not in PyMscclppCommunicator._SUPPORTED_DTYPE:
+            return False
+        if not self._is_weak_contiguous(inp):
+            return False
+        if op is not ReduceOp.SUM:
+            return False
+        if self._get_tuned_config(inp.numel() * inp.element_size()) is None:
+            return False
+        # mscclpp must not be used during any piecewise CUDA graph phase
+        # (compile, capture, or replay) as it changes the allreduce dispatch
+        # path and triggers recompilation.
+        if (
+            is_in_tc_piecewise_cuda_graph()
+            or is_in_torch_compile_warmup()
+            or get_pcg_capture_stream() is not None
+        ):
+            return False
+        return True
+
+    def dtype_to_mscclpp_dtype(self, dtype: torch.dtype):
+        if dtype == torch.float16:
+            return self.mscclpp.DataType.float16
+        elif dtype == torch.float32:
+            return self.mscclpp.DataType.float32
+        elif dtype == torch.int32:
+            return self.mscclpp.DataType.int32
+        elif dtype == torch.bfloat16:
+            return self.mscclpp.DataType.bfloat16
+        else:
+            raise ValueError(f"Unknown data type: {dtype}")
+
+    def all_reduce(
+        self,
+        tensor: torch.Tensor,
+        op: ReduceOp = ReduceOp.SUM,
+        stream: torch.cuda.Stream = None,
+    ):
+        assert op == torch.distributed.ReduceOp.SUM
+        nbytes = tensor.numel() * tensor.element_size()
+        algo, nblocks, nthreads = self._get_tuned_config(nbytes)
+        self._run_algo(algo, tensor, nbytes, nblocks, nthreads, self.symm_mem_enabled)
+        return tensor
+
+    @contextmanager
+    def change_state(
+        self,
+        enable: Optional[bool] = None,
+    ):
+        if enable is None or self.available is False:
+            # guess a default value when not specified
+            # DO: Decided if raise an exception here or not
+            enable = self.available
+
+        old_disable = self.disabled
+        self.disabled = not enable
+
+        yield
+
+        self.disabled = old_disable

python/sglang/srt/distributed/device_communicators/pynccl.py

@@ -0,0 +1,385 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/distributed/device_communicators/pynccl.py
+
+import logging
+from contextlib import contextmanager
+from typing import Optional, Union
+
+# ===================== import region =====================
+import torch
+import torch.distributed as dist
+from torch.distributed import ProcessGroup, ReduceOp
+
+from sglang.srt.distributed.device_communicators.pynccl_wrapper import (
+    NCCLLibrary,
+    buffer_type,
+    cudaStream_t,
+    ncclComm_t,
+    ncclDataTypeEnum,
+    ncclRedOpTypeEnum,
+    ncclUniqueId,
+)
+from sglang.srt.distributed.utils import StatelessProcessGroup
+from sglang.srt.utils.common import get_current_device_stream_fast
+
+logger = logging.getLogger(__name__)
+
+
+class PyNcclCommunicator:
+
+    def __init__(
+        self,
+        group: Union[ProcessGroup, StatelessProcessGroup],
+        device: Union[int, str, torch.device],
+        library_path: Optional[str] = None,
+    ):
+        """
+        Args:
+            group: the process group to work on. If None, it will use the
+                default process group.
+            device: the device to bind the PyNcclCommunicator to. If None,
+                it will be bind to f"cuda:{local_rank}".
+            library_path: the path to the NCCL library. If None, it will
+                use the default library path.
+        It is the caller's responsibility to make sure each communicator
+        is bind to a unique device.
+        """
+        if not isinstance(group, StatelessProcessGroup):
+            assert dist.is_initialized()
+            assert (
+                dist.get_backend(group) != dist.Backend.NCCL
+            ), "PyNcclCommunicator should be attached to a non-NCCL group."
+            # note: this rank is the rank in the group
+            self.rank = dist.get_rank(group)
+            self.world_size = dist.get_world_size(group)
+        else:
+            self.rank = group.rank
+            self.world_size = group.world_size
+
+        self.group = group
+
+        # if world_size == 1, no need to create communicator
+        if self.world_size == 1:
+            self.available = False
+            self.disabled = True
+            return
+        try:
+            self.nccl = NCCLLibrary(library_path)
+        except Exception:
+            # disable because of missing NCCL library
+            # e.g. in a non-GPU environment
+            self.available = False
+            self.disabled = True
+            return
+
+        self.available = True
+        self.disabled = False
+
+        self.nccl_version = self.nccl.ncclGetRawVersion()
+        if self.rank == 0:
+            logger.info("sglang is using nccl==%s", self.nccl.ncclGetVersion())
+
+        if self.rank == 0:
+            # get the unique id from NCCL
+            self.unique_id = self.nccl.ncclGetUniqueId()
+        else:
+            # construct an empty unique id
+            self.unique_id = ncclUniqueId()
+
+        if not isinstance(group, StatelessProcessGroup):
+            tensor = torch.ByteTensor(list(self.unique_id.internal))
+            ranks = dist.get_process_group_ranks(group)
+            # arg `src` in `broadcast` is the global rank
+            dist.broadcast(tensor, src=ranks[0], group=group)
+            byte_list = tensor.tolist()
+            for i, byte in enumerate(byte_list):
+                self.unique_id.internal[i] = byte
+        else:
+            self.unique_id = group.broadcast_obj(self.unique_id, src=0)
+        if isinstance(device, int):
+            device = torch.device(f"cuda:{device}")
+        elif isinstance(device, str):
+            device = torch.device(device)
+        # now `device` is a `torch.device` object
+        assert isinstance(device, torch.device)
+        self.device = device
+        # nccl communicator and stream will use this device
+        # `torch.cuda.device` is a context manager that changes the
+        # current cuda device to the specified one
+        with torch.cuda.device(device):
+            self.comm: ncclComm_t = self.nccl.ncclCommInitRank(
+                self.world_size, self.unique_id, self.rank
+            )
+            warmup_stream = torch.cuda.Stream()
+
+            # A small all_reduce for warmup.
+            with torch.cuda.stream(warmup_stream):
+                data = torch.zeros(1, device=device)
+                self.all_reduce(data)
+            warmup_stream.synchronize()
+            del data
+
+        # by default it is disabled, e.g. in profiling models and prefill phase.
+        # to use it, use under `with obj.change_state(enable=True)`, usually
+        # when we are using CUDA graph.
+        self.disabled = True
+
+    def _resolve_stream(self) -> torch.cuda.Stream:
+        """Return the current device stream used for NCCL calls."""
+        return get_current_device_stream_fast()
+
+    def all_reduce(self, tensor: torch.Tensor, op: ReduceOp = ReduceOp.SUM):
+        if self.disabled:
+            return
+        # nccl communicator created on a specific device
+        # will only work on tensors on the same device
+        # otherwise it will cause "illegal memory access"
+        assert tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the input tensor is on {tensor.device}"
+        )
+        stream = self._resolve_stream()
+        self.nccl.ncclAllReduce(
+            buffer_type(tensor.data_ptr()),
+            buffer_type(tensor.data_ptr()),
+            tensor.numel(),
+            ncclDataTypeEnum.from_torch(tensor.dtype),
+            ncclRedOpTypeEnum.from_torch(op),
+            self.comm,
+            cudaStream_t(stream.cuda_stream),
+        )
+
+    def outplace_all_reduce(
+        self,
+        in_tensor: torch.Tensor,
+        out_tensor: Optional[torch.Tensor] = None,
+        op: ReduceOp = ReduceOp.SUM,
+    ) -> Optional[torch.Tensor]:
+        if self.disabled:
+            return None
+        assert in_tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the input tensor is on {in_tensor.device}"
+        )
+
+        if out_tensor is None:
+            out_tensor = torch.empty_like(in_tensor)
+
+        stream = self._resolve_stream()
+        self.nccl.ncclAllReduce(
+            buffer_type(in_tensor.data_ptr()),  # sendbuff
+            buffer_type(out_tensor.data_ptr()),  # recvbuff - DIFFERENT pointer
+            in_tensor.numel(),
+            ncclDataTypeEnum.from_torch(in_tensor.dtype),
+            ncclRedOpTypeEnum.from_torch(op),
+            self.comm,
+            cudaStream_t(stream.cuda_stream),
+        )
+        return out_tensor
+
+    def all_gather(
+        self,
+        output_tensor: torch.Tensor,
+        input_tensor: torch.Tensor,
+        sizes: Optional[list[int]] = None,
+    ):
+        if self.disabled:
+            return
+        # nccl communicator created on a specific device
+        # will only work on tensors on the same device
+        # otherwise it will cause "illegal memory access"
+        assert input_tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the input tensor is on {input_tensor.device}"
+        )
+        stream = self._resolve_stream()
+
+        if sizes is not None:
+            split_offset = 0
+
+            self.nccl.ncclGroupStart()
+            for root, split_size in enumerate(sizes):
+                dst_slice = output_tensor[split_offset : split_offset + split_size]
+                self.nccl.ncclBroadcast(
+                    buffer_type(input_tensor.data_ptr()),
+                    buffer_type(dst_slice.data_ptr()),
+                    dst_slice.numel(),
+                    ncclDataTypeEnum.from_torch(input_tensor.dtype),
+                    root,
+                    self.comm,
+                    cudaStream_t(stream.cuda_stream),
+                )
+                split_offset += split_size
+            self.nccl.ncclGroupEnd()
+        else:
+            self.nccl.ncclAllGather(
+                buffer_type(input_tensor.data_ptr()),
+                buffer_type(output_tensor.data_ptr()),
+                input_tensor.numel(),
+                ncclDataTypeEnum.from_torch(input_tensor.dtype),
+                self.comm,
+                cudaStream_t(stream.cuda_stream),
+            )
+
+    def cp_all_gather_into_tensor(
+        self,
+        output_tensor: torch.Tensor,
+        input_tensor: torch.Tensor,
+        stream: torch.cuda.Stream,
+        sizes: Optional[list[int]] = None,
+    ):
+        """
+        Currently, it is mainly used in context parallelism,
+        primarily leveraging pynccl to implement non-blocking allgather communication.
+        """
+        # nccl communicator created on a specific device
+        # will only work on tensors on the same device
+        # otherwise it will cause "illegal memory access"
+        assert input_tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the input tensor is on {input_tensor.device}"
+        )
+        self.nccl.ncclAllGather(
+            buffer_type(input_tensor.data_ptr()),
+            buffer_type(output_tensor.data_ptr()),
+            input_tensor.numel(),
+            ncclDataTypeEnum.from_torch(input_tensor.dtype),
+            self.comm,
+            cudaStream_t(stream.cuda_stream),
+        )
+
+    def reduce_scatter(
+        self,
+        output_tensor: torch.Tensor,
+        input_tensor: torch.Tensor,
+        op: ReduceOp = ReduceOp.SUM,
+        sizes: Optional[list[int]] = None,
+    ):
+        if self.disabled:
+            return
+        # nccl communicator created on a specific device
+        # will only work on tensors on the same device
+        # otherwise it will cause "illegal memory access"
+        assert input_tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the input tensor is on {input_tensor.device}"
+        )
+        stream = self._resolve_stream()
+
+        if sizes is not None:
+            split_offset = 0
+            self.nccl.ncclGroupStart()
+            for root, split_size in enumerate(sizes):
+                chunk = input_tensor[split_offset : split_offset + split_size, ...]
+
+                self.nccl.ncclReduce(
+                    buffer_type(chunk.data_ptr()),
+                    buffer_type(output_tensor.data_ptr()),
+                    chunk.numel(),
+                    ncclDataTypeEnum.from_torch(input_tensor.dtype),
+                    ncclRedOpTypeEnum.from_torch(op),
+                    root,
+                    self.comm,
+                    cudaStream_t(stream.cuda_stream),
+                )
+                split_offset += split_size
+            self.nccl.ncclGroupEnd()
+        else:
+            self.nccl.ncclReduceScatter(
+                buffer_type(input_tensor.data_ptr()),
+                buffer_type(output_tensor.data_ptr()),
+                output_tensor.numel(),
+                ncclDataTypeEnum.from_torch(input_tensor.dtype),
+                ncclRedOpTypeEnum.from_torch(op),
+                self.comm,
+                cudaStream_t(stream.cuda_stream),
+            )
+
+    def send(self, tensor: torch.Tensor, dst: int):
+        if self.disabled:
+            return
+        assert tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the input tensor is on {tensor.device}"
+        )
+        stream = self._resolve_stream()
+        self.nccl.ncclSend(
+            buffer_type(tensor.data_ptr()),
+            tensor.numel(),
+            ncclDataTypeEnum.from_torch(tensor.dtype),
+            dst,
+            self.comm,
+            cudaStream_t(stream.cuda_stream),
+        )
+
+    def recv(self, tensor: torch.Tensor, src: int):
+        if self.disabled:
+            return
+        assert tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the input tensor is on {tensor.device}"
+        )
+        stream = self._resolve_stream()
+        self.nccl.ncclRecv(
+            buffer_type(tensor.data_ptr()),
+            tensor.numel(),
+            ncclDataTypeEnum.from_torch(tensor.dtype),
+            src,
+            self.comm,
+            cudaStream_t(stream.cuda_stream),
+        )
+
+    def broadcast(self, tensor: torch.Tensor, src: int):
+        if self.disabled:
+            return
+        assert tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the input tensor is on {tensor.device}"
+        )
+        stream = self._resolve_stream()
+
+        if src == self.rank:
+            sendbuff = buffer_type(tensor.data_ptr())
+            # NCCL requires the sender also to have a receive buffer
+            recvbuff = buffer_type(tensor.data_ptr())
+        else:
+            sendbuff = buffer_type()
+            recvbuff = buffer_type(tensor.data_ptr())
+        self.nccl.ncclBroadcast(
+            sendbuff,
+            recvbuff,
+            tensor.numel(),
+            ncclDataTypeEnum.from_torch(tensor.dtype),
+            src,
+            self.comm,
+            cudaStream_t(stream.cuda_stream),
+        )
+
+    def register_comm_window_raw(self, ptr: int, size: int):
+        return self.nccl.ncclCommWindowRegister(self.comm, buffer_type(ptr), size, 1)
+
+    def deregister_comm_window(self, window):
+        return self.nccl.ncclCommWindowDeregister(self.comm, window)
+
+    def group_start(self):
+        self.nccl.ncclGroupStart()
+
+    def group_end(self):
+        self.nccl.ncclGroupEnd()
+
+    @contextmanager
+    def change_state(self, enable: Optional[bool] = None):
+        """
+        A context manager to change the enabled state of the communicator.
+        """
+        if enable is None:
+            # guess a default value when not specified
+            enable = self.available
+
+        old_disable = self.disabled
+        self.disabled = not enable
+        try:
+            yield
+        finally:
+            self.disabled = old_disable

python/sglang/srt/distributed/device_communicators/pynccl_allocator.py

@@ -0,0 +1,406 @@
+import ctypes
+import logging
+import os
+import tempfile
+import traceback
+from contextlib import nullcontext
+
+import torch
+from torch.cuda.memory import (
+    CUDAPluggableAllocator,
+    _cuda_beginAllocateCurrentThreadToPool,
+    _cuda_endAllocateToPool,
+    _cuda_releasePool,
+)
+
+from sglang.srt.distributed.parallel_state import GroupCoordinator
+from sglang.srt.environ import envs
+from sglang.srt.runtime_context import get_server_args
+from sglang.srt.utils.common import torch_release
+
+after_2_8_0 = torch_release >= (2, 8)
+
+# C++ source for the NCCL allocator plugin
+# Key design:
+# 1. nccl_alloc_plug: Allocates memory via ncclMemAlloc and TRACKS the segment
+#    (ptr, size). Does NOT register with any comm at allocation time.
+# 2. nccl_free_plug: Frees memory via ncclMemFree and UNTRACKS the segment.
+#    Each segment is tracked only during its lifetime (from alloc to free).
+# 3. Segment tracking uses thread-safe std::vector + unordered_map for O(1) operations.
+# 4. Registration via nccl_allocator_register_segments_with_comm: Registers all
+#    tracked segments with a given comm, using index-based tracking to avoid
+#    re-registration. Registration state is maintained per-communicator in C++.
+nccl_allocator_source = """
+
+#include <cuda_runtime.h>
+#include <mutex>
+#include <vector>
+#include <unordered_map>
+#include <utility>
+
+extern "C" {
+
+// copy from https://github.com/NVIDIA/nccl/blob/master/src/nccl.h.in
+typedef enum { ncclSuccess                 =  0,
+               ncclUnhandledCudaError      =  1,
+               ncclSystemError             =  2,
+               ncclInternalError           =  3,
+               ncclInvalidArgument         =  4,
+               ncclInvalidUsage            =  5,
+               ncclRemoteError             =  6,
+               ncclInProgress              =  7,
+               ncclNumResults              =  8 } ncclResult_t;
+
+// NCCL symmetric memory window flags
+#define NCCL_WIN_COLL_SYMMETRIC 0x01
+
+typedef struct ncclComm* ncclComm_t;
+typedef struct ncclWindow_vidmem* ncclWindow_t;
+
+ncclResult_t  ncclMemAlloc(void** ptr, size_t size);
+ncclResult_t  ncclMemFree(void *ptr);
+ncclResult_t  ncclCommWindowRegister(ncclComm_t comm, void* buff, size_t size, ncclWindow_t* win, int winFlags);
+const char*  ncclGetErrorString(ncclResult_t result);
+
+#define NCCLCHECK(cmd) do {                                               \
+  ncclResult_t res = cmd;                                                 \
+  if (res != ncclSuccess) {                                               \
+    fprintf(stderr, "ERROR: NCCL symmetric memory allocation failed. Most likely out of device memory. '%s'\\n", \
+           ncclGetErrorString(res));                       \
+    return NULL;                                                        \
+  }                                                                       \
+} while(0)
+
+// Segment information structure
+struct Segment {
+    void* ptr;
+    size_t size;
+    Segment(void* p, size_t s) : ptr(p), size(s) {}
+};
+
+// Thread-safe segment tracking
+// Segment tracking using std::vector for FIFO order.
+// g_segments is maintained in insertion order (oldest first).
+static std::vector<Segment> g_segments;
+static std::mutex g_segment_mutex;
+
+// Track which segments have been registered with each communicator.
+// Key: comm_ptr, Value: the next segment index to register for this comm.
+static std::unordered_map<uintptr_t, size_t> g_comm_registration_index;
+
+// Add a segment to the tracking (appends to end, maintaining FIFO order)
+static void track_segment(void* ptr, size_t size) {
+    std::lock_guard<std::mutex> lock(g_segment_mutex);
+    g_segments.emplace_back(ptr, size);
+}
+
+void* nccl_alloc_plug(size_t size, int device, void* stream) {
+    void* ptr;
+    NCCLCHECK(ncclMemAlloc(&ptr, size));
+
+    // Track the segment but do NOT register with any comm
+    // Registration will be done at context exit via register_segments_with_comm
+    track_segment(ptr, size);
+
+    return ptr;
+}
+
+void nccl_free_plug(void* ptr, size_t size, int device, void* stream) {
+    ncclResult_t err = ncclMemFree(ptr);
+    // NOTE: We assume that no individual allocation will be freed until the
+    // entire memory pool is destroyed. If this assumption does not hold,
+    // we will encounter asymmetry issues between GPUs. For now, we clear
+    // all tracking state when the pool is destroyed.
+    std::lock_guard<std::mutex> lock(g_segment_mutex);
+    g_segments = std::vector<Segment>();
+    g_comm_registration_index = std::unordered_map<uintptr_t, size_t>();
+}
+
+// Register all tracked segments with a communicator.
+// Uses an index-based approach to avoid re-registering already-registered segments.
+// Returns 0 on success, non-zero on failure.
+int nccl_allocator_register_segments_with_comm(uintptr_t comm_ptr) {
+    std::lock_guard<std::mutex> lock(g_segment_mutex);
+
+    ncclComm_t comm = reinterpret_cast<ncclComm_t>(comm_ptr);
+
+    // Get the starting index for this communicator
+    size_t start_index = g_comm_registration_index[comm_ptr];
+
+    // Register all segments from start_index to the current end
+    for (size_t i = start_index; i < g_segments.size(); ++i) {
+        const Segment& seg = g_segments[i];
+        ncclWindow_t win;
+        ncclResult_t res = ncclCommWindowRegister(comm, seg.ptr, seg.size, &win, NCCL_WIN_COLL_SYMMETRIC);
+        if (res != ncclSuccess) {
+            fprintf(stderr, "ERROR: NCCL symmetric memory registration failed. '%s'\\n", ncclGetErrorString(res));
+            return res;
+        }
+    }
+
+    // Update the registration index for this communicator
+    g_comm_registration_index[comm_ptr] = g_segments.size();
+
+    return ncclSuccess;
+}
+
+}
+"""
+
+_allocator = None
+_mem_pool = None
+_graph_pool_id = None
+_cur_device = None
+_active_symmetric_memory_context = None
+
+# Reference to the C registration function (with arg types set)
+_register_func = None
+
+
+def is_symmetric_memory_enabled():
+    try:
+        return get_server_args().enable_symm_mem
+    except ValueError:
+        return False
+
+
+def set_graph_pool_id(graph_pool_id):
+    global _graph_pool_id
+    _graph_pool_id = graph_pool_id
+
+
+def disable_symmetric_memory_context():
+    if _active_symmetric_memory_context is None:
+        return None
+    saved_context = _active_symmetric_memory_context
+    saved_context.__exit__(None, None, None)
+    return saved_context
+
+
+def restore_symmetric_memory_context(saved_context):
+    if saved_context is not None:
+        saved_context.__enter__()
+
+
+def get_nccl_mem_pool() -> torch.cuda.MemPool:
+    """
+    Get the shared MemPool for all groups.
+
+    All groups share the same pool to avoid memory fragmentation.
+    Comm registration is handled at context exit time.
+    """
+    global _allocator, _mem_pool, _cur_device, _register_func
+    if _allocator is None:
+        import torch.utils.cpp_extension
+
+        out_dir = os.path.join(tempfile.gettempdir(), "symm_allocator")
+        os.makedirs(out_dir, exist_ok=True)
+        # Make sure to clean up leftover pytorch lock files
+        # from previous runs and synchronize across processes
+        # right after
+        try:
+            os.remove(os.path.join(out_dir, "lock"))
+        except FileNotFoundError:
+            pass
+        torch.distributed.barrier()
+
+        nccl_allocator_libname = "nccl_allocator"
+        lib_path = torch.utils.cpp_extension.load_inline(
+            name=nccl_allocator_libname,
+            cpp_sources=nccl_allocator_source,
+            with_cuda=True,
+            extra_ldflags=["-lnccl"],
+            verbose=True,
+            is_python_module=False,
+            build_directory=out_dir,
+        )
+        nccl_allocator_lib = ctypes.CDLL(lib_path)
+        _allocator = CUDAPluggableAllocator(
+            f"{out_dir}/{nccl_allocator_libname}.so",
+            "nccl_alloc_plug",
+            "nccl_free_plug",
+        ).allocator()
+        _mem_pool = torch.cuda.MemPool(_allocator)
+        _cur_device = torch.cuda.current_device()
+
+        # Setup the C function for registration with correct arg types
+        _register_func = nccl_allocator_lib.nccl_allocator_register_segments_with_comm
+        _register_func.restype = ctypes.c_int
+        _register_func.argtypes = [ctypes.c_uint64]
+
+    return _mem_pool
+
+
+class SymmetricMemoryContext:
+    """
+    Context manager for using symmetric memory with pynccl.
+
+    To Utilize the symmetric memory feature in NCCL, the buffers need to be allocated
+    by `ncclMemAlloc` and registered by `ncclCommWindowRegister`. Due to this, we introduce
+    this context manager. All tensors created under this context will be correctly
+    allocated and registered with a custom allocator.
+
+    Key design:
+    - All groups share a single MemPool to avoid memory fragmentation.
+    - At allocation time, ptrs are tracked but NOT registered with any comm.
+    - At context exit time, nccl_allocator_register_segments_with_comm is called
+      to register all tracked segments with the current comm. The C++ layer
+      tracks per-comm registration state using index-based tracking to avoid
+      re-registration of already-registered segments.
+    """
+
+    def __init__(
+        self,
+        group_coordinator: GroupCoordinator,
+    ):
+        self.group_coordinator = group_coordinator
+        self._pool_id = get_nccl_mem_pool().id
+        self._device_index = torch.cuda.current_device()
+        self.is_graph_capture = torch.cuda.is_current_stream_capturing()
+
+        # Get comm ptr for tracking registrations
+        # Use the comm pointer value as unique identifier
+        self._comm_ptr = self.group_coordinator.pynccl_comm.comm.value
+
+    def __enter__(self):
+        assert (
+            self.group_coordinator.pynccl_comm is not None
+        ), f"Symmetric memory requires pynccl to be enabled in group '{self.group_coordinator.unique_name}'"
+
+        if self.is_graph_capture:
+            assert (
+                _graph_pool_id is not None
+            ), "graph_pool_id is not set under graph capture"
+            # Pause graph memory pool to use symmetric memory with cuda graph
+            if after_2_8_0:
+                torch._C._cuda_endAllocateToPool(_cur_device, _graph_pool_id)
+            else:
+                torch._C._cuda_endAllocateCurrentStreamToPool(
+                    _cur_device, _graph_pool_id
+                )
+
+        _cuda_beginAllocateCurrentThreadToPool(self._device_index, self._pool_id)
+
+        global _active_symmetric_memory_context
+        _active_symmetric_memory_context = self
+
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        _cuda_endAllocateToPool(self._device_index, self._pool_id)
+        _cuda_releasePool(self._device_index, self._pool_id)
+        # Register all unregistered segments
+        # with the current comm
+        self._register_segments_for_comm()
+
+        if self.is_graph_capture:
+            if after_2_8_0:
+                torch._C._cuda_beginAllocateCurrentThreadToPool(
+                    _cur_device, _graph_pool_id
+                )
+            else:
+                torch._C._cuda_beginAllocateToPool(_cur_device, _graph_pool_id)
+
+        global _active_symmetric_memory_context
+        _active_symmetric_memory_context = None
+
+    def _register_segments_for_comm(self):
+        """
+        Register all tracked segments with the current comm.
+
+        Delegates to C++ layer which handles:
+        1. Tracking which segments have been registered with each comm
+        2. Only registering new segments (avoiding re-registration)
+        3. Thread-safe access to the segment registry
+        """
+
+        # Call C++ API to register all segments with this comm
+        # C++ layer tracks per-comm registration state internally
+        result = _register_func(self._comm_ptr)
+        assert (
+            result == 0
+        ), f"nccl_allocator_register_segments_with_comm failed with return code: {result}"
+
+
+def use_symmetric_memory(group_coordinator: GroupCoordinator, disabled: bool = False):
+    disabled = (
+        not is_symmetric_memory_enabled()
+        or disabled
+        or group_coordinator.world_size == 1
+    )
+    return SymmetricMemoryContext(group_coordinator) if not disabled else nullcontext()
+
+
+# --- Debug mode for symmetric memory validation ---
+
+_symm_mem_logger = logging.getLogger(__name__)
+_debug_seen_traces: set = set()
+
+
+def is_tensor_in_symmetric_mempool(tensor: torch.Tensor) -> bool:
+    """Check if a tensor's storage is allocated in the NCCL symmetric memory pool."""
+
+    if _mem_pool is None:
+        return False  # Pool not initialized
+
+    data_ptr = tensor.untyped_storage().data_ptr()
+
+    for segment in _mem_pool.snapshot():
+        for block in segment["blocks"]:
+            if block["address"] == data_ptr:
+                return True
+    return False
+
+
+def debug_check_symmetric_mempool(
+    group_coordinator: GroupCoordinator,
+    tensors: dict,
+    op_name: str,
+) -> None:
+    """
+    Debug check: verify that tensors passed to communication ops are allocated
+    in the NCCL symmetric memory pool.
+
+    Enabled by setting SGLANG_DEBUG_SYMM_MEM=1.
+    Only prints warnings on rank 0 and deduplicates identical stack traces.
+
+    Args:
+        tensors: dict mapping argument name to tensor
+                 (e.g. {"input": t1, "output": t2})
+        op_name: name of the communication operation being checked
+    """
+    if not envs.SGLANG_DEBUG_SYMM_MEM.get() or not is_symmetric_memory_enabled():
+        return
+
+    # Only print on rank 0
+    if not group_coordinator.is_first_rank:
+        return
+
+    bad_names = []
+    bad_details = []
+    for name, tensor in tensors.items():
+        if not is_tensor_in_symmetric_mempool(tensor):
+            bad_names.append(name)
+            bad_details.append(
+                f"  - '{name}' (data_ptr=0x{tensor.storage().data_ptr():x}, "
+                f"shape={list(tensor.shape)}, dtype={tensor.dtype})"
+            )
+
+    if bad_names:
+        traces = traceback.format_stack()
+        # Skip autotune stack traces
+        if any("_flashinfer_autotune" in trace for trace in traces):
+            return
+        stack = "".join(traces[:-1])
+        trace_key = f"{op_name}:{','.join(bad_names)}:{stack}"
+        if trace_key not in _debug_seen_traces:
+            _debug_seen_traces.add(trace_key)
+            _symm_mem_logger.warning(
+                "[SymmMem Debug] %s: %d tensor(s) are NOT in the "
+                "NCCL symmetric memory pool:\n%s\n"
+                "Stack trace:\n%s",
+                op_name,
+                len(bad_names),
+                "\n".join(bad_details),
+                stack,
+            )

python/sglang/srt/distributed/device_communicators/pynccl_wrapper.py

@@ -0,0 +1,567 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/distributed/device_communicators/pynccl.py
+
+# This file is a pure Python wrapper for the NCCL library.
+# The main purpose is to use NCCL combined with CUDA graph.
+# Before writing this script, we tried the following approach:
+# 1. We tried to use `cupy`, it calls NCCL correctly, but `cupy` itself
+#  often gets stuck when initializing the NCCL communicator.
+# 2. We tried to use `torch.distributed`, but `torch.distributed.all_reduce`
+#  contains many other potential cuda APIs, that are not allowed during
+#  capturing the CUDA graph. For further details, please check
+# https://discuss.pytorch.org/t/pytorch-cudagraph-with-nccl-operation-failed/ .
+#
+# Another rejected idea is to write a C/C++ binding for NCCL. It is usually
+# doable, but we often encounter issues related with nccl versions, and need
+# to switch between different versions of NCCL. See
+# https://github.com/NVIDIA/nccl/issues/1234 for more details.
+# A C/C++ binding is not flexible enough to handle this. It requires
+# recompilation of the code every time we want to switch between different
+# versions. This current implementation, with a **pure** Python wrapper, is
+# more flexible. We can easily switch between different versions of NCCL by
+# changing the environment variable `SGLANG_NCCL_SO_PATH`, or the `so_file`
+# variable in the code.
+
+import ctypes
+import logging
+import os
+import platform
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional
+
+import torch
+from torch.distributed import ReduceOp
+
+logger = logging.getLogger(__name__)
+
+
+def find_nccl_library() -> str:
+    """
+    We either use the library file specified by the `SGLANG_NCCL_SO_PATH`
+    environment variable, or we find the library file brought by PyTorch.
+    After importing `torch`, `libnccl.so.2`, `librccl.so.1` or `libmccl.so.2`
+    can be found by `ctypes` automatically.
+    """
+
+    # so_file can be set to None in sglang
+    so_file = os.environ.get("SGLANG_NCCL_SO_PATH", None)
+
+    # manually load the nccl library
+    if so_file:
+        logger.info(
+            "Found nccl from environment variable SGLANG_NCCL_SO_PATH=%s", so_file
+        )
+    else:
+        if torch.version.cuda is not None:
+            so_file = "libnccl.so.2"
+        elif torch.version.hip is not None:
+            so_file = "librccl.so.1"
+        elif hasattr(torch.version, "musa") and torch.version.musa is not None:
+            so_file = "libmccl.so.2"
+        else:
+            raise ValueError("NCCL only supports CUDA, ROCm and MUSA backends.")
+        logger.debug("Found nccl from library %s", so_file)
+    return so_file
+
+
+# === export types and functions from nccl to Python ===
+# for the original nccl definition, please check
+# https://github.com/NVIDIA/nccl/blob/master/src/nccl.h.in
+
+ncclResult_t = ctypes.c_int
+ncclComm_t = ctypes.c_void_p
+ncclWindow_t = ctypes.c_void_p
+
+
+class ncclUniqueId(ctypes.Structure):
+    _fields_ = [("internal", ctypes.c_byte * 128)]
+
+
+cudaStream_t = ctypes.c_void_p
+buffer_type = ctypes.c_void_p
+
+ncclDataType_t = ctypes.c_int
+
+
+class ncclDataTypeEnum:
+    ncclInt8 = 0
+    ncclChar = 0
+    ncclUint8 = 1
+    ncclInt32 = 2
+    ncclInt = 2
+    ncclUint32 = 3
+    ncclInt64 = 4
+    ncclUint64 = 5
+    ncclFloat16 = 6
+    ncclHalf = 6
+    ncclFloat32 = 7
+    ncclFloat = 7
+    ncclFloat64 = 8
+    ncclDouble = 8
+    ncclBfloat16 = 9
+    ncclNumTypes = 10
+
+    @classmethod
+    def from_torch(cls, dtype: torch.dtype) -> int:
+        if dtype == torch.int8:
+            return cls.ncclInt8
+        if dtype == torch.uint8:
+            return cls.ncclUint8
+        if dtype == torch.int32:
+            return cls.ncclInt32
+        if dtype == torch.int64:
+            return cls.ncclInt64
+        if dtype == torch.float16:
+            return cls.ncclFloat16
+        if dtype == torch.float32:
+            return cls.ncclFloat32
+        if dtype == torch.float64:
+            return cls.ncclFloat64
+        if dtype == torch.bfloat16:
+            return cls.ncclBfloat16
+        raise ValueError(f"Unsupported dtype: {dtype}")
+
+
+ncclRedOp_t = ctypes.c_int
+
+
+class ncclRedOpTypeEnum:
+    ncclSum = 0
+    ncclProd = 1
+    ncclMax = 2
+    ncclMin = 3
+    ncclAvg = 4
+    ncclNumOps = 5
+
+    @classmethod
+    def from_torch(cls, op: ReduceOp) -> int:
+        if op == ReduceOp.SUM:
+            return cls.ncclSum
+        if op == ReduceOp.PRODUCT:
+            return cls.ncclProd
+        if op == ReduceOp.MAX:
+            return cls.ncclMax
+        if op == ReduceOp.MIN:
+            return cls.ncclMin
+        if op == ReduceOp.AVG:
+            return cls.ncclAvg
+        raise ValueError(f"Unsupported op: {op}")
+
+
+@dataclass
+class Function:
+    name: str
+    restype: Any
+    argtypes: List[Any]
+
+
+class NCCLLibrary:
+    exported_functions = [
+        # const char* ncclGetErrorString(ncclResult_t result)
+        Function("ncclGetErrorString", ctypes.c_char_p, [ncclResult_t]),
+        # ncclResult_t  ncclGetVersion(int *version);
+        Function("ncclGetVersion", ncclResult_t, [ctypes.POINTER(ctypes.c_int)]),
+        # ncclResult_t ncclGetUniqueId(ncclUniqueId* uniqueId);
+        Function("ncclGetUniqueId", ncclResult_t, [ctypes.POINTER(ncclUniqueId)]),
+        # ncclResult_t  ncclCommInitRank(
+        #   ncclComm_t* comm, int nranks, ncclUniqueId commId, int rank);
+        # note that ncclComm_t is a pointer type, so the first argument
+        # is a pointer to a pointer
+        Function(
+            "ncclCommInitRank",
+            ncclResult_t,
+            [ctypes.POINTER(ncclComm_t), ctypes.c_int, ncclUniqueId, ctypes.c_int],
+        ),
+        # ncclResult_t  ncclAllReduce(
+        #   const void* sendbuff, void* recvbuff, size_t count,
+        #   ncclDataType_t datatype, ncclRedOp_t op, ncclComm_t comm,
+        #   cudaStream_t stream);
+        # note that cudaStream_t is a pointer type, so the last argument
+        # is a pointer
+        Function(
+            "ncclAllReduce",
+            ncclResult_t,
+            [
+                buffer_type,
+                buffer_type,
+                ctypes.c_size_t,
+                ncclDataType_t,
+                ncclRedOp_t,
+                ncclComm_t,
+                cudaStream_t,
+            ],
+        ),
+        # ncclResult_t  ncclAllGather(
+        #   const void* sendbuff, void* recvbuff, size_t count,
+        #   ncclDataType_t datatype, ncclComm_t comm,
+        #   cudaStream_t stream);
+        # note that cudaStream_t is a pointer type, so the last argument
+        # is a pointer
+        Function(
+            "ncclAllGather",
+            ncclResult_t,
+            [
+                buffer_type,
+                buffer_type,
+                ctypes.c_size_t,
+                ncclDataType_t,
+                ncclComm_t,
+                cudaStream_t,
+            ],
+        ),
+        # ncclResult_t  ncclReduce(
+        #   const void* sendbuff, void* recvbuff, size_t count,
+        #   ncclDataType_t datatype, ncclRedOp_t op, int root,
+        #   ncclComm_t comm,  cudaStream_t stream);
+        # note that cudaStream_t is a pointer type, so the last argument
+        # is a pointer
+        Function(
+            "ncclReduce",
+            ncclResult_t,
+            [
+                buffer_type,
+                buffer_type,
+                ctypes.c_size_t,
+                ncclDataType_t,
+                ncclRedOp_t,
+                ctypes.c_int,
+                ncclComm_t,
+                cudaStream_t,
+            ],
+        ),
+        # ncclResult_t  ncclReduceScatter(
+        #   const void* sendbuff, void* recvbuff, size_t count,
+        #   ncclDataType_t datatype, ncclRedOp_t op, ncclComm_t comm,
+        #   cudaStream_t stream);
+        # note that cudaStream_t is a pointer type, so the last argument
+        # is a pointer
+        Function(
+            "ncclReduceScatter",
+            ncclResult_t,
+            [
+                buffer_type,
+                buffer_type,
+                ctypes.c_size_t,
+                ncclDataType_t,
+                ncclRedOp_t,
+                ncclComm_t,
+                cudaStream_t,
+            ],
+        ),
+        # ncclResult_t  ncclSend(
+        #   const void* sendbuff, size_t count, ncclDataType_t datatype,
+        #   int dest, ncclComm_t comm, cudaStream_t stream);
+        Function(
+            "ncclSend",
+            ncclResult_t,
+            [
+                buffer_type,
+                ctypes.c_size_t,
+                ncclDataType_t,
+                ctypes.c_int,
+                ncclComm_t,
+                cudaStream_t,
+            ],
+        ),
+        # ncclResult_t  ncclRecv(
+        #   void* recvbuff, size_t count, ncclDataType_t datatype,
+        #   int src, ncclComm_t comm, cudaStream_t stream);
+        Function(
+            "ncclRecv",
+            ncclResult_t,
+            [
+                buffer_type,
+                ctypes.c_size_t,
+                ncclDataType_t,
+                ctypes.c_int,
+                ncclComm_t,
+                cudaStream_t,
+            ],
+        ),
+        # ncclResult_t ncclBroadcast(
+        #   const void* sendbuff, void* recvbuff, size_t count,
+        #   ncclDataType_t datatype, int root, ncclComm_t comm,
+        #   cudaStream_t stream);
+        Function(
+            "ncclBroadcast",
+            ncclResult_t,
+            [
+                buffer_type,
+                buffer_type,
+                ctypes.c_size_t,
+                ncclDataType_t,
+                ctypes.c_int,
+                ncclComm_t,
+                cudaStream_t,
+            ],
+        ),
+        # be cautious! this is a collective call, it will block until all
+        # processes in the communicator have called this function.
+        # because Python object destruction can happen in random order,
+        # it is better not to call it at all.
+        # ncclResult_t  ncclCommDestroy(ncclComm_t comm);
+        Function("ncclCommDestroy", ncclResult_t, [ncclComm_t]),
+        # ncclResult_t ncclGroupStart();
+        Function("ncclGroupStart", ncclResult_t, []),
+        # ncclResult_t ncclGroupEnd();
+        Function("ncclGroupEnd", ncclResult_t, []),
+    ]
+
+    exported_functions_symm_mem = [
+        # ncclResult_t ncclCommWindowRegister(ncclComm_t comm, void* buff, size_t size, ncclWindow_t* win, int winFlags);
+        Function(
+            "ncclCommWindowRegister",
+            ncclResult_t,
+            [
+                ncclComm_t,
+                buffer_type,
+                ctypes.c_size_t,
+                ctypes.POINTER(ncclWindow_t),
+                ctypes.c_int,
+            ],
+        ),
+        # ncclResult_t ncclCommWindowDeregister(ncclComm_t comm, ncclWindow_t win);
+        Function("ncclCommWindowDeregister", ncclResult_t, [ncclComm_t, ncclWindow_t]),
+    ]
+
+    # class attribute to store the mapping from the path to the library
+    # to avoid loading the same library multiple times
+    path_to_library_cache: Dict[str, Any] = {}
+
+    # class attribute to store the mapping from library path
+    #  to the corresponding dictionary
+    path_to_dict_mapping: Dict[str, Dict[str, Any]] = {}
+
+    def __init__(self, so_file: Optional[str] = None):
+
+        so_file = so_file or find_nccl_library()
+
+        try:
+            if so_file not in NCCLLibrary.path_to_dict_mapping:
+                lib = ctypes.CDLL(so_file)
+                NCCLLibrary.path_to_library_cache[so_file] = lib
+            self.lib = NCCLLibrary.path_to_library_cache[so_file]
+        except Exception as e:
+            logger.error(
+                "Failed to load NCCL library from %s . "
+                "It is expected if you are not running on NVIDIA/AMD/MTHREADS GPUs. "
+                "Otherwise, the nccl library might not exist, be corrupted "
+                "or it does not support the current platform %s. "
+                "If you already have the library, please set the "
+                "environment variable SGLANG_NCCL_SO_PATH"
+                " to point to the correct nccl library path.",
+                so_file,
+                platform.platform(),
+            )
+            raise e
+
+        if so_file not in NCCLLibrary.path_to_dict_mapping:
+            _funcs: Dict[str, Any] = {}
+            exported_functions = NCCLLibrary.exported_functions
+            if hasattr(self.lib, "ncclCommWindowRegister"):
+                exported_functions.extend(NCCLLibrary.exported_functions_symm_mem)
+            for func in exported_functions:
+                f = getattr(self.lib, func.name)
+                f.restype = func.restype
+                f.argtypes = func.argtypes
+                _funcs[func.name] = f
+            NCCLLibrary.path_to_dict_mapping[so_file] = _funcs
+        self._funcs = NCCLLibrary.path_to_dict_mapping[so_file]
+
+    def ncclGetErrorString(self, result: ncclResult_t) -> str:
+        return self._funcs["ncclGetErrorString"](result).decode("utf-8")
+
+    def NCCL_CHECK(self, result: ncclResult_t) -> None:
+        if result != 0:
+            error_str = self.ncclGetErrorString(result)
+            raise RuntimeError(f"NCCL error: {error_str}")
+
+    def ncclGetRawVersion(self) -> int:
+        version = ctypes.c_int()
+        self.NCCL_CHECK(self._funcs["ncclGetVersion"](ctypes.byref(version)))
+        # something like 21903
+        return version.value
+
+    def ncclGetVersion(self) -> str:
+        version_str = str(self.ncclGetRawVersion())
+        # something like 21903 --> "2.19.3"
+        major = version_str[0].lstrip("0")
+        minor = version_str[1:3].lstrip("0")
+        patch = version_str[3:].lstrip("0")
+        return f"{major}.{minor}.{patch}"
+
+    def ncclGetUniqueId(self) -> ncclUniqueId:
+        unique_id = ncclUniqueId()
+        self.NCCL_CHECK(self._funcs["ncclGetUniqueId"](ctypes.byref(unique_id)))
+        return unique_id
+
+    def ncclCommInitRank(
+        self, world_size: int, unique_id: ncclUniqueId, rank: int
+    ) -> ncclComm_t:
+        comm = ncclComm_t()
+        self.NCCL_CHECK(
+            self._funcs["ncclCommInitRank"](
+                ctypes.byref(comm), world_size, unique_id, rank
+            )
+        )
+        return comm
+
+    def ncclAllReduce(
+        self,
+        sendbuff: buffer_type,
+        recvbuff: buffer_type,
+        count: int,
+        datatype: int,
+        op: int,
+        comm: ncclComm_t,
+        stream: cudaStream_t,
+    ) -> None:
+        # `datatype` actually should be `ncclDataType_t`
+        # and `op` should be `ncclRedOp_t`
+        # both are aliases of `ctypes.c_int`
+        # when we pass int to a function, it will be converted to `ctypes.c_int`
+        # by ctypes automatically
+        self.NCCL_CHECK(
+            self._funcs["ncclAllReduce"](
+                sendbuff, recvbuff, count, datatype, op, comm, stream
+            )
+        )
+
+    def ncclReduce(
+        self,
+        sendbuff: buffer_type,
+        recvbuff: buffer_type,
+        count: int,
+        datatype: int,
+        op: int,
+        root: int,
+        comm: ncclComm_t,
+        stream: cudaStream_t,
+    ) -> None:
+        # `datatype` actually should be `ncclDataType_t`
+        # and `op` should be `ncclRedOp_t`
+        # both are aliases of `ctypes.c_int`
+        # when we pass int to a function, it will be converted to `ctypes.c_int`
+        # by ctypes automatically
+        self.NCCL_CHECK(
+            self._funcs["ncclReduce"](
+                sendbuff, recvbuff, count, datatype, op, root, comm, stream
+            )
+        )
+
+    def ncclReduceScatter(
+        self,
+        sendbuff: buffer_type,
+        recvbuff: buffer_type,
+        count: int,
+        datatype: int,
+        op: int,
+        comm: ncclComm_t,
+        stream: cudaStream_t,
+    ) -> None:
+        # `datatype` actually should be `ncclDataType_t`
+        # and `op` should be `ncclRedOp_t`
+        # both are aliases of `ctypes.c_int`
+        # when we pass int to a function, it will be converted to `ctypes.c_int`
+        # by ctypes automatically
+        self.NCCL_CHECK(
+            self._funcs["ncclReduceScatter"](
+                sendbuff, recvbuff, count, datatype, op, comm, stream
+            )
+        )
+
+    def ncclAllGather(
+        self,
+        sendbuff: buffer_type,
+        recvbuff: buffer_type,
+        count: int,
+        datatype: int,
+        comm: ncclComm_t,
+        stream: cudaStream_t,
+    ) -> None:
+        # `datatype` actually should be `ncclDataType_t`
+        # which is an aliases of `ctypes.c_int`
+        # when we pass int to a function, it will be converted to `ctypes.c_int`
+        # by ctypes automatically
+        self.NCCL_CHECK(
+            self._funcs["ncclAllGather"](
+                sendbuff, recvbuff, count, datatype, comm, stream
+            )
+        )
+
+    def ncclSend(
+        self,
+        sendbuff: buffer_type,
+        count: int,
+        datatype: int,
+        dest: int,
+        comm: ncclComm_t,
+        stream: cudaStream_t,
+    ) -> None:
+        self.NCCL_CHECK(
+            self._funcs["ncclSend"](sendbuff, count, datatype, dest, comm, stream)
+        )
+
+    def ncclRecv(
+        self,
+        recvbuff: buffer_type,
+        count: int,
+        datatype: int,
+        src: int,
+        comm: ncclComm_t,
+        stream: cudaStream_t,
+    ) -> None:
+        self.NCCL_CHECK(
+            self._funcs["ncclRecv"](recvbuff, count, datatype, src, comm, stream)
+        )
+
+    def ncclBroadcast(
+        self,
+        sendbuff: buffer_type,
+        recvbuff: buffer_type,
+        count: int,
+        datatype: int,
+        root: int,
+        comm: ncclComm_t,
+        stream: cudaStream_t,
+    ) -> None:
+        self.NCCL_CHECK(
+            self._funcs["ncclBroadcast"](
+                sendbuff, recvbuff, count, datatype, root, comm, stream
+            )
+        )
+
+    def ncclCommDestroy(self, comm: ncclComm_t) -> None:
+        self.NCCL_CHECK(self._funcs["ncclCommDestroy"](comm))
+
+    def ncclCommWindowRegister(
+        self, comm: ncclComm_t, buff: buffer_type, size: int, win_flags: int
+    ) -> ncclWindow_t:
+        window = ncclWindow_t()
+        self.NCCL_CHECK(
+            self._funcs["ncclCommWindowRegister"](
+                comm, buff, size, ctypes.byref(window), win_flags
+            )
+        )
+        return window
+
+    def ncclCommWindowDeregister(self, comm: ncclComm_t, window: ncclWindow_t) -> None:
+        self.NCCL_CHECK(self._funcs["ncclCommWindowDeregister"](comm, window))
+
+    def ncclGroupStart(self) -> None:
+        self.NCCL_CHECK(self._funcs["ncclGroupStart"]())
+
+    def ncclGroupEnd(self) -> None:
+        self.NCCL_CHECK(self._funcs["ncclGroupEnd"]())
+
+
+__all__ = [
+    "NCCLLibrary",
+    "ncclDataTypeEnum",
+    "ncclRedOpTypeEnum",
+    "ncclUniqueId",
+    "ncclComm_t",
+    "cudaStream_t",
+    "buffer_type",
+]

python/sglang/srt/distributed/device_communicators/quick_all_reduce.py

@@ -0,0 +1,267 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+import os
+from enum import Enum
+from functools import cache
+from typing import Union
+
+import torch
+import torch.distributed as dist
+from torch.distributed import ProcessGroup
+
+import sglang.srt.distributed.device_communicators.custom_all_reduce_ops as ops
+from sglang.srt.distributed.device_communicators.custom_all_reduce_utils import (
+    is_full_nvlink,
+    is_weak_contiguous,
+)
+from sglang.srt.distributed.parallel_state import in_the_same_node_as
+from sglang.srt.utils import is_cuda, is_hip
+
+logger = logging.getLogger(__name__)
+
+_is_cuda = is_cuda()
+_is_hip = is_hip()
+
+
+@cache
+def qr_rocm_arch_available():
+    if not _is_hip:
+        return False
+    try:
+        props = torch.cuda.get_device_properties(0)
+        gcn_arch = getattr(props, "gcnArchName", "")
+        supported_archs = ["gfx94", "gfx95"]
+        return any(gfx in gcn_arch for gfx in supported_archs)
+    except Exception as e:
+        logger.warning("Failed to determine ROCm for quick allreduce: %s", e)
+        return False
+
+
+class QuickReduceRegime(Enum):
+    FP = 0
+    INT8 = 1
+    INT6 = 2
+    INT4 = 3
+    NONE = 4
+
+
+MB = 1024 * 1024
+
+
+class QuickAllReduce:
+
+    _SUPPORTED_WORLD_SIZES = [2, 4, 8]
+    _SUPPORTED_DTYPES = [torch.float16, torch.bfloat16]
+    # The following data is based on kernel tests.
+    # In this order [FP, INT8, INT6, INT4].
+    _QR_MIN_SIZE = {
+        (torch.float16, 2): [1 * MB, 2 * MB, 2 * MB, 1 * MB],
+        (torch.float16, 4): [1 * MB, 16 * MB, 4 * MB, 2 * MB],
+        (torch.float16, 8): [16 * MB, 4 * MB, 4 * MB, 2 * MB],
+        (torch.bfloat16, 2): [2 * MB, 8 * MB, 8 * MB, 8 * MB],
+        (torch.bfloat16, 4): [8 * MB, 64 * MB, 64 * MB, 16 * MB],
+        (torch.bfloat16, 8): [16 * MB, 2048 * MB, 2048 * MB, 2048 * MB],
+    }
+
+    def __init__(
+        self, group: ProcessGroup, device: Union[int, str, torch.device]
+    ) -> None:
+        """
+        Custom allreduce provides non-destructive acceleration and is
+        available for CUDA and ROCm MI300 series.
+        Custom quick allreduce leverages quantization for further
+        acceleration on ROCm. It currently supports Q8, Q6, and Q4
+        quantization formats and FP(float16, bfloat16).
+        Quick allreduce is designed as a complement to custom allreduce.
+        Its initialization requires even stricter conditions.
+        Only the ROCm MI300 series is supported for quick allreduce at
+        this time.
+        Args:
+            group: the process group to work on. If None, it will use the
+                default process group.
+            device: the device to bind the CustomAllreduce to. If None,
+                it will be bind to f"cuda:{local_rank}".
+        It is the caller's responsibility to make sure each communicator
+        is bind to a unique device, and all communicators in this group
+        are in the same node.
+        """
+        self.disabled = True
+        if not qr_rocm_arch_available():
+            logger.debug(
+                "Custom quick allreduce is only supported on ROCm MI300 series."
+            )
+            return
+
+        if not ops.IS_QUICK_AR_AVAILABLE:
+            # disable because of missing quick reduce library
+            # e.g. in a cuda environment
+            logger.info(
+                "Custom quick allreduce is disabled because "
+                "of missing custom quick allreduce library"
+            )
+            return
+
+        self.group = group
+        assert (
+            dist.get_backend(group) != dist.Backend.NCCL
+        ), "Custom quick allreduce should be attached to a non-NCCL group."
+        if not all(in_the_same_node_as(group, source_rank=0)):
+            # No need to initialize custom quick allreduce for
+            # multi-node case.
+            logger.warning(
+                "Custom quick allreduce is disabled because this "
+                "process group spans across nodes."
+            )
+            return
+        rank = dist.get_rank(group=self.group)
+        world_size = dist.get_world_size(group=self.group)
+        self.rank = rank
+        self.world_size = world_size
+        if world_size == 1:
+            # No need to initialize QuickReduce for single GPU case.
+            return
+
+        if world_size not in QuickAllReduce._SUPPORTED_WORLD_SIZES:
+            logger.warning(
+                "Custom quick allreduce is disabled due to an "
+                "unsupported world size: %d. Supported world sizes: %s.",
+                world_size,
+                str(QuickAllReduce._SUPPORTED_WORLD_SIZES),
+            )
+            return
+
+        if isinstance(device, int):
+            device = torch.device(f"cuda:{device}")
+        elif isinstance(device, str):
+            device = torch.device(device)
+        assert isinstance(device, torch.device)
+        self.device = device
+
+        cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES", None)
+        if cuda_visible_devices:
+            device_ids = list(map(int, cuda_visible_devices.split(",")))
+        else:
+            device_ids = list(range(torch.cuda.device_count()))
+        physical_device_id = device_ids[device.index]
+        tensor = torch.tensor([physical_device_id], dtype=torch.int, device="cpu")
+        gather_list = [
+            torch.tensor([0], dtype=torch.int, device="cpu")
+            for _ in range(self.world_size)
+        ]
+        dist.all_gather(gather_list, tensor, group=self.group)
+        physical_device_ids = [t.item() for t in gather_list]
+
+        # test nvlink first, this will filter out most of the cases
+        # where custom quick allreduce is not supported
+        # this checks hardware and driver support for NVLink
+        if _is_cuda or _is_hip:
+            self.fully_connected = is_full_nvlink(physical_device_ids, self.world_size)
+        if self.world_size > 2 and not self.fully_connected:
+            logger.debug(
+                "Custom quick allreduce is disabled because it's not supported "
+                "on more than two PCIe-only GPUs. "
+            )
+            return
+
+        self.init_quick_all_reduce()
+
+    def init_quick_all_reduce(self):
+        # On RocM, bfloat16 kernels are slower than fp16
+        # due to slower match operations
+        # If environment variable is set to 1, we convert input to fp16
+        self.use_fp16_kernels = int(
+            os.environ.get("ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16", 1)
+        )
+        regime_str = os.environ.get("ROCM_QUICK_REDUCE_QUANTIZATION", "NONE")
+        if regime_str not in QuickReduceRegime.__members__:
+            logger.warning(
+                "Custom quick allreduce:",
+                f"Invalid quantization level: {regime_str}. "
+                "Supported levels: "
+                f"{list(QuickReduceRegime.__members__.keys())}",
+            )
+            return
+
+        if regime_str == "NONE":
+            logger.debug(
+                "Custom quick allreduce is disabled based "
+                "on env variable "
+                "ROCM_QUICK_REDUCE_QUANTIZATION='NONE'"
+            )
+            return
+        self.qr_quant_level = QuickReduceRegime[regime_str]
+
+        # TODO: If the dtype is not bfloat16 or then float16,
+        # quickallreduce should not be created.
+
+        # ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB is specified in MB
+        qr_max_size = int(os.environ.get("ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB", 0))
+        if qr_max_size > 0:
+            if qr_max_size < 1:
+                logger.info(
+                    "You should not set a max_size smaller than 1MB, which can "
+                    "lead to error or degradation to custom allreduce or rccl."
+                )
+            qr_max_size = qr_max_size * MB
+        # If qr_max_size is None, then 2GB is used by default.
+        self._ptr = ops.init_custom_qr(self.rank, self.world_size, qr_max_size)
+        self.qr_max_size = qr_max_size if qr_max_size > 0 else ops.qr_max_size()
+        self.create_shared_buffer()
+        self.disabled = False
+
+    def create_shared_buffer(self):
+        """
+        Creates a shared buffer for quickreduce.
+        Has to be called after init_custom_qr
+        """
+        handle = ops.qr_get_handle(self._ptr)
+        world_size = dist.get_world_size(group=self.group)
+        handles = [None] * world_size
+        dist.all_gather_object(handles, handle, group=self.group)
+        ops.qr_open_handles(self._ptr, handles)
+
+    def should_quick_allreduce(self, inp: torch.Tensor):
+        """
+        Check if quickreduce is available
+        """
+        if self.disabled:
+            return False
+        if inp.dtype not in self._SUPPORTED_DTYPES:
+            return False
+        inp_size = inp.numel() * inp.element_size()
+        # custom quick allreduce requires input byte size to be
+        # multiples of 16
+        if inp_size % 16 != 0:
+            return False
+        if not is_weak_contiguous(inp):
+            return False
+        dtype = inp.dtype
+        if self.use_fp16_kernels:
+            dtype = torch.float16
+        return (
+            inp_size <= self.qr_max_size
+            and inp_size
+            >= self._QR_MIN_SIZE[(dtype, self.world_size)][self.qr_quant_level.value]
+        )
+
+    def quick_all_reduce(self, inp: torch.Tensor, *, out: torch.Tensor = None):
+        """Performs an out-of-place custom quick all reduce."""
+        # quick allreduce doesn't require a separate graph mode,
+        # as QR uses static IPC buffer.
+        if out is None:
+            out = torch.empty_like(inp)
+        ops.qr_all_reduce(
+            self._ptr, inp, out, self.qr_quant_level.value, self.use_fp16_kernels
+        )
+        return out
+
+    def close(self):
+        if not self.disabled and getattr(self, "_ptr", None):
+            if ops is not None:
+                ops.qr_destroy(self._ptr)
+            self._ptr = 0
+            self.disabled = True
+
+    def __del__(self):
+        self.close()

python/sglang/srt/distributed/device_communicators/shm_broadcast.py

@@ -0,0 +1,519 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/distributed/device_communicators/shm_broadcast.py
+
+import logging
+import os
+import pickle
+import time
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from multiprocessing import shared_memory
+from typing import List, Optional
+from unittest.mock import patch
+
+import torch
+import torch.distributed as dist
+from torch.distributed import ProcessGroup
+from zmq import IPV6  # type: ignore
+from zmq import SUB, SUBSCRIBE, XPUB, XPUB_VERBOSE, Context  # type: ignore
+
+from sglang.srt.environ import envs
+from sglang.srt.utils.network import NetworkAddress, get_local_ip_auto, get_open_port
+from sglang.srt.utils.stale_shm_cleanup import make_shm_name
+
+# SGLANG_RINGBUFFER_WARNING_INTERVAL can be set to 60
+SGLANG_RINGBUFFER_WARNING_INTERVAL = int(
+    os.environ.get("SGLANG_RINGBUFFER_WARNING_INTERVAL", "60")
+)
+
+logger = logging.getLogger(__name__)
+
+
+class ShmRingBuffer:
+
+    def __init__(
+        self,
+        n_reader: int,
+        max_chunk_bytes: int,
+        max_chunks: int,
+        name: Optional[str] = None,
+    ):
+        """
+        A shared memory ring buffer implementation for broadcast communication.
+        Essentially, it is a queue where only one will `enqueue` and multiple
+        will `dequeue`. The max size of each item, together with the max number
+        of items that can be stored in the buffer are known in advance.
+        In this case, we don't need to synchronize the access to
+         the buffer.
+
+        Buffer memory layout:
+                  data                                 metadata
+                    |                                      |
+                    | (current_idx)                        | (current_idx)
+                    v                                      v
+        +-------------------------------+----------------------------------------+
+        | chunk0 | chunk1 | ... | chunk | metadata0 | metadata1 | ... | metadata |
+        +-------------------------------+----------------------------------------+
+        | max_chunks x max_chunk_bytes  | max_chunks x (1 + n_reader) bytes      |
+
+        metadata memory layout: each byte is a flag, the first byte is the written
+        flag, and the rest are reader flags. The flags are set to 0 by default.
+        +--------------+--------------+--------------+-----+--------------+
+        | written_flag | reader0_flag | reader1_flag | ... | readerN_flag |
+        +--------------+--------------+--------------+-----+--------------+
+
+        The state of metadata is as follows:
+
+        (case 1) 0???...???: the block is not written yet, cannot read, can write
+        (case 2) 1000...000: the block is just written, can read, cannot write
+        (case 3) 1???...???: the block is written and read by some readers, can read if not read, cannot write
+        (case 4) 1111...111: the block is written and read by all readers, cannot read, can write
+
+        State transition for readers:
+
+        When a reader finds a block that it can read (case 2 or 3), it can yield the block for caller to read.
+        Only after the caller finishes reading the block, the reader can mark the block as read.
+        Readers only mark the block as read (from 0 to 1), the writer marks the block as ready to read (from 1 to 0).
+
+        State transition for writer:
+
+        When the writer writes to a block (case 1 or 4), it first resets the written flag to 0, converting either case
+        to case 1. Then it can yield the block for caller to write. After the caller finishes writing the block, the writer
+        can reset the reader flags to 0, and mark the block as written (from 0 to 1).
+        NOTE: the order is important here, first reset the reader flags (so that we are still in case 1), then mark the block as written. The state transition is atomic. If we do it in the reverse order, it will go through case 3 and then back to case 2, and readers might read the intermediate case 3, which is not correct.
+
+        During creation, `name` is None and the buffer is created. We can pass the
+        created object to other processes by pickling it. The other processes will
+        get the name of the shared memory and open it, so that they can access the
+        same shared memory buffer.
+        """  # noqa
+        self.n_reader = n_reader
+        self.metadata_size = 1 + n_reader
+        self.max_chunk_bytes = max_chunk_bytes
+        self.max_chunks = max_chunks
+        self.total_bytes_of_buffer = (
+            self.max_chunk_bytes + self.metadata_size
+        ) * self.max_chunks
+        self.data_offset = 0
+        self.metadata_offset = self.max_chunk_bytes * self.max_chunks
+
+        if name is None:
+            # we are creating a buffer
+            self.is_creator = True
+            self.shared_memory = shared_memory.SharedMemory(
+                create=True,
+                size=self.total_bytes_of_buffer,
+                name=make_shm_name("mq"),
+            )
+            # initialize the metadata section to 0
+            with memoryview(
+                self.shared_memory.buf[self.metadata_offset :]
+            ) as metadata_buffer:
+                torch.frombuffer(metadata_buffer, dtype=torch.uint8).fill_(0)
+        else:
+            # we are opening an existing buffer
+            self.is_creator = False
+            # fix to https://stackoverflow.com/q/62748654/9191338
+            # Python incorrectly tracks shared memory even if it is not
+            # created by the process. The following patch is a workaround.
+            with patch(
+                "multiprocessing.resource_tracker.register",
+                lambda *args, **kwargs: None,
+            ):
+                try:
+                    self.shared_memory = shared_memory.SharedMemory(name=name)
+                    assert self.shared_memory.size == self.total_bytes_of_buffer
+                except FileNotFoundError:
+                    # we might deserialize the object in a different node
+                    # in this case, this object is not used,
+                    # and we should suppress the error
+                    pass
+
+    def __reduce__(self):
+        return (
+            self.__class__,
+            (
+                self.n_reader,
+                self.max_chunk_bytes,
+                self.max_chunks,
+                self.shared_memory.name,
+            ),
+        )
+
+    def __del__(self):
+        if hasattr(self, "shared_memory"):
+            self.shared_memory.close()
+            if self.is_creator:
+                self.shared_memory.unlink()
+
+    @contextmanager
+    def get_data(self, current_idx: int):
+        start = self.data_offset + current_idx * self.max_chunk_bytes
+        end = start + self.max_chunk_bytes
+        with memoryview(self.shared_memory.buf[start:end]) as buf:
+            yield buf
+
+    @contextmanager
+    def get_metadata(self, current_idx: int):
+        start = self.metadata_offset + current_idx * self.metadata_size
+        end = start + self.metadata_size
+        with memoryview(self.shared_memory.buf[start:end]) as buf:
+            yield buf
+
+
+@dataclass
+class Handle:
+    connect_ip: str
+    local_reader_ranks: List[int] = field(default_factory=list)
+
+    buffer: Optional[ShmRingBuffer] = None
+    local_subscribe_port: Optional[int] = None
+    remote_subscribe_port: Optional[int] = None
+
+
+class MessageQueue:
+
+    def __init__(
+        self,
+        n_reader,  # number of all readers
+        n_local_reader,  # number of local readers through shared memory
+        local_reader_ranks: Optional[List[int]] = None,
+        max_chunk_bytes: int = 1024 * 1024 * 10,
+        max_chunks: int = 10,
+        connect_ip: Optional[str] = None,
+    ):
+        if local_reader_ranks is None:
+            local_reader_ranks = list(range(n_local_reader))
+        else:
+            assert len(local_reader_ranks) == n_local_reader
+        self.n_local_reader = n_local_reader
+        n_remote_reader = n_reader - n_local_reader
+        self.n_remote_reader = n_remote_reader
+
+        if connect_ip is None:
+            connect_ip = (
+                get_local_ip_auto("0.0.0.0") if n_remote_reader > 0 else "127.0.0.1"
+            )
+
+        context = Context()
+
+        if n_local_reader > 0:
+            # for local readers, we will:
+            # 1. create a shared memory ring buffer to communicate small data
+            # 2. create a publish-subscribe socket to communicate large data
+            self.buffer = ShmRingBuffer(n_local_reader, max_chunk_bytes, max_chunks)
+
+            # XPUB is very similar to PUB,
+            # except that it can receive subscription messages
+            # to confirm the number of subscribers
+            self.local_socket = context.socket(XPUB)
+            # set the verbose option so that we can receive every subscription
+            # message. otherwise, we will only receive the first subscription
+            # see http://api.zeromq.org/3-3:zmq-setsockopt for more details
+            self.local_socket.setsockopt(XPUB_VERBOSE, True)
+            # Bind atomically to avoid get_open_port()'s check-then-bind race;
+            # search from SGLANG_PORT to keep the existing port range.
+            sglang_port = envs.SGLANG_PORT.get()
+            if sglang_port is not None:
+                local_subscribe_port = self.local_socket.bind_to_random_port(
+                    "tcp://127.0.0.1", min_port=sglang_port, max_port=sglang_port + 8
+                )
+            else:
+                local_subscribe_port = self.local_socket.bind_to_random_port(
+                    "tcp://127.0.0.1"
+                )
+            logger.debug("Bound to tcp://127.0.0.1:%d", local_subscribe_port)
+            self.current_idx = 0
+
+        else:
+            self.buffer = None  # type: ignore
+            local_subscribe_port = None
+            self.local_socket = None
+            self.current_idx = -1
+
+        if n_remote_reader > 0:
+            # for remote readers, we will:
+            # create a publish-subscribe socket to communicate large data
+            self.remote_socket = context.socket(XPUB)
+            self.remote_socket.setsockopt(XPUB_VERBOSE, True)
+            remote_subscribe_port = get_open_port()
+            na = NetworkAddress(connect_ip, remote_subscribe_port)
+            if na.is_ipv6:
+                self.remote_socket.setsockopt(IPV6, 1)
+            address = na.to_tcp()
+            logger.debug(f"class MessageQueue: Binding remote socket to {address=}")
+            self.remote_socket.bind(address)
+
+        else:
+            remote_subscribe_port = None
+            self.remote_socket = None
+
+        self._is_writer = True
+        self._is_local_reader = False
+        self.local_reader_rank = -1
+        # rank does not matter for remote readers
+        self._is_remote_reader = False
+
+        self.handle = Handle(
+            connect_ip=connect_ip,
+            local_reader_ranks=local_reader_ranks,
+            buffer=self.buffer,
+            local_subscribe_port=local_subscribe_port,
+            remote_subscribe_port=remote_subscribe_port,
+        )
+
+        logger.debug("Message queue communication handle: %s", self.handle)
+
+    def export_handle(self) -> Handle:
+        return self.handle
+
+    @staticmethod
+    def create_from_handle(handle: Handle, rank) -> "MessageQueue":
+        self = MessageQueue.__new__(MessageQueue)
+        self.handle = handle
+        self._is_writer = False
+
+        context = Context()
+
+        if rank in handle.local_reader_ranks:
+            assert handle.buffer is not None
+            self.buffer = handle.buffer
+            self.current_idx = 0
+            self.local_reader_rank = handle.local_reader_ranks.index(rank)
+            self._is_local_reader = True
+            self._is_remote_reader = False
+
+            self.local_socket = context.socket(SUB)
+            self.local_socket.setsockopt_string(SUBSCRIBE, "")
+            socket_addr = f"tcp://127.0.0.1:{handle.local_subscribe_port}"
+            logger.debug("Connecting to %s", socket_addr)
+            self.local_socket.connect(socket_addr)
+
+            self.remote_socket = None
+        else:
+            self.buffer = None  # type: ignore
+            self.current_idx = -1
+            self.local_reader_rank = -1
+            self._is_local_reader = False
+            self._is_remote_reader = True
+
+            self.local_socket = None
+
+            self.remote_socket = context.socket(SUB)
+            self.remote_socket.setsockopt_string(SUBSCRIBE, "")
+            na = NetworkAddress(handle.connect_ip, handle.remote_subscribe_port)
+            if na.is_ipv6:
+                self.remote_socket.setsockopt(IPV6, 1)
+            socket_addr = na.to_tcp()
+            logger.debug("Connecting to %s", socket_addr)
+            self.remote_socket.connect(socket_addr)
+
+        return self
+
+    def wait_until_ready(self):
+        """This is a collective operation. All processes (including the
+        readers and the writer) should call this function.
+        """
+        if self._is_writer:
+            # wait for all readers to connect
+
+            # local readers
+            for i in range(self.n_local_reader):
+                # wait for subscription messages from all local readers
+                self.local_socket.recv()
+            if self.n_local_reader > 0:
+                # send a message to all local readers
+                # to make sure the publish channel is working
+                self.local_socket.send(b"READY")
+
+            # remote readers
+            for i in range(self.n_remote_reader):
+                # wait for subscription messages from all remote readers
+                self.remote_socket.recv()
+            if self.n_remote_reader > 0:
+                # send a message to all remote readers
+                # to make sure the publish channel is working
+                self.remote_socket.send(b"READY")
+        elif self._is_local_reader:
+            # wait for the writer to send a message
+            recv = self.local_socket.recv()
+            assert recv == b"READY"
+        elif self._is_remote_reader:
+            # wait for the writer to send a message
+            recv = self.remote_socket.recv()
+            assert recv == b"READY"
+
+    @contextmanager
+    def acquire_write(self):
+        assert self._is_writer, "Only writers can acquire write"
+        start_time = time.monotonic()
+        n_warning = 1
+        while True:
+            with self.buffer.get_metadata(self.current_idx) as metadata_buffer:
+                read_count = sum(metadata_buffer[1:])
+                written_flag = metadata_buffer[0]
+                if written_flag and read_count != self.buffer.n_reader:
+                    # this block is written and not read by all readers
+                    # for writers, `self.current_idx` is the next block to write
+                    # if this block is not ready to write,
+                    # we need to wait until it is read by all readers
+
+                    # Release the processor to other threads
+                    os.sched_yield()
+
+                    # if we wait for a long time, we should warn the user
+                    if (
+                        time.monotonic() - start_time
+                        > SGLANG_RINGBUFFER_WARNING_INTERVAL * n_warning
+                    ):
+                        logger.warning(
+                            "No available block found in %s second. ",
+                            SGLANG_RINGBUFFER_WARNING_INTERVAL,
+                        )
+                        n_warning += 1
+
+                    continue
+                # found a block that is either
+                # (1) not written
+                # (2) read by all readers
+
+                # mark the block as not written
+                metadata_buffer[0] = 0
+                # let caller write to the buffer
+                with self.buffer.get_data(self.current_idx) as buf:
+                    yield buf
+
+                # caller has written to the buffer
+                # NOTE: order is important here
+                # first set the read flags to 0
+                # then set the written flag to 1
+                # otherwise, the readers may think they already read the block
+                for i in range(1, self.buffer.n_reader + 1):
+                    # set read flag to 0, meaning it is not read yet
+                    metadata_buffer[i] = 0
+                # mark the block as written
+                metadata_buffer[0] = 1
+                self.current_idx = (self.current_idx + 1) % self.buffer.max_chunks
+                break
+
+    @contextmanager
+    def acquire_read(self):
+        assert self._is_local_reader, "Only readers can acquire read"
+        start_time = time.monotonic()
+        n_warning = 1
+        while True:
+            with self.buffer.get_metadata(self.current_idx) as metadata_buffer:
+                read_flag = metadata_buffer[self.local_reader_rank + 1]
+                written_flag = metadata_buffer[0]
+                if not written_flag or read_flag:
+                    # this block is either
+                    # (1) not written
+                    # (2) already read by this reader
+
+                    # for readers, `self.current_idx` is the next block to read
+                    # if this block is not ready,
+                    # we need to wait until it is written
+
+                    # Release the processor to other threads
+                    os.sched_yield()
+
+                    # if we wait for a long time, we should warn the user
+                    if (
+                        time.monotonic() - start_time
+                        > SGLANG_RINGBUFFER_WARNING_INTERVAL * n_warning
+                    ):
+                        logger.warning(
+                            "No available block found in %s second. ",
+                            SGLANG_RINGBUFFER_WARNING_INTERVAL,
+                        )
+                        n_warning += 1
+
+                    continue
+                # found a block that is not read by this reader
+                # let caller read from the buffer
+                with self.buffer.get_data(self.current_idx) as buf:
+                    yield buf
+
+                # caller has read from the buffer
+                # set the read flag
+                metadata_buffer[self.local_reader_rank + 1] = 1
+                self.current_idx = (self.current_idx + 1) % self.buffer.max_chunks
+                break
+
+    def enqueue(self, obj):
+        assert self._is_writer, "Only writers can enqueue"
+        serialized_obj = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL)
+        if self.n_local_reader > 0:
+            if len(serialized_obj) >= self.buffer.max_chunk_bytes:
+                with self.acquire_write() as buf:
+                    buf[0] = 1  # overflow
+                self.local_socket.send(serialized_obj)
+            else:
+                with self.acquire_write() as buf:
+                    buf[0] = 0  # not overflow
+                    buf[1 : len(serialized_obj) + 1] = serialized_obj
+        if self.n_remote_reader > 0:
+            self.remote_socket.send(serialized_obj)
+
+    def dequeue(self):
+        if self._is_local_reader:
+            with self.acquire_read() as buf:
+                overflow = buf[0] == 1
+                if not overflow:
+                    # no need to know the size of serialized object
+                    # pickle format contains the size information internally
+                    # see https://docs.python.org/3/library/pickle.html
+                    obj = pickle.loads(buf[1:])
+            if overflow:
+                recv = self.local_socket.recv()
+                obj = pickle.loads(recv)
+        elif self._is_remote_reader:
+            recv = self.remote_socket.recv()
+            obj = pickle.loads(recv)
+        else:
+            raise RuntimeError("Only readers can dequeue")
+        return obj
+
+    def broadcast_object(self, obj=None):
+        if self._is_writer:
+            self.enqueue(obj)
+            return obj
+        else:
+            return self.dequeue()
+
+    @staticmethod
+    def create_from_process_group(
+        pg: ProcessGroup, max_chunk_bytes, max_chunks, writer_rank=0
+    ) -> "MessageQueue":
+        group_rank = dist.get_rank(pg)
+        group_world_size = dist.get_world_size(pg)
+        global_ranks = dist.get_process_group_ranks(pg)
+
+        from sglang.srt.distributed.parallel_state import in_the_same_node_as
+
+        status = in_the_same_node_as(pg, source_rank=writer_rank)
+        same_node_ranks = [i for i, s in enumerate(status) if s]
+        n_reader = group_world_size - 1
+        n_local_reader = len(same_node_ranks) - 1
+        local_reader_ranks = [i for i in same_node_ranks if i != writer_rank]
+        buffer_io: MessageQueue
+        if group_rank == writer_rank:
+            buffer_io = MessageQueue(
+                n_reader=n_reader,
+                n_local_reader=n_local_reader,
+                local_reader_ranks=local_reader_ranks,
+                max_chunk_bytes=max_chunk_bytes,
+                max_chunks=max_chunks,
+            )
+            handle = buffer_io.export_handle()
+            dist.broadcast_object_list(
+                [handle], src=global_ranks[writer_rank], group=pg
+            )
+        else:
+            recv = [None]
+            dist.broadcast_object_list(recv, src=global_ranks[writer_rank], group=pg)
+            handle = recv[0]  # type: ignore
+            buffer_io = MessageQueue.create_from_handle(handle, group_rank)
+        buffer_io.wait_until_ready()
+        return buffer_io

python/sglang/srt/distributed/device_communicators/torch_symm_mem.py

@@ -0,0 +1,171 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from https://github.com/vllm-project/vllm/blob/bf214ca22625e311a2c4c0dfbf7af19128f4919c/vllm/distributed/device_communicators/symm_mem.py
+import logging
+from typing import Optional, Union
+
+import torch
+import torch.distributed as dist
+from torch.distributed import ProcessGroup
+
+from sglang.srt.distributed.device_communicators.all_reduce_utils import (
+    TORCH_SYMM_MEM_ALL_REDUCE_MAX_SIZES,
+)
+from sglang.srt.utils import is_cuda, is_hip
+
+try:
+    import torch.distributed._symmetric_memory as torch_symm_mem
+
+    _is_cuda = is_cuda()
+    _is_hip = is_hip()
+
+    torch_symm_mem_available = False
+    if _is_cuda:
+        torch_symm_mem_available = True
+except ImportError:
+    torch_symm_mem_available = False
+
+
+logger = logging.getLogger(__name__)
+
+
+class TorchSymmMemCommunicator:
+    """
+    Thin wrapper around torch-symmetric-memory collectives.
+
+    This communicator:
+      - Validates device capability and world size.
+      - Allocates a shared symmetric buffer.
+      - Chooses between 'multimem' and 'two-shot' all-reduce kernels.
+      - Exposes a fast-path all_reduce() compatible with bfloat16 inputs.
+
+    If any prerequisite is not met, the instance remains disabled and will
+    decline to perform symmetric-memory all-reduce.
+    """
+
+    # Mapping: compute capability major -> supported world sizes for multimem
+    # If the current (cc_major, world_size) is not listed, we fall back
+    # to the two-shot path.
+    _WORLD_SIZES_MULTIMEM = {
+        9: [4, 6, 8],
+        10: [6, 8],
+    }
+
+    def __init__(self, group: ProcessGroup, device: Union[int, str, torch.device]):
+        """
+        Args:
+            group: Torch process group used for rendezvous and naming.
+            device: Target CUDA device (index, 'cuda:X', or torch.device).
+        """
+
+        self.disabled = True
+        self.buffer = None
+        self.max_size = 0
+
+        if not torch_symm_mem_available:
+            return
+
+        if isinstance(device, int):
+            device = torch.device(f"cuda:{device}")
+        elif isinstance(device, str):
+            device = torch.device(device)
+        torch.cuda.set_device(device)
+        self.dtype = torch.bfloat16
+        self.device = device
+        self.group = group
+        self.world_size = dist.get_world_size(self.group)
+        self.device_capability = torch.cuda.get_device_capability(device)[0]
+        supported_max_sizes = TORCH_SYMM_MEM_ALL_REDUCE_MAX_SIZES.get(
+            self.device_capability
+        )
+        if supported_max_sizes is None:
+            logger.warning(
+                "TorchSymmMemCommunicator: Device capability %s not supported, "
+                "communicator is not available.",
+                self.device_capability,
+            )
+            return
+        if self.world_size not in supported_max_sizes:
+            logger.warning(
+                "TorchSymmMemCommunicator: World size %d not supported, "
+                "communicator is not available.",
+                self.world_size,
+            )
+            return
+        self.max_size = supported_max_sizes[self.world_size]
+        self.buffer = torch_symm_mem.empty(
+            self.max_size // self.dtype.itemsize,
+            device=self.device,
+            dtype=self.dtype,
+        )
+        handle = torch_symm_mem.rendezvous(self.buffer, self.group.group_name)
+        if handle.multicast_ptr == 0:
+            logger.warning(
+                "TorchSymmMemCommunicator: torch symmetric memory "
+                "multicast operations are not supported."
+            )
+            self.buffer = None
+            self.disabled = True
+            return
+        self.disabled = False
+
+    def should_torch_symm_mem_allreduce(self, inp: torch.Tensor):
+        """
+        Fast-path eligibility check for a given tensor.
+
+        Conditions:
+          - Communicator must be enabled.
+          - dtype must be bfloat16 (matches kernel + buffer dtype).
+          - Total byte size must be 4-byte aligned (hardware requirement).
+          - Payload must be smaller than the symmetric-memory max size.
+
+        Returns:
+            True if the symmetric-memory path can handle this tensor.
+        """
+        if self.disabled:
+            return False
+        if inp.device != self.device:
+            return False
+        if inp.dtype != self.dtype:
+            return False
+        inp_size = inp.numel() * inp.element_size()
+        # enforce 4-byte alignment
+        if inp_size % 4 != 0:
+            return False
+        return inp_size < self.max_size
+
+    def all_reduce(
+        self, inp: torch.Tensor, *, out: Optional[torch.Tensor] = None
+    ) -> Optional[torch.Tensor]:
+        """
+        Perform an in-place sum all-reduce via torch symmetric memory.
+
+        Args:
+            inp: Input tensor on the target CUDA device (bfloat16).
+            out: Optional output tensor; if omitted, a new tensor is allocated.
+
+        Returns:
+            The reduced tensor (same shape as inp), or None if disabled.
+
+        Implementation details:
+            - Stages 'inp' into the symmetric buffer.
+            - Selects 'multimem' or 'two_shot' kernel based on topology.
+            - Writes the result into 'out' and returns it.
+        """
+        if not self.should_torch_symm_mem_allreduce(inp):
+            return None
+        if out is None:
+            out = torch.empty_like(inp)
+        self.buffer[: inp.numel()].copy_(inp.view(-1))
+        if self.world_size in self._WORLD_SIZES_MULTIMEM.get(
+            self.device_capability, ()
+        ):
+            torch.ops.symm_mem.multimem_all_reduce_(
+                self.buffer[: inp.numel()], "sum", self.group.group_name
+            )
+        else:
+            torch.ops.symm_mem.two_shot_all_reduce_(
+                self.buffer[: inp.numel()], "sum", self.group.group_name
+            )
+        out.copy_(self.buffer[: inp.numel()].view(out.shape))
+        return out

python/sglang/srt/distributed/device_communicators/triton_symm_mem_ag.py

@@ -0,0 +1,548 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Symmetric-memory ``multimem.st`` all-gather along the hidden (last) dim.
+
+Each rank stores its ``[T, H/TP]`` shard into a multicast buffer in one NVLink
+pass instead of an NCCL ring; ``create_state`` rendezvous once so launches are
+CUDA-graph capturable.
+"""
+
+import logging
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.distributed as dist
+import torch.distributed._symmetric_memory as symm_mem
+import triton
+import triton.language as tl
+
+logger = logging.getLogger(__name__)
+
+# Each thread moves _NUMEL_PER_THREAD bf16 via one 128-bit multimem op; the
+# grid-strided block count is tunable in [_MIN_BLOCKS, _MAX_BLOCKS].
+_BLOCK_THREADS = 1024
+_NUMEL_PER_THREAD = 8
+_MIN_BLOCKS = 4
+_MAX_BLOCKS = 32
+
+
+# ------------------------------------------------------------------------------
+# Low-level PTX helpers
+# ------------------------------------------------------------------------------
+
+
+@triton.jit
+def _multimem_st_128(multicast_ptrs, x, y, z, w, mask):
+    return tl.inline_asm_elementwise(
+        """
+        {
+            .reg .pred %p0;
+            setp.eq.s32 %p0, $6, 1;
+            @!%p0 bra end;
+            multimem.st.relaxed.sys.global.v4.f32 [$1], {$2, $3, $4, $5};
+            end:
+        }
+        """,
+        "=r,l,r,r,r,r,r",
+        args=[multicast_ptrs, x, y, z, w, mask.to(tl.int32)],
+        dtype=(tl.uint32),
+        is_pure=False,
+        pack=1,
+    )
+
+
+@triton.jit
+def _local_ld_128(in_ptr, mask):
+    return tl.inline_asm_elementwise(
+        """
+        {
+            .reg .pred %p0;
+            setp.eq.s32 %p0, $5, 1;
+            @!%p0 bra end;
+            ld.relaxed.sys.global.v4.b32 {$0, $1, $2, $3}, [$4];
+            end:
+        }
+        """,
+        "=r,=r,=r,=r,l,r",
+        args=[in_ptr, mask.to(tl.int32)],
+        dtype=(tl.uint32, tl.uint32, tl.uint32, tl.uint32),
+        is_pure=True,
+        pack=1,
+    )
+
+
+@triton.jit
+def _get_tid():
+    return tl.inline_asm_elementwise(
+        """
+        mov.u32 $0, %tid.x;
+        mov.u32 $1, %tid.y;
+        mov.u32 $2, %tid.z;
+        """,
+        "=r,=r,=r",
+        [],
+        dtype=(tl.uint32, tl.uint32, tl.uint32),
+        is_pure=True,
+        pack=1,
+    )
+
+
+@triton.jit
+def _get_ntid():
+    return tl.inline_asm_elementwise(
+        """
+        mov.u32 $0, %ntid.x;
+        mov.u32 $1, %ntid.y;
+        mov.u32 $2, %ntid.z;
+        """,
+        "=r,=r,=r",
+        [],
+        dtype=(tl.uint32, tl.uint32, tl.uint32),
+        is_pure=True,
+        pack=1,
+    )
+
+
+@triton.jit
+def _get_flat_tid():
+    tid_x, tid_y, tid_z = _get_tid()
+    ntid_x, ntid_y, _ = _get_ntid()
+    return tid_z * ntid_y * ntid_x + tid_y * ntid_x + tid_x
+
+
+@triton.jit
+def _sync_threads():
+    tl.inline_asm_elementwise(
+        "bar.sync 0;", "=r", [], dtype=tl.int32, is_pure=False, pack=1
+    )
+
+
+@triton.jit
+def _send_signal(addrs):
+    tl.inline_asm_elementwise(
+        """
+        {
+            .reg .u32   %tmp32_<1>;
+            .reg .pred  %p<1>;
+
+            send_signal:
+                atom.global.relaxed.sys.cas.b32 %tmp32_0, [$1], 0, 1;
+                setp.eq.u32 %p0, %tmp32_0, 0;
+                @!%p0 bra send_signal;
+        }
+        """,
+        "=r, l",
+        [addrs],
+        dtype=tl.int32,
+        is_pure=False,
+        pack=1,
+    )
+
+
+@triton.jit
+def _send_signal_release(addrs):
+    tl.inline_asm_elementwise(
+        """
+        {
+            .reg .u32   %tmp32_<1>;
+            .reg .pred  %p<1>;
+
+            send_signal:
+                atom.global.release.sys.cas.b32 %tmp32_0, [$1], 0, 1;
+                setp.eq.u32 %p0, %tmp32_0, 0;
+                @!%p0 bra send_signal;
+        }
+        """,
+        "=r, l",
+        [addrs],
+        dtype=tl.int32,
+        is_pure=False,
+        pack=1,
+    )
+
+
+@triton.jit
+def _wait_signal(addrs):
+    tl.inline_asm_elementwise(
+        """
+        {
+            .reg .u32   %tmp32_<1>;
+            .reg .pred  %p<1>;
+
+            wait_signal:
+                atom.global.sys.relaxed.cas.b32 %tmp32_0, [$1], 1, 0;
+                setp.eq.u32 %p0, %tmp32_0, 1;
+                @!%p0 bra wait_signal;
+        }
+        """,
+        "=r, l",
+        [addrs],
+        dtype=tl.int32,
+        is_pure=False,
+        pack=1,
+    )
+
+
+@triton.jit
+def _wait_signal_acquire(addrs):
+    tl.inline_asm_elementwise(
+        """
+        {
+            .reg .u32   %tmp32_<1>;
+            .reg .pred  %p<1>;
+
+            wait_signal:
+                atom.global.sys.acquire.cas.b32 %tmp32_0, [$1], 1, 0;
+                setp.eq.u32 %p0, %tmp32_0, 1;
+                @!%p0 bra wait_signal;
+        }
+        """,
+        "=r, l",
+        [addrs],
+        dtype=tl.int32,
+        is_pure=False,
+        pack=1,
+    )
+
+
+@triton.jit
+def _blockwise_barrier(
+    signal_pad_ptrs,
+    rank: tl.constexpr,
+    world_size: tl.constexpr,
+    sem: tl.constexpr,
+):
+    block_id = (
+        tl.program_id(2) * tl.num_programs(1) * tl.num_programs(0)
+        + tl.program_id(1) * tl.num_programs(0)
+        + tl.program_id(0)
+    )
+    flat_tid = _get_flat_tid()
+
+    remote_ranks = tl.arange(0, world_size)
+    signal_pad_ptrs = signal_pad_ptrs.to(tl.pointer_type(tl.uint64))
+    remote_signal_pad_addrs = tl.load(signal_pad_ptrs + remote_ranks).to(
+        tl.pointer_type(tl.uint32)
+    )
+    send_addrs = remote_signal_pad_addrs + block_id * world_size + rank
+
+    local_signal_pad_addr = tl.load(signal_pad_ptrs + rank).to(
+        tl.pointer_type(tl.uint32)
+    )
+    wait_addrs = local_signal_pad_addr + block_id * world_size + remote_ranks
+
+    if flat_tid < world_size:
+        if sem == "relaxed":
+            _send_signal(send_addrs)
+            _wait_signal(wait_addrs)
+        else:
+            _send_signal_release(send_addrs)
+            _wait_signal_acquire(wait_addrs)
+
+
+@triton.jit
+def _all_gather_kernel_inner(
+    input_ptr,
+    multicast_ptr,
+    signal_pad_ptr,
+    total_tokens,
+    hidden_offset,
+    LOCAL_HIDDEN: tl.constexpr,
+    TOTAL_HIDDEN: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    NUMEL_PER_THREAD: tl.constexpr,
+    RANK: tl.constexpr,
+    WORLD_SIZE: tl.constexpr,
+    SKIP_ENTRY_SYNC: tl.constexpr,
+) -> None:
+    if SKIP_ENTRY_SYNC == 0:
+        _blockwise_barrier(signal_pad_ptr, RANK, WORLD_SIZE, sem="relaxed")
+        _sync_threads()
+
+    chunks_per_row: tl.constexpr = LOCAL_HIDDEN // NUMEL_PER_THREAD
+    total_hidden_chunks: tl.constexpr = TOTAL_HIDDEN // NUMEL_PER_THREAD
+    hidden_offset_chunks = hidden_offset // NUMEL_PER_THREAD
+    total_chunks = total_tokens * chunks_per_row
+
+    pid = tl.program_id(axis=0)
+    tid = _get_flat_tid()
+    block_start = pid * BLOCK_SIZE
+
+    while block_start < total_chunks:
+        chunk = block_start + tid
+        mask = chunk < total_chunks
+        row = chunk // chunks_per_row
+        col_chunk = chunk % chunks_per_row
+
+        in_ptr = input_ptr.to(tl.pointer_type(tl.uint64)) + chunk * 2
+        out_chunk = row * total_hidden_chunks + hidden_offset_chunks + col_chunk
+        out_ptr = (
+            multicast_ptr.to(tl.int64).to(tl.pointer_type(tl.uint64)) + out_chunk * 2
+        )
+        x, y, z, w = _local_ld_128(in_ptr, mask)
+        _multimem_st_128(out_ptr, x, y, z, w, mask)
+        block_start += tl.num_programs(axis=0) * BLOCK_SIZE
+
+    _sync_threads()
+    _blockwise_barrier(signal_pad_ptr, RANK, WORLD_SIZE, sem="acq_rel")
+
+
+# ------------------------------------------------------------------------------
+# Public API
+# ------------------------------------------------------------------------------
+
+
+@dataclass
+class MultimemAllGatherState:
+    group: dist.ProcessGroup
+    rank_in_group: int
+    world_size: int
+    device: torch.device
+    max_token_num: int
+    hidden_dim: int
+    comm_buff: torch.Tensor
+    # Rendezvous handle; stable for the buffer's lifetime, resolved once.
+    symm_mem_hdl: Any
+
+
+def create_state(
+    group: dist.ProcessGroup,
+    rank_in_group: int,
+    max_tokens: int,
+    hidden_size: int,
+    device: torch.device | None = None,
+) -> MultimemAllGatherState:
+    """Allocate and rendezvous the symmetric-memory buffer. Collective: call
+    once outside CUDA-graph capture with identical args on every rank."""
+    assert type(group) is dist.ProcessGroup, f"Expected ProcessGroup, got {type(group)}"
+    assert hidden_size % _NUMEL_PER_THREAD == 0, (
+        f"hidden_size={hidden_size} must be a multiple of {_NUMEL_PER_THREAD} "
+        f"bf16 for 16-byte multimem.st row alignment"
+    )
+    device = device or torch.device(f"cuda:{torch.cuda.current_device()}")
+
+    # Pad holds _MAX_BLOCKS * world_size uint32 slots; max() never shrinks it.
+    pad_bytes = _MAX_BLOCKS * group.size() * 4
+    symm_mem.set_signal_pad_size(max(symm_mem.get_signal_pad_size(), pad_bytes))
+    with torch.inference_mode(False), torch.no_grad():
+        comm_buff = symm_mem.empty(
+            (max_tokens, hidden_size), dtype=torch.bfloat16, device=device
+        )
+    hdl = symm_mem.rendezvous(comm_buff, group=group)
+    assert hdl.rank == rank_in_group, (
+        f"symm_mem handle rank {hdl.rank} != rank_in_group {rank_in_group}; the "
+        f"hidden-shard offset would be wrong"
+    )
+    return MultimemAllGatherState(
+        group=group,
+        rank_in_group=rank_in_group,
+        world_size=group.size(),
+        device=device,
+        max_token_num=max_tokens,
+        hidden_dim=hidden_size,
+        comm_buff=comm_buff,
+        symm_mem_hdl=hdl,
+    )
+
+
+def _launch_config(local_numel: int):
+    assert local_numel % _NUMEL_PER_THREAD == 0
+    return _MIN_BLOCKS, _BLOCK_THREADS, _BLOCK_THREADS // 32, _NUMEL_PER_THREAD
+
+
+def all_gather_inner(
+    state: MultimemAllGatherState,
+    hidden_states: torch.Tensor,
+    tp_hidden_dim: int,
+    skip_entry_sync: bool = False,
+    safe: bool = True,
+) -> torch.Tensor:
+    """Gather ``[T, H/TP]`` shards into ``[T, H]`` along the hidden dim.
+
+    ``tp_hidden_dim`` is the gathered width ``H``. Returns a clone when ``safe``,
+    else a view into the symmetric buffer (valid until the next collective)."""
+    world_size = state.world_size
+    assert hidden_states.dtype == torch.bfloat16, "Only bfloat16 is supported"
+    assert hidden_states.is_contiguous(), "hidden_states must be contiguous"
+    assert hidden_states.data_ptr() % 16 == 0, (
+        f"hidden_states.data_ptr()={hex(hidden_states.data_ptr())} must be "
+        f"16-byte aligned for 128-bit multimem.st"
+    )
+    assert (
+        tp_hidden_dim % world_size == 0
+    ), f"tp_hidden_dim={tp_hidden_dim} must be divisible by world_size={world_size}"
+    local_hidden = tp_hidden_dim // world_size
+    assert local_hidden % _NUMEL_PER_THREAD == 0, (
+        f"per-rank hidden shard ({local_hidden}) must be a multiple of "
+        f"{_NUMEL_PER_THREAD} bf16"
+    )
+    assert tp_hidden_dim <= state.hidden_dim, (
+        f"comm buffer too narrow: tp_hidden_dim={tp_hidden_dim} > "
+        f"state.hidden_dim={state.hidden_dim}"
+    )
+    total_tokens, in_hidden = hidden_states.shape
+    assert (
+        in_hidden == local_hidden
+    ), f"input hidden ({in_hidden}) != this rank's shard ({local_hidden})"
+    assert (
+        total_tokens <= state.max_token_num
+    ), f"total_tokens={total_tokens} exceeds max_token_num={state.max_token_num}"
+
+    hidden_offset = local_hidden * state.rank_in_group
+    symm_mem_hdl = state.symm_mem_hdl
+    num_blocks, block_size, num_warps, numel_per_thread = _launch_config(
+        total_tokens * local_hidden
+    )
+    grid = (num_blocks, 1, 1)
+    _all_gather_kernel_inner[grid](
+        input_ptr=hidden_states,
+        multicast_ptr=symm_mem_hdl.multicast_ptr,
+        signal_pad_ptr=symm_mem_hdl.signal_pad_ptrs_dev,
+        total_tokens=total_tokens,
+        hidden_offset=hidden_offset,
+        LOCAL_HIDDEN=local_hidden,
+        TOTAL_HIDDEN=state.hidden_dim,
+        BLOCK_SIZE=block_size,
+        NUMEL_PER_THREAD=numel_per_thread,
+        RANK=symm_mem_hdl.rank,
+        WORLD_SIZE=symm_mem_hdl.world_size,
+        SKIP_ENTRY_SYNC=1 if skip_entry_sync else 0,
+        num_warps=num_warps,
+    )
+    output = state.comm_buff[:total_tokens, :tp_hidden_dim]
+    return output.clone() if safe else output
+
+
+# ------------------------------------------------------------------------------
+# Guarded wrapper
+# ------------------------------------------------------------------------------
+
+
+def recommended_max_tokens(include_prefill: bool, floor: int = 0) -> int:
+    """Largest batch (tokens) to keep on the fast path; bigger falls back to
+    NCCL. Covers the spec-decode batch plus, if ``include_prefill``, a prefill
+    chunk. Returns ``floor`` if server args are unavailable."""
+    try:
+        from sglang.srt.runtime_context import get_server_args
+
+        sa = get_server_args()
+
+        def g(name: str) -> int:
+            v = getattr(sa, name, 0)
+            return v if isinstance(v, int) and v > 0 else 0
+
+        tokens = g("max_running_requests") * max(
+            g("speculative_num_draft_tokens"), g("speculative_eagle_topk"), 1
+        )
+        if include_prefill:
+            tokens = max(tokens, g("chunked_prefill_size"), g("max_prefill_tokens"))
+        return max(tokens, floor)
+    except Exception:
+        return floor
+
+
+class MultimemAllGatherer:
+    """Guarded multimem all-gather (last dim) with NCCL fallback; the single
+    entry point for every caller. Owns one symmetric buffer built lazily on the
+    first eager call, and uses the kernel only when the input fits its
+    dtype/shape/alignment contract. Guards use TP-replicated quantities so all
+    ranks pick the same path. ``skip_entry_sync=True`` drops the entry barrier;
+    only safe when a cross-rank sync sits between consecutive calls."""
+
+    _UNINIT = object()
+
+    def __init__(
+        self,
+        max_tokens: int,
+        *,
+        enabled: bool = True,
+        skip_entry_sync: bool = False,
+    ):
+        self._max_tokens = int(max_tokens)
+        self._skip_entry_sync = skip_entry_sync
+        # None => always NCCL; _UNINIT => build on first eager call.
+        self._state = self._UNINIT if enabled else None
+        if self._state is self._UNINIT:
+            # Lazy import avoids a module-load dependency on the distributed facade.
+            from sglang.srt.distributed import get_tp_group
+            from sglang.srt.distributed.parallel_state import in_the_same_node_as
+            from sglang.srt.runtime_context import get_server_args
+
+            tp_group = get_tp_group()
+            # Only probe node topology when the deployment can actually span
+            # nodes. Check world_size first so a TP=1 gatherer short-circuits
+            # before reading server args (which may be unpublished on offline
+            # paths). On a single node every TP rank is co-located, so skip the
+            # in_the_same_node_as() all-reduce, which can segfault under some
+            # EP/mooncake setups, and keep multimem enabled.
+            if (
+                tp_group.world_size > 1
+                and get_server_args().nnodes > 1
+                and not all(in_the_same_node_as(tp_group.cpu_group, source_rank=0))
+            ):
+                logger.warning(
+                    "multimem all-gather disabled because the TP group spans "
+                    "across nodes."
+                )
+                self._state = None
+
+    def __call__(self, x: torch.Tensor) -> torch.Tensor:
+        state = self._state
+        if state is self._UNINIT:
+            state = self._build(x)
+            if state is not self._UNINIT:
+                self._state = state
+        if (
+            state is not None
+            and state is not self._UNINIT
+            and x.dtype == torch.bfloat16
+            and x.dim() == 2
+            and x.is_contiguous()
+            and 0 < x.shape[0] <= state.max_token_num
+            and x.data_ptr() % 16 == 0
+            and x.shape[-1] * state.world_size <= state.hidden_dim
+        ):
+            return all_gather_inner(
+                state,
+                x,
+                tp_hidden_dim=x.shape[-1] * state.world_size,
+                skip_entry_sync=self._skip_entry_sync,
+                safe=False,
+            )
+        # Lazy import avoids a module-load dependency on the distributed facade.
+        from sglang.srt.distributed import tensor_model_parallel_all_gather
+
+        return tensor_model_parallel_all_gather(x, dim=-1)
+
+    def _build(self, x: torch.Tensor):
+        if x.dim() != 2 or x.dtype != torch.bfloat16:
+            return None
+        if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing():
+            # Can't allocate under capture; retry later.
+            return self._UNINIT
+        if x.shape[-1] % _NUMEL_PER_THREAD != 0:
+            return None
+        try:
+            from sglang.srt.distributed import get_tp_group
+
+            tp_group = get_tp_group()
+            if tp_group.world_size <= 1:
+                return None
+            state = create_state(
+                group=tp_group.device_group,
+                rank_in_group=tp_group.rank_in_group,
+                max_tokens=self._max_tokens,
+                hidden_size=x.shape[-1] * tp_group.world_size,
+            )
+            if state.symm_mem_hdl.multicast_ptr == 0:
+                # No multicast for this world size / arch; multimem.st would
+                # write nowhere. Fall back to NCCL.
+                logger.warning(
+                    "multimem all-gather disabled (no multicast for world_size=%d)",
+                    tp_group.world_size,
+                )
+                return None
+            return state
+        except Exception as e:
+            logger.warning("multimem all-gather disabled (%s)", e)
+            return None

python/sglang/srt/distributed/device_communicators/vmm_utils.py

@@ -0,0 +1,654 @@
+import logging
+import os
+import struct
+import time
+from typing import Any, List, Optional
+
+import torch
+import torch.distributed as dist
+from torch.distributed import ProcessGroup
+
+from sglang.srt.utils import log_info_on_rank0
+
+logger = logging.getLogger(__name__)
+
+_drv = None
+_FD_HEADER_BYTES = 24
+_FD_SEND_TIMEOUT_S = 120.0
+
+
+def _get_cuda_driver():
+    """Lazily import cuda.bindings.driver (cached after first call)."""
+    global _drv
+    if _drv is None:
+        from cuda.bindings import driver
+
+        _drv = driver
+    return _drv
+
+
+def check_drv(result_tuple, label):
+    """Check a cuda.bindings driver call result and return the value."""
+    if not isinstance(result_tuple, tuple):
+        result_tuple = (result_tuple,)
+    err = result_tuple[0]
+    drv = _get_cuda_driver()
+    if err != drv.CUresult.CUDA_SUCCESS:
+        raise RuntimeError(f"{label}: {err}")
+    return result_tuple[1] if len(result_tuple) > 1 else None
+
+
+def is_vmm_pointer(ptr: int) -> bool:
+    """Check if a device pointer is VMM-backed (cuMemCreate/cuMemMap).
+
+    cuMemRetainAllocationHandle succeeds only on pointers from cuMemCreate;
+    it fails on cudaMalloc pointers.
+    """
+    drv = _get_cuda_driver()
+    err, handle = drv.cuMemRetainAllocationHandle(ptr)
+    if err == drv.CUresult.CUDA_SUCCESS:
+        drv.cuMemRelease(handle)
+        return True
+    return False
+
+
+def make_rw_access_desc(device_id: int):
+    """A read-write, device-local ``CUmemAccessDesc`` for ``device_id``."""
+    drv = _get_cuda_driver()
+    desc = drv.CUmemAccessDesc()
+    desc.location.type = drv.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE
+    desc.location.id = device_id
+    desc.flags = drv.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE
+    return desc
+
+
+def all_ranks_ok(group: ProcessGroup, ok: bool) -> bool:
+    """True iff ``ok`` holds on every rank in ``group`` (BAND all-reduce)."""
+    flag = torch.tensor([1 if ok else 0], dtype=torch.int32)
+    dist.all_reduce(flag, op=dist.ReduceOp.BAND, group=group)
+    return flag.item() == 1
+
+
+def release_mappings(mappings) -> None:
+    """Unmap + address-free each ``(va, span_size, [(rel, size), ...])`` mapping.
+
+    Pops from ``mappings`` so a partially-released list is safe to retry.
+    """
+    drv = _get_cuda_driver()
+    while mappings:
+        va, span_size, mapped_chunks = mappings.pop()
+        for rel, size in mapped_chunks:
+            check_drv(drv.cuMemUnmap(int(va) + int(rel), int(size)), "cuMemUnmap")
+        check_drv(drv.cuMemAddressFree(int(va), int(span_size)), "cuMemAddressFree")
+
+
+def _send_fd(sock, fd: int, src_rank: int, base_idx: int) -> None:
+    import array
+    import socket
+
+    fds = array.array("i", [int(fd)])
+    header = struct.pack("<QQQ", int(src_rank), int(base_idx), 1)
+    sent = sock.sendmsg(
+        [header],
+        [(socket.SOL_SOCKET, socket.SCM_RIGHTS, fds.tobytes())],
+    )
+    if sent != len(header):
+        raise RuntimeError(f"sendmsg sent {sent} bytes, expected {len(header)}")
+
+
+def _recv_fd(sock):
+    import array
+    import socket
+
+    fd_item_size = array.array("i").itemsize
+    data, ancdata, _, _ = sock.recvmsg(
+        _FD_HEADER_BYTES, socket.CMSG_SPACE(fd_item_size)
+    )
+    if not data:
+        return None
+    if len(data) != _FD_HEADER_BYTES:
+        raise RuntimeError(
+            f"received truncated fd header: {len(data)} < {_FD_HEADER_BYTES}"
+        )
+    src_rank, base_idx, fd_count = struct.unpack("<QQQ", data)
+    fds = array.array("i")
+    for level, cmsg_type, cmsg_data in ancdata:
+        if level == socket.SOL_SOCKET and cmsg_type == socket.SCM_RIGHTS:
+            fds.frombytes(cmsg_data[: len(cmsg_data) - (len(cmsg_data) % fd_item_size)])
+    if fd_count != 1 or len(fds) != 1:
+        for fd in fds:
+            os.close(fd)
+        raise RuntimeError(
+            f"expected one fd, got header={fd_count}, ancillary={len(fds)}"
+        )
+    return int(src_rank), int(base_idx), int(fds[0])
+
+
+def export_shareable_handles(retained_handles, group: ProcessGroup, rank: int):
+    """Export retained VMM handles, preferring FABRIC and falling back to POSIX fds.
+
+    FABRIC is used only if every rank can export it; otherwise all ranks use POSIX
+    fds. Returns ``(fabric_handles, posix_fds, use_fabric)`` (one list populated);
+    raises if both fail on any rank. Caller owns the returned ``posix_fds``.
+    """
+    drv = _get_cuda_driver()
+    FABRIC = drv.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC
+    POSIX_FD = drv.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR
+
+    fabric_handles: List[bytes] = []
+    fabric_error: Optional[Exception] = None
+    try:
+        for alloc_h in retained_handles:
+            fabric_h = check_drv(
+                drv.cuMemExportToShareableHandle(alloc_h, FABRIC, 0),
+                "cuMemExportToShareableHandle(FABRIC)",
+            )
+            fabric_handles.append(bytes(fabric_h.data))
+        fabric_ok = True
+    except Exception as e:
+        fabric_error = e
+        fabric_ok = False
+        fabric_handles = []
+        logger.info(
+            "FABRIC handle export failed on rank %s; falling back to "
+            "POSIX fd transport: %s",
+            rank,
+            e,
+        )
+
+    if all_ranks_ok(group, fabric_ok):
+        return fabric_handles, [], True
+
+    posix_fds: List[int] = []
+    posix_error: Optional[Exception] = None
+    try:
+        for alloc_h in retained_handles:
+            fd = check_drv(
+                drv.cuMemExportToShareableHandle(alloc_h, POSIX_FD, 0),
+                "cuMemExportToShareableHandle(POSIX_FD)",
+            )
+            posix_fds.append(int(fd))
+        posix_ok = True
+    except Exception as e:
+        posix_error = e
+        posix_ok = False
+        for fd in posix_fds:
+            try:
+                os.close(fd)
+            except OSError:
+                pass
+        posix_fds = []
+
+    if not all_ranks_ok(group, posix_ok):
+        cause = posix_error or fabric_error
+        message = (
+            "VMM handle export failed: FABRIC export failed on at least one "
+            "rank and POSIX fd export failed on at least one rank"
+        )
+        if cause is not None:
+            message += f"; local rank {rank} error: {cause}"
+        raise RuntimeError(message) from posix_error
+
+    return [], posix_fds, False
+
+
+def exchange_posix_fds(
+    group: ProcessGroup,
+    rank: int,
+    world_size: int,
+    local_fds: List[int],
+    peer_base_counts: List[int],
+):
+    """Exchange POSIX file descriptors across ranks via SCM_RIGHTS over a UNIX
+    socket. Returns ``{(src_rank, base_idx): fd}`` for every peer. The caller
+    owns the received fds and must close them.
+    """
+    import socket
+    import tempfile
+    import threading
+
+    sock_kind = getattr(socket, "SOCK_SEQPACKET", socket.SOCK_STREAM)
+    sock_dir = tempfile.mkdtemp(prefix="sgl_ar_fd_")
+    sock_path = os.path.join(sock_dir, f"rank_{rank}.sock")
+    server = socket.socket(socket.AF_UNIX, sock_kind)
+    server.settimeout(_FD_SEND_TIMEOUT_S)
+    received_fds = {}
+    errors = []
+
+    def recv_loop():
+        try:
+            for _ in range(world_size - 1):
+                conn, _ = server.accept()
+                with conn:
+                    conn.settimeout(_FD_SEND_TIMEOUT_S)
+                    while True:
+                        packet = _recv_fd(conn)
+                        if packet is None:
+                            break
+                        src_rank, base_idx, fd = packet
+                        key = (src_rank, base_idx)
+                        if key in received_fds:
+                            os.close(fd)
+                            raise RuntimeError(f"duplicate fd for {key}")
+                        received_fds[key] = fd
+        except BaseException as e:
+            errors.append(e)
+
+    try:
+        server.bind(sock_path)
+        server.listen(world_size)
+        paths = [None] * world_size
+        dist.all_gather_object(paths, sock_path, group=group)
+
+        thread = threading.Thread(target=recv_loop, daemon=True)
+        thread.start()
+        try:
+            for peer_rank, peer_path in enumerate(paths):
+                if peer_rank == rank:
+                    continue
+                with socket.socket(socket.AF_UNIX, sock_kind) as sock:
+                    sock.settimeout(_FD_SEND_TIMEOUT_S)
+                    sock.connect(peer_path)
+                    for base_idx, fd in enumerate(local_fds):
+                        _send_fd(sock, fd, rank, base_idx)
+        finally:
+            thread.join(_FD_SEND_TIMEOUT_S)
+
+        if thread.is_alive():
+            raise RuntimeError("timed out waiting for POSIX fd exchange")
+        if errors:
+            raise RuntimeError("POSIX fd exchange receive failed") from errors[0]
+
+        expected = {
+            (src_rank, base_idx)
+            for src_rank, count in enumerate(peer_base_counts)
+            if src_rank != rank
+            for base_idx in range(count)
+        }
+        missing = expected.difference(received_fds)
+        extra = set(received_fds).difference(expected)
+        if missing or extra:
+            for fd in received_fds.values():
+                os.close(fd)
+            raise RuntimeError(
+                "POSIX fd exchange mismatch: "
+                f"missing={sorted(missing)[:8]}, extra={sorted(extra)[:8]}"
+            )
+        return received_fds
+    finally:
+        server.close()
+        try:
+            os.unlink(sock_path)
+        except FileNotFoundError:
+            pass
+        try:
+            os.rmdir(sock_dir)
+        except OSError:
+            pass
+
+
+def import_peer_handle(fabric_handle, fd, *, use_fabric: bool, peer_rank: int):
+    """Import a peer allocation handle (FABRIC or POSIX fd). Returns the handle.
+
+    For POSIX the fd is duped before import so the caller keeps ownership of the
+    original.
+    """
+    drv = _get_cuda_driver()
+    if use_fabric:
+        FABRIC = drv.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC
+        return check_drv(
+            drv.cuMemImportFromShareableHandle(fabric_handle, FABRIC),
+            f"cuMemImportFromShareableHandle(rank={peer_rank})",
+        )
+    POSIX_FD = drv.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR
+    dup_fd = os.dup(fd)
+    try:
+        return check_drv(
+            drv.cuMemImportFromShareableHandle(dup_fd, POSIX_FD),
+            f"cuMemImportFromShareableHandle(rank={peer_rank}, POSIX_FD)",
+        )
+    finally:
+        try:
+            os.close(dup_fd)
+        except OSError:
+            pass
+
+
+def import_and_map_alloc(
+    fabric_handle,
+    fd,
+    alloc_size: int,
+    device_id: int,
+    *,
+    use_fabric: bool,
+    peer_rank: int,
+) -> int:
+    """Import a peer allocation, map it at a freshly reserved VA, return the VA."""
+    drv = _get_cuda_driver()
+    imp_h = import_peer_handle(
+        fabric_handle, fd, use_fabric=use_fabric, peer_rank=peer_rank
+    )
+    prop = check_drv(
+        drv.cuMemGetAllocationPropertiesFromHandle(imp_h),
+        "cuMemGetAllocationPropertiesFromHandle",
+    )
+    gran = check_drv(
+        drv.cuMemGetAllocationGranularity(
+            prop,
+            drv.CUmemAllocationGranularity_flags.CU_MEM_ALLOC_GRANULARITY_RECOMMENDED,
+        ),
+        "cuMemGetAllocationGranularity",
+    )
+    va = check_drv(
+        drv.cuMemAddressReserve(alloc_size, int(gran), 0, 0), "cuMemAddressReserve"
+    )
+    check_drv(drv.cuMemMap(int(va), alloc_size, 0, imp_h, 0), "cuMemMap")
+    access = make_rw_access_desc(device_id)
+    check_drv(drv.cuMemSetAccess(int(va), alloc_size, [access], 1), "cuMemSetAccess")
+    check_drv(drv.cuMemRelease(imp_h), "cuMemRelease(peer)")
+    return int(va)
+
+
+def map_chunk_into_span(
+    fabric_handle,
+    fd,
+    span_va: int,
+    rel: int,
+    alloc_size: int,
+    device_id: int,
+    *,
+    use_fabric: bool,
+    peer_rank: int,
+) -> None:
+    """Import + map a peer chunk into a caller-reserved span at ``span_va + rel``."""
+    drv = _get_cuda_driver()
+    imp_h = import_peer_handle(
+        fabric_handle, fd, use_fabric=use_fabric, peer_rank=peer_rank
+    )
+    check_drv(
+        drv.cuMemMap(int(span_va) + rel, int(alloc_size), 0, imp_h, 0),
+        "cuMemMap(span)",
+    )
+    access = make_rw_access_desc(device_id)
+    check_drv(
+        drv.cuMemSetAccess(int(span_va) + rel, int(alloc_size), [access], 1),
+        "cuMemSetAccess(span)",
+    )
+    check_drv(drv.cuMemRelease(imp_h), "cuMemRelease(span)")
+
+
+class VmmGraphInputManager:
+    def __init__(
+        self,
+        obj: Any,
+        group: ProcessGroup,
+        rank: int,
+        world_size: int,
+    ) -> None:
+        self.obj = obj
+        self.group = group
+        self.rank = rank
+        self.world_size = world_size
+        self._peer_mappings = []
+
+    def register_graph_inputs(self):
+        """Register graph capture inputs via VMM handle exchange.
+
+        VMM-compatible path for expandable_segments. The C++ side deduplicates
+        graph capture pointers into unique base allocations via cuMemGetAddressRange.
+        Python exports handles for each unique base, imports + cuMemMaps peer
+        allocations, then registers the peer VAs. FABRIC handles are preferred;
+        POSIX file descriptors are used when FABRIC is unavailable.
+        """
+        FABRIC_HANDLE_BYTES = 64
+        MAX_VMM_BASES = 4096
+        MAX_CHUNKS_PER_INPUT = 16
+
+        t0 = time.perf_counter()
+
+        bases_info, input_chunk_indices, input_offsets = (
+            self.obj.get_graph_capture_bases()
+        )
+        if not bases_info:
+            return
+        new_count = len(input_chunk_indices)
+        num_bases = len(bases_info)
+        device_id = torch.cuda.current_device()
+
+        if num_bases > MAX_VMM_BASES:
+            raise RuntimeError(
+                f"Too many VMM bases to share: {num_bases} > {MAX_VMM_BASES}"
+            )
+
+        drv = _get_cuda_driver()
+        local_posix_fds: List[int] = []
+        retained_handles = []
+        try:
+            for base_ptr, _ in bases_info:
+                alloc_h = check_drv(
+                    drv.cuMemRetainAllocationHandle(base_ptr),
+                    "cuMemRetainAllocationHandle",
+                )
+                retained_handles.append(alloc_h)
+
+            local_fabric_handles, local_posix_fds, use_fabric = (
+                export_shareable_handles(retained_handles, self.group, self.rank)
+            )
+
+            local_input_chunks = [
+                [int(idx) for idx in indices] for indices in input_chunk_indices
+            ]
+            for chunks in local_input_chunks:
+                if len(chunks) > MAX_CHUNKS_PER_INPUT:
+                    raise RuntimeError(
+                        "Too many VMM chunks for graph input: "
+                        f"{len(chunks)} > {MAX_CHUNKS_PER_INPUT}"
+                    )
+
+            # All-gather base metadata and per-input VMM spans. A captured tensor
+            # can cross expandable-segment allocation boundaries, so peer mappings
+            # must preserve each input's contiguous virtual-address span. FABRIC
+            # handles are inline metadata; POSIX fds are exchanged separately via
+            # SCM_RIGHTS because fd integers are process-local.
+            header_struct = struct.Struct("<QQ")
+            base_struct = struct.Struct(
+                f"<QQ{FABRIC_HANDLE_BYTES}s" if use_fabric else "<QQ"
+            )
+            input_struct = struct.Struct(f"<QQ{MAX_CHUNKS_PER_INPUT}Q")
+            base_offset = header_struct.size
+            input_offset = base_offset + MAX_VMM_BASES * base_struct.size
+            payload_size = input_offset + new_count * input_struct.size
+            local_payload = bytearray(payload_size)
+
+            header_struct.pack_into(local_payload, 0, num_bases, new_count)
+            for i, (base_ptr, alloc_size) in enumerate(bases_info):
+                if use_fabric:
+                    base_struct.pack_into(
+                        local_payload,
+                        base_offset + i * base_struct.size,
+                        int(base_ptr),
+                        int(alloc_size),
+                        local_fabric_handles[i],
+                    )
+                else:
+                    base_struct.pack_into(
+                        local_payload,
+                        base_offset + i * base_struct.size,
+                        int(base_ptr),
+                        int(alloc_size),
+                    )
+            for i, (chunks, offset) in enumerate(
+                zip(local_input_chunks, input_offsets)
+            ):
+                padded_chunks = chunks + [0] * (MAX_CHUNKS_PER_INPUT - len(chunks))
+                input_struct.pack_into(
+                    local_payload,
+                    input_offset + i * input_struct.size,
+                    int(offset),
+                    len(chunks),
+                    *padded_chunks,
+                )
+
+            in_buf = torch.frombuffer(local_payload, dtype=torch.uint8).clone()
+            gather_list = [torch.empty_like(in_buf) for _ in range(self.world_size)]
+            dist.all_gather(gather_list, in_buf, group=self.group)
+
+            all_base_payload = []
+            all_input_chunks = []
+            all_input_offsets = []
+            for rank, gathered in enumerate(gather_list):
+                payload = gathered.numpy().tobytes()
+                peer_num_bases, peer_new_count = header_struct.unpack_from(payload, 0)
+                if peer_new_count != new_count:
+                    raise RuntimeError(
+                        "Mismatched graph input count across ranks: "
+                        f"rank {rank} has {peer_new_count}, expected {new_count}"
+                    )
+
+                peer_bases = []
+                for i in range(peer_num_bases):
+                    if use_fabric:
+                        base_ptr, alloc_size, fabric_handle = base_struct.unpack_from(
+                            payload, base_offset + i * base_struct.size
+                        )
+                    else:
+                        base_ptr, alloc_size = base_struct.unpack_from(
+                            payload, base_offset + i * base_struct.size
+                        )
+                        fabric_handle = None
+                    peer_bases.append((base_ptr, fabric_handle, alloc_size))
+
+                peer_chunks = []
+                peer_offsets = []
+                for i in range(new_count):
+                    unpacked = input_struct.unpack_from(
+                        payload, input_offset + i * input_struct.size
+                    )
+                    offset, chunk_count, *chunks = unpacked
+                    peer_offsets.append(offset)
+                    peer_chunks.append(list(chunks[:chunk_count]))
+
+                all_base_payload.append(peer_bases)
+                all_input_chunks.append(peer_chunks)
+                all_input_offsets.append(peer_offsets)
+
+            posix_peer_fds = {}
+            if not use_fabric:
+                posix_peer_fds = exchange_posix_fds(
+                    self.group,
+                    self.rank,
+                    self.world_size,
+                    local_posix_fds,
+                    [len(peer_bases) for peer_bases in all_base_payload],
+                )
+
+            # Import + map peer allocations. Individual base mappings are kept for
+            # single-chunk inputs; span mappings reserve a contiguous VA range and
+            # map each chunk at its original relative offset.
+            peer_base_va = {}  # (rank, base_idx) -> local VA
+            peer_span_va = {}  # (rank, chunk_indices...) -> (local VA, peer base)
+            new_mappings = []
+
+            try:
+                for peer_rank in range(self.world_size):
+                    if peer_rank == self.rank:
+                        for idx, (bp, _) in enumerate(bases_info):
+                            peer_base_va[(peer_rank, idx)] = int(bp)
+                        continue
+
+                    peer_bases = all_base_payload[peer_rank]
+                    for idx, (_, fb, alloc_size) in enumerate(peer_bases):
+                        fd = None if use_fabric else posix_peer_fds[(peer_rank, idx)]
+                        va = import_and_map_alloc(
+                            fb,
+                            fd,
+                            alloc_size,
+                            device_id,
+                            use_fabric=use_fabric,
+                            peer_rank=peer_rank,
+                        )
+                        peer_base_va[(peer_rank, idx)] = va
+                        new_mappings.append((va, alloc_size, [(0, alloc_size)]))
+
+                # Build per-input peer VA lists and register.
+                peer_ptrs = []
+                for j in range(new_count):
+                    ptrs_j = []
+                    for rank in range(self.world_size):
+                        chunks = all_input_chunks[rank][j]
+                        off = all_input_offsets[rank][j]
+                        if len(chunks) == 1:
+                            ptrs_j.append(peer_base_va[(rank, chunks[0])] + off)
+                            continue
+
+                        span_key = (rank, *chunks)
+                        if span_key not in peer_span_va:
+                            peer_bases = all_base_payload[rank]
+                            first_base = peer_bases[chunks[0]][0]
+                            last_base, _, last_size = peer_bases[chunks[-1]]
+                            span_size = (
+                                int(last_base) + int(last_size) - int(first_base)
+                            )
+                            if rank == self.rank:
+                                span_va = int(first_base)
+                            else:
+                                span_va = check_drv(
+                                    drv.cuMemAddressReserve(span_size, 0, 0, 0),
+                                    "cuMemAddressReserve(span)",
+                                )
+                                mapped_chunks = []
+                                for chunk_idx in chunks:
+                                    base_ptr, fb, alloc_size = peer_bases[chunk_idx]
+                                    rel = int(base_ptr) - int(first_base)
+                                    fd = (
+                                        None
+                                        if use_fabric
+                                        else posix_peer_fds[(rank, chunk_idx)]
+                                    )
+                                    map_chunk_into_span(
+                                        fb,
+                                        fd,
+                                        span_va,
+                                        rel,
+                                        int(alloc_size),
+                                        device_id,
+                                        use_fabric=use_fabric,
+                                        peer_rank=rank,
+                                    )
+                                    mapped_chunks.append((rel, int(alloc_size)))
+                                new_mappings.append(
+                                    (int(span_va), span_size, mapped_chunks)
+                                )
+                            peer_span_va[span_key] = (int(span_va), int(first_base))
+
+                        span_va, _ = peer_span_va[span_key]
+                        ptrs_j.append(span_va + off)
+                    peer_ptrs.append(ptrs_j)
+
+                self.obj.register_peer_mapped_inputs(peer_ptrs)
+                self._peer_mappings.extend(new_mappings)
+            except Exception:
+                release_mappings(new_mappings)
+                raise
+            finally:
+                for fd in posix_peer_fds.values():
+                    os.close(fd)
+
+            elapsed_ms = (time.perf_counter() - t0) * 1000
+            transport = "FABRIC" if use_fabric else "POSIX fd"
+            log_info_on_rank0(
+                logger,
+                f"Registered {new_count} cuda graph addresses via "
+                f"{transport} handles ({num_bases} unique allocations) "
+                f"in {elapsed_ms:.1f} ms",
+            )
+        finally:
+            for fd in local_posix_fds:
+                os.close(fd)
+            for h in retained_handles:
+                check_drv(drv.cuMemRelease(h), "cuMemRelease(retained)")
+
+    def close(self):
+        if not self._peer_mappings:
+            return
+        release_mappings(self._peer_mappings)

python/sglang/srt/distributed/device_communicators/xpu_communicator.py

@@ -0,0 +1,50 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/distributed/device_communicators/xpu_communicator.py
+
+import torch
+import torch.distributed as dist
+from torch.distributed import ProcessGroup
+
+from sglang.srt.utils import is_xpu
+
+
+class XpuCommunicator:
+
+    def __init__(self, group: ProcessGroup):
+        if not is_xpu():
+            self.disabled = True
+            return
+        self.disabled = False
+        self.group = group
+        self.world_size = dist.get_world_size(self.group)
+
+    def all_reduce(self, x: torch.Tensor) -> torch.Tensor:
+        dist.all_reduce(x, group=self.group)
+        return x
+
+    def gather(
+        self, input_: torch.Tensor, rank_in_group: int, dst: int = 0, dim: int = -1
+    ):
+        # For xpu path, gather doesn't work properly together with ray
+        # cluster so we use all_gather instead for now.
+        input_size = input_.size()
+        # Allocate output tensor.
+        output_tensor = torch.empty(
+            (self.world_size,) + input_size, dtype=input_.dtype, device=input_.device
+        )
+        # All-gather.
+        torch.distributed.all_gather_into_tensor(
+            output_tensor, input_, group=self.group
+        )
+        if rank_in_group == dst:
+            # Reshape
+            output_tensor = output_tensor.movedim(0, dim)
+            output_tensor = output_tensor.reshape(
+                input_size[:dim]
+                + (self.world_size * input_size[dim],)
+                + input_size[dim + 1 :]
+            )
+        else:
+            output_tensor = None
+        return output_tensor

python/sglang/srt/distributed/naive_distributed.py

@@ -0,0 +1,113 @@
+import pickle
+import time
+from pathlib import Path
+from typing import Any, List, Optional
+
+import pybase64
+import torch
+
+from sglang.srt.utils import MultiprocessingSerializer
+
+
+class NaiveDistributed:
+    def __init__(self, rank: int, world_size: int, rendezvous: str):
+        self._rank = rank
+        self._world_size = world_size
+        self._operation_index = 0
+        self._directory = Path(rendezvous)
+        self._directory.mkdir(parents=True, exist_ok=True)
+        assert 0 <= rank < world_size
+
+        # both barrier to be safe, and as a sanity check
+        self.barrier()
+
+    def get_rank(self):
+        return self._rank
+
+    def get_world_size(self):
+        return self._world_size
+
+    def scatter(
+        self, tensor: torch.Tensor, scatter_list: List[torch.Tensor], src: int = 0
+    ):
+        if self._rank == src:
+            assert len(scatter_list) == self._world_size
+        else:
+            assert scatter_list is None
+
+        gathered_objects = self.all_gather_object(
+            dict(
+                serialized_scatter_list=[
+                    (
+                        None
+                        if item_rank == src
+                        else MultiprocessingSerializer.serialize(item)
+                    )
+                    for item_rank, item in enumerate(scatter_list)
+                ]
+            )
+            if self._rank == src
+            else dict()
+        )
+
+        remote_serialized_tensor = gathered_objects[src]["serialized_scatter_list"][
+            self._rank
+        ]
+        if self._rank == src:
+            assert remote_serialized_tensor is None
+            remote_tensor = scatter_list[self._rank]
+        else:
+            remote_tensor = MultiprocessingSerializer.deserialize(
+                remote_serialized_tensor
+            )
+        tensor.copy_(remote_tensor)
+
+        # avoid src tensor be deleted too early
+        self.barrier()
+
+    def all_gather_object(self, obj: Any) -> List[Any]:
+        self._operation_index += 1
+
+        text_postfix = "\n"
+
+        def _get_path(interesting_rank: int):
+            return (
+                self._directory
+                / f"rank{interesting_rank}_op{self._operation_index}.txt"
+            )
+
+        _get_path(self._rank).write_text(
+            pybase64.b64encode(pickle.dumps(obj)).decode("utf-8") + text_postfix
+        )
+
+        def _read_one(interesting_rank: int):
+            p = _get_path(interesting_rank)
+            while True:
+                if p.exists() and (text := p.read_text()).endswith(text_postfix):
+                    return pickle.loads(
+                        pybase64.b64decode(text[: -len(text_postfix)], validate=True)
+                    )
+                time.sleep(0.001)
+
+        return [
+            _read_one(interesting_rank) for interesting_rank in range(self._world_size)
+        ]
+
+    def barrier(self):
+        actual_objs = self.all_gather_object(self._rank)
+        assert actual_objs == list(range(self._world_size)), f"{actual_objs=}"
+
+
+# Can have multi instances if needed
+_instance: Optional[NaiveDistributed] = None
+
+
+def get_naive_distributed():
+    assert _instance is not None
+    return _instance
+
+
+def set_naive_distributed(instance: NaiveDistributed):
+    global _instance
+    assert _instance is None
+    _instance = instance

python/sglang/srt/distributed/parallel_state_wrapper.py

@@ -0,0 +1,23 @@
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class ParallelState:
+    tp_rank: int
+    tp_size: int
+    pp_rank: int
+    pp_size: int
+    dp_rank: Optional[int]
+    dp_size: int
+    attn_tp_rank: int
+    attn_tp_size: int
+    attn_cp_rank: int
+    attn_cp_size: int
+    attn_dp_rank: int
+    attn_dp_size: int
+    moe_ep_rank: int
+    moe_ep_size: int
+    moe_dp_rank: Optional[int]
+    moe_dp_size: int
+    gpu_id: int

python/sglang/srt/distributed/utils.py

@@ -0,0 +1,262 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/distributed/utils.py
+
+# Copyright 2023 The vLLM team.
+# Adapted from
+# https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/tensor_parallel/utils.py
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+import dataclasses
+import logging
+import os
+import pickle
+import time
+from collections import deque
+from typing import Any, Deque, Dict, Optional, Sequence, Tuple
+
+import torch
+from torch.distributed import TCPStore
+
+logger = logging.getLogger(__name__)
+
+
+def set_global_tcp_store(store: TCPStore) -> None:
+    """Install the shared TCPStore created during distributed initialization;
+    the handle lives on ``ctx.resources``."""
+    from sglang.srt.runtime_context import get_resources
+
+    get_resources().tcp_store = store
+    logger.info("Global TCPStore has been set")
+
+
+def get_global_tcp_store() -> Optional[TCPStore]:
+    """Get the existing global TCPStore.
+
+    This function provides access to the shared TCPStore instance that was
+    created during distributed initialization. All components (like NIXL buffers)
+    should use this same store for coordination.
+
+    Returns:
+        The global TCPStore instance, or None if not initialized yet.
+    """
+    from sglang.srt.runtime_context import get_resources
+
+    store = get_resources().tcp_store
+    if store is None:
+        logger.warning(
+            "Global TCPStore not found. Make sure init_distributed_environment "
+            "was called with a tcp:// init method."
+        )
+    return store
+
+
+def ensure_divisibility(numerator, denominator):
+    """Ensure that numerator is divisible by the denominator."""
+    assert numerator % denominator == 0, "{} is not divisible by {}".format(
+        numerator, denominator
+    )
+
+
+def divide(numerator, denominator):
+    """Ensure that numerator is divisible by the denominator and return
+    the division value."""
+    ensure_divisibility(numerator, denominator)
+    return numerator // denominator
+
+
+def split_tensor_along_last_dim(
+    tensor: torch.Tensor,
+    num_partitions: int,
+    contiguous_split_chunks: bool = False,
+) -> Sequence[torch.Tensor]:
+    """Split a tensor along its last dimension.
+
+    Arguments:
+        tensor: input tensor.
+        num_partitions: number of partitions to split the tensor
+        contiguous_split_chunks: If True, make each chunk contiguous
+                                 in memory.
+
+    Returns:
+        A list of Tensors
+    """
+    # Get the size and dimension.
+    last_dim = tensor.dim() - 1
+    last_dim_size = divide(tensor.size()[last_dim], num_partitions)
+    # Split.
+    tensor_list = torch.split(tensor, last_dim_size, dim=last_dim)
+    # NOTE: torch.split does not create contiguous tensors by default.
+    if contiguous_split_chunks:
+        return tuple(chunk.contiguous() for chunk in tensor_list)
+
+    return tensor_list
+
+
+def get_pp_indices(
+    num_hidden_layers: int, pp_rank: int, pp_size: int
+) -> Tuple[int, int]:
+    """Try to evenly distribute layers across partitions.
+    If the number of layers is not divisible by the number of partitions,
+    the last N partitions will have one extra layer, where N = remainder.
+    """
+    # partition_list_str can be set to None in sglang
+    partition_list_str = os.getenv("SGLANG_PP_LAYER_PARTITION", None)
+    if partition_list_str is not None:
+        try:
+            partitions = [int(layer) for layer in partition_list_str.split(",")]
+        except ValueError as err:
+            raise ValueError(
+                "Invalid partition string: {}".format(partition_list_str)
+            ) from err
+        if len(partitions) != pp_size:
+            raise ValueError(f"{len(partitions)=} does not match {pp_size=}.")
+        if sum(partitions) != num_hidden_layers:
+            raise ValueError(f"{sum(partitions)=} does not match {num_hidden_layers=}.")
+        start_layer = sum(partitions[:pp_rank])
+        end_layer = start_layer + partitions[pp_rank]
+    else:
+        base_layers = num_hidden_layers // pp_size
+        remainder = num_hidden_layers % pp_size
+        # Distribute the extra layers to the last 'remainder' partitions
+        if pp_rank >= pp_size - remainder:
+            partitions_without_extra_layer = pp_size - remainder
+            # This partition gets one extra layer
+            start_layer = pp_rank * (base_layers + 1) - partitions_without_extra_layer
+            end_layer = start_layer + (base_layers + 1)
+        else:
+            # This partition gets only base layers
+            start_layer = pp_rank * base_layers
+            end_layer = start_layer + base_layers
+
+    return (start_layer, end_layer)
+
+
+@dataclasses.dataclass
+class StatelessProcessGroup:
+    """A dataclass to hold a metadata store, and the rank, world_size of the
+    group. Only use it to communicate metadata between processes.
+    For data-plane communication, create NCCL-related objects.
+    """
+
+    rank: int
+    world_size: int
+    store: torch._C._distributed_c10d.Store
+    data_expiration_seconds: int = 3600  # 1 hour
+
+    # dst rank -> counter
+    send_dst_counter: Dict[int, int] = dataclasses.field(default_factory=dict)
+    # src rank -> counter
+    recv_src_counter: Dict[int, int] = dataclasses.field(default_factory=dict)
+    broadcast_send_counter: int = 0
+    broadcast_recv_src_counter: Dict[int, int] = dataclasses.field(default_factory=dict)
+
+    # A deque to store the data entries, with key and timestamp.
+    entries: Deque[Tuple[str, float]] = dataclasses.field(default_factory=deque)
+
+    def __post_init__(self):
+        assert self.rank < self.world_size
+        self.send_dst_counter = {i: 0 for i in range(self.world_size)}
+        self.recv_src_counter = {i: 0 for i in range(self.world_size)}
+        self.broadcast_recv_src_counter = {i: 0 for i in range(self.world_size)}
+
+    def send_obj(self, obj: Any, dst: int):
+        """Send an object to a destination rank."""
+        self.expire_data()
+        key = f"send_to/{dst}/{self.send_dst_counter[dst]}"
+        self.store.set(key, pickle.dumps(obj))
+        self.send_dst_counter[dst] += 1
+        self.entries.append((key, time.perf_counter()))
+
+    def expire_data(self):
+        """Expire data that is older than `data_expiration_seconds` seconds."""
+        while self.entries:
+            # check the oldest entry
+            key, timestamp = self.entries[0]
+            if time.perf_counter() - timestamp > self.data_expiration_seconds:
+                self.store.delete_key(key)
+                self.entries.popleft()
+            else:
+                break
+
+    def recv_obj(self, src: int) -> Any:
+        """Receive an object from a source rank."""
+        obj = pickle.loads(
+            self.store.get(f"send_to/{self.rank}/{self.recv_src_counter[src]}")
+        )
+        self.recv_src_counter[src] += 1
+        return obj
+
+    def broadcast_obj(self, obj: Optional[Any], src: int) -> Any:
+        """Broadcast an object from a source rank to all other ranks.
+        It does not clean up after all ranks have received the object.
+        Use it for limited times, e.g., for initialization.
+        """
+        if self.rank == src:
+            self.expire_data()
+            key = f"broadcast_from/{src}/" f"{self.broadcast_send_counter}"
+            self.store.set(key, pickle.dumps(obj))
+            self.broadcast_send_counter += 1
+            self.entries.append((key, time.perf_counter()))
+            return obj
+        else:
+            key = f"broadcast_from/{src}/" f"{self.broadcast_recv_src_counter[src]}"
+            recv_obj = pickle.loads(self.store.get(key))
+            self.broadcast_recv_src_counter[src] += 1
+            return recv_obj
+
+    def all_gather_obj(self, obj: Any) -> list[Any]:
+        """All gather an object from all ranks."""
+        gathered_objs = []
+        for i in range(self.world_size):
+            if i == self.rank:
+                gathered_objs.append(obj)
+                self.broadcast_obj(obj, src=self.rank)
+            else:
+                recv_obj = self.broadcast_obj(None, src=i)
+                gathered_objs.append(recv_obj)
+        return gathered_objs
+
+    def barrier(self):
+        """A barrier to synchronize all ranks."""
+        for i in range(self.world_size):
+            if i == self.rank:
+                self.broadcast_obj(None, src=self.rank)
+            else:
+                self.broadcast_obj(None, src=i)
+
+    @staticmethod
+    def create(
+        host: str,
+        port: int,
+        rank: int,
+        world_size: int,
+        data_expiration_seconds: int = 3600,
+    ) -> "StatelessProcessGroup":
+        """A replacement for `torch.distributed.init_process_group` that does not
+        pollute the global state.
+
+        If we have process A and process B called `torch.distributed.init_process_group`
+        to form a group, and then we want to form another group with process A, B, C,
+        D, it is not possible in PyTorch, because process A and process B have already
+        formed a group, and process C and process D cannot join that group. This
+        function is a workaround for this issue.
+
+        `torch.distributed.init_process_group` is a global call, while this function
+        is a stateless call. It will return a `StatelessProcessGroup` object that can be
+        used for exchanging metadata. With this function, process A and process B
+        can call `StatelessProcessGroup.create` to form a group, and then process A, B,
+        C, and D can call `StatelessProcessGroup.create` to form another group.
+        """  # noqa
+        store = TCPStore(
+            host_name=host,
+            port=port,
+            world_size=world_size,
+            is_master=(rank == 0),
+        )
+
+        return StatelessProcessGroup(
+            rank=rank,
+            world_size=world_size,
+            store=store,
+            data_expiration_seconds=data_expiration_seconds,
+        )