From 9cd12afddd06281c0f3ccffaa62b5283dd363b85 Mon Sep 17 00:00:00 2001 From: Summer Yang Date: Fri, 21 Aug 2026 13:55:37 -0700 Subject: [PATCH] [Perf] Tune FlashInfer all-reduce selection on SM103 Prefer standalone FlashInfer all-reduce ahead of NCCL symmetric memory when the existing FlashInfer gate is enabled. Use topology-keyed cutoffs for SM103 MNNVL and size the shared workspace accordingly. Co-authored-by: OpenAI Codex Signed-off-by: Summer Yang --- tests/distributed/test_comm_ops.py | 84 +++++++++++++++++++ .../device_communicators/all_reduce_utils.py | 8 ++ .../device_communicators/cuda_communicator.py | 29 ++++--- .../flashinfer_all_reduce.py | 64 ++++++++++++-- 4 files changed, 164 insertions(+), 21 deletions(-) diff --git a/tests/distributed/test_comm_ops.py b/tests/distributed/test_comm_ops.py index 19a095c7c6f7..deac28cde1ee 100644 --- a/tests/distributed/test_comm_ops.py +++ b/tests/distributed/test_comm_ops.py @@ -318,6 +318,90 @@ def test_cuda_communicator_checkpoints_flashinfer_workspaces( workspace.checkpoint_restore.assert_called_once_with(group) +@pytest.mark.parametrize( + ("backend", "capability", "world_size", "nodes", "expected"), + [ + ("mnnvl", 103, 4, 1, 80 * flashinfer_all_reduce.MiB - 1), + ("mnnvl", 103, 8, 2, 64 * flashinfer_all_reduce.MiB - 1), + ("mnnvl", 103, 16, 4, 8 * flashinfer_all_reduce.MiB - 1), + ("mnnvl", 103, 2, 1, None), + ("mnnvl", 103, 12, 3, None), + ("mnnvl", 103, 8, 1, None), + ("trtllm", 103, 8, 2, None), + ("mnnvl", 90, 8, 2, None), + ], +) +def test_flashinfer_standalone_size_tuning( + monkeypatch: pytest.MonkeyPatch, + backend: str, + capability: int, + world_size: int, + nodes: int, + expected: int | None, +) -> None: + monkeypatch.setattr( + flashinfer_all_reduce, + "current_platform", + Mock(get_device_capability=lambda: Mock(to_int=lambda: capability)), + ) + monkeypatch.setattr(flashinfer_all_reduce, "_node_count", lambda _: nodes) + + assert ( + flashinfer_all_reduce._get_tuned_standalone_max_size( + world_size, backend, Mock() + ) + == expected + ) + + +@pytest.mark.parametrize(("enabled", "expected"), [(True, 4681), (False, 128)]) +def test_flashinfer_standalone_workspace_size( + monkeypatch: pytest.MonkeyPatch, enabled: bool, expected: int +) -> None: + create_workspace = Mock(return_value=Mock(backend="mnnvl")) + monkeypatch.setattr( + flashinfer_all_reduce.envs, "VLLM_ALLREDUCE_USE_FLASHINFER", enabled + ) + monkeypatch.setattr(flashinfer_all_reduce, "_fi_ar_workspace", None) + monkeypatch.setattr(flashinfer_all_reduce, "_fi_ar_quant_workspace", None) + monkeypatch.setattr( + flashinfer_all_reduce, + "_resolve_fi_ar_backend", + Mock(return_value=("mnnvl", False)), + ) + monkeypatch.setattr(flashinfer_all_reduce, "get_node_count", lambda: 2) + monkeypatch.setattr( + flashinfer_all_reduce, + "_get_tuned_standalone_max_size", + Mock(return_value=64 * flashinfer_all_reduce.MiB - 1), + ) + monkeypatch.setattr(flashinfer_all_reduce, "_create_workspace", create_workspace) + + flashinfer_all_reduce.get_fi_ar_workspace(8, 0, 128, 7168, torch.bfloat16, Mock()) + + assert create_workspace.call_args.args[3] == expected + + +def test_flashinfer_all_reduce_precedes_nccl(monkeypatch: pytest.MonkeyPatch) -> None: + output = torch.empty(2) + fi_ar_comm = Mock(disabled=False) + fi_ar_comm.should_use_fi_ar.return_value = True + fi_ar_comm.all_reduce.return_value = output + communicator = CudaCommunicator.__new__(CudaCommunicator) + communicator.fi_ar_comm = fi_ar_comm + communicator.pynccl_comm = Mock(world_size=8) + communicator.qr_comm = None + nccl_selector = Mock(return_value=True) + monkeypatch.setattr( + "vllm.distributed.device_communicators.cuda_communicator." + "should_nccl_symm_mem_allreduce", + nccl_selector, + ) + + assert communicator.all_reduce(torch.empty(1)) is output + nccl_selector.assert_not_called() + + def test_async_intermediate_tensors_lazy_wait() -> None: work = _DummyWork() post_calls = {"n": 0} diff --git a/vllm/distributed/device_communicators/all_reduce_utils.py b/vllm/distributed/device_communicators/all_reduce_utils.py index 86b709ed41c7..b6da29099b0d 100644 --- a/vllm/distributed/device_communicators/all_reduce_utils.py +++ b/vllm/distributed/device_communicators/all_reduce_utils.py @@ -82,6 +82,14 @@ }, } +# Per-rank input limit for standalone FlashInfer MNNVL all-reduce. +# The key is (compute capability, world size, node count). +FI_MNNVL_ALLREDUCE_MAX_SIZE_MB: dict[tuple[int, int, int], float] = { + (103, 4, 1): 80, + (103, 8, 2): 64, + (103, 16, 4): 8, +} + # NCCL symmetric memory allreduce configuration based on H100 and GB200 benchmarks. # PyNCCL-symm outperforms custom_AR for small and large tensor sizes, # while custom_AR wins for mid-range sizes. diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index fa3a702fb0af..c0f2e0a4106d 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -220,15 +220,17 @@ def _log_all_reduce_backend_selection(self) -> None: depends on the input tensor. """ all_potential_ar_backends = [ + "FLASHINFER", "NCCL_SYMM_MEM", "QUICK_REDUCE", - "FLASHINFER", "AITER_CUSTOM", "CUSTOM", "SYMM_MEM", "PYNCCL", ] enabled_ar_backends: list[str] = [] + if self.fi_ar_comm is not None and not self.fi_ar_comm.disabled: + enabled_ar_backends.append("FLASHINFER") # Mirror the static preconditions of `should_nccl_symm_mem_allreduce`: # VLLM_BATCH_INVARIANT off, NCCL symm mem enabled, world_size meets # min_world_size, and world_size either has a tuned entry in @@ -255,8 +257,6 @@ def _log_all_reduce_backend_selection(self) -> None: enabled_ar_backends.append("NCCL_SYMM_MEM") if self.qr_comm is not None and not self.qr_comm.disabled: enabled_ar_backends.append("QUICK_REDUCE") - if self.fi_ar_comm is not None and not self.fi_ar_comm.disabled: - enabled_ar_backends.append("FLASHINFER") if self.aiter_ar_comm is not None and not self.aiter_ar_comm.disabled: enabled_ar_backends.append("AITER_CUSTOM") if self.ca_comm is not None and not self.ca_comm.disabled: @@ -276,16 +276,23 @@ def _log_all_reduce_backend_selection(self) -> None: ) def all_reduce(self, input_): + fi_ar_comm = self.fi_ar_comm + use_fi_ar = ( + fi_ar_comm is not None + and not fi_ar_comm.disabled + and fi_ar_comm.should_use_fi_ar(input_) + ) + # since currently we perform copy input -> symm_input -> out-of-place AR # return symm_output, we don't need to check if input is symmetric - if self.pynccl_comm is not None and should_nccl_symm_mem_allreduce( - self.pynccl_comm.world_size, input_ + if ( + self.pynccl_comm is not None + and not use_fi_ar + and should_nccl_symm_mem_allreduce(self.pynccl_comm.world_size, input_) ): out = torch.ops.vllm.all_reduce_symmetric_with_copy(input_) if out is not None: return out - # always try quick reduce first, then flashinfer, then the AITER or vLLM - # custom allreduce, and then pynccl. (quick reduce just for ROCM MI3*) qr_comm = self.qr_comm if ( qr_comm is not None @@ -295,12 +302,8 @@ def all_reduce(self, input_): out = qr_comm.quick_all_reduce(input_) assert out is not None return out - fi_ar_comm = self.fi_ar_comm - if ( - fi_ar_comm is not None - and not fi_ar_comm.disabled - and fi_ar_comm.should_use_fi_ar(input_) - ): + if use_fi_ar: + assert fi_ar_comm is not None out = fi_ar_comm.all_reduce(input_) assert out is not None return out diff --git a/vllm/distributed/device_communicators/flashinfer_all_reduce.py b/vllm/distributed/device_communicators/flashinfer_all_reduce.py index b7fe09081ac6..9cd3902dd809 100644 --- a/vllm/distributed/device_communicators/flashinfer_all_reduce.py +++ b/vllm/distributed/device_communicators/flashinfer_all_reduce.py @@ -14,7 +14,10 @@ import vllm.envs as envs from vllm.config.compilation import PassConfig -from vllm.distributed.parallel_state import get_node_count +from vllm.distributed.device_communicators.all_reduce_utils import ( + FI_MNNVL_ALLREDUCE_MAX_SIZE_MB, +) +from vllm.distributed.parallel_state import _node_count, get_node_count from vllm.logger import init_logger from vllm.platforms import current_platform @@ -23,6 +26,8 @@ # The empirical value for small batch PDL_ADVANCE_LAUNCH_TOKENS = 16 +MiB = 1024 * 1024 + fi_ar_available = False try: import flashinfer.comm as flashinfer_comm # type: ignore[no-redef] @@ -43,6 +48,23 @@ _fi_ar_workspace_groups: dict[int, ProcessGroup] = {} +def _get_tuned_standalone_max_size( + world_size: int, + backend: str, + group: ProcessGroup, +) -> int | None: + if backend != "mnnvl": + return None + capability = current_platform.get_device_capability() + if capability is None: + return None + max_size_mb = FI_MNNVL_ALLREDUCE_MAX_SIZE_MB.get( + (capability.to_int(), world_size, _node_count(group)) + ) + # Tuned cutoffs are exclusive; store the largest accepted size. + return None if max_size_mb is None else int(max_size_mb * MiB) - 1 + + def _create_workspace( backend: str, world_size: int, @@ -165,6 +187,14 @@ def get_fi_ar_workspace( "'trtllm' backend. Please use 'mnnvl' backend instead." ) + if ( + envs.VLLM_ALLREDUCE_USE_FLASHINFER + and (max_size := _get_tuned_standalone_max_size(world_size, backend, group)) + is not None + ): + element_size = torch.empty((), dtype=dtype, device="cpu").element_size() + max_token_num = max(max_token_num, max_size // (hidden_dim * element_size)) + def _get_or_create(be: str): # Reuse the quant workspace if it was already created with the same backend if _fi_ar_quant_workspace is not None and _fi_ar_quant_workspace.backend == be: @@ -346,20 +376,28 @@ def __init__( if self.world_size == 1: return - # Use the same threshold as the allreduce-rms fusion pass - # TODO: tune the threshold - MiB = 1024 * 1024 - max_workspace_size = PassConfig.default_fi_allreduce_fusion_max_size_mb().get( - self.world_size, None + default_max_size_mb = PassConfig.default_fi_allreduce_fusion_max_size_mb().get( + self.world_size ) - if not max_workspace_size: + if not default_max_size_mb: logger.warning( "FlashInfer All Reduce is disabled because it " "is not supported for world_size=%d.", self.world_size, ) return - self.max_workspace_size = max_workspace_size * MiB + + backend, _ = _resolve_fi_ar_backend() + tuned_max_size = _get_tuned_standalone_max_size( + self.world_size, + backend, + self.group, + ) + self.max_workspace_size = ( + tuned_max_size + if tuned_max_size is not None + else int(default_max_size_mb * MiB) + ) self.max_num_tokens = 0 self.disabled = False @@ -394,6 +432,16 @@ def should_use_fi_ar(self, input_tensor: torch.Tensor) -> bool: if len(input_tensor.shape) != 2: return False + if input_tensor.dtype not in ( + torch.float16, + torch.bfloat16, + torch.float32, + ): + return False + + if input_tensor.nbytes > self.max_workspace_size: + return False + num_tokens, hidden_dim = input_tensor.shape if not self.max_num_tokens: element_size = torch.tensor([], dtype=input_tensor.dtype).element_size()