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

[Fix] Select custom all-reduce v2 by topology capability (#35061)

Browse Source 
Co-authored-by: xingyuliu <xingyuliu@fb.com>
This commit is contained in:
Xingyu Liu
2026-08-18 10:30:10 -07:00
committed by GitHub
co-authored by xingyuliu
parent 480033def0
commit 7dcaf11987
8 changed files with 268 additions and 236 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/distributed/device_communicators/configs/custom_all_reduce_v2.py
+126 -67
View File
@@ -112,65 +112,137 @@ def _pack_heuristic(*args) -> Heuristic:
return Heuristic(*arg_list)
def _sm100_config(world_size: int, num_sm: int) -> AllReduceConfig:
# 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),
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)),
16: (0.250 * MB, 0.250 * MB, 128.0 * MB, Range(256 * KB, 128 * MB)),
# SM100 (Blackwell, B200/B300/GB200). Tuned on B200 (148 SMs); world 16 on GB200.
@cache
def _sm100_configs(num_sm: int) -> dict[int, AllReduceConfig]:
mc_blocks = {5: 64, 6: 48, 7: 48, 8: 32, 16: 32}
def config(world_size: int, *, graph: tuple, eager: tuple) -> AllReduceConfig:
return AllReduceConfig(
graph=_pack_heuristic(*graph),
eager=_pack_heuristic(*eager),
num_push_blocks=num_sm,
num_pull_blocks=num_sm if world_size == 2 else 96,
num_mc_blocks=mc_blocks.get(world_size),
)
return {
2: config(
2,
graph=(8.000 * MB, 32.00 * MB, 128.0 * MB),
eager=(16.00 * MB, 128.0 * MB, 128.0 * MB),
),
3: config(
3,
graph=(4.000 * MB, 4.000 * MB, 128.0 * MB),
eager=(8.000 * MB, 8.000 * MB, 32.00 * MB),
),
4: config(
4,
graph=(2.250 * MB, 2.250 * MB, 128.0 * MB),
eager=(3.000 * MB, 3.000 * MB, 32.00 * MB),
),
5: config(
5,
graph=(1.500 * MB, 1.500 * MB, 128.0 * MB),
eager=(2.000 * MB, 2.000 * MB, 32.00 * MB, Range(0, 32 * MB)),
),
6: config(
6,
graph=(1.000 * MB, 1.000 * MB, 128.0 * MB),
eager=(1.250 * MB, 1.250 * MB, 64.00 * MB, Range(0, 64 * MB)),
),
7: config(
7,
graph=(0.625 * MB, 0.625 * MB, 128.0 * MB),
eager=(1.000 * MB, 1.000 * MB, 64.00 * MB, Range(0, 64 * MB)),
),
8: config(
8,
graph=(0.500 * MB, 0.500 * MB, 128.0 * MB, Range(8 * MB, 128 * MB)),
eager=(0.750 * MB, 0.750 * MB, 128.0 * MB, Range(0, 128 * MB)),
),
16: config(
16,
graph=(0.250 * MB, 0.250 * MB, 128.0 * MB, Range(256 * KB, 128 * MB)),
eager=(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),
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)),
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, 16: 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)),
# SM90 (Hopper, H100/H200). Tuned on H200.
@cache
def _sm90_configs(num_sm: int) -> dict[int, AllReduceConfig]:
def config(world_size: int, *, graph: tuple, eager: tuple) -> AllReduceConfig:
return AllReduceConfig(
graph=_pack_heuristic(*graph),
eager=_pack_heuristic(*eager),
num_push_blocks=num_sm,
num_pull_blocks=64,
num_mc_blocks=None if world_size < 4 else 128 // world_size,
)
return {
2: config(
2,
graph=(16.00 * MB, 128.0 * MB, 128.0 * MB),
eager=(32.00 * MB, 128.0 * MB, 128.0 * MB),
),
3: config(
3,
graph=(1.250 * MB, 1.250 * MB, 128.0 * MB),
eager=(3.000 * MB, 3.000 * MB, 16.00 * MB),
),
4: config(
4,
graph=(384.0 * KB, 384.0 * KB, 128.0 * MB),
eager=(896.0 * KB, 896.0 * KB, 32.00 * MB, Range(0, 32 * MB)),
),
5: config(
5,
graph=(192.0 * KB, 192.0 * KB, 32.00 * MB),
eager=(384.0 * KB, 384.0 * KB, 32.00 * MB, Range(0, 32 * MB)),
),
6: config(
6,
graph=(128.0 * KB, 128.0 * KB, 32.00 * MB, Range(8 * MB, 32 * MB)),
eager=(192.0 * KB, 192.0 * KB, 32.00 * MB, Range(0, 32 * MB)),
),
7: config(
7,
graph=(128.0 * KB, 128.0 * KB, 32.00 * MB, Range(1 * MB, 32 * MB)),
eager=(128.0 * KB, 128.0 * KB, 32.00 * MB, Range(0, 32 * MB)),
),
8: config(
8,
graph=(128.0 * KB, 128.0 * KB, 32.00 * MB, Range(512 * KB, 128 * MB)),
eager=(128.0 * KB, 128.0 * KB, 128.0 * MB, Range(0, 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]),
@cache
def _get_all_reduce_configs() -> dict[int, AllReduceConfig]:
cuda_major, _ = torch.cuda.get_device_capability()
num_sm = torch.cuda.get_device_properties().multi_processor_count
if cuda_major == 9:
return _sm90_configs(num_sm)
if cuda_major == 10:
return _sm100_configs(num_sm)
default = AllReduceConfig(
graph=Heuristic(1 * MB, 1 * MB, 16 * MB),
eager=Heuristic(1 * MB, 1 * MB, 16 * MB),
num_push_blocks=num_sm,
num_pull_blocks=64,
num_mc_blocks=None if world_size < 4 else 128 // world_size,
num_pull_blocks=num_sm,
num_mc_blocks=None,
)
return {world_size: default for world_size in range(2, 17)}
@cache
def get_supported_world_sizes() -> tuple[int, ...]:
return tuple(_get_all_reduce_configs())
@cache
@@ -180,17 +252,4 @@ def get_all_reduce_config(world_size: int) -> AllReduceConfig:
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,
)
return _get_all_reduce_configs()[world_size]
python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py
+7 -24
View File
@@ -42,7 +42,10 @@ from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph impo
is_in_tc_piecewise_cuda_graph,
)
from .configs.custom_all_reduce_v2 import get_all_reduce_config
from .configs.custom_all_reduce_v2 import (
get_all_reduce_config,
get_supported_world_sizes,
)
from .custom_all_reduce_utils import (
can_use_custom_all_reduce_with_nvlink,
is_one_nvlink_clique,
@@ -447,8 +450,7 @@ class CustomAllReduceV2:
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."""
"""Check whether expandable-segments VMM backs the caching allocator."""
probe = torch.empty(1, dtype=torch.uint8, device=device)
return is_vmm_pointer(probe.data_ptr())
@@ -457,34 +459,15 @@ def can_use_custom_all_reduce_v2(
group: ProcessGroup,
device: torch.device,
) -> bool:
# Multi-node (MNNVL): the node-local NVLink/P2P topology checks below
# are meaningless across nodes; the torch symm-mem rendezvous (fabric
# handles) is the real capability gate there.
if envs.SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE.get() and not all(
in_the_same_node_as(group, source_rank=0)
):
world_size = dist.get_world_size(group=group)
if world_size in range(2, 9):
logger.warning(
"CustomAllReduceV2 enabled on a multi-node group "
"(world_size=%d); graph zero-copy is disabled.",
world_size,
)
return True
return False
supported = list(range(2, 17))
supported = get_supported_world_sizes()
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=supported,
supported_world_size=list(supported),
cls_name="CustomAllReduceV2",
)
return full_nvlink is True
python/sglang/srt/distributed/parallel_state.py
+3 -12
View File
@@ -724,14 +724,8 @@ class GroupCoordinator:
self.pymscclpp_comm is not None
and self.pymscclpp_comm.should_mscclpp_allreduce(input_)
)
# With the MNNVL opt-in, let CustomAllReduceV2 take eligible (small)
# inputs ahead of the symm-mem pynccl fast path; otherwise pynccl
# would absorb every all-reduce whenever --enable-symm-mem is on and
# v2 never runs. Large inputs fail should_custom_ar and still go to
# the symm-mem path below.
_ca_takes_input = (
_CA_V2_MULTINODE
and self.ca_comm is not None
should_use_custom_allreduce = (
self.ca_comm is not None
and not self.ca_comm.disabled
and self.ca_comm.should_custom_ar(input_)
)
@@ -739,7 +733,7 @@ class GroupCoordinator:
self.pynccl_comm is not None
and self.is_symmetric_memory_enabled()
and not should_use_pymscclpp_allreduce
and not _ca_takes_input
and not should_use_custom_allreduce
):
self.debug_check_symmetric_mempool(self, {"input": input_}, "all_reduce")
with self.pynccl_comm.change_state(enable=True):
@@ -2104,9 +2098,6 @@ logger = logging.getLogger(__name__)
_ENABLE_CUSTOM_ALL_REDUCE = True
_ENABLE_MSCCLPP_ALL_REDUCE = False
_ENABLE_TORCH_SYMM_MEM_ALL_REDUCE = False
# Read once at import: whether CustomAllReduceV2 is opted in on a multi-node
# (MNNVL) group. Used on the all_reduce hot path (see GroupCoordinator).
_CA_V2_MULTINODE = envs.SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE.get()
_ENABLE_FLASHINFER_ALLREDUCE_ONLY = False
python/sglang/srt/environ.py
-8
View File
@@ -1104,14 +1104,6 @@ class Envs:
SGLANG_CUSTOM_ALL_REDUCE_V2_MAX_SIZE_KB = EnvInt(16 * 1024)
SGLANG_FORCE_CUSTOM_ALL_REDUCE_V2_PULL_SIZE_KB = EnvInt(None)
SGLANG_FORCE_CUSTOM_ALL_REDUCE_V2_PUSH_SIZE_KB = EnvInt(None)
# Allow CustomAllReduceV2 on a process group that spans nodes (MNNVL
# fabric). Requires torch symmetric memory to rendezvous across nodes
# (fabric handles + IMEX). Graph zero-copy input registration is not
# supported in this mode and is disabled; all-reduce inside CUDA graphs
# falls back to eager pull from the symm workspace. Auto-enabled on
# MNNVL-fabric devices (GB200/GB300) when nnodes > 1; set 0/1 to
# override in either direction.
SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE = EnvBool(False)
# ===================================================================
# RoPE cache
python/sglang/srt/server_args.py
-36
View File
@@ -8026,41 +8026,6 @@ class ServerArgs:
"1" if requested_transport == "cuda_ipc" else "0"
)
def _handle_custom_all_reduce_v2_multinode(self):
# Custom all-reduce v2's graph zero-copy path uses IPC handles and is
# intra-node only. On MNNVL-fabric devices (GB200/GB300) the eager pull
# path works across nodes via the symm-mem workspace, so opt into the
# multinode mode automatically (a failed fabric rendezvous falls back
# to the legacy path at init). Elsewhere force-disable v2 on
# multi-node so the dispatch falls back to the legacy CustomAllreduce
# path, unless the MNNVL opt-in is set explicitly.
if self.nnodes <= 1 or not envs.SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2.get():
return
if (
not envs.SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE.is_set()
and is_mnnvl_fabric_device()
# CustomAllReduceV2 supports world sizes 2..8 only
# (can_use_custom_all_reduce_v2 rejects larger groups); don't
# auto-opt-in a TP16+ launch just to fall back downstream.
and self.tp_size <= 8
):
logger.info(
"MNNVL fabric device detected with nnodes=%d: enabling "
"custom all-reduce v2 multinode mode "
"(SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE=1; set it "
"to 0 to opt out).",
self.nnodes,
)
envs.SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE.set("1")
if not envs.SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE.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")
def _handle_environment_variables(self):
self._handle_multimodal_feature_transport()
envs.SGLANG_ENABLE_TORCH_COMPILE.set("1" if self.enable_torch_compile else "0")
@@ -8072,7 +8037,6 @@ class ServerArgs:
envs.SGLANG_ENABLE_DETERMINISTIC_INFERENCE.set(
"1" if self.enable_deterministic_inference else "0"
)
self._handle_custom_all_reduce_v2_multinode()
if self.enable_deterministic_inference:
envs.SGLANG_FLASHINFER_MOE_FUSED_FINALIZE.set("0")
if self.debug_cuda_graph: