diff --git a/python/sglang/kernels/jit/csrc/distributed/custom_all_reduce.cuh b/python/sglang/kernels/jit/csrc/distributed/custom_all_reduce.cuh index 8412eea81..ff49c3e3e 100644 --- a/python/sglang/kernels/jit/csrc/distributed/custom_all_reduce.cuh +++ b/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 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 + uint8_t* pull_workspaces[kWorldSize]; // must be symmetric memory + uint8_t* push_workspaces[kWorldSize]; // 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 @@ -228,12 +229,12 @@ struct AllReducePushImpl { using vec_t = device::AlignedVector; static_assert(kWorldSize <= kMaxWorldSize); - static SGL_DEVICE bool sync_enter_push(const AllReduceParams& params) { + static SGL_DEVICE bool sync_enter_push(const AllReduceParams& params) { device::PDLWaitPrimary(); return (params.push_counter[blockIdx.x].get() % 2) != 0; } - static SGL_DEVICE void sync_exit_push(const AllReduceParams& params) { + static SGL_DEVICE void sync_exit_push(const AllReduceParams& params) { device::PDLTriggerSecondary(); __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& 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 - static SGL_DEVICE uint32_t sync_enter_pull(const AllReduceParams& params) { + 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(); @@ -372,7 +373,7 @@ struct AllReducePullImpl { } template - static SGL_DEVICE void sync_exit_pull(const AllReduceParams& params, uint32_t current) { + static SGL_DEVICE void sync_exit_pull(const AllReduceParams& params, uint32_t current) { device::PDLTriggerSecondary(); __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& 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(params, counter); } - static SGL_DEVICE void forward_2shot(const AllReduceParams& params) { + 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; @@ -481,9 +482,9 @@ struct AllReducePullImpl { } }; -template +template __global__ __launch_bounds__(1024, 1) // - void all_reduce_kernel(const __grid_constant__ AllReduceParams params) { + 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); @@ -528,9 +529,10 @@ struct AllReduceKernel { using Tensor = tvm::ffi::Tensor; using TensorView = tvm::ffi::TensorView; template - static constexpr auto kernel_pull = all_reduce_kernel, kShot>; + static constexpr auto kernel_pull = + all_reduce_kernel, kWorldSize, kShot>; template - static constexpr auto kernel_push = all_reduce_kernel, kShot>; + static constexpr auto kernel_push = all_reduce_kernel, kWorldSize, kShot>; public: static Tensor run(CommunicatorRef ref, Tensor in_, std::string algo, std::variant pull_arg) { @@ -549,7 +551,7 @@ struct AllReduceKernel { const auto graph_ptr = use_graph ? std::get(pull_arg).data_ptr() : nullptr; const bool inplace = use_graph && algo == "2shot_pull"; Tensor out = inplace ? in_ : ffi::empty_like(in_); - AllReduceParams params{ + AllReduceParams params{ .input = in_.data_ptr(), .output = out.data_ptr(), .num_elements = num_elems, diff --git a/python/sglang/kernels/jit/include/sgl_kernel/distributed/communicator.cuh b/python/sglang/kernels/jit/include/sgl_kernel/distributed/communicator.cuh index f7a2f9332..99475c623 100644 --- a/python/sglang/kernels/jit/include/sgl_kernel/distributed/communicator.cuh +++ b/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: diff --git a/python/sglang/srt/distributed/device_communicators/configs/custom_all_reduce_v2.py b/python/sglang/srt/distributed/device_communicators/configs/custom_all_reduce_v2.py index 8751df57e..3d0d2556f 100644 --- a/python/sglang/srt/distributed/device_communicators/configs/custom_all_reduce_v2.py +++ b/python/sglang/srt/distributed/device_communicators/configs/custom_all_reduce_v2.py @@ -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]), diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py index e98584634..55c292cde 100644 --- a/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py +++ b/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 - nnodes > 1 since custom AR is intra-node only. + Multi-node v2 is admitted only for a single NVLink clique (see + 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 ( diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_utils.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_utils.py index 0e4bc9d13..3bfd85ae1 100644 --- a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_utils.py +++ b/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() diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py index 047926280..0875ee20e 100644 --- a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py +++ b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py @@ -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 diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 0887fc71a..474984554 100644 --- a/python/sglang/srt/server_args.py +++ b/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( diff --git a/test/registered/kernels/benchmark/communication/bench_custom_all_reduce.py b/test/registered/kernels/benchmark/communication/bench_custom_all_reduce.py index aa918ef69..f0a4bb5df 100644 --- a/test/registered/kernels/benchmark/communication/bench_custom_all_reduce.py +++ b/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) diff --git a/test/registered/kernels/ops/communication/test_custom_all_reduce.py b/test/registered/kernels/ops/communication/test_custom_all_reduce.py index 894b16334..ad507b172 100644 --- a/test/registered/kernels/ops/communication/test_custom_all_reduce.py +++ b/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 - # count in 2..8 serially overruns the per-file CI time budget, and 3/5/6/7 - # are rare in practice. Use --num-gpu to exercise them explicitly. + # Only sweep the common world sizes (2, 4, 8, 16) by default: testing every + # count in 2..16 serially overruns the per-file CI time budget, and numbers + # 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, )