[Kernel] Rewrite JIT custom all-reduce (v2) with a decoupled kernel/storage design (#31049)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: root <root@GPUC5A6.maas>

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:
+    def algo_name(self) -> str:
-        return 2 if self == AllReduceAlgo.TWO_SHOT_PULL else 1
+        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]
+    # (cudaIpcMemHandle bytes, offset-in-allocation) for one device pointer
-    CUSTOM_AR_PAIR = Tuple[int, CUSTOM_AR_HANDLE]
+    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.
+def _init_communicator() -> None:
-            :param world_size: The total number of processes in the group.
+    module = load_jit(
-            :param pull_buffer_bytes: The size of the buffer (in bytes) used for pull-based all-reduce.
+        "communicator",
-            :param push_buffer_bytes: The size of the buffer (in bytes) used for push-based all-reduce.
+        cuda_files=["distributed/communicator.cuh"],
-            :param graph_input_count: The maximum number of inputs in all CUDA graphs.
+        cuda_wrappers=[("register_once", "register_communicator")],
-            :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.
+    module.register_once()
-            :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.
-            """
 
-        @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.
+@lazy_register_class("sgl.Communicator", _init_communicator)
-            :param num_threads: The maximum number of threads per block. -1 means no limit.
+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.cuh"],
-        cuda_files=["distributed/custom_all_reduce_pull.cuh"],
         cuda_wrappers=[("all_reduce", f"custom_all_reduce<{args}>")],
     )
 
 
-@cache_once
+@debug_kernel_api
-def _jit_custom_all_reduce_push_module(dtype: torch.dtype, world_size: int) -> Module:
+def custom_all_reduce(
-    args = make_cpp_args(dtype, world_size, is_arch_support_pdl())
+    comm: Communicator,
-    return load_jit(
+    input: torch.Tensor,
-        "custom_all_reduce_push",
+    algo: AllReduceAlgo,
-        *args,
+    pull_arg: PullArg,
-        extra_ldflags=["-lcuda"],
+) -> tvm_ffi.Tensor:
-        cuda_files=["distributed/custom_all_reduce_push.cuh"],
+    module = get_all_reduce_module(input.dtype, comm.world_size)
-        cuda_wrappers=[("all_reduce", f"custom_all_reduce<{args}>")],
+    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 = get_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.fused_parallel_qknorm(comm, q, k, q_weight, k_weight, eps)

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

python/sglang/jit_kernel/csrc/distributed/communicator.cuh

@@ -0,0 +1,113 @@
+#include <sgl_kernel/ffi.h>
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/distributed/communicator.cuh>
+
+#include <tvm/ffi/extra/stl.h>
+#include <tvm/ffi/reflection/registry.h>
+
+#include <cstdint>
+#include <map>
+#include <optional>
+#include <string>
+#include <vector>
+
+namespace host::distributed {
+
+inline CommunicatorObj::CommunicatorObj(
+    const uint32_t rank,
+    const uint32_t world_size,
+    std::vector<TensorView> push_workspaces,
+    std::vector<TensorView> pull_workspaces,
+    std::vector<TensorView> pull_semaphores,
+    TensorView push_counter,
+    const std::optional<int64_t> 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<kDLCUDA>();
+  for (uint32_t i = 0; i < world_size; ++i) {
+    TensorMatcher({2 * world_size, push_bytes}).with_dtype<uint8_t>().with_device(device).verify(push_workspaces[i]);
+    TensorMatcher({pull_bytes})  //
+        .with_dtype<uint8_t>()
+        .with_device(device)
+        .verify(pull_workspaces[i]);
+    TensorMatcher({num_pull_blocks, static_cast<int64_t>(sizeof(Semaphore))})
+        .with_dtype<uint8_t>()
+        .with_device(device)
+        .verify(pull_semaphores[i]);
+    this->push_workspaces[i] = static_cast<uint8_t*>(push_workspaces[i].data_ptr());
+    this->pull_workspaces[i] = static_cast<uint8_t*>(pull_workspaces[i].data_ptr());
+    this->pull_semaphores[i] = static_cast<Semaphore*>(pull_semaphores[i].data_ptr());
+  }
+  TensorMatcher({num_push_blocks, static_cast<int64_t>(sizeof(Counter))})
+      .with_dtype<uint8_t>()
+      .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<uint8_t*>(static_cast<uintptr_t>(pull_mc_workspace_ptr.value()));
+  } else {
+    this->pull_mc_workspace = nullptr;
+  }
+
+  // push config
+  this->push_counter = static_cast<Counter*>(push_counter.data_ptr());
+  this->push_bytes = push_bytes.unwrap();
+  this->num_push_blocks = static_cast<uint32_t>(num_push_blocks.unwrap());
+
+  // pull config
+  this->pull_bytes = pull_bytes.unwrap();
+  this->num_pull_blocks = static_cast<uint32_t>(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<std::string, uint32_t> 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<Class>()
+      .def(
+          refl::init<
+              uint32_t,
+              uint32_t,
+              std::vector<TensorView>,
+              std::vector<TensorView>,
+              std::vector<TensorView>,
+              TensorView,
+              std::optional<int64_t>>(),
+          "__init__")
+      .def_ro("world_size", &Class::world_size)
+      .def_ro("rank", &Class::rank)
+      .def("_config", &Class::config);
+}

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 <sgl_kernel/ffi.h>
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+
+#include <sgl_kernel/distributed/communicator.cuh>
+
+#include <tvm/ffi/extra/stl.h>
+
+#include <algorithm>
+#include <bit>
+#include <cstdint>
+#include <cstring>
+#include <string>
+#include <variant>
+
+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 <typename T>
+struct fp_trait {};
+
+template <>
+struct fp_trait<bf16_t> {
+  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<fp16_t> {
+  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<float> {
+  using type = uint32_t;
+  [[maybe_unused]]
+  static constexpr uint32_t pos_zero = 0x00000000u;
+  [[maybe_unused]]
+  static constexpr uint32_t neg_zero = 0x80000000u;
+};
+
+template <typename DType>
+SGL_DEVICE void clear_pos_zero(DType& val) {
+  using Trait = fp_trait<DType>;
+  const auto ptr = reinterpret_cast<typename Trait::type*>(&val);
+  if (*ptr == Trait::pos_zero) *ptr = Trait::neg_zero;
+}
+
+template <typename DType>
+SGL_DEVICE bool is_pos_zero(const DType& val) {
+  using Trait = fp_trait<DType>;
+  const auto ptr = reinterpret_cast<const typename Trait::type*>(&val);
+  return *ptr == Trait::pos_zero;
+}
+
+template <typename DType>
+SGL_DEVICE DType get_pos_zero() {
+  using Trait = fp_trait<DType>;
+  const auto value = Trait::pos_zero;
+  return *reinterpret_cast<const DType*>(&value);
+}
+
+template <typename T2, size_t N, size_t M>
+SGL_DEVICE auto reduce(device::AlignedVector<T2, N> (&vec)[M]) -> device::AlignedVector<T2, N> {
+  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<fp32x2_t>(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<T2, N> out_vec;
+#pragma unroll
+  for (size_t j = 0; j < N; ++j) {
+    out_vec[j] = device::cast<T2>(acc_vec[j]);
+  }
+  return out_vec;
+}
+
+template <typename V>
+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<const uint8_t*>(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<const V*>(&val);
+}
+
+template <typename V>
+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<const uint4*>(&x);
+  addr = static_cast<uint8_t*>(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 <typename V>
+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<const uint8_t*>(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<const V*>(&val);
+}
+
+template <typename V>
+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<const uint4*>(&x);
+  addr = static_cast<uint8_t*>(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 <typename V>
+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<const uint8_t*>(mc_addr) + vec_offset * 16;
+  if constexpr (std::is_same_v<V, device::AlignedVector<fp32x2_t, 2>>) {
+    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<const V*>(&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<V, device::AlignedVector<fp16x2_t, 4>>) {
+      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<V, device::AlignedVector<bf16x2_t, 4>>);  // 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<const V*>(&val);
+  }
+#else
+  assert(false && "multimem load is only supported on Hopper or later architecture");
+#endif
+}
+
+template <typename V>
+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<const float4*>(&x);
+  mc_addr = static_cast<uint8_t*>(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 <typename T, uint32_t kWorldSize, bool kUsePDL>
+struct AllReducePushImpl {
+ private:
+  using T2 = packed_t<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<T2, kVecSize>;
+  static_assert(kWorldSize <= kMaxWorldSize);
+
+  static SGL_DEVICE bool sync_enter_push(const AllReduceParams& params) {
+    device::PDLWaitPrimary<kUsePDL>();
+    return (params.push_counter[blockIdx.x].get() % 2) != 0;
+  }
+
+  static SGL_DEVICE void sync_exit_push(const AllReduceParams& params) {
+    device::PDLTriggerSecondary<kUsePDL>();
+    __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<T>();
+#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 <typename T, uint32_t kWorldSize, PullMode kMode, bool kUsePDL>
+struct AllReducePullImpl {
+ private:
+  using T2 = packed_t<T>;
+  static constexpr uint32_t kVecSize = 16 / sizeof(T2);
+  static constexpr uint32_t kElemsPerVec = 16 / sizeof(T);
+  using vec_t = device::AlignedVector<T2, kVecSize>;
+  static_assert(kWorldSize <= kMaxWorldSize);
+
+  template <bool kFence>
+  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<kUsePDL>();
+      const auto bx = blockIdx.x;
+      const auto semaphore = &params.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 <bool kFence>
+  static SGL_DEVICE void sync_exit_pull(const AllReduceParams& params, uint32_t current) {
+    device::PDLTriggerSecondary<kUsePDL>();
+    __syncthreads();
+    if (const auto tx = threadIdx.x; tx < kWorldSize) {
+      const auto bx = blockIdx.x;
+      const auto semaphore = &params.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 <bool kIs2shot>
+  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<false>(params);
+    reduce_impl<false>(total_num_vecs, data, params.output, params.pull_mc_workspace);
+    sync_exit_pull<false>(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<vec_t*>(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<vec_t*>(params.pull_workspaces[i]) + local_vec_bias;
+      }
+    }
+    const auto counter = sync_enter_pull<false>(params);
+    const auto mc_addr = reinterpret_cast<vec_t*>(params.pull_mc_workspace) + local_vec_bias;
+    reduce_impl<true>(local_num_vecs, data, params.output, mc_addr);
+    sync_exit_pull<true>(params, counter);
+  }
+};
+
+template <typename Impl, int kShot>
+__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 <uint32_t N>
+__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<uint32_t, N / 4>;
+  const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
+  device::PDLWaitPrimary<true>();
+  device::PDLTriggerSecondary<true>();
+  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<uint32_t>(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 <typename T, uint32_t kWorldSize, bool kUsePDL>
+struct AllReduceKernel {
+ private:
+  using Tensor = tvm::ffi::Tensor;
+  using TensorView = tvm::ffi::TensorView;
+  template <int kShot, PullMode kPullMode>
+  static constexpr auto kernel_pull = all_reduce_kernel<AllReducePullImpl<T, kWorldSize, kPullMode, kUsePDL>, kShot>;
+  template <int kShot>
+  static constexpr auto kernel_push = all_reduce_kernel<AllReducePushImpl<T, kWorldSize, kUsePDL>, kShot>;
+
+ public:
+  static Tensor run(CommunicatorRef ref, Tensor in_, std::string algo, std::variant<TensorView, bool> 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<T>(in_.dtype()), "Input dtype mismatch");
+    RuntimeCheck(in_.device().device_type == kDLCUDA, "Only CUDA device is supported");
+    RuntimeCheck(std::bit_cast<intptr_t>(in_.data_ptr()) % 16 == 0, "Input pointer is not properly aligned");
+    const auto num_elems_int64 = in_.numel();
+    const auto num_elems = static_cast<uint32_t>(num_elems_int64);
+    RuntimeCheck(static_cast<int64_t>(num_elems) == num_elems_int64, "Number of items exceeds 4G limit");
+    const bool use_graph = std::holds_alternative<TensorView>(pull_arg);
+    const auto graph_ptr = use_graph ? std::get<TensorView>(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<void* const*>(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<bool>(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<device::kMaxVecBytes>;
+          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 <typename T, uint32_t kWorldSize, bool kUsePDL>
+tvm::ffi::Tensor custom_all_reduce(
+    CommunicatorRef comm, tvm::ffi::Tensor input, std::string algo, std::variant<tvm::ffi::TensorView, bool> pull_arg) {
+  return AllReduceKernel<T, kWorldSize, kUsePDL>::run(comm, input, algo, pull_arg);
+}
+
+}  // namespace

python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_base.cuh

@@ -1,30 +0,0 @@
-#include <sgl_kernel/ffi.h>
-#include <sgl_kernel/tensor.h>
-#include <sgl_kernel/utils.h>
-
-#include <sgl_kernel/type.cuh>
-#include <sgl_kernel/utils.cuh>
-#include <sgl_kernel/vec.cuh>
-
-#include <sgl_kernel/distributed/custom_all_reduce.cuh>
-
-#include <cstdint>
-#include <cstring>
-
-inline void register_custom_all_reduce() {
-  namespace refl = tvm::ffi::reflection;
-  using Class = host::distributed::CustomAllReduceBase;
-  refl::ObjectDef<Class>()
-      .def(refl::init<uint32_t, uint32_t, uint32_t, uint32_t, int64_t, int64_t, int64_t>(), "__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);
-}

python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_pull.cuh

@@ -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 <sgl_kernel/ffi.h>
-#include <sgl_kernel/tensor.h>
-#include <sgl_kernel/utils.h>
-
-#include <sgl_kernel/type.cuh>
-#include <sgl_kernel/utils.cuh>
-#include <sgl_kernel/vec.cuh>
-
-#include <sgl_kernel/distributed/common.cuh>
-#include <sgl_kernel/distributed/custom_all_reduce.cuh>
-
-#include <bit>
-#include <cstdint>
-#include <cstring>
-
-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 <bool kBroadcast, typename DType, uint32_t kNumGPU>
-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<DType>;
-  using Storage = AlignedVector<DType2, kVecSize>;
-  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 <typename DType, uint32_t kNumGPU, bool kUsePDL>
-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<DType*>(data->input[i]);
-  device::PDLWaitPrimary<kUsePDL>();
-
-  ctrl.sync</*kFence=*/0, /*kStart=*/1>(params.rank, kNumGPU);
-  all_reduce_impl</*kBroadcast=*/false>(params, input);
-
-  device::PDLTriggerSecondary<kUsePDL>();
-  ctrl.sync</*kFence=*/0, /*kStart=*/0>(params.rank, kNumGPU);
-}
-
-template <typename DType, uint32_t kNumGPU, bool kUsePDL>
-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<DType*>(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<kUsePDL>();
-
-  ctrl.sync</*kFence=*/0, /*kStart=*/1>(params.rank, kNumGPU);
-  all_reduce_impl</*kBroadcast=*/true>(local_params, input);
-
-  device::PDLTriggerSecondary<kUsePDL>();
-  ctrl.sync</*kFence=*/1, /*kStart=*/0>(params.rank, kNumGPU);
-}
-
-template <typename DType, uint32_t kNumGPU, bool kUsePDL>
-struct CustomAllReducePull : public CustomAllReduceBase {
-  static constexpr uint32_t kVecSize = 16 / (sizeof(DType) * 2);
-  static constexpr auto one_shot_kernel = all_reduce_one_shot_kernel<DType, kNumGPU, kUsePDL>;
-  static constexpr auto two_shot_kernel = all_reduce_two_shot_kernel<DType, kNumGPU, kUsePDL>;
-  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<uint32_t>(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<DType>(input.dtype()), "Input dtype mismatch");
-    RuntimeCheck(std::bit_cast<intptr_t>(input_ptr) % 16 == 0, "Input pointer is not properly aligned");
-    RuntimeCheck(m_pull_ctrl.has_value(), "Controller is not initialized");
-    RuntimeCheck(static_cast<int64_t>(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<int64_t>(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 <typename DType, uint32_t kNumGPU, bool kUsePDL>
-tvm::ffi::Tensor custom_all_reduce(CustomAllReduceRef obj, tvm::ffi::Tensor input, int shot) {
-  using Impl = CustomAllReducePull<DType, kNumGPU, kUsePDL>;
-  return static_cast<Impl&>(*obj.get()).all_reduce(input, shot);
-}
-
-}  // namespace

python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_push.cuh

@@ -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 <sgl_kernel/ffi.h>
-#include <sgl_kernel/tensor.h>
-#include <sgl_kernel/utils.h>
-
-#include <sgl_kernel/type.cuh>
-#include <sgl_kernel/utils.cuh>
-#include <sgl_kernel/vec.cuh>
-
-#include <sgl_kernel/distributed/common.cuh>
-#include <sgl_kernel/distributed/custom_all_reduce.cuh>
-
-#include <cstdint>
-#include <cstring>
-
-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 <typename T>
-struct fp_trait {};
-
-// TODO: support more dtypes
-template <>
-struct fp_trait<bf16_t> {
-  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<fp16_t> {
-  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<float> {
-  using type = uint32_t;
-  [[maybe_unused]]
-  static constexpr uint32_t pos_zero = 0x00000000u;
-  [[maybe_unused]]
-  static constexpr uint32_t neg_zero = 0x80000000u;
-};
-
-template <typename DType>
-SGL_DEVICE void clear_pos_zero(DType& val) {
-  using Trait = fp_trait<DType>;
-  const auto ptr = reinterpret_cast<typename Trait::type*>(&val);
-  if (*ptr == Trait::pos_zero) *ptr = Trait::neg_zero;
-}
-
-template <typename DType>
-SGL_DEVICE bool is_pos_zero(const DType& val) {
-  using Trait = fp_trait<DType>;
-  const auto ptr = reinterpret_cast<const typename Trait::type*>(&val);
-  return *ptr == Trait::pos_zero;
-}
-
-template <typename DType>
-SGL_DEVICE DType get_pos_zero() {
-  using Trait = fp_trait<DType>;
-  const auto value = Trait::pos_zero;
-  return *reinterpret_cast<const DType*>(&value);
-}
-
-template <typename T>
-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<T>(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<const T*>(&val);
-}
-
-template <typename T>
-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<const uint4*>(&x);
-  addr = device::pointer::offset<T>(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 <typename DType, uint32_t kNumGPU>
-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<packed_t<DType>, 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 <typename DType, uint32_t kNumGPU>
-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<packed_t<DType>, 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<DType>(), get_pos_zero<DType>()});
-#pragma unroll
-    for (uint32_t i = 0; i < kNumGPU; ++i) {
-      pos_zeros.store(poll_buf[i], offset);
-    }
-  }
-}
-
-template <typename DType, uint32_t kNumGPU, bool kUsePDL>
-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<kUsePDL>();
-
-  // 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<DType*>(pointer::offset(buffer[i], rank * buffer_bytes, epoch_offset));
-  }
-  push_impl(push_buf, input, num_items);
-
-  PDLTriggerSecondary<kUsePDL>();
-
-  // Phase 2: Poll local data
-  DType* poll_buf[kNumGPU];
-#pragma unroll
-  for (uint32_t i = 0; i < kNumGPU; ++i) {
-    poll_buf[i] = static_cast<DType*>(pointer::offset(buffer[rank], i * buffer_bytes, epoch_offset));
-  }
-  poll_impl(poll_buf, output, num_items);
-  ctrl.exit();
-}
-
-template <typename DType, uint32_t kNumGPU, bool kUsePDL>
-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<uint32_t>(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<DType>(input.dtype()), "Input dtype mismatch");
-    RuntimeCheck(std::bit_cast<intptr_t>(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<int64_t>(num_items) == num_items_int64, "Number of items exceeds 4G limit");
-
-    const auto input_bytes = static_cast<int64_t>(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<DType, kNumGPU, kUsePDL>;
-    LaunchKernel(num_blocks, num_threads, device)  //
-        .enable_pdl(kUsePDL)(kernel, params, *m_push_ctrl);
-    return output;
-  }
-};
-
-template <typename DType, uint32_t kNumGPU, bool kUsePDL>
-tvm::ffi::Tensor custom_all_reduce(CustomAllReduceRef obj, tvm::ffi::Tensor input, int shot) {
-  using Impl = CustomAllReducePush<DType, kNumGPU, kUsePDL>;
-  return static_cast<Impl&>(*obj.get()).all_reduce(input, shot);
-}
-
-}  // namespace

python/sglang/jit_kernel/csrc/distributed/ipc.cuh

@@ -0,0 +1,194 @@
+#include <sgl_kernel/ffi.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/utils.cuh>
+
+#include <tvm/ffi/container/array.h>
+#include <tvm/ffi/container/tuple.h>
+#include <tvm/ffi/reflection/registry.h>
+
+#include <array>
+#include <cstddef>
+#include <cstdint>
+#include <cstring>
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <functional>
+#include <map>
+#include <string_view>
+#include <unordered_map>
+#include <unordered_set>
+
+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<uintptr_t>(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<char, sizeof(cudaIpcMemHandle_t)>;
+  using FFIHandle = tvm::ffi::Array<char>;
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("sgl.IPCManager", IPCManager, tvm::ffi::Object);
+  static constexpr bool _type_mutable = true;
+
+  using BatchGetResult = tvm::ffi::Array<tvm::ffi::Tuple<FFIHandle, size_t>>;
+  using BatchGetInputs = tvm::ffi::Array<uintptr_t>;
+
+  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<void*>(base_addr)));
+    }
+    m_handle2ptr_cache.clear();
+    m_ptr2handle_cache.clear();
+    m_local_handles.clear();
+  }
+
+  BatchGetInputs batch_open_handles(BatchGetResult handles) {
+    tvm::ffi::Array<uintptr_t> 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<FFIHandle, size_t>;
+    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<char>(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<uint8_t>(ipc_handle[i]));
+    }
+    return ffi_handle;
+  }
+
+  std::pair<IPCHandle, size_t> 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<void*>(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<uintptr_t>(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<std::string_view>{}(sv);
+    }
+  };
+
+  std::unordered_map<IPCHandle, uintptr_t, HashCUDAIPC, EqualCUDAIPC> m_handle2ptr_cache;
+  std::unordered_set<IPCHandle, HashCUDAIPC, EqualCUDAIPC> m_local_handles;
+  std::map<uintptr_t, std::pair<IPCHandle, size_t>> m_ptr2handle_cache;
+};
+
+}  // namespace host::distributed
+
+inline void register_ipc_manager() {
+  namespace refl = tvm::ffi::reflection;
+  using Class = host::distributed::IPCManager;
+  refl::ObjectDef<Class>()
+      .def(refl::init<>(), "__init__")
+      .def("batch_get_handles", &Class::batch_get_handles)
+      .def("batch_open_handles", &Class::batch_open_handles)
+      .def("destroy", &Class::destroy);
+}

python/sglang/jit_kernel/csrc/distributed/tp_qknorm.cuh

@@ -11,19 +11,18 @@
 #include <sgl_kernel/vec.cuh>
 #include <sgl_kernel/warp.cuh>
 
-#include <sgl_kernel/distributed/common.cuh>
+#include <sgl_kernel/distributed/communicator.cuh>
-#include <sgl_kernel/distributed/custom_all_reduce.cuh>
 
 #include <cstdint>
 #include <cstring>
 
 namespace {
 
-using device::distributed::PushController;
+using device::distributed::Counter;
-using host::distributed::CustomAllReduceBase, host::distributed::CustomAllReduceRef;
+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 <typename Trait>
 __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 <typename DType, uint32_t kNumGPU, int64_t kQDim, int64_t kKDim, bool kUsePDL>
-struct FusedParallelQKNormAcrossHead : public CustomAllReduceBase {
+struct FusedParallelQKNormAcrossHead {
   using Trait = KernelTrait<DType, kNumGPU, kQDim, kKDim, kUsePDL>;
   static constexpr auto kernel = parallel_qknorm_across_head<Trait>;
-  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 max_num_blocks = comm.num_push_blocks;
-    const auto num_clean = need_clean ? (m_max_num_cta_push - num_tokens) : 0;
+    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<uint32_t>(comm.push_bytes);
     params.num_clean_up_count = num_clean;
 
     const auto needed_buffer_bytes = static_cast<int64_t>(num_tokens) * 2 * sizeof(float);
-    RuntimeCheck(m_num_gpu == kNumGPU, "Number of GPUs mismatch");
+    RuntimeCheck(comm.world_size == kNumGPU, "Number of GPUs mismatch");
-    RuntimeCheck(m_push_ctrl.has_value(), "Controller is not initialized");
     RuntimeCheck(std::bit_cast<intptr_t>(params.q_ptr) % 16 == 0, "q pointer is not properly aligned");
     RuntimeCheck(std::bit_cast<intptr_t>(params.k_ptr) % 16 == 0, "k pointer is not properly aligned");
     RuntimeCheck(std::bit_cast<intptr_t>(params.q_weight) % 16 == 0, "q_weight pointer is not properly aligned");
     RuntimeCheck(std::bit_cast<intptr_t>(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 _run(*comm.get(), q, k, q_weight, k_weight, eps);
-    return static_cast<Self*>(obj.get())->_run(q, k, q_weight, k_weight, eps);
   }
 };
 

python/sglang/jit_kernel/include/sgl_kernel/distributed/common.cuh

@@ -1,120 +0,0 @@
-#pragma once
-#include <sgl_kernel/utils.cuh>
-
-namespace device::distributed {
-
-inline constexpr uint32_t kMaxNumGPU = 8;
-
-struct alignas(128) Semaphore {
- public:
-  constexpr Semaphore() : m_flag(0), m_counter(0) {}
-
-  template <bool kFence>
-  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 <bool kFence>
-  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<Semaphore*>(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 <bool kFence, bool kStart>
-  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<kFence>(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<kFence>() - 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<SignalType*>(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

python/sglang/jit_kernel/include/sgl_kernel/distributed/communicator.cuh

@@ -0,0 +1,120 @@
+#pragma once
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/utils.cuh>
+
+#include <tvm/ffi/container/tensor.h>
+#include <tvm/ffi/object.h>
+
+#include <array>
+#include <cstdint>
+#include <map>
+#include <optional>
+#include <string>
+#include <vector>
+
+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<TensorView> push_workspaces,
+      std::vector<TensorView> pull_workspaces,
+      std::vector<TensorView> pull_semaphores,
+      TensorView push_counter,
+      std::optional<int64_t> pull_mc_workspace_ptr);
+
+  void config(std::map<std::string, uint32_t> 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<uint8_t*, kMaxWorldSize> pull_workspaces;    // symmetric memory
+  std::array<uint8_t*, kMaxWorldSize> push_workspaces;    // symmetric memory
+  std::array<Semaphore*, kMaxWorldSize> 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

python/sglang/jit_kernel/include/sgl_kernel/distributed/custom_all_reduce.cuh

@@ -1,446 +0,0 @@
-#pragma once
-#include <sgl_kernel/utils.h>
-
-#include <sgl_kernel/utils.cuh>
-#include <sgl_kernel/vec.cuh>
-
-#include <sgl_kernel/distributed/common.cuh>
-
-#include <tvm/ffi/container/array.h>
-#include <tvm/ffi/container/tuple.h>
-#include <tvm/ffi/reflection/registry.h>
-
-#include <algorithm>
-#include <array>
-#include <cstdint>
-#include <cstring>
-#include <functional>
-#include <numeric>
-#include <optional>
-#include <span>
-#include <unordered_map>
-#include <vector>
-
-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<char>;
-
-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<std::string_view>{}({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<int64_t, ExternHandle>;  // (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<int64_t>(std::numeric_limits<uint32_t>::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<InputPair> share_graph_inputs() {
-    tvm::ffi::Array<InputPair> 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<void*, ExternHandle> 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<char*>(ptr) - reinterpret_cast<char*>(base_ptr);
-      result.push_back(InputPair{offset, get_handle(base_ptr)});
-    }
-    return result;
-  }
-
-  void post_init(tvm::ffi::Array<ExternHandle> 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<tvm::ffi::Array<InputPair>> 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<AllReduceData> 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<int64_t> get_graph_capture_ptrs() {
-    tvm::ffi::Array<int64_t> 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<int64_t>(ptr));
-    }
-    return result;
-  }
-
-  using BaseInfo = tvm::ffi::Tuple<int64_t, int64_t>;  // (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<BaseInfo>, tvm::ffi::Array<tvm::ffi::Array<int64_t>>, tvm::ffi::Array<int64_t>>
-  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<uintptr_t, int64_t> base_to_idx;
-    tvm::ffi::Array<BaseInfo> bases;
-    tvm::ffi::Array<tvm::ffi::Array<int64_t>> input_indices;
-    tvm::ffi::Array<int64_t> 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<CUdeviceptr>(ptr);
-      CUdeviceptr first_base = 0;
-      tvm::ffi::Array<int64_t> 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<int64_t>(cursor - base);
-        RuntimeCheck(
-            byte_offset >= 0 && static_cast<size_t>(byte_offset) < size,
-            "graph capture input at ",
-            reinterpret_cast<uintptr_t>(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<int64_t>(base), static_cast<int64_t>(size)});
-        }
-        chunks.push_back(it->second);
-
-        const auto available = static_cast<int64_t>(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<CUdeviceptr>(ptr) - first_base);
-    }
-    using Result =
-        tvm::ffi::Tuple<tvm::ffi::Array<BaseInfo>, tvm::ffi::Array<tvm::ffi::Array<int64_t>>, tvm::ffi::Array<int64_t>>;
-    return Result(bases, input_indices, offsets);
-  }
-
-  void register_peer_mapped_inputs(tvm::ffi::Array<tvm::ffi::Array<int64_t>> 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<AllReduceData> 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<void*>(static_cast<int64_t>(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<uint32_t, uint32_t> 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<uint32_t, uint32_t>{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<AllReduceData*>(start) + (1 + which);
-  }
-  int64_t registered_count() const {
-    return static_cast<int64_t>(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<PullController> m_pull_ctrl;
-  std::optional<PushController> m_push_ctrl;
-  void* m_storage = nullptr;
-  std::vector<void*> m_graph_capture_inputs;
-  std::vector<int64_t> m_graph_capture_input_bytes;
-  std::vector<void*> m_peer_storage;
-  std::unordered_map<cudaIpcMemHandle_t, void*, HandleHash, HandleEqual> 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 <typename DType2, size_t N, uint32_t M>
-SGL_DEVICE auto reduce_impl(AlignedVector<DType2, N> (&storage)[M]) -> AlignedVector<DType2, N> {
-  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<fp32x2_t>(storage[i][j]);
-      auto& [x_acc, y_acc] = acc[j];
-      x_acc += x;
-      y_acc += y;
-    }
-  }
-
-  AlignedVector<DType2, N> result;
-#pragma unroll
-  for (uint32_t j = 0; j < N; ++j) {
-    result[j] = cast<DType2>(acc[j]);
-  }
-
-  return result;
-}
-
-}  // namespace device::distributed

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

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}));"

python/sglang/srt/distributed/device_communicators/configs/__init__.py

@@ -0,0 +1 @@
+"""Tuned dispatch configs for device communicators."""

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,
+    )

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 List, Optional, Tuple
-from typing import Dict, List, Optional, TypeVar
 
 import torch
 import torch.distributed as dist
 from torch.distributed import ProcessGroup
 
-from sglang.jit_kernel.all_reduce import AllReduceAlgo, get_custom_all_reduce_cls
+from sglang.jit_kernel.all_reduce import (
-from sglang.srt.distributed.device_communicators.custom_all_reduce_utils import (
+    AllReduceAlgo,
-    can_use_custom_all_reduce_with_nvlink,
+    Communicator,
-    is_weak_contiguous,
+    IPCManager,
-)
+    custom_all_reduce,
-from sglang.srt.distributed.device_communicators.vmm_utils import (
-    VmmGraphInputManager,
-    is_vmm_pointer,
 )
 from sglang.srt.environ import envs
 from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
     is_in_tc_piecewise_cuda_graph,
 )
-from sglang.srt.utils import is_sm100_supported
+
+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 _PullMode(enum.Enum):
-class ModeConfig:
+    EAGER = enum.auto()  # pull_arg = False (also used for 1shot_push)
-    one_shot_push_threshold: int  # below this, use one-shot push
+    MULTICAST = enum.auto()  # pull_arg = True
-    one_shot_pull_threshold: int  # below this, use one-shot pull
+    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
+        base_config = get_all_reduce_config(self.world_size)
-            max_pull_size = 16 * 1024 * 1024
+        if max_pull_size is None:
-        if max_push_size is None:  # default to recommended size
+            max_pull_size = min(base_config.max_pull_bytes, max_size)
-            config = THRESHOLD_2_SHOT_MAP[self.world_size]
+        if max_push_size is None:
-            max_push_size = config.one_shot_push_threshold
+            max_push_size = min(base_config.max_push_bytes, max_size)
-        self.max_pull_size = max_pull_size
+        # a minimal workspace keeps the Communicator valid even when a caller
-        self.max_push_size = max_push_size
+        # only uses one direction (e.g. push-only fused qk-norm instances)
-        self.max_size = max(max_pull_size, max_push_size)
+        self.max_pull_size = _ceil_align(max(max_pull_size, _ALIGN_BYTES), _ALIGN_BYTES)
-        self.override_shot(None)  # set default config based on world size
+        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,
+        # device-side pointer table: one row of world_size pointers per
-            world_size=self.world_size,
+        # graph-captured all-reduce input (at most 8 MB at world_size = 8)
-            pull_buffer_bytes=self.max_pull_size,
+        self.graph_params = torch.zeros(
-            push_buffer_bytes=self.max_push_size,
+            (_MAX_GRAPH_INPUTS, self.world_size),
-            graph_input_count=131072,
+            dtype=torch.uint64,
-            max_pull_blocks=max_pull_blocks,
+            device=self.device,
-            max_push_blocks=max_push_blocks,
         )
+        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):
+    def _init_workspace(self) -> None:
-        if shot is None:
+        """Slice one symmetric-memory allocation into all shared buffers.
-            config = THRESHOLD_2_SHOT_MAP[self.world_size]
+
-        else:
+        Layout per rank: ``[2 * world_size push buffers | pull buffer |
-            assert shot in (1, 2)
+        pull semaphores]``. The push counter is rank-local, so it lives in a
-            threshold = INF if shot == 1 else 0
+        plain CUDA tensor instead.
-            config = replace(self.config, one_shot_pull_threshold=threshold)
+        """
-        # need to clip the config thresholds to max sizes to avoid invalid config
+        cfg = self.config
-        push_threshold = min(config.one_shot_push_threshold, self.max_push_size)
+        push_num_bufs = 2 * self.world_size  # 2 phases x world_size peers
-        pull_threshold = min(config.one_shot_pull_threshold, self.max_pull_size)
+        push_ws_bytes = push_num_bufs * self.max_push_size
-        self.config: ModeConfig = replace(
+        pull_ws_bytes = self.max_pull_size
-            config,
+        pull_sem_bytes = _SEMAPHORE_BYTES * cfg.num_pull_blocks
-            one_shot_push_threshold=push_threshold,
+        total_bytes = push_ws_bytes + pull_ws_bytes + pull_sem_bytes
-            one_shot_pull_threshold=pull_threshold,
+        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
+        multicast_ptr = int(symm_mem.multicast_ptr)
-    def capture(self):
+        can_multicast = multicast_ptr != 0
-        if self.disabled:
+        pull_mc_workspace = multicast_ptr + pull_ws_offset if can_multicast else None
-            yield
+        if not can_multicast or cfg.num_mc_blocks is None:
-            return
+            self.config = self.config._replace(num_mc_blocks=None)
-        try:
-            self.obj.set_cuda_graph_capture(not self.tms_cudagraph)
-            yield
-        finally:
-            self.obj.set_cuda_graph_capture(False)
-        assert (
-            not torch.cuda.is_current_stream_capturing()
-        ), "Cannot register graph inputs while capturing CUDA graph"
-        raw_ptrs = self.obj.get_graph_capture_ptrs()
-        if raw_ptrs and is_vmm_pointer(raw_ptrs[0]):
-            self._vmm_graph_input_manager.register_graph_inputs()
-        else:
-            self._register_graph_inputs_ipc()
 
-    def _register_graph_inputs_ipc(self):
+        self.obj = Communicator(
-        """Register graph capture inputs via cudaIpcGetMemHandle.
+            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]
+        def uncap(heuristic):
-        offsets = [offset for offset, _ in pairs]
+            return heuristic._replace(two_shot_pull_threshold=self.max_pull_size)
-        handles_all = self._share_list(handles)
+
-        offsets_all = self._share_list(offsets)
+        self.config = self.config._replace(
-        result = [list(zip(o, h)) for o, h in zip(offsets_all, handles_all)]
+            graph=uncap(self.config.graph),
-        self.obj.register_inputs(result)
+            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
+        nbytes = input.numel() * input.element_size()
-            try:
+        can_use_graph = self._can_use_graph()
-                self.obj.set_cuda_graph_capture(False)
+        if self.override_algo is not None:
-                return self._all_reduce(input)
+            algo = self.override_algo
-            finally:
+            use_graph = can_use_graph and not algo.is_push()
-                self.obj.set_cuda_graph_capture(not self.tms_cudagraph)
+            mode = _PullMode.GRAPH if use_graph else _PullMode.EAGER
-        return self._all_reduce(input)
+        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] = {}

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

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)

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 = DTYPE.itemsize
-DTYPE_ITEMSIZE = torch.tensor([], dtype=DTYPE).element_size()
+MESSAGE_SIZES_KB = [2**x for x in range(2, 17)]
-MESSAGE_SIZES_BYTES = [
+MESSAGE_SIZES_KB += [192, 384, 640, 768, 896, 1536, 3072]
-    4 * 1024,  # 4K
+MESSAGE_SIZES_KB.sort()
-    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
-]
 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
