diff --git a/docs/design/fusions.md b/docs/design/fusions.md index 002984c113a1..a9ca2c0a2528 100644 --- a/docs/design/fusions.md +++ b/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: diff --git a/tests/distributed/test_comm_ops.py b/tests/distributed/test_comm_ops.py index 253ec8a188af..35843de7885b 100644 --- a/tests/distributed/test_comm_ops.py +++ b/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) diff --git a/tests/test_config.py b/tests/test_config.py index a2e75b10725d..18a534d8f04d 100644 --- a/tests/test_config.py +++ b/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( diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index 7df8f4e0a4ed..d9b4f5880b6d 100644 --- a/vllm/config/parallel.py +++ b/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.""" diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 32b49f42725e..690b9b740ad4 100644 --- a/vllm/config/vllm.py +++ b/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 diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index 99ebcd717ddf..cd29d9bb6db0 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/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.