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

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
@@ -112,65 +112,137 @@ def _pack_heuristic(*args) -> Heuristic:
return Heuristic(*arg_list) 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.
# SM100 (Blackwell, B200/B300/GB200). Tuned on B200 (148 SMs); world 16 on GB200. @cache
graph_map = { def _sm100_configs(num_sm: int) -> dict[int, AllReduceConfig]:
2: (8.000 * MB, 32.00 * MB, 128.0 * MB), mc_blocks = {5: 64, 6: 48, 7: 48, 8: 32, 16: 32}
3: (4.000 * MB, 4.000 * MB, 128.0 * MB),
4: (2.250 * MB, 2.250 * MB, 128.0 * MB), def config(world_size: int, *, graph: tuple, eager: tuple) -> AllReduceConfig:
5: (1.500 * MB, 1.500 * MB, 128.0 * MB), return AllReduceConfig(
6: (1.000 * MB, 1.000 * MB, 128.0 * MB), graph=_pack_heuristic(*graph),
7: (0.625 * MB, 0.625 * MB, 128.0 * MB), eager=_pack_heuristic(*eager),
8: (0.500 * MB, 0.500 * MB, 128.0 * MB, Range(8 * MB, 128 * MB)), num_push_blocks=num_sm,
16: (0.250 * MB, 0.250 * MB, 128.0 * MB, Range(256 * KB, 128 * MB)), 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.
# SM90 (Hopper, H100/H200). Tuned on H200. @cache
graph_map = { def _sm90_configs(num_sm: int) -> dict[int, AllReduceConfig]:
2: (16.00 * MB, 128.0 * MB, 128.0 * MB), def config(world_size: int, *, graph: tuple, eager: tuple) -> AllReduceConfig:
3: (1.250 * MB, 1.250 * MB, 128.0 * MB), return AllReduceConfig(
4: (384.0 * KB, 384.0 * KB, 128.0 * MB), graph=_pack_heuristic(*graph),
5: (192.0 * KB, 192.0 * KB, 32.00 * MB), eager=_pack_heuristic(*eager),
6: (128.0 * KB, 128.0 * KB, 32.00 * MB, Range(8 * MB, 32 * MB)), num_push_blocks=num_sm,
7: (128.0 * KB, 128.0 * KB, 32.00 * MB, Range(1 * MB, 32 * MB)), num_pull_blocks=64,
8: (128.0 * KB, 128.0 * KB, 32.00 * MB, Range(512 * KB, 128 * MB)), 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), @cache
4: (896.0 * KB, 896.0 * KB, 32.00 * MB, Range(0, 32 * MB)), def _get_all_reduce_configs() -> dict[int, AllReduceConfig]:
5: (384.0 * KB, 384.0 * KB, 32.00 * MB, Range(0, 32 * MB)), cuda_major, _ = torch.cuda.get_device_capability()
6: (192.0 * KB, 192.0 * KB, 32.00 * MB, Range(0, 32 * MB)), num_sm = torch.cuda.get_device_properties().multi_processor_count
7: (128.0 * KB, 128.0 * KB, 32.00 * MB, Range(0, 32 * MB)), if cuda_major == 9:
8: (128.0 * KB, 128.0 * KB, 128.0 * MB, Range(0, 128 * MB)), return _sm90_configs(num_sm)
} if cuda_major == 10:
return AllReduceConfig( return _sm100_configs(num_sm)
graph=_pack_heuristic(*graph_map[world_size]),
eager=_pack_heuristic(*eager_map[world_size]), default = AllReduceConfig(
graph=Heuristic(1 * MB, 1 * MB, 16 * MB),
eager=Heuristic(1 * MB, 1 * MB, 16 * MB),
num_push_blocks=num_sm, num_push_blocks=num_sm,
num_pull_blocks=64, num_pull_blocks=num_sm,
num_mc_blocks=None if world_size < 4 else 128 // world_size, 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 @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 Only SM90 and SM100 are benchmarked so far; other archs get a
conservative default (1 MB one-shot crossovers, no multicast). conservative default (1 MB one-shot crossovers, no multicast).
""" """
cuda_major, _ = torch.cuda.get_device_capability() return _get_all_reduce_configs()[world_size]
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,
)
@@ -42,7 +42,10 @@ from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph impo
is_in_tc_piecewise_cuda_graph, 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 ( from .custom_all_reduce_utils import (
can_use_custom_all_reduce_with_nvlink, can_use_custom_all_reduce_with_nvlink,
is_one_nvlink_clique, is_one_nvlink_clique,
@@ -447,8 +450,7 @@ class CustomAllReduceV2:
def _is_vmm_backed_allocator(device: torch.device) -> bool: def _is_vmm_backed_allocator(device: torch.device) -> bool:
"""True iff the caching allocator is VMM-backed (expandable_segments). Uniform """Check whether expandable-segments VMM backs the caching allocator."""
launch, so the local probe reflects every rank."""
probe = torch.empty(1, dtype=torch.uint8, device=device) probe = torch.empty(1, dtype=torch.uint8, device=device)
return is_vmm_pointer(probe.data_ptr()) return is_vmm_pointer(probe.data_ptr())
@@ -457,34 +459,15 @@ def can_use_custom_all_reduce_v2(
group: ProcessGroup, group: ProcessGroup,
device: torch.device, device: torch.device,
) -> bool: ) -> bool:
# Multi-node (MNNVL): the node-local NVLink/P2P topology checks below supported = get_supported_world_sizes()
# 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))
if dist.get_world_size(group=group) not in supported: if dist.get_world_size(group=group) not in supported:
return False 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)): 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) return is_one_nvlink_clique(group, device) and _is_vmm_backed_allocator(device)
full_nvlink = can_use_custom_all_reduce_with_nvlink( full_nvlink = can_use_custom_all_reduce_with_nvlink(
group=group, group=group,
device=device, device=device,
supported_world_size=supported, supported_world_size=list(supported),
cls_name="CustomAllReduceV2", cls_name="CustomAllReduceV2",
) )
return full_nvlink is True return full_nvlink is True
@@ -724,14 +724,8 @@ class GroupCoordinator:
self.pymscclpp_comm is not None self.pymscclpp_comm is not None
and self.pymscclpp_comm.should_mscclpp_allreduce(input_) and self.pymscclpp_comm.should_mscclpp_allreduce(input_)
) )
# With the MNNVL opt-in, let CustomAllReduceV2 take eligible (small) should_use_custom_allreduce = (
# inputs ahead of the symm-mem pynccl fast path; otherwise pynccl self.ca_comm is not None
# 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
and not self.ca_comm.disabled and not self.ca_comm.disabled
and self.ca_comm.should_custom_ar(input_) and self.ca_comm.should_custom_ar(input_)
) )
@@ -739,7 +733,7 @@ class GroupCoordinator:
self.pynccl_comm is not None self.pynccl_comm is not None
and self.is_symmetric_memory_enabled() and self.is_symmetric_memory_enabled()
and not should_use_pymscclpp_allreduce 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") self.debug_check_symmetric_mempool(self, {"input": input_}, "all_reduce")
with self.pynccl_comm.change_state(enable=True): with self.pynccl_comm.change_state(enable=True):
@@ -2104,9 +2098,6 @@ logger = logging.getLogger(__name__)
_ENABLE_CUSTOM_ALL_REDUCE = True _ENABLE_CUSTOM_ALL_REDUCE = True
_ENABLE_MSCCLPP_ALL_REDUCE = False _ENABLE_MSCCLPP_ALL_REDUCE = False
_ENABLE_TORCH_SYMM_MEM_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 _ENABLE_FLASHINFER_ALLREDUCE_ONLY = False
-8
View File
@@ -1104,14 +1104,6 @@ class Envs:
SGLANG_CUSTOM_ALL_REDUCE_V2_MAX_SIZE_KB = EnvInt(16 * 1024) 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_PULL_SIZE_KB = EnvInt(None)
SGLANG_FORCE_CUSTOM_ALL_REDUCE_V2_PUSH_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 # RoPE cache
-36
View File
@@ -8026,41 +8026,6 @@ class ServerArgs:
"1" if requested_transport == "cuda_ipc" else "0" "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): def _handle_environment_variables(self):
self._handle_multimodal_feature_transport() self._handle_multimodal_feature_transport()
envs.SGLANG_ENABLE_TORCH_COMPILE.set("1" if self.enable_torch_compile else "0") 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( envs.SGLANG_ENABLE_DETERMINISTIC_INFERENCE.set(
"1" if self.enable_deterministic_inference else "0" "1" if self.enable_deterministic_inference else "0"
) )
self._handle_custom_all_reduce_v2_multinode()
if self.enable_deterministic_inference: if self.enable_deterministic_inference:
envs.SGLANG_FLASHINFER_MOE_FUSED_FINALIZE.set("0") envs.SGLANG_FLASHINFER_MOE_FUSED_FINALIZE.set("0")
if self.debug_cuda_graph: if self.debug_cuda_graph:
@@ -0,0 +1,92 @@
from unittest.mock import Mock
import pytest
from sglang.srt.distributed.device_communicators import custom_all_reduce_v2
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
def _patch_group(monkeypatch, *, world_size, same_node):
group = object()
device = object()
monkeypatch.setattr(
custom_all_reduce_v2.dist,
"get_world_size",
lambda group: world_size,
)
monkeypatch.setattr(
custom_all_reduce_v2,
"get_supported_world_sizes",
lambda: (world_size,),
)
monkeypatch.setattr(
custom_all_reduce_v2,
"in_the_same_node_as",
lambda group, source_rank: [same_node] * world_size,
)
return group, device
@pytest.mark.parametrize(
("same_node", "has_fabric_clique", "uses_vmm", "expected"),
[
(False, True, True, True),
(False, False, True, False),
(False, True, False, False),
(True, None, None, True),
],
)
def test_topology_capability(
monkeypatch, same_node, has_fabric_clique, uses_vmm, expected
):
world_size = 8 if same_node else 16
group, device = _patch_group(
monkeypatch,
world_size=world_size,
same_node=same_node,
)
def is_one_clique(group, device):
if same_node:
pytest.fail("intra-node groups do not need a fabric clique")
return has_fabric_clique
def is_vmm_backed(device):
if same_node:
pytest.fail("intra-node groups do not need VMM")
return uses_vmm
intra_node_capability = Mock(return_value=True)
monkeypatch.setattr(
custom_all_reduce_v2,
"is_one_nvlink_clique",
is_one_clique,
)
monkeypatch.setattr(
custom_all_reduce_v2,
"_is_vmm_backed_allocator",
is_vmm_backed,
)
monkeypatch.setattr(
custom_all_reduce_v2,
"can_use_custom_all_reduce_with_nvlink",
intra_node_capability,
)
assert custom_all_reduce_v2.can_use_custom_all_reduce_v2(group, device) is expected
if same_node:
intra_node_capability.assert_called_once_with(
group=group,
device=device,
supported_world_size=[world_size],
cls_name="CustomAllReduceV2",
)
else:
intra_node_capability.assert_not_called()
if __name__ == "__main__":
raise SystemExit(pytest.main([__file__, "-v"]))
@@ -37,6 +37,7 @@ not the per-rank group membership logic.
from __future__ import annotations from __future__ import annotations
import sys import sys
from contextlib import nullcontext
from unittest.mock import Mock, patch from unittest.mock import Mock, patch
import pytest import pytest
@@ -49,6 +50,45 @@ register_cpu_ci(est_time=8, suite="base-a-test-cpu")
parallel_state = pytest.importorskip("sglang.srt.distributed.parallel_state") parallel_state = pytest.importorskip("sglang.srt.distributed.parallel_state")
def test_custom_allreduce_precedes_symmetric_memory_pynccl():
coordinator = parallel_state.GroupCoordinator.__new__(
parallel_state.GroupCoordinator
)
coordinator.world_size = 2
coordinator.unique_name = "test"
coordinator.hpu_communicator = None
coordinator.xpu_communicator = None
coordinator.npu_communicator = None
coordinator.qr_comm = None
coordinator.pymscclpp_comm = None
coordinator.torch_symm_mem_comm = None
coordinator._fi_workspace_hint = None
coordinator.ca_comm = Mock(disabled=False)
coordinator.ca_comm.should_custom_ar.return_value = True
coordinator.pynccl_comm = Mock()
coordinator.pynccl_comm.change_state.return_value = nullcontext()
coordinator.is_symmetric_memory_enabled = Mock(return_value=True)
coordinator.debug_check_symmetric_mempool = Mock()
input_ = Mock(is_cpu=False)
custom_output = object()
with (
patch.object(parallel_state.torch.compiler, "is_compiling", return_value=False),
patch.object(
parallel_state, "outplace_all_reduce", return_value=custom_output
) as outplace_all_reduce,
):
output = coordinator.all_reduce(input_)
assert output is custom_output
outplace_all_reduce.assert_called_once_with(
input_,
group_name="test",
outplace_all_reduce_method="ca",
)
coordinator.pynccl_comm.all_reduce.assert_not_called()
def test_parallel_group_construction_tp8_attn_cp2(): def test_parallel_group_construction_tp8_attn_cp2():
""" """
Test parallel group construction for 8 GPU configuration with: Test parallel group construction for 8 GPU configuration with:
@@ -1,89 +0,0 @@
"""Unit tests for the MNNVL auto-inference gate.
The TP8 best-throughput launch used to require exporting
``SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE=1`` by hand. It is now
capability-inferred; these cases pin the negative-branch contracts so a
refactor cannot silently turn the predicate into always-true (engaging fabric
paths on non-fabric clusters) or drop the explicit-off override.
"""
import unittest
from types import SimpleNamespace
from unittest.mock import patch
from sglang.srt.environ import envs
from sglang.srt.server_args import ServerArgs
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
_HANDLE = ServerArgs._handle_custom_all_reduce_v2_multinode
def _cleared(*fields):
"""Context helper: run with the given env fields unset, restore after."""
import contextlib
import os
@contextlib.contextmanager
def ctx():
backup = {f.name: os.environ.pop(f.name, None) for f in fields}
try:
yield
finally:
for name, val in backup.items():
if val is None:
os.environ.pop(name, None)
else:
os.environ[name] = val
return ctx()
class TestCaV2MultinodeAuto(CustomTestCase):
def test_fabric_multinode_auto_enables(self):
"""GB200/GB300 + nnodes>1 + unset opt-in -> multinode mode on, v2 kept."""
with _cleared(
envs.SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE,
envs.SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2,
), patch("sglang.srt.server_args.is_mnnvl_fabric_device", return_value=True):
_HANDLE(SimpleNamespace(nnodes=2, tp_size=8))
self.assertTrue(envs.SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE.get())
self.assertTrue(envs.SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2.get())
def test_non_fabric_multinode_still_disables_v2(self):
"""Non-fabric multi-node keeps the legacy force-disable (the predicate
must not degrade to always-true)."""
with _cleared(
envs.SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE,
envs.SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2,
), patch("sglang.srt.server_args.is_mnnvl_fabric_device", return_value=False):
_HANDLE(SimpleNamespace(nnodes=2, tp_size=8))
self.assertFalse(envs.SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE.get())
self.assertFalse(envs.SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2.get())
def test_explicit_off_wins_over_fabric(self):
"""SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE=0 on a fabric device
must still force-disable v2 (explicit off beats auto-detection)."""
with _cleared(envs.SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2), patch(
"sglang.srt.server_args.is_mnnvl_fabric_device", return_value=True
), envs.SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE.override("0"):
_HANDLE(SimpleNamespace(nnodes=2, tp_size=8))
self.assertFalse(envs.SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE.get())
self.assertFalse(envs.SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2.get())
def test_tp16_not_auto_opted_in(self):
"""CustomAllReduceV2 supports world sizes 2..8 only; a TP16 fabric
launch must not auto-set the multinode opt-in (it would log
'enabling' and then silently fall back downstream)."""
with _cleared(
envs.SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE,
envs.SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2,
), patch("sglang.srt.server_args.is_mnnvl_fabric_device", return_value=True):
_HANDLE(SimpleNamespace(nnodes=2, tp_size=16))
self.assertFalse(envs.SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE.is_set())
if __name__ == "__main__":
unittest.main()