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 3d0d2556f..152066168 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 @@ -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] 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 cfda063b3..77833836a 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 @@ -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 diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index b383c5404..b131a48e6 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/python/sglang/srt/distributed/parallel_state.py @@ -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 diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index bc0e6004c..3abbabe13 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -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 diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index c5b5c84ab..b6956a95d 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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: diff --git a/test/registered/unit/distributed/test_custom_all_reduce_v2_capability.py b/test/registered/unit/distributed/test_custom_all_reduce_v2_capability.py new file mode 100644 index 000000000..579781465 --- /dev/null +++ b/test/registered/unit/distributed/test_custom_all_reduce_v2_capability.py @@ -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"])) diff --git a/test/registered/unit/distributed/test_parallel_state.py b/test/registered/unit/distributed/test_parallel_state.py index a4395344f..8043b3d40 100644 --- a/test/registered/unit/distributed/test_parallel_state.py +++ b/test/registered/unit/distributed/test_parallel_state.py @@ -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: diff --git a/test/registered/unit/server_args/test_mnnvl_auto_inference.py b/test/registered/unit/server_args/test_mnnvl_auto_inference.py deleted file mode 100644 index cd0afda09..000000000 --- a/test/registered/unit/server_args/test_mnnvl_auto_inference.py +++ /dev/null @@ -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()