From 6c2758c9b361b3e17a0dcf750444e6cc76d6bef2 Mon Sep 17 00:00:00 2001 From: Nicholaus Shupe Date: Mon, 6 Jul 2026 06:34:52 -0700 Subject: [PATCH] [Bugfix] Disable custom all-reduce on consumer Blackwell (sm_12x) The custom all-reduce CUDA kernels are only built and tuned for sm_90 and sm_100 (see CUSTOM_ALL_REDUCE_MAX_SIZES, which lists 9.0/10.0/10.3 but no 12.x). On consumer Blackwell (sm_120 / sm_121, e.g. RTX PRO 6000 and the RTX 50-series) the kernel's IPC / peer-mapping path aborts with `custom_all_reduce.cuh 'invalid argument'` during cudagraph capture, killing the TP worker. The existing topology guards do not catch this: they only disable custom all-reduce for more than two PCIe-only GPUs, so a 2-GPU tensor-parallel setup passes the checks, enables the kernel, and then crashes at capture. Disable custom all-reduce for the whole sm_12x family via current_platform.is_device_capability_family(120) (extracted into CustomAllreduce._is_unsupported_arch for testability) and fall back to NCCL all-reduce, so users no longer have to pass disable_custom_all_reduce=True manually. The early return leaves the communicator in its default-disabled state (self.disabled = True) before any buffer allocation or ops.init_custom_ar call. Add a unit test covering the sm_12x family (disabled), other CUDA archs (enabled), and non-CUDA platforms (not gated). Co-Authored-By: Claude Opus 4.8 Signed-off-by: Nicholaus Shupe --- tests/distributed/test_custom_all_reduce.py | 25 +++++++++++++++++++ .../device_communicators/custom_all_reduce.py | 21 ++++++++++++++++ 2 files changed, 46 insertions(+) diff --git a/tests/distributed/test_custom_all_reduce.py b/tests/distributed/test_custom_all_reduce.py index edddb6ec8455..2cd84bad4b28 100644 --- a/tests/distributed/test_custom_all_reduce.py +++ b/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 diff --git a/vllm/distributed/device_communicators/custom_all_reduce.py b/vllm/distributed/device_communicators/custom_all_reduce.py index 86180a51a71e..98ecd06d5e31 100644 --- a/vllm/distributed/device_communicators/custom_all_reduce.py +++ b/vllm/distributed/device_communicators/custom_all_reduce.py @@ -66,6 +66,19 @@ class CustomAllreduce: _DEFAULT_REDUCE_SCATTER_MAX_SIZE = 16 * 1024 * 1024 _DEFAULT_MNNVL_REDUCE_SCATTER_MAX_SIZE = 16 * 1024 * 1024 + @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, @@ -147,6 +160,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