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
test/registered/unit/distributed/test_custom_all_reduce_v2_capability.py
+92
View File
@@ -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"]))
test/registered/unit/distributed/test_parallel_state.py
+40
View File
@@ -37,6 +37,7 @@ not the per-rank group membership logic.
from __future__ import annotations
import sys
from contextlib import nullcontext
from unittest.mock import Mock, patch
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")
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():
"""
Test parallel group construction for 8 GPU configuration with:
test/registered/unit/server_args/test_mnnvl_auto_inference.py
-89
View File
@@ -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()