[Bugfix][Distributed] Keep cross-node MNNVL groups off CUDA IPC all-reduce

tests/distributed/test_custom_all_reduce.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import random
+from unittest.mock import Mock
 
 import pytest
 import ray
@@ -25,6 +26,18 @@
     test_sizes[i] -= v % 8
 
 
+@pytest.fixture
+def custom_allreduce():
+    communicator = car.CustomAllreduce.__new__(car.CustomAllreduce)
+    communicator.disabled = False
+    communicator.world_size = 2
+    communicator.max_size = 1024
+    communicator._ptr = 0
+    communicator.mnnvl_only = False
+    communicator.fully_connected = False
+    return communicator
+
+
 @ray.remote(num_gpus=1, max_calls=1)
 def _all_reduce_mhc(monkeypatch, tp_size, pp_size, rank, distributed_init_port):
     from vllm.model_executor.kernels.mhc.tilelang import (
@@ -146,16 +159,72 @@ def test_all_reduce_mhc_preserves_bf16_boundaries_and_graph_replay(monkeypatch):
     ],
 )
 def test_custom_allreduce_filters_dtype(
+    custom_allreduce,
     dtype: torch.dtype,
     expected: bool,
 ) -> None:
-    communicator = car.CustomAllreduce.__new__(car.CustomAllreduce)
-    communicator.disabled = False
-    communicator.world_size = 2
-    communicator.max_size = 1024
-    communicator._ptr = 0
+    assert custom_allreduce.should_custom_ar(torch.empty(16, dtype=dtype)) is expected
+
+
+@pytest.mark.parametrize("capturing", [False, True])
+def test_cross_node_custom_allreduce_declines_without_launching_ipc(
+    custom_allreduce, monkeypatch, capturing
+):
+    custom_allreduce.mnnvl_only = True
+    custom_allreduce._IS_CAPTURING = capturing
+    custom_allreduce.all_reduce = Mock()
+    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: capturing)
 
-    assert communicator.should_custom_ar(torch.empty(16, dtype=dtype)) is expected
+    assert custom_allreduce.custom_all_reduce(torch.empty(16)) is None
+    custom_allreduce.all_reduce.assert_not_called()
+
+
+@pytest.mark.parametrize("mnnvl_only", [False, True])
+def test_graph_capture_registers_ipc_buffers_only_on_same_node(
+    custom_allreduce, monkeypatch, mnnvl_only
+):
+    custom_allreduce.mnnvl_only = mnnvl_only
+    custom_allreduce.group = object()
+    custom_allreduce.rank = 0
+    get_meta = Mock(return_value=([1], [0]))
+    register = Mock()
+
+    def share_meta(data, **kwargs):
+        data[:] = [[1], [0]]
+
+    broadcast = Mock(side_effect=share_meta)
+    monkeypatch.setattr(car.ops, "get_graph_buffer_ipc_meta", get_meta)
+    monkeypatch.setattr(car.ops, "register_graph_buffers", register)
+    monkeypatch.setattr(car.dist, "get_world_size", lambda **kwargs: 2)
+    monkeypatch.setattr(car.dist, "get_process_group_ranks", lambda **kwargs: [0, 1])
+    monkeypatch.setattr(car.dist, "broadcast_object_list", broadcast)
+
+    with custom_allreduce.capture():
+        assert custom_allreduce._IS_CAPTURING
+
+    assert not custom_allreduce._IS_CAPTURING
+    if mnnvl_only:
+        get_meta.assert_not_called()
+        broadcast.assert_not_called()
+        register.assert_not_called()
+    else:
+        get_meta.assert_called_once_with(0)
+        assert broadcast.call_count == 2
+        register.assert_called_once_with(0, [[1], [1]], [[0], [0]])
+
+
+def test_cross_node_mnnvl_all_gather_and_reduce_scatter_remain_eligible(
+    custom_allreduce, monkeypatch
+):
+    custom_allreduce.mnnvl_only = True
+    custom_allreduce.mnnvl_multicast_ptr = 1
+    custom_allreduce.max_mnnvl_all_gather_size = 1024
+    custom_allreduce.max_mnnvl_reduce_scatter_size = 1024
+    monkeypatch.setattr(car.current_platform, "is_cuda", lambda: True)
+    inp = torch.empty(16, 8)
+
+    assert custom_allreduce.should_custom_all_gather(inp)
+    assert custom_allreduce.should_custom_reduce_scatter(inp)
 
 
 @pytest.mark.parametrize(

vllm/distributed/device_communicators/custom_all_reduce.py

@@ -135,8 +135,9 @@ def __init__(
             device: the device to bind the CustomAllreduce to. If None,
                 it will be bound to f"cuda:{local_rank}".
         It is the caller's responsibility to make sure each communicator
-        is bind to a unique device, and all communicators in this group
-        are in the same node.
+        is bound to a unique device. CUDA IPC collectives require all ranks
+        to be on the same node; cross-node groups only support MNNVL
+        all-gather and reduce-scatter.
 
         """
         self._IS_CAPTURING = False
@@ -383,6 +384,9 @@ def capture(self):
                 self.register_graph_buffers()
 
     def register_graph_buffers(self):
+        if self.mnnvl_only:
+            # MNNVL collectives use pre-registered symmetric memory, not CUDA IPC.
+            return
         handle, offset = ops.get_graph_buffer_ipc_meta(self._ptr)
         logger.debug("Registering %d cuda graph addresses", len(offset))
         # We cannot directly use `dist.all_gather_object` here
@@ -402,7 +406,7 @@ def register_graph_buffers(self):
         ops.register_graph_buffers(self._ptr, handles, offsets)
 
     def should_custom_ar(self, inp: torch.Tensor):
-        if self.disabled or self.world_size > 8:
+        if self.disabled or self.mnnvl_only or self.world_size > 8:
             return False
         if inp.dtype not in (torch.float32, torch.float16, torch.bfloat16):
             return False