[Distributed] Make --disable-custom-all-reduce fall back to NCCL

docs/design/fusions.md

@@ -110,6 +110,7 @@ Fields set explicitly by the user always take precedence over optimization-level
 RMSNorm, and optionally a quantization step into a single FlashInfer / TRT-LLM communication kernel.
 This fusion is only profitable for small `num_tokens`,
 so the fusion is only performed in the lower compiled range.
+It is off by default with `--disable-custom-all-reduce`, which makes all-reduce fall back to NCCL.
 
 Patterns covered:
 

tests/distributed/test_comm_ops.py

@@ -24,6 +24,7 @@
 from vllm.distributed.device_communicators import flashinfer_all_reduce
 from vllm.distributed.device_communicators.cuda_communicator import CudaCommunicator
 from vllm.distributed.parallel_state import GroupCoordinator, TensorMetadata
+from vllm.platforms import current_platform
 from vllm.v1.worker.gpu_worker import AsyncIntermediateTensors
 
 from ..utils import (
@@ -466,6 +467,49 @@ def test_flashinfer_all_reduce_precedes_nccl(monkeypatch: pytest.MonkeyPatch) ->
     nccl_selector.assert_not_called()
 
 
+def test_disable_custom_all_reduce_falls_back_to_nccl(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """--disable-custom-all-reduce sends TP all-reduce to NCCL, not to the
+    FlashInfer or symmetric-memory all-reduce."""
+    backends = {}
+    for module, cls in (
+        ("custom_all_reduce", "CustomAllreduce"),
+        ("flashinfer_all_reduce", "FlashInferAllReduce"),
+        ("flashinfer_pcie_ipc_all_reduce", "FlashInferPcieIpcAllReduce"),
+        ("pynccl", "PyNcclCommunicator"),
+        ("symm_mem", "SymmMemCommunicator"),
+    ):
+        backends[cls] = Mock(disabled=False, world_size=2)
+        monkeypatch.setattr(
+            f"vllm.distributed.device_communicators.{module}.{cls}",
+            Mock(return_value=backends[cls]),
+        )
+    # Symmetric-memory all-reduce is only created on CUDA.
+    monkeypatch.setattr(current_platform, "is_cuda", lambda: True)
+    monkeypatch.setenv("VLLM_ALLREDUCE_USE_FLASHINFER_PCIE_IPC", "1")
+    monkeypatch.setattr(
+        "vllm.distributed.parallel_state._ENABLE_CUSTOM_ALL_REDUCE", False
+    )
+
+    # Not registered with torch.distributed, so it is used as a stateless group.
+    group = Mock(rank=Mock(return_value=0), size=Mock(return_value=2))
+    communicator = CudaCommunicator(
+        group,
+        device_group=Mock(),
+        unique_name="tp:0",
+        global_ranks=[0, 1],
+        global_world_size=2,
+    )
+    input_ = torch.ones(4)
+    communicator.all_reduce(input_)
+
+    backends["PyNcclCommunicator"].all_reduce.assert_called_once_with(input_)
+    backends["FlashInferAllReduce"].all_reduce.assert_not_called()
+    backends["FlashInferPcieIpcAllReduce"].all_reduce.assert_not_called()
+    backends["SymmMemCommunicator"].all_reduce.assert_not_called()
+
+
 def test_aiter_all_gather_precedes_pynccl(monkeypatch: pytest.MonkeyPatch) -> None:
     """Ensure can use aiter all_gather if available even if pynccl is not available."""
     output = torch.empty(2)

tests/test_config.py

@@ -3169,6 +3169,30 @@ def test_vllm_config_defaults_are_none():
                 assert getattr(config.compilation_config, k) is None
 
 
+@pytest.mark.parametrize("disable_custom_all_reduce", [False, True])
+def test_allreduce_rms_fusion_default_follows_disable_custom_all_reduce(
+    monkeypatch, disable_custom_all_reduce
+):
+    """--disable-custom-all-reduce also turns off the FlashInfer all-reduce +
+    RMSNorm fusion by default, so TP all-reduce falls back to NCCL."""
+    monkeypatch.setattr(current_platform, "is_rocm", lambda: False)
+    monkeypatch.setattr(current_platform, "is_cuda", lambda: True)
+    monkeypatch.setattr(
+        current_platform, "is_device_capability_family", lambda capability: False
+    )
+    monkeypatch.setattr(
+        current_platform, "is_device_capability", lambda capability: capability == 90
+    )
+    monkeypatch.setattr("vllm.utils.flashinfer.has_flashinfer", lambda: True)
+    config = object.__new__(VllmConfig)
+    config.parallel_config = SimpleNamespace(
+        tensor_parallel_size=2, disable_custom_all_reduce=disable_custom_all_reduce
+    )
+
+    enabled = vllm_config_module.enable_allreduce_rms_fusion(config)
+    assert enabled is not disable_custom_all_reduce
+
+
 def test_validate_mamba_align_subblock_prefill():
     """Align mode permits configured prefill chunks smaller than a block."""
     config = SimpleNamespace(

vllm/config/parallel.py

@@ -234,7 +234,10 @@ class ParallelConfig:
     parallel and large models."""
 
     disable_custom_all_reduce: bool = False
-    """Disable the custom all-reduce kernel and fall back to NCCL."""
+    """Disable the custom all-reduce kernel and fall back to NCCL. This also
+    disables the FlashInfer and symmetric-memory all-reduce for tensor
+    parallelism, and the all-reduce + RMSNorm fusion unless
+    `pass_config.fuse_allreduce_rms` is set explicitly."""
 
     enable_elastic_ep: bool = False
     """Enable elastic expert parallelism with stateless NCCL groups for DP/EP."""

vllm/config/vllm.py

@@ -184,6 +184,10 @@ def enable_allreduce_rms_fusion(cfg: "VllmConfig") -> bool:
     if envs.VLLM_BATCH_INVARIANT:
         return False
 
+    # --disable-custom-all-reduce falls back to NCCL
+    if cfg.parallel_config.disable_custom_all_reduce:
+        return False
+
     if current_platform.is_rocm():
         from vllm._aiter_ops import rocm_aiter_ops
 

vllm/distributed/device_communicators/cuda_communicator.py

@@ -60,13 +60,20 @@ def __init__(
             from vllm.distributed.parallel_state import _ENABLE_CUSTOM_ALL_REDUCE
 
             use_custom_allreduce = _ENABLE_CUSTOM_ALL_REDUCE
-            use_torch_symm_mem = envs.VLLM_ALLREDUCE_USE_SYMM_MEM
+            # --disable-custom-all-reduce falls back to NCCL, so it also turns
+            # off the FlashInfer and symmetric-memory all-reduce.
+            use_torch_symm_mem = (
+                use_custom_allreduce and envs.VLLM_ALLREDUCE_USE_SYMM_MEM
+            )
             # FlashInfer all-reduce does not provide a fixed reduction order.
             use_flashinfer_allreduce = (
-                envs.VLLM_ALLREDUCE_USE_FLASHINFER and not envs.VLLM_BATCH_INVARIANT
+                use_custom_allreduce
+                and envs.VLLM_ALLREDUCE_USE_FLASHINFER
+                and not envs.VLLM_BATCH_INVARIANT
             )
             use_flashinfer_pcie_ipc_allreduce = (
-                envs.VLLM_ALLREDUCE_USE_FLASHINFER_PCIE_IPC
+                use_custom_allreduce
+                and envs.VLLM_ALLREDUCE_USE_FLASHINFER_PCIE_IPC
                 and not envs.VLLM_BATCH_INVARIANT
             )
             # Neither AITER nor QuickReduce all-reduce has a fixed reduction order.