[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>

tests/distributed/test_custom_all_reduce.py

@@ -9,6 +9,7 @@ import torch
 import torch.distributed as dist
 
 from vllm.distributed.communication_op import tensor_model_parallel_all_reduce  # noqa
+from vllm.distributed.device_communicators import custom_all_reduce as car
 from vllm.distributed.parallel_state import get_tp_group, graph_capture
 
 from ..utils import (
@@ -23,6 +24,67 @@ for i, v in enumerate(test_sizes):
     test_sizes[i] -= v % 8
 
 
+@pytest.mark.parametrize(
+    ("major", "local_multicast", "expected"),
+    [
+        (8, True, False),
+        (9, True, False),
+        (10, False, False),
+        (10, True, True),
+    ],
+)
+def test_cross_node_mnnvl_gate_checks_generation_and_multicast(
+    monkeypatch,
+    major,
+    local_multicast,
+    expected,
+):
+    def has_device_capability(capability, device_id):
+        assert capability == 100
+        assert device_id == 3
+        return major >= 10
+
+    monkeypatch.setattr(
+        car.current_platform,
+        "has_device_capability",
+        has_device_capability,
+    )
+    monkeypatch.setattr(
+        car,
+        "_has_local_multicast_support",
+        lambda _device: local_multicast,
+    )
+    monkeypatch.setattr(car.dist, "all_reduce", lambda *_args, **_kwargs: None)
+
+    assert car._group_can_attempt_mnnvl(object(), torch.device("cuda:3")) is expected
+
+
+def test_cross_node_mnnvl_gate_requires_support_on_every_rank(monkeypatch):
+    monkeypatch.setattr(
+        car.current_platform,
+        "has_device_capability",
+        lambda *_args: True,
+    )
+    monkeypatch.setattr(
+        car,
+        "_has_local_multicast_support",
+        lambda _device: True,
+    )
+
+    def report_unsupported_peer(support, **_kwargs):
+        support.zero_()
+
+    monkeypatch.setattr(car.dist, "all_reduce", report_unsupported_peer)
+
+    assert not car._group_can_attempt_mnnvl(object(), torch.device("cuda:0"))
+
+
+def test_local_multicast_support_rejects_non_cuda(monkeypatch):
+    monkeypatch.setattr(car.current_platform, "is_cuda", lambda: False)
+
+    assert not car._has_local_multicast_support(torch.device("cuda:0"))
+
+
 @ray.remote(num_gpus=1, max_calls=1)
 def graph_allreduce(
     monkeypatch: pytest.MonkeyPatch,

vllm/distributed/device_communicators/custom_all_reduce.py

@@ -33,6 +33,57 @@ except ImportError:
 logger = init_logger(__name__)
 
 
+def _has_local_multicast_support(device: torch.device) -> bool:
+    """Return whether this CUDA device can allocate multicast memory."""
+    if torch_symm_mem is None or not current_platform.is_cuda():
+        return False
+    if device.index is None:
+        return False
+
+    try:
+        from torch._C._autograd import DeviceType
+        from torch._C._distributed_c10d import _SymmetricMemory
+
+        return _SymmetricMemory.has_multicast_support(
+            DeviceType.CUDA,
+            device.index,
+        )
+    except Exception as error:
+        logger.debug("MNNVL capability probe failed: %s", error)
+        return False
+
+
+def _group_can_attempt_mnnvl(
+    group: ProcessGroup,
+    device: torch.device,
+) -> bool:
+    """Return whether every rank can enter the cross-node MNNVL path.
+
+    MNNVL is available only on Blackwell-class GPUs. Local multicast support
+    is necessary but does not establish that the process group spans an MNNVL
+    domain; the symmetric-memory rendezvous below performs that group-level
+    check. The CPU all-reduce keeps every rank on the same control-flow path
+    when a heterogeneous or partially configured group is encountered.
+    """
+    device_index = device.index
+    local_support = (
+        device_index is not None
+        and current_platform.has_device_capability(100, device_index)
+        and _has_local_multicast_support(device)
+    )
+    group_support = torch.tensor(
+        int(local_support),
+        dtype=torch.int32,
+        device="cpu",
+    )
+    dist.all_reduce(
+        group_support,
+        op=dist.ReduceOp.MIN,
+        group=group,
+    )
+    return bool(group_support.item())
+
+
 def _can_p2p(rank: int, world_size: int) -> bool:
     for i in range(world_size):
         if i == rank:
@@ -146,6 +197,12 @@ class CustomAllreduce:
         # now `device` is a `torch.device` object
         assert isinstance(device, torch.device)
         self.device = device
+        if not same_node and not _group_can_attempt_mnnvl(group, device):
+            logger.warning(
+                "Custom collectives are disabled because this multi-node "
+                "group does not have MNNVL-capable GPUs on every rank."
+            )
+            return
         device_capability = current_platform.get_device_capability()
         if (
             current_platform.is_cuda()