From 132ade55cd65c3679a60adec563c8e6ca81238c5 Mon Sep 17 00:00:00 2001 From: DarkSharpness <76582120+DarkSharpness@users.noreply.github.com> Date: Fri, 17 Jul 2026 18:37:22 +0800 Subject: [PATCH] [Kernel] Rewrite JIT custom all-reduce (v2) with a decoupled kernel/storage design (#31049) Co-authored-by: Claude Fable 5 Co-authored-by: root --- python/sglang/jit_kernel/all_reduce.py | 288 ++++---- python/sglang/jit_kernel/benchmark/utils.py | 2 +- .../csrc/distributed/communicator.cuh | 113 +++ .../csrc/distributed/custom_all_reduce.cuh | 650 ++++++++++++++++++ .../distributed/custom_all_reduce_base.cuh | 30 - .../distributed/custom_all_reduce_pull.cuh | 205 ------ .../distributed/custom_all_reduce_push.cuh | 253 ------- .../jit_kernel/csrc/distributed/ipc.cuh | 194 ++++++ .../jit_kernel/csrc/distributed/tp_qknorm.cuh | 54 +- .../include/sgl_kernel/distributed/common.cuh | 120 ---- .../sgl_kernel/distributed/communicator.cuh | 120 ++++ .../distributed/custom_all_reduce.cuh | 446 ------------ python/sglang/jit_kernel/tests/utils.py | 2 +- python/sglang/jit_kernel/utils.py | 29 + .../device_communicators/configs/__init__.py | 1 + .../configs/custom_all_reduce_v2.py | 194 ++++++ .../custom_all_reduce_v2.py | 490 ++++++++----- .../device_communicators/vmm_utils.py | 50 ++ python/sglang/srt/environ.py | 3 + .../jit/benchmark/bench_custom_all_reduce.py | 90 +-- .../jit/benchmark/bench_tp_qknorm.py | 13 +- test/registered/jit/test_custom_all_reduce.py | 32 +- test/registered/jit/test_tp_qknorm.py | 24 +- 23 files changed, 1917 insertions(+), 1486 deletions(-) create mode 100644 python/sglang/jit_kernel/csrc/distributed/communicator.cuh create mode 100644 python/sglang/jit_kernel/csrc/distributed/custom_all_reduce.cuh delete mode 100644 python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_base.cuh delete mode 100644 python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_pull.cuh delete mode 100644 python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_push.cuh create mode 100644 python/sglang/jit_kernel/csrc/distributed/ipc.cuh delete mode 100644 python/sglang/jit_kernel/include/sgl_kernel/distributed/common.cuh create mode 100644 python/sglang/jit_kernel/include/sgl_kernel/distributed/communicator.cuh delete mode 100644 python/sglang/jit_kernel/include/sgl_kernel/distributed/custom_all_reduce.cuh create mode 100644 python/sglang/srt/distributed/device_communicators/configs/__init__.py create mode 100644 python/sglang/srt/distributed/device_communicators/configs/custom_all_reduce_v2.py diff --git a/python/sglang/jit_kernel/all_reduce.py b/python/sglang/jit_kernel/all_reduce.py index 05b31127a..ec8dd12f4 100644 --- a/python/sglang/jit_kernel/all_reduce.py +++ b/python/sglang/jit_kernel/all_reduce.py @@ -1,7 +1,7 @@ from __future__ import annotations import enum -from typing import TYPE_CHECKING, List, NamedTuple, Optional, Tuple, cast +from typing import TYPE_CHECKING, List, Tuple, Union import torch import tvm_ffi @@ -10,17 +10,13 @@ from tvm_ffi import Module from sglang.jit_kernel.utils import ( cache_once, is_arch_support_pdl, + lazy_register_class, load_jit, make_cpp_args, ) from sglang.kernel_api_logging import debug_kernel_api -class ConfigResult(NamedTuple): - num_blocks: int - num_threads: int - - class AllReduceAlgo(enum.Enum): ONE_SHOT_PUSH = enum.auto() ONE_SHOT_PULL = enum.auto() @@ -30,97 +26,142 @@ class AllReduceAlgo(enum.Enum): return self == AllReduceAlgo.ONE_SHOT_PUSH @property - def shot(self) -> int: - return 2 if self == AllReduceAlgo.TWO_SHOT_PULL else 1 + def algo_name(self) -> str: + return _ALGO_NAMES[self] +_ALGO_NAMES = { + AllReduceAlgo.ONE_SHOT_PUSH: "1shot_push", + AllReduceAlgo.ONE_SHOT_PULL: "1shot_pull", + AllReduceAlgo.TWO_SHOT_PULL: "2shot_pull", +} + +# ``pull_arg`` of the all-reduce kernel: a row of the graph-params pointer +# table selects graph mode; a plain bool selects multicast (True) / eager. +PullArg = Union[torch.Tensor, bool] + if TYPE_CHECKING: - CUSTOM_AR_HANDLE = List[int] - CUSTOM_AR_PAIR = Tuple[int, CUSTOM_AR_HANDLE] + # (cudaIpcMemHandle bytes, offset-in-allocation) for one device pointer + IPC_HANDLE_PAIR = Tuple[List[int], int] - class CustomAllReduceObj: - def __init__( - self, - rank: int, - world_size: int, - pull_buffer_bytes: int, - push_buffer_bytes: int, - graph_input_count: int, - *, - max_pull_blocks: Optional[int] = None, - max_push_blocks: Optional[int] = None, - ) -> None: - """ - Create a CustomAllReduceObj instance. - :param rank: The rank of the current process. - :param world_size: The total number of processes in the group. - :param pull_buffer_bytes: The size of the buffer (in bytes) used for pull-based all-reduce. - :param push_buffer_bytes: The size of the buffer (in bytes) used for push-based all-reduce. - :param graph_input_count: The maximum number of inputs in all CUDA graphs. - :param max_pull_blocks: The maximum number of thread blocks to launch for pull-based all-reduce. - If None, it will be determined by the implementation. - :param max_push_blocks: The maximum number of thread blocks to launch for push-based all-reduce. - If None, it will be determined by the implementation. - """ +def _init_communicator() -> None: + module = load_jit( + "communicator", + cuda_files=["distributed/communicator.cuh"], + cuda_wrappers=[("register_once", "register_communicator")], + ) + module.register_once() - @property - def world_size(self) -> int: ... - def share_storage(self) -> CUSTOM_AR_HANDLE: ... - def share_graph_inputs(self) -> List[CUSTOM_AR_PAIR]: ... - def post_init(self, handles: List[CUSTOM_AR_HANDLE]) -> None: ... - def register_inputs(self, handles: List[List[CUSTOM_AR_PAIR]]) -> None: ... - def set_cuda_graph_capture(self, is_capturing: bool) -> None: ... - def get_graph_capture_bases( - self, - ) -> Tuple[List[Tuple[int, int]], List[List[int]], List[int]]: ... - def free(self, tp_cpu_group: torch.distributed.ProcessGroup) -> None: ... - def all_reduce( - self, input: torch.Tensor, algo: AllReduceAlgo - ) -> tvm_ffi.Tensor: ... - def config_pull( - self, num_blocks: int = -1, num_threads: int = -1 - ) -> ConfigResult: - """ - Configure the CUDA kernel's grid and block dimensions. - This provides only the upper bound of the configuration, - and the actual launch configuration may be determined by implementation. - Note that push-based all-reduce can not be configured currently. - :param num_blocks: The maximum number of thread blocks to launch. -1 means no limit. - :param num_threads: The maximum number of threads per block. -1 means no limit. +@lazy_register_class("sgl.Communicator", _init_communicator) +class Communicator(tvm_ffi.Object): + """Storage plane of the custom all-reduce: a thin pointer holder. - :return: The previous configuration as a ConfigResult named tuple. - """ - ... + All buffers are owned by the caller (symmetric-memory tensor views plus + a local push counter); this object only validates and records them. + """ + + if TYPE_CHECKING: + # C++ interface + rank: int + world_size: int + + def _config(self, kwargs: dict) -> None: ... + + def __init__( + self, + rank: int, + world_size: int, + push_workspaces: List[torch.Tensor], + pull_workspaces: List[torch.Tensor], + pull_semaphores: List[torch.Tensor], + push_counter: torch.Tensor, + pull_mc_workspace: int | None, + ) -> None: + """ + :param push_workspaces: per-rank ``[2 * world_size, push_bytes]`` + uint8 views of symmetric memory. + :param pull_workspaces: per-rank ``[pull_bytes]`` uint8 views of + symmetric memory. + :param pull_semaphores: per-rank ``[num_pull_blocks, 128]`` uint8 + views of symmetric memory. + :param push_counter: local ``[num_push_blocks, 4]`` uint8 tensor. + :param pull_mc_workspace: multicast address of the pull workspace, + or None when multicast is unavailable. + """ + self.__ffi_init__( + rank, + world_size, + push_workspaces, + pull_workspaces, + pull_semaphores, + push_counter, + pull_mc_workspace, + ) + + def config( + self, + num_pull_blocks: int | None = None, + num_multicast_blocks: int | None = None, + ) -> Communicator: + kwargs = {} + if num_pull_blocks is not None: + kwargs["num_pull_blocks"] = num_pull_blocks + if num_multicast_blocks is not None: + kwargs["num_multicast_blocks"] = num_multicast_blocks + self._config(kwargs) + return self + + +def _init_ipc_manager() -> None: + module = load_jit( + "cuda_ipc", + extra_ldflags=["-lcuda"], + cuda_files=["distributed/ipc.cuh"], + cuda_wrappers=[("register_once", "register_ipc_manager")], + ) + module.register_once() + + +@lazy_register_class("sgl.IPCManager", _init_ipc_manager) +class IPCManager(tvm_ffi.Object): + """Batched cudaIpc handle exchange for CUDA-graph input pointers.""" + + if TYPE_CHECKING: + # C++ interface + def destroy(self) -> None: ... + def batch_get_handles(self, ptrs: List[int]) -> List[IPC_HANDLE_PAIR]: ... + def batch_open_handles(self, handles: List[IPC_HANDLE_PAIR]) -> List[int]: ... + + def __init__(self) -> None: + self.__ffi_init__() @cache_once -def _jit_custom_all_reduce_pull_module(dtype: torch.dtype, world_size: int) -> Module: +def get_all_reduce_module(dtype: torch.dtype, world_size: int) -> Module: args = make_cpp_args(dtype, world_size, is_arch_support_pdl()) return load_jit( - "custom_all_reduce_pull", + "custom_all_reduce", *args, - extra_ldflags=["-lcuda"], - cuda_files=["distributed/custom_all_reduce_pull.cuh"], + cuda_files=["distributed/custom_all_reduce.cuh"], cuda_wrappers=[("all_reduce", f"custom_all_reduce<{args}>")], ) -@cache_once -def _jit_custom_all_reduce_push_module(dtype: torch.dtype, world_size: int) -> Module: - args = make_cpp_args(dtype, world_size, is_arch_support_pdl()) - return load_jit( - "custom_all_reduce_push", - *args, - extra_ldflags=["-lcuda"], - cuda_files=["distributed/custom_all_reduce_push.cuh"], - cuda_wrappers=[("all_reduce", f"custom_all_reduce<{args}>")], - ) +@debug_kernel_api +def custom_all_reduce( + comm: Communicator, + input: torch.Tensor, + algo: AllReduceAlgo, + pull_arg: PullArg, +) -> tvm_ffi.Tensor: + module = get_all_reduce_module(input.dtype, comm.world_size) + return module.all_reduce(comm, input, algo.algo_name, pull_arg) @cache_once -def _jit_fused_parallel_qknorm_module( +def get_fused_parallel_qknorm_module( dtype: torch.dtype, world_size: int, q_dim: int, k_dim: int ) -> Module: args = make_cpp_args(dtype, world_size, q_dim, k_dim, is_arch_support_pdl()) @@ -128,7 +169,6 @@ def _jit_fused_parallel_qknorm_module( return load_jit( "tp_qknorm", *args, - extra_ldflags=["-lcuda"], cuda_files=["distributed/tp_qknorm.cuh"], cuda_wrappers=[ ("fused_parallel_qknorm", f"{cls_name}::run"), @@ -137,107 +177,23 @@ def _jit_fused_parallel_qknorm_module( ) -@cache_once -def get_custom_all_reduce_cls() -> type[CustomAllReduceObj]: - module = load_jit( - "custom_all_reduce_base", - extra_ldflags=["-lcuda"], - cuda_files=["distributed/custom_all_reduce_base.cuh"], - cuda_wrappers=[("register_once", "register_custom_all_reduce")], - ) - module.register_once() - device = torch.cuda.current_device() - props = torch.cuda.get_device_properties(device) - NUM_CTA = props.multi_processor_count - MAX_THREADS = 512 - - @tvm_ffi.register_object("sgl.CustomAllReduce") - class CustomAllReduceObjReal(tvm_ffi.Object): - __slots__ = ("__dict__",) - - def __init__( - self, - rank: int, - world_size: int, - pull_buffer_bytes: int, - push_buffer_bytes: int, - graph_input_count: int, - *, - max_pull_blocks: Optional[int] = None, - max_push_blocks: Optional[int] = None, - ) -> None: - max_pull_blocks = NUM_CTA if max_pull_blocks is None else max_pull_blocks - max_push_blocks = NUM_CTA if max_push_blocks is None else max_push_blocks - self.__ffi_init__( - rank, - world_size, - max_pull_blocks, - max_push_blocks, - pull_buffer_bytes, - push_buffer_bytes, - graph_input_count, - ) - self._world_size = world_size - self._pull_config = ConfigResult(min(NUM_CTA, max_pull_blocks), MAX_THREADS) - if max_pull_blocks > 0: # special case: cannot configure 0 blocks - self.configure_pull(*self._pull_config) # type: ignore - - @property - def world_size(self) -> int: - return self._world_size - - @debug_kernel_api - def all_reduce( - self, - input: torch.Tensor, - algo: AllReduceAlgo, - ) -> tvm_ffi.Tensor: - compile_fn = ( - _jit_custom_all_reduce_push_module - if algo.is_push() - else _jit_custom_all_reduce_pull_module - ) - module = compile_fn(input.dtype, self._world_size) - return module.all_reduce(self, input, algo.shot) - - def config_pull( - self, num_blocks: int = -1, num_threads: int = -1 - ) -> ConfigResult: - old_config = self._pull_config - num_blocks = num_blocks if num_blocks != -1 else old_config.num_blocks - num_threads = num_threads if num_threads != -1 else old_config.num_threads - new_config = ConfigResult(num_blocks, num_threads) - if new_config != old_config: - result = ConfigResult(*self.configure_pull(*new_config)) # type: ignore - assert result == self._pull_config - self._pull_config = new_config - return old_config - - def free(self, tp_cpu_group: torch.distributed.ProcessGroup) -> None: - self.free_ipc_handles() # type: ignore - torch.distributed.barrier(group=tp_cpu_group) - self.free_storage() # type: ignore - - return cast(type["CustomAllReduceObj"], CustomAllReduceObjReal) - - def get_fused_parallel_qknorm_max_occupancy( dtype: torch.dtype, world_size: int, q_dim: int, k_dim: int ) -> int: - module = _jit_fused_parallel_qknorm_module(dtype, world_size, q_dim, k_dim) + module = get_fused_parallel_qknorm_module(dtype, world_size, q_dim, k_dim) return module.get_max_occupancy() def fused_parallel_qknorm( - custom_ar: CustomAllReduceObj, + comm: Communicator, q: torch.Tensor, k: torch.Tensor, q_weight: torch.Tensor, k_weight: torch.Tensor, eps: float = 1e-6, ) -> None: - world_size = custom_ar.world_size + world_size = comm.world_size q_dim = q.shape[-1] * world_size k_dim = k.shape[-1] * world_size - module = _jit_fused_parallel_qknorm_module(q.dtype, world_size, q_dim, k_dim) - module.fused_parallel_qknorm(custom_ar, q, k, q_weight, k_weight, eps) + module = get_fused_parallel_qknorm_module(q.dtype, world_size, q_dim, k_dim) + module.fused_parallel_qknorm(comm, q, k, q_weight, k_weight, eps) diff --git a/python/sglang/jit_kernel/benchmark/utils.py b/python/sglang/jit_kernel/benchmark/utils.py index 1aa15c74f..74fed6b79 100644 --- a/python/sglang/jit_kernel/benchmark/utils.py +++ b/python/sglang/jit_kernel/benchmark/utils.py @@ -18,7 +18,7 @@ def multigpu_bench_main( pre_launch_fn: Optional[Callable[[List[int]], None]] = None, timeout: Optional[int] = None, ) -> None: - """cudalib-style multi-GPU benchmark entry point. + """Torchrun-based multi-GPU benchmark entry point. Drop this at the bottom of a benchmark file:: diff --git a/python/sglang/jit_kernel/csrc/distributed/communicator.cuh b/python/sglang/jit_kernel/csrc/distributed/communicator.cuh new file mode 100644 index 000000000..19210e449 --- /dev/null +++ b/python/sglang/jit_kernel/csrc/distributed/communicator.cuh @@ -0,0 +1,113 @@ +#include +#include +#include + +#include + +#include +#include + +#include +#include +#include +#include +#include + +namespace host::distributed { + +inline CommunicatorObj::CommunicatorObj( + const uint32_t rank, + const uint32_t world_size, + std::vector push_workspaces, + std::vector pull_workspaces, + std::vector pull_semaphores, + TensorView push_counter, + const std::optional pull_mc_workspace_ptr) { + this->rank = rank; + this->world_size = world_size; + RuntimeCheck(1 < world_size && world_size <= kMaxWorldSize, "Invalid world size: ", world_size); + RuntimeCheck(rank < world_size, "Invalid rank: ", rank); + RuntimeCheck(push_workspaces.size() == world_size, "Bad push workspace count"); + RuntimeCheck(pull_workspaces.size() == world_size, "Bad pull workspace count"); + RuntimeCheck(pull_semaphores.size() == world_size, "Bad pull semaphore count"); + // Shared symbolic sizes / device enforce consistency across ranks; the + // matchers also require contiguity (no strides given) and uint8 dtype. + auto push_bytes = SymbolicSize{"push_bytes"}; + auto pull_bytes = SymbolicSize{"pull_bytes"}; + auto num_pull_blocks = SymbolicSize{"num_pull_blocks"}; + auto num_push_blocks = SymbolicSize{"num_push_blocks"}; + auto device = SymbolicDevice{}; + device.set_options(); + for (uint32_t i = 0; i < world_size; ++i) { + TensorMatcher({2 * world_size, push_bytes}).with_dtype().with_device(device).verify(push_workspaces[i]); + TensorMatcher({pull_bytes}) // + .with_dtype() + .with_device(device) + .verify(pull_workspaces[i]); + TensorMatcher({num_pull_blocks, static_cast(sizeof(Semaphore))}) + .with_dtype() + .with_device(device) + .verify(pull_semaphores[i]); + this->push_workspaces[i] = static_cast(push_workspaces[i].data_ptr()); + this->pull_workspaces[i] = static_cast(pull_workspaces[i].data_ptr()); + this->pull_semaphores[i] = static_cast(pull_semaphores[i].data_ptr()); + } + TensorMatcher({num_push_blocks, static_cast(sizeof(Counter))}) + .with_dtype() + .with_device(device) + .verify(push_counter); + RuntimeCheck(push_bytes.unwrap() > 0 && pull_bytes.unwrap() > 0, "Workspace sizes must be positive"); + + if (pull_mc_workspace_ptr.has_value()) { + this->pull_mc_workspace = reinterpret_cast(static_cast(pull_mc_workspace_ptr.value())); + } else { + this->pull_mc_workspace = nullptr; + } + + // push config + this->push_counter = static_cast(push_counter.data_ptr()); + this->push_bytes = push_bytes.unwrap(); + this->num_push_blocks = static_cast(num_push_blocks.unwrap()); + + // pull config + this->pull_bytes = pull_bytes.unwrap(); + this->num_pull_blocks = static_cast(num_pull_blocks.unwrap()); + this->num_multicast_blocks = this->num_pull_blocks; + this->total_pull_blocks = this->num_pull_blocks; +} + +inline void CommunicatorObj::config(std::map config) { + for (const auto& [key, value] : config) { + if (key == "num_pull_blocks") { + RuntimeCheck(value > 0 && value <= total_pull_blocks, "Invalid number of pull blocks: ", value); + this->num_pull_blocks = value; + } else if (key == "num_multicast_blocks") { + RuntimeCheck(value > 0 && value <= total_pull_blocks, "Invalid number of multicast blocks: ", value); + this->num_multicast_blocks = value; + } else { + RuntimeCheck(false, "Unknown config key: ", key); + } + } +} + +} // namespace host::distributed + +inline void register_communicator() { + namespace refl = tvm::ffi::reflection; + using Class = host::distributed::CommunicatorObj; + using TensorView = tvm::ffi::TensorView; + refl::ObjectDef() + .def( + refl::init< + uint32_t, + uint32_t, + std::vector, + std::vector, + std::vector, + TensorView, + std::optional>(), + "__init__") + .def_ro("world_size", &Class::world_size) + .def_ro("rank", &Class::rank) + .def("_config", &Class::config); +} diff --git a/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce.cuh b/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce.cuh new file mode 100644 index 000000000..8412eea81 --- /dev/null +++ b/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce.cuh @@ -0,0 +1,650 @@ +// Custom all-reduce kernels over the decoupled Communicator storage plane. +// +// Three algorithms are provided behind one entry point: +// - 1shot_push: lamport-style push of local data to every peer's push +// workspace, then a local polling reduce (best at small sizes). +// - 1shot_pull: every rank reduces all peers' data (from the symmetric pull +// workspaces, a CUDA-graph pointer table, or a multicast address). +// - 2shot_pull: reduce-scatter fused with all-gather; each rank reduces its +// shard in place so every workspace ends up holding the full result. +// +// Unlike the previous implementation, the kernels carry no storage or IPC +// logic: all pointers arrive via `CommunicatorObj` (owned by Python) and the +// per-call `AllReduceParams`. +#include +#include +#include + +#include +#include +#include + +#include + +#include + +#include +#include +#include +#include +#include +#include + +namespace { + +using device::distributed::Counter, device::distributed::Semaphore; +using host::distributed::CommunicatorRef; + +inline constexpr uint32_t kMaxWorldSize = device::distributed::kMaxWorldSize; + +enum class PullMode { + Graph, + Eager, + Multicast, // also eager +}; + +template +struct fp_trait {}; + +template <> +struct fp_trait { + using type = uint16_t; + [[maybe_unused]] + static constexpr uint16_t pos_zero = 0x0000u; + [[maybe_unused]] + static constexpr uint16_t neg_zero = 0x8000u; +}; + +template <> +struct fp_trait { + using type = uint16_t; + [[maybe_unused]] + static constexpr uint16_t pos_zero = 0x0000u; + [[maybe_unused]] + static constexpr uint16_t neg_zero = 0x8000u; +}; + +template <> +struct fp_trait { + using type = uint32_t; + [[maybe_unused]] + static constexpr uint32_t pos_zero = 0x00000000u; + [[maybe_unused]] + static constexpr uint32_t neg_zero = 0x80000000u; +}; + +template +SGL_DEVICE void clear_pos_zero(DType& val) { + using Trait = fp_trait; + const auto ptr = reinterpret_cast(&val); + if (*ptr == Trait::pos_zero) *ptr = Trait::neg_zero; +} + +template +SGL_DEVICE bool is_pos_zero(const DType& val) { + using Trait = fp_trait; + const auto ptr = reinterpret_cast(&val); + return *ptr == Trait::pos_zero; +} + +template +SGL_DEVICE DType get_pos_zero() { + using Trait = fp_trait; + const auto value = Trait::pos_zero; + return *reinterpret_cast(&value); +} + +template +SGL_DEVICE auto reduce(device::AlignedVector (&vec)[M]) -> device::AlignedVector { + fp32x2_t acc_vec[N]; +#pragma unroll + for (size_t i = 0; i < M; ++i) { +#pragma unroll + for (size_t j = 0; j < N; ++j) { + const auto [x, y] = device::cast(vec[i][j]); + auto& [acc_x, acc_y] = acc_vec[j]; + acc_x = i == 0 ? x : acc_x + x; + acc_y = i == 0 ? y : acc_y + y; + } + } + device::AlignedVector out_vec; +#pragma unroll + for (size_t j = 0; j < N; ++j) { + out_vec[j] = device::cast(acc_vec[j]); + } + return out_vec; +} + +template +SGL_DEVICE void ld_global_16B(V& x, const void* addr, int64_t vec_offset) { + static_assert(alignof(V) == 16 && sizeof(V) == 16); + addr = static_cast(addr) + vec_offset * sizeof(V); + uint4 val; + asm volatile("ld.global.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(val.x), "=r"(val.y), "=r"(val.z), "=r"(val.w) + : "l"(addr)); + x = *reinterpret_cast(&val); +} + +template +SGL_DEVICE void st_global_16B(const V& x, void* addr, int64_t vec_offset) { + static_assert(alignof(V) == 16 && sizeof(V) == 16); + const uint4 val = *reinterpret_cast(&x); + addr = static_cast(addr) + vec_offset * sizeof(V); + asm volatile("st.global.v4.b32 [%4], {%0, %1, %2, %3};" + : // + : "r"(val.x), "r"(val.y), "r"(val.z), "r"(val.w), "l"(addr)); +} + +template +SGL_DEVICE void ld_relaxed_16B(V& x, const void* addr, int64_t vec_offset) { + static_assert(alignof(V) == 16 && sizeof(V) == 16); + addr = static_cast(addr) + vec_offset * sizeof(V); + uint4 val; + asm volatile("ld.relaxed.sys.global.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(val.x), "=r"(val.y), "=r"(val.z), "=r"(val.w) + : "l"(addr)); + x = *reinterpret_cast(&val); +} + +template +SGL_DEVICE void st_relaxed_16B(const V& x, void* addr, int64_t vec_offset) { + static_assert(alignof(V) == 16 && sizeof(V) == 16); + const uint4 val = *reinterpret_cast(&x); + addr = static_cast(addr) + vec_offset * sizeof(V); + asm volatile("st.relaxed.sys.global.v4.b32 [%4], {%0, %1, %2, %3};" + : // + : "r"(val.x), "r"(val.y), "r"(val.z), "r"(val.w), "l"(addr)); +} + +template +SGL_DEVICE void ld_multimem_16B(V& x, const void* mc_addr, int64_t vec_offset) { +#if SGL_ARCH_HOPPER_OR_GREATER + static_assert(alignof(V) == 16 && sizeof(V) == 16); + mc_addr = static_cast(mc_addr) + vec_offset * 16; + if constexpr (std::is_same_v>) { + float4 val; + asm volatile("multimem.ld_reduce.weak.add.v4.f32 {%0, %1, %2, %3}, [%4];" + : "=f"(val.x), "=f"(val.y), "=f"(val.z), "=f"(val.w) + : "l"(mc_addr)); + x = *reinterpret_cast(&val); + } else { + // Packed f16x2/bf16x2 results live in b32 registers ("=r"); .acc::f32 only + // raises the accumulation precision, not the result register type — ptxas + // rejects .f32 ("=f") destinations with "Arguments mismatch". + uint4 val; + if constexpr (std::is_same_v>) { + asm volatile("multimem.ld_reduce.weak.add.acc::f32.v4.f16x2 {%0, %1, %2, %3}, [%4];" + : "=r"(val.x), "=r"(val.y), "=r"(val.z), "=r"(val.w) + : "l"(mc_addr)); + } else { + static_assert(std::is_same_v>); // 4x bf16x2 + asm volatile("multimem.ld_reduce.weak.add.acc::f32.v4.bf16x2 {%0, %1, %2, %3}, [%4];" + : "=r"(val.x), "=r"(val.y), "=r"(val.z), "=r"(val.w) + : "l"(mc_addr)); + } + x = *reinterpret_cast(&val); + } +#else + assert(false && "multimem load is only supported on Hopper or later architecture"); +#endif +} + +template +SGL_DEVICE void st_multimem_16B(const V& x, void* mc_addr, int64_t vec_offset) { +#if SGL_ARCH_HOPPER_OR_GREATER + static_assert(alignof(V) == 16 && sizeof(V) == 16); + const auto val = *reinterpret_cast(&x); + mc_addr = static_cast(mc_addr) + vec_offset * 16; + asm volatile("multimem.st.weak.v4.f32 [%4], {%0, %1, %2, %3};" + : + : "f"(val.x), "f"(val.y), "f"(val.z), "f"(val.w), "l"(mc_addr)); +#else + assert(false && "multimem store is only supported on Hopper or later architecture"); +#endif +} + +struct AllReduceParams { + const void* __restrict__ input; + void* __restrict__ output; + uint32_t num_elements; + uint32_t rank; + void* const* __restrict__ graph_params; + uint8_t* pull_workspaces[kMaxWorldSize]; // must be symmetric memory + uint8_t* push_workspaces[kMaxWorldSize]; // must be symmetric memory + Semaphore* pull_semaphores[kMaxWorldSize]; // must be symmetric memory + Counter* push_counter; + uint8_t* pull_mc_workspace; // must be a multicast address + int64_t push_buffer_stride; // per-buffer bytes; each rank holds 2 * kMaxWorldSize buffers +}; + +template +struct AllReducePushImpl { + private: + using T2 = packed_t; + /// NOTE: force 16B load/store to reduce register pressure + static constexpr uint32_t kVecSize = 16 / sizeof(T2); + static constexpr uint32_t kElemsPerVec = 16 / sizeof(T); + using vec_t = device::AlignedVector; + static_assert(kWorldSize <= kMaxWorldSize); + + static SGL_DEVICE bool sync_enter_push(const AllReduceParams& params) { + device::PDLWaitPrimary(); + return (params.push_counter[blockIdx.x].get() % 2) != 0; + } + + static SGL_DEVICE void sync_exit_push(const AllReduceParams& params) { + device::PDLTriggerSecondary(); + __syncthreads(); + if (threadIdx.x == 0) { + params.push_counter[blockIdx.x].inc(1); // NOTE: u32 overflow is safe under mod 2 + } + } + + static SGL_DEVICE void push_impl(uint32_t num_vecs, void* (&data)[kWorldSize], const void* src) { + const auto num_threads = blockDim.x * gridDim.x; + const auto global_tid = blockIdx.x * blockDim.x + threadIdx.x; +#pragma unroll + for (auto vid = global_tid; vid < num_vecs; vid += num_threads) { + vec_t vec; + ld_global_16B(vec, src, vid); +#pragma unroll + for (uint32_t j = 0; j < kVecSize; ++j) { + clear_pos_zero(vec[j].x); + clear_pos_zero(vec[j].y); + } +#pragma unroll + for (uint32_t i = 0; i < kWorldSize; ++i) { + st_relaxed_16B(vec, data[i], vid); + } + } + } + + static SGL_DEVICE void poll_impl(uint32_t num_vecs, void* (&data)[kWorldSize], void* out) { + // need polling to ensure data is ready + const auto num_threads = blockDim.x * gridDim.x; + const auto global_tid = blockIdx.x * blockDim.x + threadIdx.x; + // pos_zero-filled vec we write back after consuming each slot, so the + // double-buffered phase comes back around with the "slot empty" marker + // re-established. + vec_t pos_zero_vec; + { + const auto z = get_pos_zero(); +#pragma unroll + for (uint32_t j = 0; j < kVecSize; ++j) { + pos_zero_vec[j].x = z; + pos_zero_vec[j].y = z; + } + } +#pragma unroll + for (auto vid = global_tid; vid < num_vecs; vid += num_threads) { + vec_t vec[kWorldSize]; + do { + bool has_zero = false; +#pragma unroll + for (uint32_t i = 0; i < kWorldSize; ++i) { + ld_relaxed_16B(vec[i], data[i], vid); + } +#pragma unroll + for (uint32_t i = 0; i < kWorldSize; ++i) { +#pragma unroll + for (uint32_t j = 0; j < kVecSize; ++j) { + has_zero |= is_pos_zero(vec[i][j].x); + has_zero |= is_pos_zero(vec[i][j].y); + } + } + if (!has_zero) break; + } while (true); + const auto out_vec = reduce(vec); + st_global_16B(out_vec, out, vid); +#pragma unroll + for (uint32_t i = 0; i < kWorldSize; ++i) { + st_global_16B(pos_zero_vec, data[i], vid); + } + } + } + + public: + static SGL_DEVICE void forward_1shot(const AllReduceParams& params) { + // push local data to peer ranks, then reduce locally + const auto phase = sync_enter_push(params); + const auto r = params.rank; + const auto num_vecs = device::div_ceil(params.num_elements, kElemsPerVec); + const auto stride_bytes = params.push_buffer_stride; + const auto phase_stride_bytes = phase * stride_bytes * kWorldSize; + + // push to peer + void* push_buf[kWorldSize]; +#pragma unroll + for (uint32_t i = 0; i < kWorldSize; ++i) { + push_buf[i] = params.push_workspaces[i] + r * stride_bytes + phase_stride_bytes; + } + push_impl(num_vecs, push_buf, params.input); + + // poll from local + void* poll_buf[kWorldSize]; +#pragma unroll + for (uint32_t i = 0; i < kWorldSize; ++i) { + poll_buf[i] = params.push_workspaces[r] + i * stride_bytes + phase_stride_bytes; + } + poll_impl(num_vecs, poll_buf, params.output); + + sync_exit_push(params); + } +}; + +template +struct AllReducePullImpl { + private: + using T2 = packed_t; + static constexpr uint32_t kVecSize = 16 / sizeof(T2); + static constexpr uint32_t kElemsPerVec = 16 / sizeof(T); + using vec_t = device::AlignedVector; + static_assert(kWorldSize <= kMaxWorldSize); + + template + static SGL_DEVICE uint32_t sync_enter_pull(const AllReduceParams& params) { + uint32_t current_counter_val = 0; + if (const auto tx = threadIdx.x; tx < kWorldSize) { + device::PDLWaitPrimary(); + const auto bx = blockIdx.x; + const auto semaphore = ¶ms.pull_semaphores[tx][bx]; + const auto counter = semaphore->counter_ptr(); + const auto current = tx == params.rank ? counter->inc(2 * kWorldSize) : 0; + current_counter_val = current; + if constexpr (kFence) { + semaphore->put_release(); + } else { + semaphore->put_relaxed(); + } + if (tx == params.rank) { + if constexpr (kFence) { + while (semaphore->get_acquire() - current < kWorldSize) + ; + } else { + while (semaphore->get_relaxed() - current < kWorldSize) + ; + } + } + } + __syncthreads(); + return current_counter_val + kWorldSize; + } + + template + static SGL_DEVICE void sync_exit_pull(const AllReduceParams& params, uint32_t current) { + device::PDLTriggerSecondary(); + __syncthreads(); + if (const auto tx = threadIdx.x; tx < kWorldSize) { + const auto bx = blockIdx.x; + const auto semaphore = ¶ms.pull_semaphores[tx][bx]; + if constexpr (kFence) { + semaphore->put_release(); + } else { + semaphore->put_relaxed(); + } + if (tx == params.rank) { + if constexpr (kFence) { + while (semaphore->get_acquire() - current < kWorldSize) + ; + } else { + while (semaphore->get_relaxed() - current < kWorldSize) + ; + } + } + } + } + + template + static SGL_DEVICE void reduce_impl( + uint32_t num_vecs, // + [[maybe_unused]] void* (&data)[kWorldSize], + [[maybe_unused]] void* out, + [[maybe_unused]] void* mc_addr) { + const auto num_threads = blockDim.x * gridDim.x; + const auto global_tid = blockIdx.x * blockDim.x + threadIdx.x; + for (auto vid = global_tid; vid < num_vecs; vid += num_threads) { + if constexpr (kMode == PullMode::Multicast) { + vec_t out_vec; + ld_multimem_16B(out_vec, mc_addr, vid); + if constexpr (kIs2shot) { + // inplace write to workspace for 2-shot all reduce + st_multimem_16B(out_vec, mc_addr, vid); + } else { + // write to output for 1-shot all reduce + out_vec.store(out, vid); + } + } else { + vec_t vec[kWorldSize]; +#pragma unroll + for (uint32_t i = 0; i < kWorldSize; ++i) { + vec[i].load(data[i], vid); + } + const auto out_vec = reduce(vec); + if constexpr (kIs2shot) { + // inplace write to buffer for 2-shot all reduce +#pragma unroll + for (uint32_t i = 0; i < kWorldSize; ++i) { + out_vec.store(data[i], vid); + } + } else { + // write to output for 1-shot all reduce + out_vec.store(out, vid); + } + } + } + } + + public: + static SGL_DEVICE void forward_1shot(const AllReduceParams& params) { + const auto total_num_vecs = device::div_ceil(params.num_elements, kElemsPerVec); + void* data[kWorldSize]; + if constexpr (kMode == PullMode::Graph) { +#pragma unroll + for (uint32_t i = 0; i < kWorldSize; ++i) { + data[i] = params.graph_params[i]; + } + } else if constexpr (kMode == PullMode::Eager) { +#pragma unroll + for (uint32_t i = 0; i < kWorldSize; ++i) { + data[i] = params.pull_workspaces[i]; + } + } + const auto counter = sync_enter_pull(params); + reduce_impl(total_num_vecs, data, params.output, params.pull_mc_workspace); + sync_exit_pull(params, counter); + } + + static SGL_DEVICE void forward_2shot(const AllReduceParams& params) { + const auto total_num_vecs = device::div_ceil(params.num_elements, kElemsPerVec); + const auto avg_vecs = total_num_vecs / kWorldSize; + const auto rem_vecs = total_num_vecs % kWorldSize; + // usually, hidden size is a multiple of 1024, so 1024 / 8 = 128 is typically 128-bytes aligned + const auto local_vec_bias = avg_vecs * params.rank + min(params.rank, rem_vecs); + const auto local_num_vecs = avg_vecs + (params.rank < rem_vecs ? 1 : 0); + void* data[kWorldSize]; + if constexpr (kMode == PullMode::Graph) { +#pragma unroll + for (uint32_t i = 0; i < kWorldSize; ++i) { + data[i] = reinterpret_cast(params.graph_params[i]) + local_vec_bias; + } + } else if constexpr (kMode == PullMode::Eager) { +#pragma unroll + for (uint32_t i = 0; i < kWorldSize; ++i) { + data[i] = reinterpret_cast(params.pull_workspaces[i]) + local_vec_bias; + } + } + const auto counter = sync_enter_pull(params); + const auto mc_addr = reinterpret_cast(params.pull_mc_workspace) + local_vec_bias; + reduce_impl(local_num_vecs, data, params.output, mc_addr); + sync_exit_pull(params, counter); + } +}; + +template +__global__ __launch_bounds__(1024, 1) // + void all_reduce_kernel(const __grid_constant__ AllReduceParams params) { + static_assert(kShot == 1 || kShot == 2, "invalid shot"); + if constexpr (kShot == 1) { + return Impl::forward_1shot(params); + } else { + return Impl::forward_2shot(params); + } +} + +template +__global__ void memcpy_kernel(void* __restrict__ dst, const void* __restrict__ src, uint32_t num_vecs) { + static_assert(N % 4 == 0, "at least 4-bytes aligned for uint32_t load/store"); + using vec_t = device::AlignedVector; + const auto tid = blockIdx.x * blockDim.x + threadIdx.x; + device::PDLWaitPrimary(); + device::PDLTriggerSecondary(); + if (tid < num_vecs) { + vec_t vec; + vec.load(src, tid); + vec.store(dst, tid); + } +} + +// Pick the smallest block size whose grid still fits in one wave; the kernels +// are grid-stride so any choice is correct, this only tunes occupancy. +[[maybe_unused]] +uint32_t choose_block_size(uint32_t num_threads) { + static const uint32_t kNumSM = [] { + int device = 0, sm_count = 0; + host::RuntimeDeviceCheck(cudaGetDevice(&device)); + host::RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, device)); + return static_cast(sm_count); + }(); + for (const uint32_t block_size : {128u, 256u, 512u}) { + if (host::div_ceil(num_threads, block_size) <= kNumSM) return block_size; + } + return 1024u; +} + +template +struct AllReduceKernel { + private: + using Tensor = tvm::ffi::Tensor; + using TensorView = tvm::ffi::TensorView; + template + static constexpr auto kernel_pull = all_reduce_kernel, kShot>; + template + static constexpr auto kernel_push = all_reduce_kernel, kShot>; + + public: + static Tensor run(CommunicatorRef ref, Tensor in_, std::string algo, std::variant pull_arg) { + using namespace host; + const auto& data = *ref.get(); + RuntimeCheck(algo == "1shot_pull" || algo == "2shot_pull" || algo == "1shot_push", "Invalid algo: ", algo); + RuntimeCheck(data.world_size == kWorldSize, "Mismatch world size"); + RuntimeCheck(in_.IsContiguous(), "Input tensor must be contiguous"); + RuntimeCheck(is_type(in_.dtype()), "Input dtype mismatch"); + RuntimeCheck(in_.device().device_type == kDLCUDA, "Only CUDA device is supported"); + RuntimeCheck(std::bit_cast(in_.data_ptr()) % 16 == 0, "Input pointer is not properly aligned"); + const auto num_elems_int64 = in_.numel(); + const auto num_elems = static_cast(num_elems_int64); + RuntimeCheck(static_cast(num_elems) == num_elems_int64, "Number of items exceeds 4G limit"); + const bool use_graph = std::holds_alternative(pull_arg); + const auto graph_ptr = use_graph ? std::get(pull_arg).data_ptr() : nullptr; + const bool inplace = use_graph && algo == "2shot_pull"; + Tensor out = inplace ? in_ : ffi::empty_like(in_); + AllReduceParams params{ + .input = in_.data_ptr(), + .output = out.data_ptr(), + .num_elements = num_elems, + .rank = data.rank, + .graph_params = static_cast(graph_ptr), + .pull_workspaces = {}, + .push_workspaces = {}, + .pull_semaphores = {}, + .push_counter = data.push_counter, + .pull_mc_workspace = data.pull_mc_workspace, + .push_buffer_stride = data.push_bytes, + }; + for (uint32_t i = 0; i < kWorldSize; ++i) { + params.pull_workspaces[i] = data.pull_workspaces[i]; + params.push_workspaces[i] = data.push_workspaces[i]; + params.pull_semaphores[i] = data.pull_semaphores[i]; + } + const int64_t nbytes = num_elems_int64 * sizeof(T); + RuntimeCheck(nbytes % 16 == 0, "Input bytes must be a multiple of 16, got: ", nbytes); + const uint32_t num_vecs = num_elems / (16 / sizeof(T)); + const auto stream = LaunchKernel::resolve_device(in_.device()); + + if (algo == "1shot_push") { + RuntimeCheck(!use_graph, "Push mode doesn't have graph mode optimization"); + RuntimeCheck(nbytes <= data.push_bytes, "Input size ", nbytes, " exceeds push workspace size ", data.push_bytes); + // the grid is bound to the counter array and must stay constant + const uint32_t num_blocks = data.num_push_blocks; + LaunchKernel(num_blocks, choose_block_size(num_vecs), stream) // + .enable_pdl(kUsePDL)(kernel_push<1>, params); + return out; + } + + using enum PullMode; + RuntimeCheck(nbytes <= data.pull_bytes, "Input size ", nbytes, " exceeds pull workspace size ", data.pull_bytes); + const auto pull_mode = use_graph ? Graph : std::get(pull_arg) ? Multicast : Eager; + RuntimeCheck(pull_mode != Multicast || data.pull_mc_workspace != nullptr, "Multicast requires an mc workspace"); + + const uint32_t num_blocks = data.num_pull_blocks; + const auto cuda_memcpy = [&](void* dst, const void* src) { + if constexpr (SGL_ARCH_HOPPER_OR_GREATER) { // PDL memcpy is faster + // based on micro benchmark, only enable when batch size is small + aligned + constexpr int64_t threshold_MB = SGL_ARCH_BLACKWELL_OR_GREATER ? 1024 : 8; + if (nbytes % device::kMaxVecBytes == 0 && nbytes <= threshold_MB * 1024 * 1024) { + const auto copy_kernel = memcpy_kernel; + const uint32_t num_copy_vecs = nbytes / device::kMaxVecBytes; + const uint32_t num_copy_threads = 128u; + const uint32_t num_copy_blocks = div_ceil(num_copy_vecs, num_copy_threads); + LaunchKernel(num_copy_blocks, num_copy_threads, stream) + .enable_pdl(kUsePDL)(copy_kernel, dst, src, num_copy_vecs); + return; + } + } + // safe fallback to cudaMemcpyAsync for large size or older architecture + RuntimeDeviceCheck(cudaMemcpyAsync(dst, src, nbytes, cudaMemcpyDeviceToDevice, stream)); + }; + + const auto local_workspace = data.pull_workspaces[data.rank]; + if (algo == "1shot_pull") { + // first copy to workspace + if (!use_graph) cuda_memcpy(local_workspace, in_.data_ptr()); + const auto kernel = (pull_mode == Graph) ? kernel_pull<1, Graph> + : pull_mode == Eager ? kernel_pull<1, Eager> + : kernel_pull<1, Multicast>; + // then launch kernel to reduce and write to output + LaunchKernel(num_blocks, choose_block_size(num_vecs), stream) // + .enable_pdl(kUsePDL)(kernel, params); + } else /* 2shot_pull */ { + const uint32_t avg_vecs = div_ceil(num_vecs, kWorldSize); + // first copy to workspace + if (!use_graph) cuda_memcpy(local_workspace, in_.data_ptr()); + // then launch kernel to reduce in workspace + const auto kernel = (pull_mode == Graph) ? kernel_pull<2, Graph> + : pull_mode == Eager ? kernel_pull<2, Eager> + : kernel_pull<2, Multicast>; + if (pull_mode == Multicast) { + const auto max_blocks = data.num_multicast_blocks; + constexpr uint32_t kMulticastNumThreads = 512u; + // NOTE: too much traffic will degrade performance in multicast + LaunchKernel(std::min(num_blocks, max_blocks), kMulticastNumThreads, stream) + .enable_pdl(kUsePDL)(kernel, params); + } else { + LaunchKernel(num_blocks, choose_block_size(avg_vecs), stream) // + .enable_pdl(kUsePDL)(kernel, params); + } + // finally copy from workspace to output + if (!use_graph) cuda_memcpy(out.data_ptr(), local_workspace); + } + return out; + } +}; + +template +tvm::ffi::Tensor custom_all_reduce( + CommunicatorRef comm, tvm::ffi::Tensor input, std::string algo, std::variant pull_arg) { + return AllReduceKernel::run(comm, input, algo, pull_arg); +} + +} // namespace diff --git a/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_base.cuh b/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_base.cuh deleted file mode 100644 index 00e265513..000000000 --- a/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_base.cuh +++ /dev/null @@ -1,30 +0,0 @@ -#include -#include -#include - -#include -#include -#include - -#include - -#include -#include - -inline void register_custom_all_reduce() { - namespace refl = tvm::ffi::reflection; - using Class = host::distributed::CustomAllReduceBase; - refl::ObjectDef() - .def(refl::init(), "__init__") - .def("share_storage", &Class::share_storage) - .def("share_graph_inputs", &Class::share_graph_inputs) - .def("post_init", &Class::post_init) - .def("register_inputs", &Class::register_inputs) - .def("set_cuda_graph_capture", &Class::set_cuda_graph_capture) - .def("get_graph_capture_ptrs", &Class::get_graph_capture_ptrs) - .def("get_graph_capture_bases", &Class::get_graph_capture_bases) - .def("register_peer_mapped_inputs", &Class::register_peer_mapped_inputs) - .def("free_ipc_handles", &Class::free_ipc_handles) - .def("free_storage", &Class::free_storage) - .def("configure_pull", &Class::configure_pull); -} diff --git a/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_pull.cuh b/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_pull.cuh deleted file mode 100644 index 0dfc63ac5..000000000 --- a/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_pull.cuh +++ /dev/null @@ -1,205 +0,0 @@ -// Partially migrated from AOT kernel: -// https://github.com/sgl-project/sglang/blob/v0.5.9/sgl-kernel/csrc/allreduce/custom_all_reduce.cu -// Which was originally adapted from: -// https://github.com/vllm-project/vllm/blob/v0.8.2/csrc/custom_all_reduce.cu -// We redesign the controller interface to minimize control plane traffic, -// and fuse the reduce-scatter and broadcast in the 2-shot all reduce -#include -#include -#include - -#include -#include -#include - -#include -#include - -#include -#include -#include - -namespace { - -using device::distributed::PullController; -using host::distributed::AllReduceData; -using host::distributed::CustomAllReduceBase, host::distributed::CustomAllReduceRef; - -struct AllReduceParams { - void* __restrict__ output; - uint32_t rank; - uint32_t num_items; // NOTE: support at most 4G, but that's too much -}; - -[[maybe_unused]] -SGL_DEVICE void prefetch_uniform_ptr(const void* ptr) { - asm volatile("prefetchu.L1 [%0];" ::"l"(ptr) : "memory"); -} - -#define CUSTOM_AR_KERNEL __global__ __launch_bounds__(1024, 1) - -template -SGL_DEVICE void all_reduce_impl(const AllReduceParams& params, DType* (&input)[kNumGPU]) { - using namespace device; - - constexpr uint32_t kVecSize = 16 / (sizeof(DType) * 2); - using DType2 = packed_t; - using Storage = AlignedVector; - const auto& [output, rank, num_items] = params; - - for (auto i = blockIdx.x;; i += gridDim.x) { - const auto offset = i * blockDim.x + threadIdx.x; - if (offset * kVecSize * 2 >= num_items) break; - Storage storage[kNumGPU]; - -#pragma unroll - for (uint32_t i = 0; i < kNumGPU; ++i) { - storage[i].load(input[i], offset); - } - const Storage result = distributed::reduce_impl(storage); - if constexpr (kBroadcast) { -#pragma unroll - for (uint32_t i = 0; i < kNumGPU; ++i) { - result.store(input[i], offset); - } - } else { - result.store(output, offset); - } - } -} - -template -CUSTOM_AR_KERNEL void all_reduce_one_shot_kernel( - const AllReduceData* __restrict__ data, - const AllReduceParams __grid_constant__ params, - const PullController __grid_constant__ ctrl) { - /// NOTE: we assume the data array is ready before the previous kernel - DType* input[kNumGPU]; - prefetch_uniform_ptr(data); -#pragma unroll - for (uint32_t i = 0; i < kNumGPU; ++i) - input[i] = static_cast(data->input[i]); - device::PDLWaitPrimary(); - - ctrl.sync(params.rank, kNumGPU); - all_reduce_impl(params, input); - - device::PDLTriggerSecondary(); - ctrl.sync(params.rank, kNumGPU); -} - -template -CUSTOM_AR_KERNEL void all_reduce_two_shot_kernel( - const AllReduceData* __restrict__ data, - const AllReduceParams __grid_constant__ params, - const PullController __grid_constant__ ctrl) { - // get the range of this rank - using device::kWarpThreads, device::div_ceil; - - prefetch_uniform_ptr(data); - DType* input[kNumGPU]; -#pragma unroll - for (uint32_t i = 0; i < kNumGPU; ++i) - input[i] = static_cast(data->input[i]); - - constexpr uint32_t kVecSize = 16 / (sizeof(DType) * 2); - const uint32_t num_items = params.num_items; - const uint32_t total_vec = num_items / (kVecSize * 2); // must be divisible here - const uint32_t vec_per_rank = div_ceil(div_ceil(total_vec, kNumGPU), kWarpThreads) * kWarpThreads; - const uint32_t local_vec_start = min(params.rank * vec_per_rank, total_vec); - const uint32_t local_vec_finish = min(local_vec_start + vec_per_rank, total_vec); - const uint32_t local_start = local_vec_start * kVecSize * 2; - const uint32_t local_length = (local_vec_finish - local_vec_start) * kVecSize * 2; - const auto local_params = AllReduceParams{ - .output = nullptr, // this is not used for 2-shot all reduce - .rank = params.rank, - .num_items = local_length, - }; - -#pragma unroll - for (uint32_t i = 0; i < kNumGPU; ++i) - input[i] += local_start; - - device::PDLWaitPrimary(); - - ctrl.sync(params.rank, kNumGPU); - all_reduce_impl(local_params, input); - - device::PDLTriggerSecondary(); - ctrl.sync(params.rank, kNumGPU); -} - -template -struct CustomAllReducePull : public CustomAllReduceBase { - static constexpr uint32_t kVecSize = 16 / (sizeof(DType) * 2); - static constexpr auto one_shot_kernel = all_reduce_one_shot_kernel; - static constexpr auto two_shot_kernel = all_reduce_two_shot_kernel; - static_assert(kNumGPU <= device::distributed::kMaxNumGPU, "kNumGPU exceeds the maximum supported GPUs"); - - tvm::ffi::Tensor all_reduce(tvm::ffi::Tensor input, int shot) { - using namespace host; - const bool use_2shot = (shot == 2); - const auto device = input.device(); - const auto input_ptr = input.data_ptr(); - const auto buffer_ptr = get_pull_buffer(m_storage); - const auto num_items_int64 = input.numel(); - const auto num_items = static_cast(num_items_int64); - const auto items_per_block = m_cta_size * kVecSize * 2; - const auto needed_blocks = div_ceil(num_items, items_per_block); - const auto num_blocks = std::min(needed_blocks, m_num_cta); - const auto kernel = use_2shot ? two_shot_kernel : one_shot_kernel; - // only 1-shot + graph capture need extra output buffer - const auto output = (m_is_graph_capturing && !use_2shot) ? ffi::empty_like(input) : input; - const auto params = AllReduceParams{ - .output = use_2shot ? nullptr : output.data_ptr(), - .rank = m_rank, - .num_items = num_items, - }; - - RuntimeCheck(input.IsContiguous(), "Input tensor must be contiguous"); - RuntimeCheck(m_num_gpu == kNumGPU, "Mismatch GPU count"); - RuntimeCheck(shot == 1 || shot == 2, "Invalid shot count: ", shot); - RuntimeCheck(device.device_type == kDLCUDA, "Only CUDA device is supported"); - RuntimeCheck(is_type(input.dtype()), "Input dtype mismatch"); - RuntimeCheck(std::bit_cast(input_ptr) % 16 == 0, "Input pointer is not properly aligned"); - RuntimeCheck(m_pull_ctrl.has_value(), "Controller is not initialized"); - RuntimeCheck(static_cast(num_items) == num_items_int64, "Number of items exceeds 4G limit"); - - const auto& ctrl = *m_pull_ctrl; - const auto stream = LaunchKernel::resolve_device(device); - auto launch = LaunchKernel{num_blocks, m_cta_size, stream}; - launch.enable_pdl(kUsePDL); - const auto input_bytes = static_cast(sizeof(DType) * num_items); - RuntimeCheck(input_bytes <= m_pull_buffer_bytes, "Input is too large, num items: ", num_items); - const auto check_capturing = [&] { - if (!m_is_graph_capturing) return false; // override to avoid cudaRT call overhead - cudaStreamCaptureStatus status; - RuntimeDeviceCheck(cudaStreamIsCapturing(stream, &status)); - return status == cudaStreamCaptureStatusActive; - }; - if (check_capturing()) { - // no-op if not really capturing, we're in a dummy run - const auto data_ptr = allocate_graph_capture_input(input_ptr, input_bytes); - /// NOTE: we assume when the graph is replayed, the data_ptr should be ready - launch(kernel, data_ptr, params, ctrl); - } else { - // 1.copy the input to the buffer - RuntimeDeviceCheck(cudaMemcpyAsync(buffer_ptr, input_ptr, input_bytes, cudaMemcpyDeviceToDevice, stream)); - // 2. launch the all reduce kernel - const auto data_ptr = get_data_ptr(); // use default buffer - launch(kernel, data_ptr, params, ctrl); - if (use_2shot) { // 3. copy the reduced result back to the output, because 2-shot doesn't write to output - RuntimeDeviceCheck(cudaMemcpyAsync(input_ptr, buffer_ptr, input_bytes, cudaMemcpyDeviceToDevice, stream)); - } - } - return output; - } -}; - -template -tvm::ffi::Tensor custom_all_reduce(CustomAllReduceRef obj, tvm::ffi::Tensor input, int shot) { - using Impl = CustomAllReducePull; - return static_cast(*obj.get()).all_reduce(input, shot); -} - -} // namespace diff --git a/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_push.cuh b/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_push.cuh deleted file mode 100644 index c4523c27e..000000000 --- a/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_push.cuh +++ /dev/null @@ -1,253 +0,0 @@ -// Partially adapted from: -// https://github.com/flashinfer-ai/flashinfer/blob/v0.6.4/include/flashinfer/comm/trtllm_allreduce_fusion.cuh -// We simplify the lamport design and minimize the ring buffer count (from 3 -> 2) -#include -#include -#include - -#include -#include -#include - -#include -#include - -#include -#include - -namespace { - -using device::distributed::PushController; -using host::distributed::CustomAllReduceBase, host::distributed::CustomAllReduceRef; - -struct AllReducePushData { - void* __restrict__ buffer[device::distributed::kMaxNumGPU]; - const void* input; - void* output; - uint32_t rank; - uint32_t num_items; - uint32_t buffer_bytes; - uint32_t epoch_bytes; -}; - -#define CUSTOM_AR_KERNEL __global__ __launch_bounds__(1024, 1) - -template -struct fp_trait {}; - -// TODO: support more dtypes -template <> -struct fp_trait { - using type = uint16_t; - [[maybe_unused]] - static constexpr uint16_t pos_zero = 0x0000u; - [[maybe_unused]] - static constexpr uint16_t neg_zero = 0x8000u; -}; - -template <> -struct fp_trait { - using type = uint16_t; - [[maybe_unused]] - static constexpr uint16_t pos_zero = 0x0000u; - [[maybe_unused]] - static constexpr uint16_t neg_zero = 0x8000u; -}; - -template <> -struct fp_trait { - using type = uint32_t; - [[maybe_unused]] - static constexpr uint32_t pos_zero = 0x00000000u; - [[maybe_unused]] - static constexpr uint32_t neg_zero = 0x80000000u; -}; - -template -SGL_DEVICE void clear_pos_zero(DType& val) { - using Trait = fp_trait; - const auto ptr = reinterpret_cast(&val); - if (*ptr == Trait::pos_zero) *ptr = Trait::neg_zero; -} - -template -SGL_DEVICE bool is_pos_zero(const DType& val) { - using Trait = fp_trait; - const auto ptr = reinterpret_cast(&val); - return *ptr == Trait::pos_zero; -} - -template -SGL_DEVICE DType get_pos_zero() { - using Trait = fp_trait; - const auto value = Trait::pos_zero; - return *reinterpret_cast(&value); -} - -template -SGL_DEVICE void ld_global_volatile_16B(T& x, const void* addr, int64_t offset) { - static_assert(alignof(T) == 16 && sizeof(T) == 16); - addr = device::pointer::offset(addr, offset); - uint4 val; - asm volatile("ld.volatile.global.v4.b32 {%0, %1, %2, %3}, [%4];" - : "=r"(val.x), "=r"(val.y), "=r"(val.z), "=r"(val.w) - : "l"(addr)); - x = *reinterpret_cast(&val); -} - -template -SGL_DEVICE void st_global_volatile_16B(const T& x, void* addr, int64_t offset) { - static_assert(alignof(T) == 16 && sizeof(T) == 16); - const uint4 val = *reinterpret_cast(&x); - addr = device::pointer::offset(addr, offset); - asm volatile( - "st.volatile.global.v4.b32 [%4], {%0, %1, %2, %3};" ::"r"(val.x), "r"(val.y), "r"(val.z), "r"(val.w), "l"(addr)); -} - -template -SGL_DEVICE void push_impl(DType* (&push_buf)[kNumGPU], const void* data, uint32_t num_items) { - using namespace device; - constexpr uint32_t kVecSize = 16 / (sizeof(DType) * 2); - using Storage = AlignedVector, kVecSize>; - - for (auto i = blockIdx.x;; i += gridDim.x) { - const auto offset = i * blockDim.x + threadIdx.x; - if (offset * kVecSize * 2 >= num_items) break; - Storage vec; - vec.load(data, offset); -#pragma unroll - for (uint32_t j = 0; j < kVecSize; ++j) { - clear_pos_zero(vec[j].x); - clear_pos_zero(vec[j].y); - } -#pragma unroll - for (uint32_t i = 0; i < kNumGPU; ++i) { - st_global_volatile_16B(vec, push_buf[i], offset); - } - } -} - -template -SGL_DEVICE void poll_impl(DType* (&poll_buf)[kNumGPU], void* data, uint32_t num_items) { - using namespace device; - constexpr uint32_t kVecSize = 16 / (sizeof(DType) * 2); - using Storage = AlignedVector, kVecSize>; - - for (auto i = blockIdx.x;; i += gridDim.x) { - const auto offset = i * blockDim.x + threadIdx.x; - if (offset * kVecSize * 2 >= num_items) break; - Storage storage[kNumGPU]; - - while (true) { - bool has_pos_zero = false; -#pragma unroll - for (uint32_t i = 0; i < kNumGPU; ++i) { - ld_global_volatile_16B(storage[i], poll_buf[i], offset); -#pragma unroll - for (auto j = 0; j < kVecSize; ++j) { - has_pos_zero |= is_pos_zero(storage[i][j].x); - has_pos_zero |= is_pos_zero(storage[i][j].y); - } - } - if (!has_pos_zero) break; - } - - const Storage result = distributed::reduce_impl(storage); - result.store(data, offset); - - Storage pos_zeros; - pos_zeros.fill({get_pos_zero(), get_pos_zero()}); -#pragma unroll - for (uint32_t i = 0; i < kNumGPU; ++i) { - pos_zeros.store(poll_buf[i], offset); - } - } -} - -template -CUSTOM_AR_KERNEL void all_reduce_one_shot_push_kernel( - const AllReducePushData __grid_constant__ params, // - const PushController __grid_constant__ ctrl) { - using namespace device; - - const auto [buffer, input, output, rank, num_items, buffer_bytes, epoch_bytes] = params; - - PDLWaitPrimary(); - - // Phase 1: Push data from input to all ranks' buffers - const auto epoch_offset = ctrl.epoch() * epoch_bytes; - DType* push_buf[kNumGPU]; -#pragma unroll - for (uint32_t i = 0; i < kNumGPU; ++i) { - push_buf[i] = static_cast(pointer::offset(buffer[i], rank * buffer_bytes, epoch_offset)); - } - push_impl(push_buf, input, num_items); - - PDLTriggerSecondary(); - - // Phase 2: Poll local data - DType* poll_buf[kNumGPU]; -#pragma unroll - for (uint32_t i = 0; i < kNumGPU; ++i) { - poll_buf[i] = static_cast(pointer::offset(buffer[rank], i * buffer_bytes, epoch_offset)); - } - poll_impl(poll_buf, output, num_items); - ctrl.exit(); -} - -template -struct CustomAllReducePush : public CustomAllReduceBase { - static constexpr uint32_t kVecSize = 16 / (sizeof(DType) * 2); - static_assert(kNumGPU <= device::distributed::kMaxNumGPU, "kNumGPU exceeds the maximum supported GPUs"); - - tvm::ffi::Tensor all_reduce(tvm::ffi::Tensor input, int shot) { - using namespace host; - const auto device = input.device(); - const auto input_ptr = input.data_ptr(); - const auto num_items_int64 = input.numel(); - const auto num_items = static_cast(num_items_int64); - const auto num_blocks = m_max_num_cta_push; // must be constant to ensure correctness - const auto num_threads = [&] { - for (const auto t : {128u, 256u, 512u}) { - if (t * num_blocks * 2 * kVecSize >= num_items) return t; - } - return 1024u; - }(); - const auto output = input; - AllReducePushData params; - for (uint32_t i = 0; i < kNumGPU; ++i) { - params.buffer[i] = get_push_buffer(m_peer_storage[i]); - } - params.input = input_ptr; - params.output = input_ptr; - params.rank = m_rank; - params.num_items = num_items; - params.buffer_bytes = m_push_buffer_bytes; - params.epoch_bytes = kNumGPU * params.buffer_bytes; - - RuntimeCheck(input.IsContiguous(), "Input must be contiguous"); - RuntimeCheck(m_num_gpu == kNumGPU, "Number of GPUs mismatch"); - RuntimeCheck(device.device_type == kDLCUDA, "Only CUDA device is supported"); - RuntimeCheck(is_type(input.dtype()), "Input dtype mismatch"); - RuntimeCheck(std::bit_cast(input_ptr) % 16 == 0, "Input pointer is not properly aligned"); - RuntimeCheck(m_push_ctrl.has_value(), "Controller is not initialized"); - RuntimeCheck(shot == 1, "Push all-reduce only supports 1-shot, got: ", shot); - RuntimeCheck(static_cast(num_items) == num_items_int64, "Number of items exceeds 4G limit"); - - const auto input_bytes = static_cast(sizeof(DType) * num_items_int64); - RuntimeCheck(input_bytes <= m_push_buffer_bytes, "Input is too large, num items: ", num_items); - - const auto kernel = all_reduce_one_shot_push_kernel; - LaunchKernel(num_blocks, num_threads, device) // - .enable_pdl(kUsePDL)(kernel, params, *m_push_ctrl); - return output; - } -}; - -template -tvm::ffi::Tensor custom_all_reduce(CustomAllReduceRef obj, tvm::ffi::Tensor input, int shot) { - using Impl = CustomAllReducePush; - return static_cast(*obj.get()).all_reduce(input, shot); -} - -} // namespace diff --git a/python/sglang/jit_kernel/csrc/distributed/ipc.cuh b/python/sglang/jit_kernel/csrc/distributed/ipc.cuh new file mode 100644 index 000000000..3847047ac --- /dev/null +++ b/python/sglang/jit_kernel/csrc/distributed/ipc.cuh @@ -0,0 +1,194 @@ +#include +#include + +#include + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace host::distributed { + +struct AllocationRange { + uintptr_t base; + size_t size; + size_t offset; +}; + +inline auto get_allocation_range(uintptr_t ptr) -> AllocationRange { + CUdeviceptr base = 0; + size_t size = 0; + const CUresult res = cuMemGetAddressRange(&base, &size, ptr); + if (res != CUDA_SUCCESS) { + const char* name = nullptr; + cuGetErrorName(res, &name); + RuntimeCheck(false, "cuMemGetAddressRange failed: ", name ? name : "unknown"); + } + const auto b = static_cast(base); + return {.base = b, .size = size, .offset = ptr - b}; +} + +/** + * \brief Batched cudaIpc handle exchange for CUDA-graph input pointers. + * + * `batch_get_handles` maps local device pointers to (base allocation IPC + * handle, offset) pairs; `batch_open_handles` opens peer handles (cached per + * unique handle) and returns absolute peer pointers. Only works for + * cudaMalloc-backed pointers; VMM-backed pointers take the fabric/posix-fd + * path in Python instead. + */ +struct IPCManager : public tvm::ffi::Object { + public: + using IPCHandle = std::array; + using FFIHandle = tvm::ffi::Array; + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("sgl.IPCManager", IPCManager, tvm::ffi::Object); + static constexpr bool _type_mutable = true; + + using BatchGetResult = tvm::ffi::Array>; + using BatchGetInputs = tvm::ffi::Array; + + IPCManager() = default; + ~IPCManager() { + this->destroy(); + } + + void destroy() { + for (const auto& [handle, base_addr] : m_handle2ptr_cache) { + if (m_local_handles.count(handle)) continue; + RuntimeDeviceCheck(cudaIpcCloseMemHandle(reinterpret_cast(base_addr))); + } + m_handle2ptr_cache.clear(); + m_ptr2handle_cache.clear(); + m_local_handles.clear(); + } + + BatchGetInputs batch_open_handles(BatchGetResult handles) { + tvm::ffi::Array result; + result.reserve(handles.size()); + for (const auto& pair : handles) { + const auto ipc_handle = to_ipc_handle(get<0>(pair)); + const auto offset = get<1>(pair); + result.push_back(open_handle(ipc_handle) + offset); + } + return result; + } + + BatchGetResult batch_get_handles(const BatchGetInputs& ptrs) { + RuntimeCheck(m_ptr2handle_cache.empty(), "Internal error: stale pointer cache"); + BatchGetResult result; + result.reserve(ptrs.size()); + using Tuple = tvm::ffi::Tuple; + for (const auto& ptr : ptrs) { + const auto [ipc_handle, offset] = get_handle(ptr); + result.emplace_back(Tuple{to_ffi_handle(ipc_handle), offset}); + } + // We intentionally do NOT cache by base address across calls. The caching + // allocator (e.g. PyTorch) may free a CUDA allocation and later return a + // new allocation whose virtual range overlaps with the freed one; a cache + // keyed on base/size would then hand back a stale `cudaIpcMemHandle_t` + // that no longer maps to live memory on the peer. Re-querying is cheap. + m_ptr2handle_cache.clear(); + return result; + } + + private: + static IPCHandle to_ipc_handle(const FFIHandle& ffi_handle) { + IPCHandle ipc_handle; + RuntimeCheck(ffi_handle.size() == sizeof(cudaIpcMemHandle_t), "Invalid IPC handle size: ", ffi_handle.size()); + for (size_t i = 0; i < sizeof(cudaIpcMemHandle_t); ++i) { + ipc_handle[i] = static_cast(ffi_handle[i]); + } + return ipc_handle; + } + + static IPCHandle to_ipc_handle(const cudaIpcMemHandle_t& cuda_handle) { + IPCHandle ipc_handle; + std::memcpy(ipc_handle.data(), &cuda_handle, sizeof(cudaIpcMemHandle_t)); + return ipc_handle; + } + + static FFIHandle to_ffi_handle(const IPCHandle& ipc_handle) { + FFIHandle ffi_handle; + ffi_handle.reserve(sizeof(cudaIpcMemHandle_t)); + for (size_t i = 0; i < sizeof(cudaIpcMemHandle_t); ++i) { + ffi_handle.push_back(static_cast(ipc_handle[i])); + } + return ffi_handle; + } + + std::pair get_handle(uintptr_t ptr) { + auto it = m_ptr2handle_cache.upper_bound(ptr); + if (it != m_ptr2handle_cache.begin()) { + --it; + const auto& [cached_handle, cached_size] = it->second; + const auto offset = ptr - it->first; + if (offset < cached_size) return {cached_handle, offset}; + } + // Not found in cache, query CUDA and cache the result + const auto range = get_allocation_range(ptr); + cudaIpcMemHandle_t handle; + RuntimeDeviceCheck(cudaIpcGetMemHandle(&handle, reinterpret_cast(range.base))); + const auto ipc_handle = to_ipc_handle(handle); + const auto [_, success] = m_ptr2handle_cache.try_emplace(range.base, ipc_handle, range.size); + RuntimeCheck(success, "Internal error: base address already exists in cache"); + m_handle2ptr_cache.try_emplace(ipc_handle, range.base); + m_local_handles.insert(ipc_handle); + return {ipc_handle, range.offset}; + } + + uintptr_t open_handle(const IPCHandle& handle) { + const auto it = m_handle2ptr_cache.find(handle); + if (it != m_handle2ptr_cache.end()) { + return it->second; + } + cudaIpcMemHandle_t cuda_handle; + std::memcpy(&cuda_handle, handle.data(), sizeof(cudaIpcMemHandle_t)); + void* base_ptr = nullptr; + RuntimeDeviceCheck(cudaIpcOpenMemHandle(&base_ptr, cuda_handle, cudaIpcMemLazyEnablePeerAccess)); + const auto base_addr = reinterpret_cast(base_ptr); + const auto [_, success] = m_handle2ptr_cache.try_emplace(handle, base_addr); + RuntimeCheck(success, "Internal error: IPC handle already exists in cache"); + return base_addr; + } + + struct EqualCUDAIPC { + bool operator()(const IPCHandle& a, const IPCHandle& b) const { + return std::memcmp(a.data(), b.data(), sizeof(cudaIpcMemHandle_t)) == 0; + } + }; + + struct HashCUDAIPC { + std::size_t operator()(const IPCHandle& handle) const { + const auto sv = std::string_view{handle.data(), sizeof(cudaIpcMemHandle_t)}; + return std::hash{}(sv); + } + }; + + std::unordered_map m_handle2ptr_cache; + std::unordered_set m_local_handles; + std::map> m_ptr2handle_cache; +}; + +} // namespace host::distributed + +inline void register_ipc_manager() { + namespace refl = tvm::ffi::reflection; + using Class = host::distributed::IPCManager; + refl::ObjectDef() + .def(refl::init<>(), "__init__") + .def("batch_get_handles", &Class::batch_get_handles) + .def("batch_open_handles", &Class::batch_open_handles) + .def("destroy", &Class::destroy); +} diff --git a/python/sglang/jit_kernel/csrc/distributed/tp_qknorm.cuh b/python/sglang/jit_kernel/csrc/distributed/tp_qknorm.cuh index ca80e1efc..a337d9f43 100644 --- a/python/sglang/jit_kernel/csrc/distributed/tp_qknorm.cuh +++ b/python/sglang/jit_kernel/csrc/distributed/tp_qknorm.cuh @@ -11,19 +11,18 @@ #include #include -#include -#include +#include #include #include namespace { -using device::distributed::PushController; -using host::distributed::CustomAllReduceBase, host::distributed::CustomAllReduceRef; +using device::distributed::Counter; +using host::distributed::CommunicatorObj, host::distributed::CommunicatorRef; struct ParallelQKNormParams { - void* __restrict__ buffer[device::distributed::kMaxNumGPU]; + void* __restrict__ buffer[device::distributed::kMaxWorldSize]; void* q_ptr; void* k_ptr; const void* __restrict__ q_weight; @@ -102,7 +101,7 @@ struct KernelTrait { template __global__ __launch_bounds__(Trait::kBlockSize, Trait::kOccupancy) void parallel_qknorm_across_head( - const ParallelQKNormParams __grid_constant__ params, const PushController __grid_constant__ ctrl) { + const ParallelQKNormParams __grid_constant__ params, Counter* const __restrict__ counters) { using namespace device; // each cta will handle exactly 1 token @@ -140,10 +139,10 @@ __global__ __launch_bounds__(Trait::kBlockSize, Trait::kOccupancy) void parallel const auto start = (bx - num_tokens) * blockDim.x + threadIdx.x; const auto stride = (gridDim.x - num_tokens) * blockDim.x; for (uint32_t i = start; i < num_clean_up_count; i += stride) - ctrl.exit_unsafe(num_tokens + i); + counters[num_tokens + i].inc(1); return; } - const auto epoch_offset = ctrl.epoch() * epoch_bytes; // only for comm + const auto epoch_offset = (counters[bx].get() % 2) * epoch_bytes; // only for comm __builtin_assume(bx < num_tokens); // since we have `bx >= num_tokens` Storage next_input; @@ -226,16 +225,20 @@ __global__ __launch_bounds__(Trait::kBlockSize, Trait::kOccupancy) void parallel } input_i_ptr = input_next_ptr; } - ctrl.exit(); + __syncthreads(); + if (threadIdx.x == 0) { + counters[bx].inc(1); // NOTE: u32 overflow is safe under mod 2 + } } template -struct FusedParallelQKNormAcrossHead : public CustomAllReduceBase { +struct FusedParallelQKNormAcrossHead { using Trait = KernelTrait; static constexpr auto kernel = parallel_qknorm_across_head; - static_assert(kNumGPU <= device::distributed::kMaxNumGPU, "kNumGPU exceeds the maximum supported GPUs"); + static_assert(kNumGPU <= device::distributed::kMaxWorldSize, "kNumGPU exceeds the maximum supported GPUs"); - void _run( + static void _run( + const CommunicatorObj& comm, const tvm::ffi::Tensor q, const tvm::ffi::Tensor k, const tvm::ffi::Tensor q_weight, @@ -271,15 +274,16 @@ struct FusedParallelQKNormAcrossHead : public CustomAllReduceBase { // use at most `world_size` blocks to clean up, // this is based on the observation that occupancy is usually linear // with respect to the world size - const bool need_clean = num_tokens < m_max_num_cta_push; - const auto num_clean = need_clean ? (m_max_num_cta_push - num_tokens) : 0; + const auto max_num_blocks = comm.num_push_blocks; + const bool need_clean = num_tokens < max_num_blocks; + const auto num_clean = need_clean ? (max_num_blocks - num_tokens) : 0; const auto num_blocks = need_clean ? num_tokens + div_ceil(num_clean, Trait::kBlockSize) // - : m_max_num_cta_push; // + : max_num_blocks; // const auto num_threads = Trait::kBlockSize; - RuntimeCheck(num_blocks <= m_max_num_cta_push, "internal error"); + RuntimeCheck(num_blocks <= max_num_blocks, "internal error"); ParallelQKNormParams params; for (uint32_t i = 0; i < kNumGPU; ++i) { - params.buffer[i] = get_push_buffer(m_peer_storage[i]); + params.buffer[i] = comm.push_workspaces[i]; } params.q_ptr = q.data_ptr(); params.k_ptr = k.data_ptr(); @@ -288,22 +292,21 @@ struct FusedParallelQKNormAcrossHead : public CustomAllReduceBase { params.q_stride_bytes = q.stride(0) * sizeof(DType); params.k_stride_bytes = k.stride(0) * sizeof(DType); params.eps = eps / kNumGPU; // scale down eps by number of GPUs - params.rank = m_rank; + params.rank = comm.rank; params.num_tokens = num_tokens; - params.epoch_bytes = m_push_buffer_bytes; + params.epoch_bytes = static_cast(comm.push_bytes); params.num_clean_up_count = num_clean; const auto needed_buffer_bytes = static_cast(num_tokens) * 2 * sizeof(float); - RuntimeCheck(m_num_gpu == kNumGPU, "Number of GPUs mismatch"); - RuntimeCheck(m_push_ctrl.has_value(), "Controller is not initialized"); + RuntimeCheck(comm.world_size == kNumGPU, "Number of GPUs mismatch"); RuntimeCheck(std::bit_cast(params.q_ptr) % 16 == 0, "q pointer is not properly aligned"); RuntimeCheck(std::bit_cast(params.k_ptr) % 16 == 0, "k pointer is not properly aligned"); RuntimeCheck(std::bit_cast(params.q_weight) % 16 == 0, "q_weight pointer is not properly aligned"); RuntimeCheck(std::bit_cast(params.k_weight) % 16 == 0, "k_weight pointer is not properly aligned"); - RuntimeCheck(needed_buffer_bytes <= m_push_buffer_bytes, "Push buffer is too small"); + RuntimeCheck(needed_buffer_bytes <= comm.push_bytes, "Push buffer is too small"); LaunchKernel(num_blocks, num_threads, device) // - .enable_pdl(kUsePDL)(kernel, params, *m_push_ctrl); + .enable_pdl(kUsePDL)(kernel, params, comm.push_counter); } static uint32_t get_max_occupancy() { @@ -311,14 +314,13 @@ struct FusedParallelQKNormAcrossHead : public CustomAllReduceBase { } static void - run(CustomAllReduceRef obj, + run(CommunicatorRef comm, const tvm::ffi::Tensor q, const tvm::ffi::Tensor k, const tvm::ffi::Tensor q_weight, const tvm::ffi::Tensor k_weight, const float eps) { - using Self = FusedParallelQKNormAcrossHead; - return static_cast(obj.get())->_run(q, k, q_weight, k_weight, eps); + return _run(*comm.get(), q, k, q_weight, k_weight, eps); } }; diff --git a/python/sglang/jit_kernel/include/sgl_kernel/distributed/common.cuh b/python/sglang/jit_kernel/include/sgl_kernel/distributed/common.cuh deleted file mode 100644 index e0ce2dc08..000000000 --- a/python/sglang/jit_kernel/include/sgl_kernel/distributed/common.cuh +++ /dev/null @@ -1,120 +0,0 @@ -#pragma once -#include - -namespace device::distributed { - -inline constexpr uint32_t kMaxNumGPU = 8; - -struct alignas(128) Semaphore { - public: - constexpr Semaphore() : m_flag(0), m_counter(0) {} - - template - SGL_DEVICE uint32_t get() const { - uint32_t val; - if constexpr (kFence) { - asm volatile("ld.acquire.sys.global.u32 %0, [%1];" : "=r"(val) : "l"(&m_flag)); - } else { - asm volatile("ld.volatile.global.u32 %0, [%1];" : "=r"(val) : "l"(&m_flag)); - } - return val; - } - - template - SGL_DEVICE uint32_t add(uint32_t val) { - uint32_t old_val; - if constexpr (kFence) { - asm volatile("atom.release.sys.global.add.u32 %0, [%1], %2;" : "=r"(old_val) : "l"(&m_flag), "r"(val)); - } else { - asm volatile("atom.global.add.u32 %0, [%1], %2;" : "=r"(old_val) : "l"(&m_flag), "r"(val)); - } - return old_val; - } - - // Only called by the owning GPU - plain load is sufficient - SGL_DEVICE uint32_t get_counter() const { - return m_counter; - } - - // Only called by the owning GPU - plain store is sufficient - SGL_DEVICE void set_counter(uint32_t val) { - m_counter = val; - } - - private: - uint32_t m_flag; - uint32_t m_counter; -}; - -struct PullController { - public: - using SignalType = Semaphore; - - PullController(void** signals, uint32_t num_gpu) { - for (uint32_t i = 0; i < num_gpu; ++i) { - m_signals[i] = static_cast(signals[i]); - } - } - - /// Synchronize all GPUs. - /// When kFence is true, establishes happens-before across GPUs using - /// release/acquire semantics, ensuring prior writes are visible system-wide. - template - SGL_DEVICE void sync(uint32_t rank, uint32_t num_gpu) const { - // For fenced sync: ensure all threads in this block have completed their writes, - // so the signaling thread's release carries them transitively. - static_assert(!(kFence && kStart), "Start stage does not need to wait fence"); - if constexpr (kFence || !kStart) __syncthreads(); - constexpr auto kStage = kStart ? 1 : 2; - const auto warp_id = threadIdx.x / kWarpThreads; - const auto lane_id = threadIdx.x % kWarpThreads; - if (lane_id == 0 && warp_id < num_gpu) { - auto& signal = m_signals[warp_id][blockIdx.x]; - signal.add(1); - if (warp_id == rank) { - const auto target = num_gpu * kStage; - /// NOTE: correctness here: - /// - base is only read/updated locally by the owning GPU - const auto base = signal.get_counter(); - while (signal.get() - base < target) - ; - if constexpr (!kStart) { - signal.set_counter(base + target); - } - } - } - if constexpr (kStart) __syncthreads(); - } - - private: - Semaphore* __restrict__ m_signals[kMaxNumGPU]; -}; - -struct PushController { - public: - using SignalType = uint32_t; - static constexpr int64_t kNumStages = 2; - - PushController(void* ptr) : m_local_signal(static_cast(ptr)) {} - - SGL_DEVICE SignalType epoch() const { - return m_local_signal[blockIdx.x]; - } - - SGL_DEVICE void exit() const { - __syncthreads(); - if (threadIdx.x == 0) { - this->exit_unsafe(blockIdx.x); - } - } - - SGL_DEVICE void exit_unsafe(uint32_t which) const { - auto& signal = m_local_signal[which]; - signal = (signal + 1) % kNumStages; - } - - private: - SignalType* m_local_signal; -}; - -} // namespace device::distributed diff --git a/python/sglang/jit_kernel/include/sgl_kernel/distributed/communicator.cuh b/python/sglang/jit_kernel/include/sgl_kernel/distributed/communicator.cuh new file mode 100644 index 000000000..f7a2f9332 --- /dev/null +++ b/python/sglang/jit_kernel/include/sgl_kernel/distributed/communicator.cuh @@ -0,0 +1,120 @@ +#pragma once +#include + +#include + +#include +#include + +#include +#include +#include +#include +#include +#include + +namespace device::distributed { + +inline constexpr uint32_t kMaxWorldSize = 8; + +struct Counter { + public: + Counter(const Counter&) = delete; + SGL_DEVICE uint32_t get() const { + return m_counter; + } + SGL_DEVICE void set(uint32_t val) { + m_counter = val; + } + SGL_DEVICE uint32_t inc(uint32_t val) { + return ::atomicAdd(&m_counter, val); + } + + private: + uint32_t m_counter; +}; + +struct alignas(128) Semaphore { + public: + Semaphore(const Semaphore&) = delete; + SGL_DEVICE Counter* counter_ptr() { + return &m_counter; + } + SGL_DEVICE uint32_t get_relaxed() const { + uint32_t val; + asm volatile("ld.relaxed.sys.global.u32 %0, [%1];" : "=r"(val) : "l"(&m_flag) : "memory"); + return val; + } + SGL_DEVICE void put_relaxed() { + asm volatile("red.relaxed.sys.global.add.u32 [%0], 1;" : : "l"(&m_flag) : "memory"); + } + SGL_DEVICE uint32_t get_acquire() const { + uint32_t val; + asm volatile("ld.acquire.sys.global.u32 %0, [%1];" : "=r"(val) : "l"(&m_flag) : "memory"); + return val; + } + SGL_DEVICE void put_release() { + asm volatile("red.release.sys.global.add.u32 [%0], 1;" : : "l"(&m_flag) : "memory"); + } + + private: + uint32_t m_flag; + Counter m_counter; +}; + +} // namespace device::distributed + +namespace host::distributed { + +using device::distributed::Counter, device::distributed::Semaphore; +inline constexpr uint32_t kMaxWorldSize = device::distributed::kMaxWorldSize; + +/** + * \brief Storage plane of the custom all-reduce implementation. + * + * A thin, kernel-agnostic view over externally owned buffers: per-rank + * symmetric workspaces, synchronization primitives, and grid-size settings. + * It performs no allocation and no IPC; the Python side owns the storage + * (symmetric memory) and its lifetime. + */ +struct CommunicatorObj : public tvm::ffi::Object { + public: + using TensorView = tvm::ffi::TensorView; + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("sgl.Communicator", CommunicatorObj, tvm::ffi::Object); + static constexpr bool _type_mutable = true; // config() mutates block counts + + // Defined in csrc/distributed/communicator.cuh (only the registration + // module needs the implementation). + CommunicatorObj( + uint32_t rank, + uint32_t world_size, + std::vector push_workspaces, + std::vector pull_workspaces, + std::vector pull_semaphores, + TensorView push_counter, + std::optional pull_mc_workspace_ptr); + + void config(std::map config); + + uint32_t rank; + uint32_t world_size; + int64_t push_bytes; // per-buffer bytes; each rank holds 2 * world_size buffers + int64_t pull_bytes; + uint32_t num_push_blocks; // not configurable (bound to the counter array) + uint32_t num_pull_blocks; + uint32_t num_multicast_blocks; + std::array pull_workspaces; // symmetric memory + std::array push_workspaces; // symmetric memory + std::array pull_semaphores; // symmetric memory + Counter* push_counter; // local memory + uint8_t* pull_mc_workspace; // multicast address of the pull workspace (may be null) + + private: // upper bounds for config() + uint32_t total_pull_blocks; +}; + +struct CommunicatorRef : public tvm::ffi::ObjectRef { + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(CommunicatorRef, tvm::ffi::ObjectRef, CommunicatorObj); +}; + +} // namespace host::distributed diff --git a/python/sglang/jit_kernel/include/sgl_kernel/distributed/custom_all_reduce.cuh b/python/sglang/jit_kernel/include/sgl_kernel/distributed/custom_all_reduce.cuh deleted file mode 100644 index bffe36727..000000000 --- a/python/sglang/jit_kernel/include/sgl_kernel/distributed/custom_all_reduce.cuh +++ /dev/null @@ -1,446 +0,0 @@ -#pragma once -#include - -#include -#include - -#include - -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace host::distributed { - -using device::distributed::PullController, device::distributed::PushController; - -struct AllReduceData { - constexpr AllReduceData() {} - void* __restrict__ input[device::distributed::kMaxNumGPU]; -}; - -using ExternHandle = tvm::ffi::Array; - -inline ExternHandle to_extern_handle(void* ptr) { - ExternHandle array; - cudaIpcMemHandle_t handle; - RuntimeDeviceCheck(cudaIpcGetMemHandle(&handle, ptr)); - for (size_t i = 0; i < sizeof(handle); ++i) { - array.push_back(handle.reserved[i]); - } - return array; -} - -inline void* from_extern_handle(const ExternHandle& array) { - cudaIpcMemHandle_t handle; - RuntimeCheck(array.size() == sizeof(handle), "Invalid IPC handle size: ", array.size()); - for (size_t i = 0; i < sizeof(handle); ++i) { - handle.reserved[i] = array[i]; - } - void* ptr; - RuntimeDeviceCheck(cudaIpcOpenMemHandle(&ptr, handle, cudaIpcMemLazyEnablePeerAccess)); - return ptr; -} - -struct HandleHash { - std::size_t operator()(const cudaIpcMemHandle_t& handle) const { - return std::hash{}({handle.reserved, sizeof(handle.reserved)}); - } -}; - -struct HandleEqual { - bool operator()(const cudaIpcMemHandle_t& a, const cudaIpcMemHandle_t& b) const { - return std::memcmp(a.reserved, b.reserved, sizeof(a.reserved)) == 0; - } -}; - -/** - * \brief The control plane of the custom all-reduce implementation. - * It manages the internal state and synchronization of the participating GPUs. - */ -struct CustomAllReduceBase : public tvm::ffi::Object { - public: - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("sgl.CustomAllReduce", CustomAllReduceBase, tvm::ffi::Object); - - static constexpr bool _type_mutable = true; - using InputPair = tvm::ffi::Tuple; // (offset, ipc handle) - - CustomAllReduceBase( - uint32_t rank, - uint32_t num_gpu, - uint32_t max_num_cta_pull, - uint32_t max_num_cta_push, - int64_t pull_buffer_size, - int64_t push_buffer_size, - int64_t graph_buffer_count) - : m_pull_buffer_bytes(pull_buffer_size), - m_push_buffer_bytes(push_buffer_size), - m_graph_buffer_count(graph_buffer_count), - m_rank(rank), - m_num_gpu(num_gpu), - m_max_num_cta_pull(max_num_cta_pull), - m_max_num_cta_push(max_num_cta_push), - // default config for pull kernel, can be updated by `configure()` - m_num_cta(max_num_cta_pull), - m_cta_size(256) { - RuntimeCheck(pull_buffer_size % 128 == 0, "Pull buffer size should be aligned to 128 bytes"); - RuntimeCheck(push_buffer_size % 128 == 0, "Push buffer size should be aligned to 128 bytes"); - RuntimeCheck(rank < num_gpu, "Invalid rank: ", rank); - const int64_t kU32Max = static_cast(std::numeric_limits::max()); - const int64_t push_buffer_size_all = push_all_ranks_bytes(); - RuntimeCheck(pull_buffer_size <= kU32Max, "Pull buffer size is too large: ", pull_buffer_size); - RuntimeCheck(push_buffer_size_all <= kU32Max, "Push buffer size is too large: ", push_buffer_size_all); - RuntimeDeviceCheck(cudaMalloc(&m_storage, storage_bytes())); - } - - ExternHandle share_storage() { - return to_extern_handle(m_storage); - } - - tvm::ffi::Array share_graph_inputs() { - tvm::ffi::Array result; - const auto new_inputs_count = registered_count() - m_cum_registered_count; - RuntimeCheck(new_inputs_count >= 0, "Invalid new count: ", new_inputs_count); - result.reserve(new_inputs_count); - std::unordered_map ipc_cache; - const auto get_handle = [&](void* ptr) -> ExternHandle { - const auto it = ipc_cache.find(ptr); - if (it != ipc_cache.end()) return it->second; - const auto handle = to_extern_handle(ptr); - ipc_cache.try_emplace(ptr, handle); - return handle; - }; - for (const auto ptr : std::span(m_graph_capture_inputs).subspan(m_cum_registered_count)) { - // note: must share the base address of each allocation, or we get wrong address - void* base_ptr; - const auto cu_result = cuPointerGetAttribute(&base_ptr, CU_POINTER_ATTRIBUTE_RANGE_START_ADDR, (CUdeviceptr)ptr); - RuntimeCheck(cu_result == CUDA_SUCCESS, "failed to get pointer attr"); - const auto offset = reinterpret_cast(ptr) - reinterpret_cast(base_ptr); - result.push_back(InputPair{offset, get_handle(base_ptr)}); - } - return result; - } - - void post_init(tvm::ffi::Array ipc_storages) { - RuntimeCheck(ipc_storages.size() == m_num_gpu, "Invalid array size: ", ipc_storages.size()); - m_peer_storage.resize(m_num_gpu); - for (const auto i : irange(m_num_gpu)) { - if (i == m_rank) { - m_peer_storage[i] = m_storage; - } else { - m_peer_storage[i] = from_extern_handle(ipc_storages[i]); - } - } - - // set signal buffer to zero - const auto pull_signal = get_pull_signal(m_storage); - RuntimeDeviceCheck(cudaMemset(pull_signal, 0, pull_signal_bytes())); - - // update the pull controller and data pointer - RuntimeCheck(!m_pull_ctrl.has_value(), "Controller is already initialized"); - m_pull_ctrl.emplace(m_peer_storage.data(), m_num_gpu); - AllReduceData data; - for (const auto i : irange(m_num_gpu)) { - data.input[i] = get_pull_buffer(m_peer_storage[i]); - } - const auto default_data_ptr = get_data_ptr(); - RuntimeDeviceCheck(cudaMemcpy(default_data_ptr, &data, sizeof(AllReduceData), cudaMemcpyHostToDevice)); - - // update the push controller and data pointer - RuntimeCheck(!m_push_ctrl.has_value(), "Controller is already initialized"); - const auto push_signal = get_push_signal(m_storage); - RuntimeDeviceCheck(cudaMemset(push_signal, 0, push_signal_bytes())); - m_push_ctrl.emplace(push_signal); - const auto push_buffer = get_push_buffer(m_storage); - RuntimeDeviceCheck(cudaMemset(push_buffer, 0, push_all_ranks_bytes())); - } - - void register_inputs(tvm::ffi::Array> ipc_graph_inputs) { - RuntimeCheck(ipc_graph_inputs.size() == m_num_gpu); - const auto new_registered_count = registered_count() - m_cum_registered_count; - RuntimeCheck(new_registered_count >= 0, "Invalid registered count: ", new_registered_count); - if (new_registered_count == 0) return; // avoid `m_get_data_ptr()` out-of-bounds - std::vector data; - data.resize(new_registered_count); - const auto open_cached = [&](const ExternHandle& h) -> void* { - RuntimeCheck(h.size() == sizeof(cudaIpcMemHandle_t), "Invalid IPC handle size: ", h.size()); - cudaIpcMemHandle_t handle; - for (size_t i = 0; i < sizeof(handle); ++i) - handle.reserved[i] = h[i]; - const auto [it, success] = m_ipc_cache.try_emplace(handle, nullptr); - if (success) { - void* ptr; - RuntimeDeviceCheck(cudaIpcOpenMemHandle(&ptr, handle, cudaIpcMemLazyEnablePeerAccess)); - it->second = ptr; - } - return it->second; - }; - for (const auto i : irange(ipc_graph_inputs.size())) { - const auto& array = ipc_graph_inputs[i]; - RuntimeCheck(int64_t(array.size()) == new_registered_count); - if (i == m_rank) { - for (const auto j : irange(new_registered_count)) { - data[j].input[i] = m_graph_capture_inputs[m_cum_registered_count + j]; - } - } else { - for (const auto j : irange(new_registered_count)) { - /// NOTE: structural binding will cause intern compiler error... - const auto elem = array[j]; - const auto offset = elem.get<0>(); - const auto ipc_handle = elem.get<1>(); - data[j].input[i] = pointer::offset(open_cached(ipc_handle), offset); - } - } - } - - const auto new_registered_bytes = sizeof(AllReduceData) * new_registered_count; - const auto dst_ptr = get_data_ptr(m_cum_registered_count); - m_cum_registered_count += new_registered_count; - RuntimeDeviceCheck(cudaMemcpy(dst_ptr, data.data(), new_registered_bytes, cudaMemcpyHostToDevice)); - } - - void set_cuda_graph_capture(bool enabled) { - m_is_graph_capturing = enabled; - } - - tvm::ffi::Array get_graph_capture_ptrs() { - tvm::ffi::Array result; - const auto new_count = registered_count() - m_cum_registered_count; - result.reserve(new_count); - for (const auto ptr : std::span(m_graph_capture_inputs).subspan(m_cum_registered_count)) { - result.push_back(reinterpret_cast(ptr)); - } - return result; - } - - using BaseInfo = tvm::ffi::Tuple; // (base_ptr, size) - - /// Returns (unique_bases, per_input_base_indices, per_input_offset). - /// unique_bases[i] = (base_ptr, alloc_size) for each unique allocation. - /// per_input_base_indices[j] = indices of VMM allocations covering input j. - /// per_input_offset[j] = byte offset from the first allocation base for input j. - tvm::ffi::Tuple, tvm::ffi::Array>, tvm::ffi::Array> - get_graph_capture_bases() { - const auto new_inputs = std::span(m_graph_capture_inputs).subspan(m_cum_registered_count); - const auto new_input_bytes = std::span(m_graph_capture_input_bytes).subspan(m_cum_registered_count); - std::unordered_map base_to_idx; - tvm::ffi::Array bases; - tvm::ffi::Array> input_indices; - tvm::ffi::Array offsets; - input_indices.reserve(new_inputs.size()); - offsets.reserve(new_inputs.size()); - RuntimeCheck(new_inputs.size() == new_input_bytes.size(), "graph input metadata mismatch"); - for (const auto input_idx : irange(new_inputs.size())) { - const auto ptr = new_inputs[input_idx]; - auto remaining = new_input_bytes[input_idx]; - RuntimeCheck(remaining > 0, "Invalid graph capture input size: ", remaining); - - auto cursor = reinterpret_cast(ptr); - CUdeviceptr first_base = 0; - tvm::ffi::Array chunks; - while (remaining > 0) { - CUdeviceptr base = 0; - size_t size = 0; - const auto r = cuMemGetAddressRange(&base, &size, cursor); - RuntimeCheck(r == CUDA_SUCCESS, "cuMemGetAddressRange failed: ", r); - if (first_base == 0) first_base = base; - const auto byte_offset = static_cast(cursor - base); - RuntimeCheck( - byte_offset >= 0 && static_cast(byte_offset) < size, - "graph capture input at ", - reinterpret_cast(ptr), - " is outside VMM allocation [base=", - base, - ", size=", - size, - "]"); - - auto [it, inserted] = base_to_idx.try_emplace(base, bases.size()); - if (inserted) { - bases.push_back(BaseInfo{static_cast(base), static_cast(size)}); - } - chunks.push_back(it->second); - - const auto available = static_cast(size) - byte_offset; - const auto advance = std::min(remaining, available); - RuntimeCheck(advance > 0, "Failed to advance VMM graph capture span"); - remaining -= advance; - cursor += advance; - } - input_indices.push_back(chunks); - offsets.push_back(reinterpret_cast(ptr) - first_base); - } - using Result = - tvm::ffi::Tuple, tvm::ffi::Array>, tvm::ffi::Array>; - return Result(bases, input_indices, offsets); - } - - void register_peer_mapped_inputs(tvm::ffi::Array> peer_ptrs_per_input) { - const auto new_count = registered_count() - m_cum_registered_count; - RuntimeCheck(int64_t(peer_ptrs_per_input.size()) == new_count, "peer_ptrs count mismatch"); - if (new_count == 0) return; - std::vector data(new_count); - for (const auto j : irange(new_count)) { - const auto& ptrs = peer_ptrs_per_input[j]; - RuntimeCheck(ptrs.size() == m_num_gpu, "peer count mismatch"); - for (const auto i : irange(m_num_gpu)) { - data[j].input[i] = reinterpret_cast(static_cast(ptrs[i])); - } - } - const auto dst_ptr = get_data_ptr(m_cum_registered_count); - m_cum_registered_count += new_count; - RuntimeDeviceCheck(cudaMemcpy(dst_ptr, data.data(), sizeof(AllReduceData) * new_count, cudaMemcpyHostToDevice)); - } - - void free_ipc_handles() { - for (const auto& pair : m_ipc_cache) { - host::RuntimeDeviceCheck(cudaIpcCloseMemHandle(pair.second)); - } - m_ipc_cache.clear(); - } - - void free_storage() { - host::RuntimeDeviceCheck(cudaFree(m_storage)); - m_storage = nullptr; - } - - tvm::ffi::Tuple configure_pull(uint32_t num_cta, uint32_t cta_size) { - using host::RuntimeCheck; - const auto min_cta_size = m_num_gpu * device::kWarpThreads; - RuntimeCheck(num_cta > 0 && num_cta <= m_max_num_cta_pull, "Invalid number of CTAs: ", num_cta); - RuntimeCheck(cta_size >= min_cta_size, "Block size must be at least ", min_cta_size); - const auto old_num_cta = m_num_cta; - const auto old_block_size = m_cta_size; - m_num_cta = num_cta; - m_cta_size = cta_size; - return tvm::ffi::Tuple{old_num_cta, old_block_size}; - } - - protected: - AllReduceData* allocate_graph_capture_input(void* data_ptr, int64_t input_bytes) { - const auto count = registered_count(); - RuntimeCheck(count < m_graph_buffer_count, "Graph buffer overflow, increase `graph_buffer_count`!"); - m_graph_capture_inputs.push_back(data_ptr); - m_graph_capture_input_bytes.push_back(input_bytes); - return get_data_ptr(count); - } - AllReduceData* get_data_ptr(int64_t which = -1) { - const auto count = registered_count(); - RuntimeCheck(which >= -1 && which < count, "Invalid graph buffer index: ", which, ", count: ", count); - const auto start = get_pull_params(m_storage); - return static_cast(start) + (1 + which); - } - int64_t registered_count() const { - return static_cast(m_graph_capture_inputs.size()); - } - int64_t pull_signal_bytes() const { - return _align_bytes(sizeof(PullController::SignalType) * m_max_num_cta_pull); - } - int64_t push_signal_bytes() const { - return _align_bytes(sizeof(PushController::SignalType) * m_max_num_cta_push); - } - int64_t graph_param_bytes() const { - return _align_bytes(sizeof(AllReduceData) * (1 + m_graph_buffer_count)); // 1 for default - } - int64_t push_all_ranks_bytes() const { - return _align_bytes(PushController::kNumStages * m_num_gpu * m_push_buffer_bytes); - } - int64_t storage_bytes() const { - return _get_offset_impl(5); - } - void* get_pull_signal(void* ptr) const { - return pointer::offset(ptr, _get_offset_impl(0)); - } - void* get_push_signal(void* ptr) const { - return pointer::offset(ptr, _get_offset_impl(1)); - } - void* get_pull_params(void* ptr) const { - return pointer::offset(ptr, _get_offset_impl(2)); - } - void* get_pull_buffer(void* ptr) const { - return pointer::offset(ptr, _get_offset_impl(3)); - } - void* get_push_buffer(void* ptr) const { - return pointer::offset(ptr, _get_offset_impl(4)); - } - int64_t _get_offset_impl(int64_t which) const { - // | SignalArray (pull + push) | GraphBuffers (pull params) | Buffers (pull + push) | - const int64_t offset_map[5] = { - /*[0]=*/pull_signal_bytes(), - /*[1]=*/push_signal_bytes(), - /*[2]=*/graph_param_bytes(), - /*[3]=*/m_pull_buffer_bytes, - /*[4]=*/push_all_ranks_bytes(), - }; - RuntimeCheck(which >= 0 && which <= 5, "Invalid offset index: ", which); - return std::accumulate(offset_map, offset_map + which, int64_t(0)); - } - static int64_t _align_bytes(int64_t size) { - return div_ceil(size, 128) * 128; - } - - const int64_t m_pull_buffer_bytes; - const int64_t m_push_buffer_bytes; - const int64_t m_graph_buffer_count; - const uint32_t m_rank; - const uint32_t m_num_gpu; - const uint32_t m_max_num_cta_pull; - const uint32_t m_max_num_cta_push; - // these 2 config should only affect pull kernel - uint32_t m_num_cta; - uint32_t m_cta_size; - // other states - bool m_is_graph_capturing = false; - int64_t m_cum_registered_count = 0; - std::optional m_pull_ctrl; - std::optional m_push_ctrl; - void* m_storage = nullptr; - std::vector m_graph_capture_inputs; - std::vector m_graph_capture_input_bytes; - std::vector m_peer_storage; - std::unordered_map m_ipc_cache; -}; - -struct CustomAllReduceRef : public tvm::ffi::ObjectRef { - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(CustomAllReduceRef, tvm::ffi::ObjectRef, CustomAllReduceBase); -}; - -} // namespace host::distributed - -namespace device::distributed { - -template -SGL_DEVICE auto reduce_impl(AlignedVector (&storage)[M]) -> AlignedVector { - fp32x2_t acc[N] = {}; -#pragma unroll // unroll num gpu - for (uint32_t i = 0; i < M; ++i) { -#pragma unroll // unroll vec - for (uint32_t j = 0; j < N; ++j) { - const auto [x, y] = cast(storage[i][j]); - auto& [x_acc, y_acc] = acc[j]; - x_acc += x; - y_acc += y; - } - } - - AlignedVector result; -#pragma unroll - for (uint32_t j = 0; j < N; ++j) { - result[j] = cast(acc[j]); - } - - return result; -} - -} // namespace device::distributed diff --git a/python/sglang/jit_kernel/tests/utils.py b/python/sglang/jit_kernel/tests/utils.py index 493df2c9a..1ead20913 100644 --- a/python/sglang/jit_kernel/tests/utils.py +++ b/python/sglang/jit_kernel/tests/utils.py @@ -14,7 +14,7 @@ def multigpu_pytest_main( pre_launch_fn: Optional[Callable[[List[int]], None]] = None, timeout: Optional[int] = 600, ) -> None: - """cudalib-style multi-GPU pytest entry point. + """Torchrun-based multi-GPU pytest entry point. Drop this at the bottom of a test file:: diff --git a/python/sglang/jit_kernel/utils.py b/python/sglang/jit_kernel/utils.py index 2ad5c25c4..66ab63a58 100644 --- a/python/sglang/jit_kernel/utils.py +++ b/python/sglang/jit_kernel/utils.py @@ -62,6 +62,35 @@ def cache_once(fn: F) -> F: return wrapper # type: ignore +_REGISTERED_CLASSES: Dict[type, type] = {} +T = TypeVar("T") + + +def lazy_register_class(name: str, init_fn: Callable[[], None]) -> Callable[[T], T]: + """A decorator to lazily register a tvm-ffi object class on first use. + + `init_fn` runs once (typically JIT-compiling and registering the C++ + reflection) right before the class is registered under the FFI type key + `name`; afterwards instantiation proceeds normally. + """ + + def decorator(cls: T) -> T: + def __new__(cls, *args, **kwargs): + import tvm_ffi + + if cls not in _REGISTERED_CLASSES: + init_fn() # lazy initialization before registration once + _REGISTERED_CLASSES[cls] = tvm_ffi.register_object(name)(cls) + cls = _REGISTERED_CLASSES[cls] + return original_new(cls, *args, **kwargs) + + original_new = cls.__new__ + cls.__new__ = __new__ + return cls + + return decorator + + def _make_wrapper(tup: Tuple[str, str]) -> str: export_name, kernel_name = tup return f"TVM_FFI_DLL_EXPORT_TYPED_FUNC({export_name}, ({kernel_name}));" diff --git a/python/sglang/srt/distributed/device_communicators/configs/__init__.py b/python/sglang/srt/distributed/device_communicators/configs/__init__.py new file mode 100644 index 000000000..2d52ebd37 --- /dev/null +++ b/python/sglang/srt/distributed/device_communicators/configs/__init__.py @@ -0,0 +1 @@ +"""Tuned dispatch configs for device communicators.""" diff --git a/python/sglang/srt/distributed/device_communicators/configs/custom_all_reduce_v2.py b/python/sglang/srt/distributed/device_communicators/configs/custom_all_reduce_v2.py new file mode 100644 index 000000000..8751df57e --- /dev/null +++ b/python/sglang/srt/distributed/device_communicators/configs/custom_all_reduce_v2.py @@ -0,0 +1,194 @@ +"""Hand-tuned dispatch configs for the JIT custom all-reduce (v2). + +Thresholds and block counts come from sweeps of +``test/registered/jit/benchmark/bench_custom_all_reduce.py`` on the listed +GPUs; ``get_all_reduce_config`` picks the table for the current arch and +world size. +""" + +from functools import cache +from typing import NamedTuple, Optional + +import torch + +KB, MB = 1024, 1024 * 1024 + + +class Range(NamedTuple): + min_bytes: int + max_bytes: int + + def contains(self, nbytes: int) -> bool: + return self.min_bytes <= nbytes <= self.max_bytes + + def clip(self, max_bytes: int) -> "Range": + return Range( + min(self.min_bytes, max_bytes), + min(self.max_bytes, max_bytes), + ) + + +class Heuristic(NamedTuple): + """Self-contained algo ranges for one dispatch context (graph or eager). + + Four algos are tried in order of preference (fastest first): + 1. ``1shot_push``: nbytes <= ``one_shot_push_threshold`` + 2. ``1shot_pull``: nbytes <= ``one_shot_pull_threshold`` + 3. ``2shot_pull`` mc: nbytes in ``mc.min_bytes..mc.max_bytes`` + (only when multicast is enabled at runtime) + 4. ``2shot_pull``: nbytes <= ``two_shot_pull_threshold`` + Above all of these, the caller falls back to NCCL. + + Setting two adjacent thresholds equal effectively disables the middle + algo; leaving ``mc`` at the default disables multicast. + """ + + one_shot_push_threshold: int + one_shot_pull_threshold: int + two_shot_pull_threshold: int + mc: Range = Range(0, 0) # default: multicast disabled in this context + + @property + def max_push_bytes(self) -> int: + return self.one_shot_push_threshold + + @property + def max_pull_bytes(self) -> int: + # The pull workspace hosts every pull-variant kernel, so it has to + # fit whichever variant runs at the largest size. + return max( + self.one_shot_pull_threshold, + self.two_shot_pull_threshold, + self.mc.max_bytes, + ) + + def clip(self, *, max_push_bytes: int, max_pull_bytes: int) -> "Heuristic": + return Heuristic( + min(self.one_shot_push_threshold, max_push_bytes), + min(self.one_shot_pull_threshold, max_pull_bytes), + min(self.two_shot_pull_threshold, max_pull_bytes), + self.mc.clip(max_pull_bytes), + ) + + +class AllReduceConfig(NamedTuple): + """All tuning knobs for a single (arch, world_size). + + The two ``Heuristic`` entries describe the size crossover for each + dispatch context (CUDA-graph capture vs eager). Block-count knobs apply + to the kernel grid: + - ``num_push_blocks``: 1shot_push grid (bound to the counter array) + - ``num_pull_blocks``: 1shot_pull (any mode) and non-mc 2shot_pull + - ``num_mc_blocks`` : mc 2shot_pull; ``None`` disables multicast + """ + + graph: Heuristic + eager: Heuristic + num_push_blocks: int + num_pull_blocks: int + num_mc_blocks: Optional[int] + + @property + def max_push_bytes(self) -> int: + return max(self.graph.max_push_bytes, self.eager.max_push_bytes) + + @property + def max_pull_bytes(self) -> int: + return max(self.graph.max_pull_bytes, self.eager.max_pull_bytes) + + def clip(self, *, max_push_bytes: int, max_pull_bytes: int) -> "AllReduceConfig": + return self._replace( + graph=self.graph.clip( + max_push_bytes=max_push_bytes, max_pull_bytes=max_pull_bytes + ), + eager=self.eager.clip( + max_push_bytes=max_push_bytes, max_pull_bytes=max_pull_bytes + ), + ) + + +def _pack_heuristic(*args) -> Heuristic: + arg_list: list = [int(p) if isinstance(p, float) else p for p in args] + return Heuristic(*arg_list) + + +def _sm100_config(world_size: int, num_sm: int) -> AllReduceConfig: + # SM100 (Blackwell, B200/B300). Tuned on B200 (148 SMs). + graph_map = { + 2: (8.000 * MB, 32.00 * MB, 128.0 * MB), + 3: (4.000 * MB, 4.000 * MB, 128.0 * MB), + 4: (2.250 * MB, 2.250 * MB, 128.0 * MB), + 5: (1.500 * MB, 1.500 * MB, 128.0 * MB), + 6: (1.000 * MB, 1.000 * MB, 128.0 * MB), + 7: (0.625 * MB, 0.625 * MB, 128.0 * MB), + 8: (0.500 * MB, 0.500 * MB, 128.0 * MB, Range(8 * MB, 128 * MB)), + } + eager_map = { + 2: (16.00 * MB, 128.0 * MB, 128.0 * MB), + 3: (8.000 * MB, 8.000 * MB, 32.00 * MB), + 4: (3.000 * MB, 3.000 * MB, 32.00 * MB), + 5: (2.000 * MB, 2.000 * MB, 32.00 * MB, Range(0, 32 * MB)), + 6: (1.250 * MB, 1.250 * MB, 64.00 * MB, Range(0, 64 * MB)), + 7: (1.000 * MB, 1.000 * MB, 64.00 * MB, Range(0, 64 * MB)), + 8: (0.750 * MB, 0.750 * MB, 128.0 * MB, Range(0, 128 * MB)), + } + mc_blocks_map = {5: 64, 6: 48, 7: 48, 8: 32} + return AllReduceConfig( + graph=_pack_heuristic(*graph_map[world_size]), + eager=_pack_heuristic(*eager_map[world_size]), + num_push_blocks=num_sm, + num_pull_blocks=num_sm if world_size == 2 else 96, + num_mc_blocks=mc_blocks_map.get(world_size, None), + ) + + +def _sm90_config(world_size: int, num_sm: int) -> AllReduceConfig: + # SM90 (Hopper, H100/H200). Tuned on H200. + graph_map = { + 2: (16.00 * MB, 128.0 * MB, 128.0 * MB), + 3: (1.250 * MB, 1.250 * MB, 128.0 * MB), + 4: (384.0 * KB, 384.0 * KB, 128.0 * MB), + 5: (192.0 * KB, 192.0 * KB, 32.00 * MB), + 6: (128.0 * KB, 128.0 * KB, 32.00 * MB, Range(8 * MB, 32 * MB)), + 7: (128.0 * KB, 128.0 * KB, 32.00 * MB, Range(1 * MB, 32 * MB)), + 8: (128.0 * KB, 128.0 * KB, 32.00 * MB, Range(512 * KB, 128 * MB)), + } + eager_map = { + 2: (32.00 * MB, 128.0 * MB, 128.0 * MB), + 3: (3.000 * MB, 3.000 * MB, 16.00 * MB), + 4: (896.0 * KB, 896.0 * KB, 32.00 * MB, Range(0, 32 * MB)), + 5: (384.0 * KB, 384.0 * KB, 32.00 * MB, Range(0, 32 * MB)), + 6: (192.0 * KB, 192.0 * KB, 32.00 * MB, Range(0, 32 * MB)), + 7: (128.0 * KB, 128.0 * KB, 32.00 * MB, Range(0, 32 * MB)), + 8: (128.0 * KB, 128.0 * KB, 128.0 * MB, Range(0, 128 * MB)), + } + return AllReduceConfig( + graph=_pack_heuristic(*graph_map[world_size]), + eager=_pack_heuristic(*eager_map[world_size]), + num_push_blocks=num_sm, + num_pull_blocks=64, + num_mc_blocks=None if world_size < 4 else 128 // world_size, + ) + + +@cache +def get_all_reduce_config(world_size: int) -> AllReduceConfig: + """Tuned thresholds and block counts for the current arch / world size. + + Only SM90 and SM100 are benchmarked so far; other archs get a + conservative default (1 MB one-shot crossovers, no multicast). + """ + cuda_major, _ = torch.cuda.get_device_capability() + num_sm = torch.cuda.get_device_properties().multi_processor_count + if cuda_major == 9: + return _sm90_config(world_size, num_sm) + if cuda_major == 10: + return _sm100_config(world_size, num_sm) + default = Heuristic(1 * MB, 1 * MB, 16 * MB) + return AllReduceConfig( + graph=default, + eager=default, + num_push_blocks=num_sm, + num_pull_blocks=num_sm, + num_mc_blocks=None, + ) diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py index a98707724..031688419 100644 --- a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py +++ b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py @@ -1,38 +1,89 @@ +"""JIT custom all-reduce (v2) over a decoupled storage plane. + +The CUDA side is split into two independent pieces: + +- ``Communicator``: a thin pointer holder over symmetric-memory workspaces + (push buffers, pull buffer, semaphores) plus a local push counter. All + storage is allocated and owned here, in Python. +- the all-reduce kernel: a pure function of ``(input, Communicator, algo, + pull_arg)`` with three algorithms (1shot_push / 1shot_pull / 2shot_pull) + and three pull data sources (eager workspace / CUDA-graph pointer table / + multicast address). + +CUDA-graph inputs are exchanged from Python after capture (cudaIpc handles +for cudaMalloc-backed pointers, fabric/posix-fd VMM mapping for expandable +segments) and written into a device-side pointer table (``graph_params``); +the kernel captured in the graph dereferences its row at replay time. +""" + +import enum import logging from contextlib import contextmanager -from dataclasses import dataclass, replace -from typing import Dict, List, Optional, TypeVar +from typing import List, Optional, Tuple 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.jit_kernel.all_reduce import ( + AllReduceAlgo, + Communicator, + IPCManager, + custom_all_reduce, ) 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 + +from .configs.custom_all_reduce_v2 import get_all_reduce_config +from .custom_all_reduce_utils import ( + can_use_custom_all_reduce_with_nvlink, + is_weak_contiguous, +) +from .vmm_utils import ( + VmmGraphInputManager, + compute_graph_capture_bases, + is_vmm_pointer, +) logger = logging.getLogger(__name__) -T = TypeVar("T") +MB = 1024 * 1024 -INF = 1 << 60 +_ALIGN_BYTES = 1024 +_SEMAPHORE_BYTES = 128 +_MAX_GRAPH_INPUTS = 131072 +# resolved once at import time; explicit constructor sizes take precedence +_DEFAULT_MAX_SIZE = envs.SGLANG_CUSTOM_ALL_REDUCE_V2_MAX_SIZE_KB.get() * 1024 -@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 _PullMode(enum.Enum): + EAGER = enum.auto() # pull_arg = False (also used for 1shot_push) + MULTICAST = enum.auto() # pull_arg = True + GRAPH = enum.auto() # pull_arg = a graph_params row + + +def _ceil_align(nbytes: int, align: int) -> int: + return (nbytes + align - 1) // align * align + + +def _allocate_symmetric_memory(nbytes: int, device: torch.device, group: ProcessGroup): + from torch._C._distributed_c10d import _SymmetricMemory + + if torch.__version__ < "2.11.0": + import torch.distributed._symmetric_memory as torch_symm_mem + + torch_symm_mem.enable_symm_mem_for_group(group.group_name) + tensor = _SymmetricMemory.empty_strided_p2p( + (nbytes,), + [1], + torch.uint8, + device, + group.group_name, + ) + symm_mem = _SymmetricMemory.rendezvous(tensor) + return tensor, symm_mem class CustomAllReduceV2: @@ -40,96 +91,195 @@ class CustomAllReduceV2: self, group: ProcessGroup, device: torch.device, + max_size: int = _DEFAULT_MAX_SIZE, + *, 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() + """ + :param max_size: direction-agnostic memory cap. Each workspace is + sized to what the tuned config wants, clipped to + this bound. Defaults to + ``SGLANG_CUSTOM_ALL_REDUCE_V2_MAX_SIZE_KB`` (16 MB). + :param max_pull_size: explicit pull workspace size; overrides both + the tuned size and ``max_size``. + :param max_push_size: explicit per-buffer push workspace size; + overrides both the tuned size and ``max_size``. + """ self.disabled = True if not can_use_custom_all_reduce_v2(group=group, device=device): return self.group = group + self.device = device 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 + base_config = get_all_reduce_config(self.world_size) + if max_pull_size is None: + max_pull_size = min(base_config.max_pull_bytes, max_size) + if max_push_size is None: + max_push_size = min(base_config.max_push_bytes, max_size) + # a minimal workspace keeps the Communicator valid even when a caller + # only uses one direction (e.g. push-only fused qk-norm instances) + self.max_pull_size = _ceil_align(max(max_pull_size, _ALIGN_BYTES), _ALIGN_BYTES) + self.max_push_size = _ceil_align(max(max_push_size, _ALIGN_BYTES), _ALIGN_BYTES) + self.max_size = max(self.max_pull_size, self.max_push_size) + num_pull_blocks = base_config.num_pull_blocks + num_push_blocks = base_config.num_push_blocks + if max_pull_blocks is not None: + num_pull_blocks = max(min(num_pull_blocks, max_pull_blocks), 1) + if max_push_blocks is not None: + num_push_blocks = max(max_push_blocks, 1) + self.config = base_config.clip( + max_push_bytes=self.max_push_size, max_pull_bytes=self.max_pull_size + )._replace(num_pull_blocks=num_pull_blocks, num_push_blocks=num_push_blocks) 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, + + # device-side pointer table: one row of world_size pointers per + # graph-captured all-reduce input (at most 8 MB at world_size = 8) + self.graph_params = torch.zeros( + (_MAX_GRAPH_INPUTS, self.world_size), + dtype=torch.uint64, + device=self.device, ) + self._init_workspace() + self._ipc_manager = IPCManager() self._vmm_graph_input_manager = VmmGraphInputManager( - obj=self.obj, + obj=self, group=self.group, rank=self.rank, world_size=self.world_size, ) - self._post_init_obj() + self._graph_inputs: List[Tuple[int, int]] = [] # (data_ptr, nbytes) + self._graph_counter = 0 + self._graph_mode_allowed = False 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, + def _init_workspace(self) -> None: + """Slice one symmetric-memory allocation into all shared buffers. + + Layout per rank: ``[2 * world_size push buffers | pull buffer | + pull semaphores]``. The push counter is rank-local, so it lives in a + plain CUDA tensor instead. + """ + cfg = self.config + push_num_bufs = 2 * self.world_size # 2 phases x world_size peers + push_ws_bytes = push_num_bufs * self.max_push_size + pull_ws_bytes = self.max_pull_size + pull_sem_bytes = _SEMAPHORE_BYTES * cfg.num_pull_blocks + total_bytes = push_ws_bytes + pull_ws_bytes + pull_sem_bytes + pull_ws_offset = push_ws_bytes + pull_sem_offset = push_ws_bytes + pull_ws_bytes + + self._symm_tensor, symm_mem = _allocate_symmetric_memory( + total_bytes, device=self.device, group=self.group + ) + workspaces = [ + symm_mem.get_buffer(i, [total_bytes], torch.uint8) + for i in range(self.world_size) + ] + workspaces[self.rank].zero_() + torch.cuda.synchronize() + dist.barrier(group=self.group) + + def slice_ws(rank: int, shape: List[int], offset: int) -> torch.Tensor: + nbytes = 1 + for s in shape: + nbytes *= s + assert offset + nbytes <= total_bytes + return workspaces[rank][offset : offset + nbytes].view(shape) + + push_workspaces = [ + slice_ws(i, [push_num_bufs, self.max_push_size], 0) + for i in range(self.world_size) + ] + pull_workspaces = [ + slice_ws(i, [pull_ws_bytes], pull_ws_offset) for i in range(self.world_size) + ] + pull_semaphores = [ + slice_ws(i, [cfg.num_pull_blocks, _SEMAPHORE_BYTES], pull_sem_offset) + for i in range(self.world_size) + ] + self._push_counter = torch.zeros( + (cfg.num_push_blocks,), dtype=torch.uint32, device=self.device ) - @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() + multicast_ptr = int(symm_mem.multicast_ptr) + can_multicast = multicast_ptr != 0 + pull_mc_workspace = multicast_ptr + pull_ws_offset if can_multicast else None + if not can_multicast or cfg.num_mc_blocks is None: + self.config = self.config._replace(num_mc_blocks=None) - def _register_graph_inputs_ipc(self): - """Register graph capture inputs via cudaIpcGetMemHandle. + self.obj = Communicator( + rank=self.rank, + world_size=self.world_size, + push_workspaces=push_workspaces, + pull_workspaces=pull_workspaces, + pull_semaphores=pull_semaphores, + push_counter=self._push_counter.view(-1, 1).view(torch.uint8), + pull_mc_workspace=pull_mc_workspace, + ) + if self.config.num_mc_blocks is not None: + self.obj.config(num_multicast_blocks=self.config.num_mc_blocks) + if self.rank == 0: + logger.info( + "All Reduce config: symmetric_memory = %.2f MB, " + "local_buffer = %.2f MB, multicast = %s", + total_bytes / MB, + (self.graph_params.nbytes + self._push_counter.nbytes) / MB, + self.config.num_mc_blocks is not None, + ) + dist.barrier(group=self.group) - This is the fast path for cudaMalloc-backed allocations. Fails - on VMM pointers (expandable_segments). + # ------------------------------------------------------------------ + # Algo selection + # ------------------------------------------------------------------ + + def uncap_pull_thresholds(self) -> None: + """Raise the 2-shot ceiling to the workspace capacity. + + The tuned config caps ``2shot_pull`` at the size where NCCL takes + over; benchmarks and tests that must keep every sweep size on the + custom-AR path can lift that cap up to ``max_pull_size``. """ - 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 uncap(heuristic): + return heuristic._replace(two_shot_pull_threshold=self.max_pull_size) + + self.config = self.config._replace( + graph=uncap(self.config.graph), + eager=uncap(self.config.eager), + ) + + def _can_use_graph(self) -> bool: + # `_graph_mode_allowed` is only set inside `capture()`, so the eager + # hot path never reaches the cudart capture query. During capture, + # warm-up runs execute immediately and must not consume a + # graph_params row (it would be dereferenced before registration). + return ( + self._graph_mode_allowed + and not is_in_tc_piecewise_cuda_graph() + and torch.cuda.is_current_stream_capturing() + ) + + def _pick_algo( + self, nbytes: int, can_use_graph: bool + ) -> Tuple[Optional[AllReduceAlgo], _PullMode]: + heuristic = self.config.graph if can_use_graph else self.config.eager + default_mode = _PullMode.GRAPH if can_use_graph else _PullMode.EAGER + use_multicast = self.config.num_mc_blocks is not None + if nbytes <= heuristic.one_shot_push_threshold: + return AllReduceAlgo.ONE_SHOT_PUSH, _PullMode.EAGER + if nbytes <= heuristic.one_shot_pull_threshold: + return AllReduceAlgo.ONE_SHOT_PULL, default_mode + if use_multicast and heuristic.mc.contains(nbytes): + return AllReduceAlgo.TWO_SHOT_PULL, _PullMode.MULTICAST + if nbytes <= heuristic.two_shot_pull_threshold: + return AllReduceAlgo.TWO_SHOT_PULL, default_mode + return None, _PullMode.EAGER def should_custom_ar(self, inp: torch.Tensor) -> bool: """Check if the input tensor is suitable for custom all-reduce.""" @@ -141,100 +291,130 @@ class CustomAllReduceV2: return False if not is_weak_contiguous(inp): return False - return inp_size <= self.max_size + if self.override_algo is not None: + return inp_size <= self.max_size + algo, _ = self._pick_algo(inp_size, can_use_graph=self._can_use_graph()) + return algo is not None + + # ------------------------------------------------------------------ + # All-reduce + # ------------------------------------------------------------------ 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) + nbytes = input.numel() * input.element_size() + can_use_graph = self._can_use_graph() + if self.override_algo is not None: + algo = self.override_algo + use_graph = can_use_graph and not algo.is_push() + mode = _PullMode.GRAPH if use_graph else _PullMode.EAGER + else: + algo, mode = self._pick_algo(nbytes, can_use_graph=can_use_graph) + assert algo is not None, f"No algo for {nbytes} bytes" + if mode == _PullMode.GRAPH: + pull_arg: torch.Tensor | bool = self._allocate_graph_row(input, nbytes) + else: + pull_arg = mode == _PullMode.MULTICAST + return torch.from_dlpack(custom_all_reduce(self.obj, input, algo, pull_arg)) + + def _allocate_graph_row(self, input: torch.Tensor, nbytes: int) -> torch.Tensor: + index = self._graph_counter + len(self._graph_inputs) + assert ( + index < _MAX_GRAPH_INPUTS + ), "Graph input table overflow, increase _MAX_GRAPH_INPUTS!" + self._graph_inputs.append((input.data_ptr(), nbytes)) + return self.graph_params[index] + + # ------------------------------------------------------------------ + # CUDA-graph input registration + # ------------------------------------------------------------------ + + @contextmanager + def capture(self): + if self.disabled: + yield + return + try: + self._graph_mode_allowed = not self.tms_cudagraph + yield + finally: + self._graph_mode_allowed = False + assert ( + not torch.cuda.is_current_stream_capturing() + ), "Cannot register graph inputs while capturing CUDA graph" + self._register_graph_inputs() + + def _register_graph_inputs(self) -> None: + if not self._graph_inputs: + return + first_ptr = self._graph_inputs[0][0] + if is_vmm_pointer(first_ptr): + # calls back into get_graph_capture_bases / register_peer_mapped_inputs + self._vmm_graph_input_manager.register_graph_inputs() + else: + self._register_graph_inputs_ipc() + + def _register_graph_inputs_ipc(self) -> None: + """Register graph capture inputs via cudaIpc handles. + + This is the fast path for cudaMalloc-backed allocations. Fails on + VMM pointers (expandable_segments), which use the VMM path instead. + """ + ptrs = [ptr for ptr, _ in self._graph_inputs] + handles = self._ipc_manager.batch_get_handles(ptrs) + local = [(list(handle), int(offset)) for handle, offset in handles] + gathered: List[Optional[list]] = [None] * self.world_size + dist.all_gather_object(gathered, local, group=self.group) + ptrs_per_rank: List[List[int]] = [] + for rank, remote in enumerate(gathered): + if rank == self.rank: + ptrs_per_rank.append(ptrs) + else: + ptrs_per_rank.append(list(self._ipc_manager.batch_open_handles(remote))) + peer_ptrs = [ + [ptrs_per_rank[rank][i] for rank in range(self.world_size)] + for i in range(len(ptrs)) + ] + self.register_peer_mapped_inputs(peer_ptrs) + + def get_graph_capture_bases(self): + """VMM base allocations of pending graph inputs (VmmGraphInputManager hook).""" + return compute_graph_capture_bases(self._graph_inputs) + + def register_peer_mapped_inputs(self, peer_ptrs: List[List[int]]) -> None: + """Write per-input peer pointers into the device-side pointer table.""" + assert len(peer_ptrs) == len(self._graph_inputs) + count = len(peer_ptrs) + rows = torch.tensor(peer_ptrs, dtype=torch.uint64, device=self.device) + self.graph_params[self._graph_counter : self._graph_counter + count].copy_(rows) + # the rows must be visible before any (PDL-chained) graph replay + torch.cuda.synchronize() + self._graph_counter += count + self._graph_inputs.clear() + + # ------------------------------------------------------------------ + # Teardown + # ------------------------------------------------------------------ def close(self): if not self.disabled and hasattr(self, "obj"): - self.obj.free(self.group) + self._ipc_manager.destroy() + dist.barrier(group=self.group) + del self.obj # drop the pointer holder before the workspace tensors 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()), + supported_world_size=list(range(2, 9)), cls_name="CustomAllReduceV2", ) return full_nvlink is True - - -THRESHOLD_2_SHOT_MAP: Dict[int, ModeConfig] = {} diff --git a/python/sglang/srt/distributed/device_communicators/vmm_utils.py b/python/sglang/srt/distributed/device_communicators/vmm_utils.py index 895bd259c..a286e9401 100644 --- a/python/sglang/srt/distributed/device_communicators/vmm_utils.py +++ b/python/sglang/srt/distributed/device_communicators/vmm_utils.py @@ -52,6 +52,56 @@ def is_vmm_pointer(ptr: int) -> bool: return False +def compute_graph_capture_bases(graph_inputs: List[tuple]): + """Map graph-capture inputs onto their VMM base allocations. + + ``graph_inputs`` is a list of ``(device_ptr, nbytes)`` pairs. A captured + tensor can cross expandable-segment allocation boundaries, so each input + is walked with ``cuMemGetAddressRange`` until its byte span is covered. + + Returns ``(bases_info, input_chunk_indices, input_offsets)``: + - ``bases_info[i] = (base_ptr, alloc_size)`` per unique allocation + - ``input_chunk_indices[j]`` = indices of allocations covering input j + - ``input_offsets[j]`` = byte offset of input j from its first base + """ + drv = _get_cuda_driver() + base_to_idx = {} + bases_info: List[tuple] = [] + input_chunk_indices: List[List[int]] = [] + input_offsets: List[int] = [] + for ptr, nbytes in graph_inputs: + ptr, remaining = int(ptr), int(nbytes) + if remaining <= 0: + raise RuntimeError(f"Invalid graph capture input size: {nbytes}") + cursor = ptr + first_base = None + chunks: List[int] = [] + while remaining > 0: + err, base, size = drv.cuMemGetAddressRange(cursor) + if err != drv.CUresult.CUDA_SUCCESS: + raise RuntimeError(f"cuMemGetAddressRange: {err}") + base, size = int(base), int(size) + if first_base is None: + first_base = base + byte_offset = cursor - base + if not 0 <= byte_offset < size: + raise RuntimeError( + f"graph capture input at {ptr} is outside VMM allocation " + f"[base={base}, size={size}]" + ) + idx = base_to_idx.setdefault(base, len(bases_info)) + if idx == len(bases_info): + bases_info.append((base, size)) + chunks.append(idx) + advance = min(remaining, size - byte_offset) + assert advance > 0, "Failed to advance VMM graph capture span" + remaining -= advance + cursor += advance + input_chunk_indices.append(chunks) + input_offsets.append(ptr - first_base) + return bases_info, input_chunk_indices, input_offsets + + def make_rw_access_desc(device_id: int): """A read-write, device-local ``CUmemAccessDesc`` for ``device_id``.""" drv = _get_cuda_driver() diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index d5beef75b..337ffd8a8 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -726,6 +726,9 @@ class Envs: # Set to 0: force disable (use default Aiter AR even with --enable-deterministic-inference) SGLANG_USE_1STAGE_ALLREDUCE = EnvBool(False) SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2 = EnvBool(True) + # Default per-direction workspace cap for CustomAllReduceV2; explicit + # constructor sizes take precedence over this. + SGLANG_CUSTOM_ALL_REDUCE_V2_MAX_SIZE_KB = EnvInt(16 * 1024) SGLANG_FLASHINFER_PREFILL_SPLIT_TILE_SIZE = EnvInt(4096) SGLANG_FLASHINFER_DECODE_SPLIT_TILE_SIZE = EnvInt(2048) SGLANG_TRITON_PREFILL_TRUNCATION_ALIGN_SIZE = EnvInt(4096) diff --git a/test/registered/jit/benchmark/bench_custom_all_reduce.py b/test/registered/jit/benchmark/bench_custom_all_reduce.py index a52d443a9..e0ea022d3 100644 --- a/test/registered/jit/benchmark/bench_custom_all_reduce.py +++ b/test/registered/jit/benchmark/bench_custom_all_reduce.py @@ -1,18 +1,3 @@ -"""Benchmark JIT custom all-reduce (v2) vs NCCL, AOT custom-AR (v1), and -FlashInfer trtllm allreduce_fusion. - -Usage:: - - # Benchmark on every supported world size (2..8 GPUs): - python benchmark/bench_custom_all_reduce.py - # Pick a specific world size (or comma-separated list): - python benchmark/bench_custom_all_reduce.py --num-gpu 4 - python benchmark/bench_custom_all_reduce.py --num-gpu 2,4,8 - -The script self-relaunches under ``torchrun --nproc_per_node=N`` for each N in -``num_gpus``; results are printed on rank 0 of every run. -""" - from __future__ import annotations import atexit @@ -44,35 +29,19 @@ register_cuda_ci( # --------------------------------------------------------------------------- DTYPE = torch.bfloat16 -# torch.dtype.itemsize exists only on newer torch; element_size() is portable. -DTYPE_ITEMSIZE = torch.tensor([], dtype=DTYPE).element_size() -MESSAGE_SIZES_BYTES = [ - 4 * 1024, # 4K - 16 * 1024, # 16K - 64 * 1024, # 64K - 128 * 1024, # 128K - 3 * 64 * 1024, # 192K - 4 * 64 * 1024, # 256K - 3 * 128 * 1024, # 384K - 4 * 128 * 1024, # 512K - 5 * 128 * 1024, # 640K - 6 * 128 * 1024, # 768K - 7 * 128 * 1024, # 896K - 1 * 1024 * 1024, # 1M - 2 * 1024 * 1024, # 2M - 3 * 1024 * 1024, # 3M - 4 * 1024 * 1024, # 4M - 8 * 1024 * 1024, # 8M - 16 * 1024 * 1024, # 16M - 32 * 1024 * 1024, # 32M -] +DTYPE_ITEMSIZE = DTYPE.itemsize +MESSAGE_SIZES_KB = [2**x for x in range(2, 17)] +MESSAGE_SIZES_KB += [192, 384, 640, 768, 896, 1536, 3072] +MESSAGE_SIZES_KB.sort() WORLD_SIZES = list(range(2, 9)) -MAX_BYTES = max(MESSAGE_SIZES_BYTES) +MAX_BYTES = max(MESSAGE_SIZES_KB) * 1024 # trtllm allreduce_fusion only supports these world sizes. FI_SUPPORTED_WORLD_SIZES = (2, 4, 8) # AOT custom_all_reduce (v1) only supports these world sizes. AOT_SUPPORTED_WORLD_SIZES = (2, 4, 6, 8) -PROVIDERS = ["nccl", "aot", "jit", "fi"] +# jit-eager times the naive-loop dispatch (eager heuristics); jit-graph +# captures the calls in a CUDA graph (graph heuristics + pointer table). +PROVIDERS = ["nccl", "aot", "jit-eager", "jit-graph", "fi"] WORLD_SIZES = get_benchmark_range(WORLD_SIZES, [2, 4, 8]) # --------------------------------------------------------------------------- @@ -92,8 +61,6 @@ def _init_cpu_group() -> dist.ProcessGroup: backend="nccl", ) atexit.register(dist.destroy_process_group) - # Quieter benchmark output. - logging.disable(logging.INFO) torch.cuda.set_stream(torch.cuda.Stream()) return coord.cpu_group @@ -101,9 +68,13 @@ def _init_cpu_group() -> dist.ProcessGroup: @cache_once def _init_nccl_group() -> dist.ProcessGroup: _init_cpu_group() - coord = ps._WORLD - assert coord is not None and coord.device_group is not None - return coord.device_group + local_rank = int(os.environ["LOCAL_RANK"]) + device_group = torch.distributed.new_group( + backend="nccl", + device_id=torch.device(f"cuda:{local_rank}"), + ) + assert isinstance(device_group, dist.ProcessGroup) + return device_group # --------------------------------------------------------------------------- @@ -133,11 +104,13 @@ class JITAllReduceBackend: ) device = torch.device(f"cuda:{int(os.environ['LOCAL_RANK'])}") - self.comm = CustomAllReduceV2( - _init_cpu_group(), device, max_pull_size=MAX_BYTES - ) + # tuned workspace sizes, capped at the sweep maximum + self.comm = CustomAllReduceV2(_init_cpu_group(), device, max_size=MAX_BYTES) if self.comm.disabled: raise RuntimeError("JIT CustomAllReduceV2 is disabled on this system") + # keep the whole sweep on the custom-AR path: the tuned config would + # otherwise send the largest sizes back to NCCL + self.comm.uncap_pull_thresholds() register_comm_cleanup(self.comm) def graph_context(self): @@ -177,7 +150,7 @@ class FlashInferAllReduceBackend: world_size = dist.get_world_size(group=group) # Use the smallest message size as the inner hidden dim, so any # message in the sweep is an integer multiple of it. - hidden_dim = min(MESSAGE_SIZES_BYTES) // DTYPE_ITEMSIZE + hidden_dim = 1024 * min(MESSAGE_SIZES_KB) // DTYPE_ITEMSIZE num_tokens = MAX_BYTES // (hidden_dim * DTYPE_ITEMSIZE) self._comm = comm self._hidden_dim = hidden_dim @@ -225,7 +198,8 @@ def _init_fi_backend() -> FlashInferAllReduceBackend: BACKEND_FACTORY = { "nccl": _init_nccl_backend, - "jit": _init_jit_backend, + "jit-eager": _init_jit_backend, + "jit-graph": _init_jit_backend, "aot": _init_aot_backend, "fi": _init_fi_backend, } @@ -236,6 +210,10 @@ def _init_all_backends() -> None: """Pre-build every supported backend before any timed iteration so JIT compilation / IPC setup don't bleed into the first measured size. """ + local_rank = int(os.environ["LOCAL_RANK"]) + if local_rank == 0: # NOTE: log some verbose info on initialization + logging.basicConfig(level=logging.INFO) + world_size = dist.get_world_size(_init_cpu_group()) factories = dict(BACKEND_FACTORY) if world_size not in AOT_SUPPORTED_WORLD_SIZES: @@ -245,15 +223,18 @@ def _init_all_backends() -> None: for fn in factories.values(): fn() + # reset level to warning + logging.getLogger().setLevel(logging.WARNING) + # --------------------------------------------------------------------------- # Benchmark # --------------------------------------------------------------------------- -@marker.parametrize("message_bytes", MESSAGE_SIZES_BYTES) +@marker.parametrize("message_KB", MESSAGE_SIZES_KB) @marker.benchmark("provider", PROVIDERS) -def benchmark(message_bytes: int, provider: str): +def benchmark(message_KB: int, provider: str): cpu_group = _init_cpu_group() gpu_group = _init_nccl_group() world_size = dist.get_world_size(cpu_group) @@ -268,15 +249,18 @@ def benchmark(message_bytes: int, provider: str): ) _init_all_backends() backend = BACKEND_FACTORY[provider]() + message_bytes = message_KB * 1024 numel = message_bytes // DTYPE_ITEMSIZE - device = torch.device(f"cuda:{int(os.environ['LOCAL_RANK'])}") + device_id = int(os.environ["LOCAL_RANK"]) + device = torch.device(f"cuda:{device_id}") x = torch.randn(numel, dtype=DTYPE, device=device) + ctx_fn = backend.graph_context if not provider.endswith("eager") else None # Bandwidth-equivalent bytes moved by a ring all-reduce per rank. effective_bytes = int(x.nbytes * 2 * (world_size - 1) / world_size) return marker.do_bench( backend.all_reduce, input_args=(x,), - graph_context_fn=backend.graph_context, + graph_context_fn=ctx_fn, sync_multigpu_fn=lambda: dist.barrier(gpu_group), # all-reduce is in-place w.r.t. its argument; explicit footprint # captures the cross-GPU traffic instead. diff --git a/test/registered/jit/benchmark/bench_tp_qknorm.py b/test/registered/jit/benchmark/bench_tp_qknorm.py index f4f71bcf4..a63fdb6ba 100644 --- a/test/registered/jit/benchmark/bench_tp_qknorm.py +++ b/test/registered/jit/benchmark/bench_tp_qknorm.py @@ -24,11 +24,10 @@ import torch.distributed as dist import sglang.srt.distributed.parallel_state as ps from sglang.jit_kernel.all_reduce import ( - _jit_custom_all_reduce_pull_module, - _jit_custom_all_reduce_push_module, - _jit_fused_parallel_qknorm_module, fused_parallel_qknorm, + get_all_reduce_module, get_fused_parallel_qknorm_max_occupancy, + get_fused_parallel_qknorm_module, ) from sglang.jit_kernel.benchmark import marker from sglang.jit_kernel.benchmark.utils import multigpu_bench_main @@ -70,13 +69,11 @@ def _compile_one(world_size: int) -> None: Top-level so it survives ``spawn`` pickling. Compiled artifacts are cached on disk by ``tvm_ffi``; torchrun children will reuse them. """ - # baseline path: sum-sq -> pull-mode all-reduce -> apply - _jit_custom_all_reduce_pull_module(DTYPE, world_size) - # fused path: push-mode all-reduce - _jit_custom_all_reduce_push_module(DTYPE, world_size) + # baseline path: sum-sq -> all-reduce -> apply (also covers push mode) + get_all_reduce_module(DTYPE, world_size) # fused path: fused QKNorm kernel (one per (dtype, world_size, q_dim, k_dim)) for q_dim, k_dim in Q_K_DIMS: - _jit_fused_parallel_qknorm_module(DTYPE, world_size, q_dim, k_dim) + get_fused_parallel_qknorm_module(DTYPE, world_size, q_dim, k_dim) def _precompile_kernels(num_gpus: List[int]) -> None: diff --git a/test/registered/jit/test_custom_all_reduce.py b/test/registered/jit/test_custom_all_reduce.py index be827a660..8fd238828 100644 --- a/test/registered/jit/test_custom_all_reduce.py +++ b/test/registered/jit/test_custom_all_reduce.py @@ -11,8 +11,6 @@ Usage:: # odd / non-power-of-two counts that the default sweep skips: python tests/test_custom_all_reduce.py --num-gpu 3 python tests/test_custom_all_reduce.py --num-gpu 2,4,6,8 - # Extra pytest args (forwarded to each torchrun worker): - python tests/test_custom_all_reduce.py -k bfloat16 """ from __future__ import annotations @@ -30,11 +28,7 @@ import torch import torch.distributed as dist import sglang.srt.distributed.parallel_state as ps -from sglang.jit_kernel.all_reduce import ( - AllReduceAlgo, - _jit_custom_all_reduce_pull_module, - _jit_custom_all_reduce_push_module, -) +from sglang.jit_kernel.all_reduce import AllReduceAlgo, get_all_reduce_module from sglang.jit_kernel.mp import register_comm_cleanup from sglang.jit_kernel.tests.utils import multigpu_pytest_main from sglang.jit_kernel.utils import cache_once, get_ci_test_range @@ -90,13 +84,12 @@ TEST_DTYPES = get_ci_test_range(TEST_DTYPES, [torch.bfloat16]) def _compile_one(dtype: torch.dtype, world_size: int) -> None: - """Compile both (push, pull) variants for a single (dtype, world_size). + """Compile the all-reduce module for a single (dtype, world_size). Top-level so it survives ``spawn`` pickling. Compiled artifacts are cached on disk by ``tvm_ffi``; torchrun children will reuse them. """ - _jit_custom_all_reduce_pull_module(dtype, world_size) - _jit_custom_all_reduce_push_module(dtype, world_size) + get_all_reduce_module(dtype, world_size) def _precompile_kernels(num_gpus: List[int]) -> None: @@ -149,10 +142,16 @@ def _init_cpu_group_once() -> dist.ProcessGroup: @cache_once def _init_nccl_group_once() -> dist.ProcessGroup: + # Reference NCCL group allocated independently of the parallel_state + # world group, so the test does not couple to framework internals. _init_cpu_group_once() - coord = ps._WORLD - assert coord is not None and coord.device_group is not None - return coord.device_group + local_rank = int(os.environ["LOCAL_RANK"]) + device_group = dist.new_group( + backend="nccl", + device_id=torch.device(f"cuda:{local_rank}"), + ) + assert isinstance(device_group, dist.ProcessGroup) + return device_group @cache_once @@ -162,7 +161,9 @@ def _init_comm_once() -> CustomAllReduceV2: max_size = max(TEST_SIZES) * max( torch.tensor([], dtype=d).element_size() for d in TEST_DTYPES ) - comm = CustomAllReduceV2(cpu_group, device, max_size, max_size) + comm = CustomAllReduceV2( + cpu_group, device, max_pull_size=max_size, max_push_size=max_size + ) if comm.disabled: raise RuntimeError("JIT CustomAllReduceV2 is disabled on this system") register_comm_cleanup(comm) @@ -224,6 +225,9 @@ def test_custom_all_reduce( dist.all_reduce(out_ref, group=nccl_group) out_jit = run(inp) # Exact equality, since values are small integers within bf16 precision. + # NOTE: use torch's assert_close: it compares on device (~2 ms here), + # while triton's converts to numpy on the host (~0.6 s per 32 MB + # tensor) and would dominate the test wall time. torch.testing.assert_close(out_ref, out_jit, atol=0, rtol=0) diff --git a/test/registered/jit/test_tp_qknorm.py b/test/registered/jit/test_tp_qknorm.py index a880a0a94..6eccd8416 100644 --- a/test/registered/jit/test_tp_qknorm.py +++ b/test/registered/jit/test_tp_qknorm.py @@ -15,9 +15,9 @@ import triton import sglang.srt.distributed.parallel_state as ps from sglang.jit_kernel.all_reduce import ( - _jit_custom_all_reduce_push_module, - _jit_fused_parallel_qknorm_module, fused_parallel_qknorm, + get_all_reduce_module, + get_fused_parallel_qknorm_module, ) from sglang.jit_kernel.mp import register_comm_cleanup from sglang.jit_kernel.tests.utils import multigpu_pytest_main @@ -55,9 +55,9 @@ def _compile_one(dtype: torch.dtype, world_size: int) -> None: Top-level so it survives ``spawn`` pickling. Compiled artifacts are cached on disk by ``tvm_ffi``; torchrun children will reuse them. """ - _jit_custom_all_reduce_push_module(dtype, world_size) + get_all_reduce_module(dtype, world_size) for q_dim, k_dim in Q_K_DIMS: - _jit_fused_parallel_qknorm_module(dtype, world_size, q_dim, k_dim) + get_fused_parallel_qknorm_module(dtype, world_size, q_dim, k_dim) def _precompile_kernels(num_gpus: List[int]) -> None: @@ -102,10 +102,16 @@ def _init_cpu_group_once() -> dist.ProcessGroup: @cache_once def _init_nccl_group_once() -> dist.ProcessGroup: + # Reference NCCL group allocated independently of the parallel_state + # world group, so the test does not couple to framework internals. _init_cpu_group_once() - coord = ps._WORLD - assert coord is not None and coord.device_group is not None - return coord.device_group + local_rank = int(os.environ["LOCAL_RANK"]) + device_group = dist.new_group( + backend="nccl", + device_id=torch.device(f"cuda:{local_rank}"), + ) + assert isinstance(device_group, dist.ProcessGroup) + return device_group @cache_once @@ -114,7 +120,9 @@ def _init_comm_once() -> CustomAllReduceV2: device = torch.device(f"cuda:{int(os.environ['LOCAL_RANK'])}") max_pull_size = 0 max_push_size = 8 * max(BATCH_SIZES) - comm = CustomAllReduceV2(cpu_group, device, max_pull_size, max_push_size) + comm = CustomAllReduceV2( + cpu_group, device, max_pull_size=max_pull_size, max_push_size=max_push_size + ) if comm.disabled: raise RuntimeError("JIT CustomAllReduceV2 is disabled on this system") register_comm_cleanup(comm)