[comm] Enable multi-node custom-AR v2 on a single NVLink clique (#32339)

Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Lianmin Zheng <lianminzheng@gmail.com>

python/sglang/kernels/jit/csrc/distributed/custom_all_reduce.cuh

@@ -204,18 +204,19 @@ SGL_DEVICE void st_multimem_16B(const V& x, void* mc_addr, int64_t vec_offset) {
 #endif
 }
 
+template <uint32_t kWorldSize>
 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* pull_workspaces[kWorldSize];    // must be symmetric memory
-  uint8_t* push_workspaces[kMaxWorldSize];    // must be symmetric memory
+  uint8_t* push_workspaces[kWorldSize];    // must be symmetric memory
-  Semaphore* pull_semaphores[kMaxWorldSize];  // must be symmetric memory
+  Semaphore* pull_semaphores[kWorldSize];  // 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
+  int64_t push_buffer_stride;  // per-buffer bytes; each rank holds 2 * world_size buffers
 };
 
 template <typename T, uint32_t kWorldSize, bool kUsePDL>
@@ -228,12 +229,12 @@ struct AllReducePushImpl {
   using vec_t = device::AlignedVector<T2, kVecSize>;
   static_assert(kWorldSize <= kMaxWorldSize);
 
-  static SGL_DEVICE bool sync_enter_push(const AllReduceParams& params) {
+  static SGL_DEVICE bool sync_enter_push(const AllReduceParams<kWorldSize>& params) {
     device::PDLWaitPrimary<kUsePDL>();
     return (params.push_counter[blockIdx.x].get() % 2) != 0;
   }
 
-  static SGL_DEVICE void sync_exit_push(const AllReduceParams& params) {
+  static SGL_DEVICE void sync_exit_push(const AllReduceParams<kWorldSize>& params) {
     device::PDLTriggerSecondary<kUsePDL>();
     __syncthreads();
     if (threadIdx.x == 0) {
@@ -305,7 +306,7 @@ struct AllReducePushImpl {
   }
 
  public:
-  static SGL_DEVICE void forward_1shot(const AllReduceParams& params) {
+  static SGL_DEVICE void forward_1shot(const AllReduceParams<kWorldSize>& params) {
     // push local data to peer ranks, then reduce locally
     const auto phase = sync_enter_push(params);
     const auto r = params.rank;
@@ -343,7 +344,7 @@ struct AllReducePullImpl {
   static_assert(kWorldSize <= kMaxWorldSize);
 
   template <bool kFence>
-  static SGL_DEVICE uint32_t sync_enter_pull(const AllReduceParams& params) {
+  static SGL_DEVICE uint32_t sync_enter_pull(const AllReduceParams<kWorldSize>& params) {
     uint32_t current_counter_val = 0;
     if (const auto tx = threadIdx.x; tx < kWorldSize) {
       device::PDLWaitPrimary<kUsePDL>();
@@ -372,7 +373,7 @@ struct AllReducePullImpl {
   }
 
   template <bool kFence>
-  static SGL_DEVICE void sync_exit_pull(const AllReduceParams& params, uint32_t current) {
+  static SGL_DEVICE void sync_exit_pull(const AllReduceParams<kWorldSize>& params, uint32_t current) {
     device::PDLTriggerSecondary<kUsePDL>();
     __syncthreads();
     if (const auto tx = threadIdx.x; tx < kWorldSize) {
@@ -436,7 +437,7 @@ struct AllReducePullImpl {
   }
 
  public:
-  static SGL_DEVICE void forward_1shot(const AllReduceParams& params) {
+  static SGL_DEVICE void forward_1shot(const AllReduceParams<kWorldSize>& params) {
     const auto total_num_vecs = device::div_ceil(params.num_elements, kElemsPerVec);
     void* data[kWorldSize];
     if constexpr (kMode == PullMode::Graph) {
@@ -455,7 +456,7 @@ struct AllReducePullImpl {
     sync_exit_pull<false>(params, counter);
   }
 
-  static SGL_DEVICE void forward_2shot(const AllReduceParams& params) {
+  static SGL_DEVICE void forward_2shot(const AllReduceParams<kWorldSize>& 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;
@@ -481,9 +482,9 @@ struct AllReducePullImpl {
   }
 };
 
-template <typename Impl, int kShot>
+template <typename Impl, uint32_t kWorldSize, int kShot>
 __global__ __launch_bounds__(1024, 1)  //
-    void all_reduce_kernel(const __grid_constant__ AllReduceParams params) {
+    void all_reduce_kernel(const __grid_constant__ AllReduceParams<kWorldSize> params) {
   static_assert(kShot == 1 || kShot == 2, "invalid shot");
   if constexpr (kShot == 1) {
     return Impl::forward_1shot(params);
@@ -528,9 +529,10 @@ struct AllReduceKernel {
   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>;
+  static constexpr auto kernel_pull =
+      all_reduce_kernel<AllReducePullImpl<T, kWorldSize, kPullMode, kUsePDL>, kWorldSize, kShot>;
   template <int kShot>
-  static constexpr auto kernel_push = all_reduce_kernel<AllReducePushImpl<T, kWorldSize, kUsePDL>, kShot>;
+  static constexpr auto kernel_push = all_reduce_kernel<AllReducePushImpl<T, kWorldSize, kUsePDL>, kWorldSize, kShot>;
 
  public:
   static Tensor run(CommunicatorRef ref, Tensor in_, std::string algo, std::variant<TensorView, bool> pull_arg) {
@@ -549,7 +551,7 @@ struct AllReduceKernel {
     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{
+    AllReduceParams<kWorldSize> params{
         .input = in_.data_ptr(),
         .output = out.data_ptr(),
         .num_elements = num_elems,

python/sglang/kernels/jit/include/sgl_kernel/distributed/communicator.cuh

@@ -15,7 +15,7 @@
 
 namespace device::distributed {
 
-inline constexpr uint32_t kMaxWorldSize = 8;
+inline constexpr uint32_t kMaxWorldSize = 16;
 
 struct Counter {
  public:

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

@@ -113,7 +113,7 @@ def _pack_heuristic(*args) -> Heuristic:
 
 
 def _sm100_config(world_size: int, num_sm: int) -> AllReduceConfig:
-    # SM100 (Blackwell, B200/B300). Tuned on B200 (148 SMs).
+    # SM100 (Blackwell, B200/B300/GB200). Tuned on B200 (148 SMs); world 16 on GB200.
     graph_map = {
         2: (8.000 * MB, 32.00 * MB, 128.0 * MB),
         3: (4.000 * MB, 4.000 * MB, 128.0 * MB),
@@ -122,6 +122,7 @@ def _sm100_config(world_size: int, num_sm: int) -> AllReduceConfig:
         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)),
+        16: (0.250 * MB, 0.250 * MB, 128.0 * MB, Range(256 * KB, 128 * MB)),
     }
     eager_map = {
         2: (16.00 * MB, 128.0 * MB, 128.0 * MB),
@@ -131,8 +132,9 @@ def _sm100_config(world_size: int, num_sm: int) -> AllReduceConfig:
         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)),
+        16: (0.250 * MB, 0.250 * MB, 128.0 * MB, Range(256 * KB, 128 * MB)),
     }
-    mc_blocks_map = {5: 64, 6: 48, 7: 48, 8: 32}
+    mc_blocks_map = {5: 64, 6: 48, 7: 48, 8: 32, 16: 32}
     return AllReduceConfig(
         graph=_pack_heuristic(*graph_map[world_size]),
         eager=_pack_heuristic(*eager_map[world_size]),

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

@@ -352,8 +352,8 @@ def dispatch_custom_allreduce(
 
     On CUDA, the JIT-compiled v2 implementation is used by default.
     Set SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2=0 to fall back to the legacy CustomAllreduce.
-    Note: ServerArgs._handle_environment_variables forces this env to "0" when
+    Multi-node v2 is admitted only for a single NVLink clique (see
-    nnodes > 1 since custom AR is intra-node only.
+    can_use_custom_all_reduce_v2); other cross-node groups fall back to NCCL.
     """
     if _is_cuda and envs.SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2.get():
         from .custom_all_reduce_v2 import (

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

@@ -391,6 +391,53 @@ def is_full_nvlink(physical_device_ids: List[int], world_size: int) -> bool:
         return True
 
 
+# NVML_GPU_FABRIC_STATE_COMPLETED: the GPU has joined its NVLink fabric clique.
+_NVML_GPU_FABRIC_STATE_COMPLETED = 3
+
+
+def _gpu_fabric_clique(device: torch.device):
+    """(cluster_uuid, clique_id) of the local GPU's NVLink fabric clique, or None if
+    the GPU has not joined a fabric (single-node box / fabric init incomplete)."""
+    cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES", None)
+    if cuda_visible_devices:
+        device_ids = list(map(int, cuda_visible_devices.split(",")))
+    else:
+        device_ids = list(range(torch.cuda.device_count()))
+    handle = pynvml.nvmlDeviceGetHandleByIndex(device_ids[device.index])
+    fabric = pynvml.c_nvmlGpuFabricInfo_v3_t()
+    fabric.version = pynvml.nvmlGpuFabricInfo_v3
+    pynvml.nvmlDeviceGetGpuFabricInfoV(handle, ctypes.byref(fabric))
+    if fabric.state != _NVML_GPU_FABRIC_STATE_COMPLETED:
+        return None
+    return (bytes(fabric.clusterUuid), int(fabric.cliqueId))
+
+
+@with_nvml_context
+def is_one_nvlink_clique(
+    group: torch.distributed.ProcessGroup, device: torch.device
+) -> bool:
+    """True iff every rank's GPU is in the same NVLink fabric clique (one NVL72 /
+    MNNVL domain). Such a clique shares a single NVLink address space even across
+    nodes, so custom-AR v2's symm-mem storage + fabric peer VAs are valid group-wide."""
+    if _is_hip:
+        return False
+    try:
+        clique = _gpu_fabric_clique(device)
+    except Exception as e:
+        logger.warning(
+            "GPU fabric clique query failed (%r); custom-AR stays intra-node.", e
+        )
+        clique = None
+    # Always all-gather (every rank calls it once) so a failed query on any rank
+    # resolves to a clean False rather than a collective mismatch.
+    world_size = dist.get_world_size(group=group)
+    gathered: List[object] = [None] * world_size
+    dist.all_gather_object(gathered, clique, group=group)
+    if any(c is None for c in gathered):
+        return False
+    return len(set(gathered)) == 1
+
+
 def is_weak_contiguous(inp: torch.Tensor):
     return inp.is_contiguous() or (
         inp.storage().nbytes() - inp.storage_offset() * inp.element_size()

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

@@ -31,6 +31,7 @@ from sglang.kernels.ops.communication.all_reduce import (
     IPCManager,
     custom_all_reduce,
 )
+from sglang.srt.distributed.parallel_state import in_the_same_node_as
 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,
@@ -39,6 +40,7 @@ from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph impo
 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_one_nvlink_clique,
     is_weak_contiguous,
 )
 from .vmm_utils import (
@@ -139,7 +141,7 @@ class CustomAllReduceV2:
         self.tms_cudagraph = envs.SGLANG_MEMORY_SAVER_CUDA_GRAPH.get()
 
         # device-side pointer table: one row of world_size pointers per
-        # graph-captured all-reduce input (at most 8 MB at world_size = 8)
+        # graph-captured all-reduce input
         self.graph_params = torch.zeros(
             (_MAX_GRAPH_INPUTS, self.world_size),
             dtype=torch.uint64,
@@ -407,14 +409,29 @@ class CustomAllReduceV2:
         self.close()
 
 
+def _is_vmm_backed_allocator(device: torch.device) -> bool:
+    """True iff the caching allocator is VMM-backed (expandable_segments). Uniform
+    launch, so the local probe reflects every rank."""
+    probe = torch.empty(1, dtype=torch.uint8, device=device)
+    return is_vmm_pointer(probe.data_ptr())
+
+
 def can_use_custom_all_reduce_v2(
     group: ProcessGroup,
     device: torch.device,
 ) -> bool:
+    supported = list(range(2, 17))
+    if dist.get_world_size(group=group) not in supported:
+        return False
+    # Multi-node needs a single NVLink clique (one NVL72 / MNNVL domain) whose
+    # allocator is VMM-backed: graph inputs cross nodes via FABRIC / POSIX-fd VMM
+    # handles, not cudaIpc (intra-node only). Else use the intra-node nvlink check.
+    if not all(in_the_same_node_as(group, source_rank=0)):
+        return is_one_nvlink_clique(group, device) and _is_vmm_backed_allocator(device)
     full_nvlink = can_use_custom_all_reduce_with_nvlink(
         group=group,
         device=device,
-        supported_world_size=list(range(2, 9)),
+        supported_world_size=supported,
         cls_name="CustomAllReduceV2",
     )
     return full_nvlink is True

python/sglang/srt/server_args.py

@@ -7317,16 +7317,6 @@ class ServerArgs:
         envs.SGLANG_ENABLE_DETERMINISTIC_INFERENCE.set(
             "1" if self.enable_deterministic_inference else "0"
         )
-        # Custom all-reduce v2 uses IPC handles and is intra-node only. Force-disable
-        # on multi-node so the dispatch falls back to the legacy CustomAllreduce path.
-        if self.nnodes > 1 and envs.SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2.get():
-            if envs.SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2.is_set():
-                logger.warning(
-                    "Disabling SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2 because nnodes=%d "
-                    "(custom all-reduce v2 is intra-node only).",
-                    self.nnodes,
-                )
-            envs.SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2.set("0")
         if self.debug_cuda_graph:
             if not (is_cuda() or is_hip()):
                 logger.warning(

test/registered/kernels/benchmark/communication/bench_custom_all_reduce.py

@@ -34,7 +34,7 @@ DTYPE_ITEMSIZE = DTYPE.itemsize
 MESSAGE_SIZES_KB = [2**x for x in range(2, 17)]
 MESSAGE_SIZES_KB += [192, 384, 640, 768, 896, 1536, 3072]
 MESSAGE_SIZES_KB.sort()
-WORLD_SIZES = list(range(2, 9))
+WORLD_SIZES = list(range(2, 9)) + [16]
 MAX_BYTES = max(MESSAGE_SIZES_KB) * 1024
 # trtllm allreduce_fusion only supports these world sizes.
 FI_SUPPORTED_WORLD_SIZES = (2, 4, 8)

test/registered/kernels/ops/communication/test_custom_all_reduce.py

@@ -235,12 +235,13 @@ def test_custom_all_reduce(
 
 
 if __name__ == "__main__":
-    # Only sweep the common world sizes (2, 4, 8) by default: testing every
+    # Only sweep the common world sizes (2, 4, 8, 16) by default: testing every
-    # count in 2..8 serially overruns the per-file CI time budget, and 3/5/6/7
+    # count in 2..16 serially overruns the per-file CI time budget, and numbers
-    # are rare in practice. Use --num-gpu to exercise them explicitly.
+    # in the middle are rare in practice. Use --num-gpu to exercise them
+    # explicitly.
     multigpu_pytest_main(
         __name__,
         __file__,
-        num_gpus=(2, 4, 8),
+        num_gpus=(2, 4, 8, 16),
         pre_launch_fn=_precompile_kernels,
     )