minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Lianmin Zheng <lianminzheng@gmail.com>
This commit is contained in:
Ming Yang
2026-07-25 17:25:18 -07:00
committed by GitHub
co-authored by Claude Lianmin Zheng
parent 2c63a2f12b
commit 55c4853487
9 changed files with 97 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/jit/csrc/distributed/custom_all_reduce.cuh
+18 -16
View File
@@ -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* 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 <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
+1 -1
View File
@@ -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
+4 -2
View File
@@ -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
+2 -2
View File
@@ -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 (
python/sglang/srt/distributed/device_communicators/custom_all_reduce_utils.py
+47
View File
@@ -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
+19 -2
View File
@@ -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
-10
View File
@@ -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
+1 -1
View File
@@ -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
+5 -4
View File
@@ -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,
)