+    local_rank = int(os.environ["LOCAL_RANK"])
-    assert coord is not None and coord.device_group is not None
+    device_group = torch.distributed.new_group(
-    return coord.device_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(
+        # tuned workspace sizes, capped at the sweep maximum
-            _init_cpu_group(), device, max_pull_size=MAX_BYTES
+        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.

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
+    # baseline path: sum-sq -> all-reduce -> apply (also covers push mode)
-    _jit_custom_all_reduce_pull_module(DTYPE, world_size)
+    get_all_reduce_module(DTYPE, world_size)
-    # fused path: push-mode all-reduce
-    _jit_custom_all_reduce_push_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:

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 (
+from sglang.jit_kernel.all_reduce import AllReduceAlgo, get_all_reduce_module
-    AllReduceAlgo,
-    _jit_custom_all_reduce_pull_module,
-    _jit_custom_all_reduce_push_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)
+    get_all_reduce_module(dtype, world_size)
-    _jit_custom_all_reduce_push_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
+    local_rank = int(os.environ["LOCAL_RANK"])
-    assert coord is not None and coord.device_group is not None
+    device_group = dist.new_group(
-    return coord.device_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)
 
 

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
+    local_rank = int(os.environ["LOCAL_RANK"])
-    assert coord is not None and coord.device_group is not None
+    device_group = dist.new_group(
-    return coord.device_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)