[Bugfix][Distributed] Gate cross-node MNNVL custom all-reduce by group capability (#53253)
Signed-off-by: JulianZJN <148066842+JulianZJN@users.noreply.github.com> Signed-off-by: Summer Yang <girasoleyang@gmail.com> Co-authored-by: Codex <codex@openai.com> Co-authored-by: Summer Yang <girasoleyang@gmail.com>
Julian committed
Aug 29, 2026 at 08:49 UTC
99013d77d332a2d21d7214b57fa495f2bad2b448
2 files changed
+119
tests/distributed/test_custom_all_reduce.py
+62
@@ -9,6 +9,7 @@ import torch
9
import torch.distributed as dist
10
11
from vllm.distributed.communication_op import tensor_model_parallel_all_reduce # noqa
12
+from vllm.distributed.device_communicators import custom_all_reduce as car
13
from vllm.distributed.parallel_state import get_tp_group, graph_capture
14
15
from ..utils import (
@@ -23,6 +24,67 @@ for i, v in enumerate(test_sizes):
24
test_sizes[i] -= v % 8
25
26
27
+@pytest.mark.parametrize(
28
+ ("major", "local_multicast", "expected"),
29
+ [
30
+ (8, True, False),
31
+ (9, True, False),
32
+ (10, False, False),
33
+ (10, True, True),
34
+ ],
35
+)
36
+def test_cross_node_mnnvl_gate_checks_generation_and_multicast(
37
+ monkeypatch,
38
+ major,
39
+ local_multicast,
40
+ expected,
41
+):
42
+ def has_device_capability(capability, device_id):
43
+ assert capability == 100
44
+ assert device_id == 3
45
+ return major >= 10
46
+
47
+ monkeypatch.setattr(
48
+ car.current_platform,
49
+ "has_device_capability",
50
+ has_device_capability,
51
+ )
52
+ monkeypatch.setattr(
53
+ car,
54
+ "_has_local_multicast_support",
55
+ lambda _device: local_multicast,
56
+ )
57
+ monkeypatch.setattr(car.dist, "all_reduce", lambda *_args, **_kwargs: None)
58
+
59
+ assert car._group_can_attempt_mnnvl(object(), torch.device("cuda:3")) is expected
60
+
61
+
62
+def test_cross_node_mnnvl_gate_requires_support_on_every_rank(monkeypatch):
63
+ monkeypatch.setattr(
64
+ car.current_platform,
65
+ "has_device_capability",
66
+ lambda *_args: True,
67
+ )
68
+ monkeypatch.setattr(
69
+ car,
70
+ "_has_local_multicast_support",
71
+ lambda _device: True,
72
+ )
73
+
74
+ def report_unsupported_peer(support, **_kwargs):
75
+ support.zero_()
76
+
77
+ monkeypatch.setattr(car.dist, "all_reduce", report_unsupported_peer)
78
+
79
+ assert not car._group_can_attempt_mnnvl(object(), torch.device("cuda:0"))
80
+
81
+
82
+def test_local_multicast_support_rejects_non_cuda(monkeypatch):
83
+ monkeypatch.setattr(car.current_platform, "is_cuda", lambda: False)
84
+
85
+ assert not car._has_local_multicast_support(torch.device("cuda:0"))
86
+
87
+
88
@ray.remote(num_gpus=1, max_calls=1)
89
def graph_allreduce(
90
monkeypatch: pytest.MonkeyPatch,
vllm/distributed/device_communicators/custom_all_reduce.py
+57
@@ -33,6 +33,57 @@ except ImportError:
33
logger = init_logger(__name__)
34
35
36
+def _has_local_multicast_support(device: torch.device) -> bool:
37
+ """Return whether this CUDA device can allocate multicast memory."""
38
+ if torch_symm_mem is None or not current_platform.is_cuda():
39
+ return False
40
+ if device.index is None:
41
+ return False
42
+
43
+ try:
44
+ from torch._C._autograd import DeviceType
45
+ from torch._C._distributed_c10d import _SymmetricMemory
46
+
47
+ return _SymmetricMemory.has_multicast_support(
48
+ DeviceType.CUDA,
49
+ device.index,
50
+ )
51
+ except Exception as error:
52
+ logger.debug("MNNVL capability probe failed: %s", error)
53
+ return False
54
+
55
+
56
+def _group_can_attempt_mnnvl(
57
+ group: ProcessGroup,
58
+ device: torch.device,
59
+) -> bool:
60
+ """Return whether every rank can enter the cross-node MNNVL path.
61
+
62
+ MNNVL is available only on Blackwell-class GPUs. Local multicast support
63
+ is necessary but does not establish that the process group spans an MNNVL
64
+ domain; the symmetric-memory rendezvous below performs that group-level
65
+ check. The CPU all-reduce keeps every rank on the same control-flow path
66
+ when a heterogeneous or partially configured group is encountered.
67
+ """
68
+ device_index = device.index
69
+ local_support = (
70
+ device_index is not None
71
+ and current_platform.has_device_capability(100, device_index)
72
+ and _has_local_multicast_support(device)
73
+ )
74
+ group_support = torch.tensor(
75
+ int(local_support),
76
+ dtype=torch.int32,
77
+ device="cpu",
78
+ )
79
+ dist.all_reduce(
80
+ group_support,
81
+ op=dist.ReduceOp.MIN,
82
+ group=group,
83
+ )
84
+ return bool(group_support.item())
85
+
86
+
87
def _can_p2p(rank: int, world_size: int) -> bool:
88
for i in range(world_size):
89
if i == rank:
@@ -146,6 +197,12 @@ class CustomAllreduce:
197
# now `device` is a `torch.device` object
198
assert isinstance(device, torch.device)
199
self.device = device
200
+ if not same_node and not _group_can_attempt_mnnvl(group, device):
201
+ logger.warning(
202
+ "Custom collectives are disabled because this multi-node "
203
+ "group does not have MNNVL-capable GPUs on every rank."
204
+ )
205
+ return
206
device_capability = current_platform.get_device_capability()
207
if (
208
current_platform.is_cuda()