[Bugfix] Disable custom all-reduce on consumer Blackwell (sm_12x)

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 patch
 
 import pytest
 import ray
@@ -130,3 +131,27 @@ def test_custom_allreduce(
     if world_size > torch.accelerator.device_count():
         pytest.skip("Not enough GPUs to run the test.")
     multi_process_parallel(monkeypatch, tp_size, pipeline_parallel_size, test_target)
+
+
+@pytest.mark.parametrize(
+    "is_cuda,is_sm12x,expected",
+    [
+        (True, True, True),  # consumer Blackwell (sm_120 / sm_121): unsupported
+        (True, False, False),  # other CUDA arch (sm_90 / sm_100 / ...): supported
+        (False, True, False),  # non-CUDA platform (e.g. ROCm): not gated
+    ],
+)
+def test_is_unsupported_arch(is_cuda, is_sm12x, expected):
+    from vllm.distributed.device_communicators import custom_all_reduce
+
+    with (
+        patch.object(
+            custom_all_reduce.current_platform, "is_cuda", return_value=is_cuda
+        ),
+        patch.object(
+            custom_all_reduce.current_platform,
+            "is_device_capability_family",
+            return_value=is_sm12x,
+        ),
+    ):
+        assert custom_all_reduce.CustomAllreduce._is_unsupported_arch() is expected

vllm/distributed/device_communicators/custom_all_reduce.py

@@ -51,6 +51,19 @@ def _can_p2p(rank: int, world_size: int) -> bool:
 class CustomAllreduce:
     _SUPPORTED_WORLD_SIZES = [2, 4, 6, 8]
 
+    @staticmethod
+    def _is_unsupported_arch() -> bool:
+        # The custom all-reduce kernels are only built and tuned for sm_90 and
+        # sm_100 (see CUSTOM_ALL_REDUCE_MAX_SIZES). On consumer Blackwell
+        # (sm_12x, e.g. RTX PRO 6000 / RTX 50-series) the kernel's IPC /
+        # peer-mapping path aborts with `custom_all_reduce.cuh 'invalid
+        # argument'` during cudagraph capture, so the entire 12.x family is
+        # unsupported and callers must fall back to NCCL all-reduce.
+        return (
+            current_platform.is_cuda()
+            and current_platform.is_device_capability_family(120)
+        )
+
     # max_size: max supported allreduce size
     def __init__(
         self,
@@ -120,6 +133,14 @@ def __init__(
         assert isinstance(device, torch.device)
         self.device = device
         device_capability = current_platform.get_device_capability()
+        if self._is_unsupported_arch():
+            logger.warning(
+                "Custom allreduce is disabled because it is not supported on "
+                "consumer Blackwell (sm_12x, e.g. RTX PRO 6000 / RTX 50-series);"
+                " falling back to NCCL all-reduce. To silence this warning, "
+                "specify disable_custom_all_reduce=True explicitly."
+            )
+            return
         if (
             current_platform.is_cuda()
             and symm_mem_enabled