[Perf] Tune FlashInfer all-reduce selection on SM103 TP8

Prefer standalone FlashInfer all-reduce ahead of NCCL symmetric memory when the existing FlashInfer gate is enabled. Use a topology-keyed strict 64 MiB cutoff for two-node SM103 TP8 MNNVL and size the shared workspace accordingly.

Co-authored-by: OpenAI Codex <codex@openai.com>
Signed-off-by: Summer Yang <girasoleyang@gmail.com>

Author: GirasoleY

tests/distributed/test_flashinfer_all_reduce.py

@@ -0,0 +1,122 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from vllm.distributed.device_communicators import (
+    cuda_communicator as cuda_communicator_module,
+)
+from vllm.distributed.device_communicators import flashinfer_all_reduce
+from vllm.distributed.device_communicators.cuda_communicator import CudaCommunicator
+
+
+def _mock_sm103_tp8_mnnvl(
+    monkeypatch: pytest.MonkeyPatch, tp_group_node_count: int
+) -> object:
+    capability = Mock()
+    capability.to_int.return_value = 103
+    group = object()
+    monkeypatch.setattr(flashinfer_all_reduce, "fi_ar_available", True)
+    monkeypatch.setattr(flashinfer_all_reduce.current_platform, "is_cuda", lambda: True)
+    monkeypatch.setattr(
+        flashinfer_all_reduce.current_platform,
+        "get_device_capability",
+        lambda: capability,
+    )
+    monkeypatch.setattr(flashinfer_all_reduce.dist, "get_world_size", lambda _: 8)
+    monkeypatch.setattr(flashinfer_all_reduce.dist, "get_rank", lambda _: 0)
+    monkeypatch.setattr(
+        flashinfer_all_reduce, "get_node_count", lambda: tp_group_node_count
+    )
+    monkeypatch.setattr(
+        flashinfer_all_reduce, "_node_count", lambda _: tp_group_node_count
+    )
+    monkeypatch.setattr(
+        flashinfer_all_reduce, "_resolve_fi_ar_backend", lambda: ("mnnvl", False)
+    )
+    monkeypatch.setattr(
+        flashinfer_all_reduce.PassConfig,
+        "default_fi_allreduce_fusion_max_size_mb",
+        lambda: {8: 2},
+    )
+    return group
+
+
+@pytest.mark.parametrize(
+    ("tp_group_node_count", "expected_size_mb", "expected_exclusive"),
+    [(2, 64, True), (1, 2, False), (4, 2, False)],
+)
+def test_standalone_threshold_uses_tp_group_topology(
+    monkeypatch: pytest.MonkeyPatch,
+    tp_group_node_count: int,
+    expected_size_mb: int,
+    expected_exclusive: bool,
+) -> None:
+    group = _mock_sm103_tp8_mnnvl(monkeypatch, tp_group_node_count)
+
+    communicator = flashinfer_all_reduce.FlashInferAllReduce(
+        group=group,  # type: ignore[arg-type]
+        device="cuda",
+    )
+
+    assert (
+        communicator.max_workspace_size == expected_size_mb * flashinfer_all_reduce.MiB
+    )
+    assert communicator.max_size_is_exclusive is expected_exclusive
+
+
+@pytest.mark.parametrize(
+    ("standalone_enabled", "expected_max_token_num"),
+    [(True, 4681), (False, 128)],
+)
+def test_workspace_promotion_respects_standalone_flag(
+    monkeypatch: pytest.MonkeyPatch,
+    standalone_enabled: bool,
+    expected_max_token_num: int,
+) -> None:
+    group = _mock_sm103_tp8_mnnvl(monkeypatch, tp_group_node_count=2)
+    monkeypatch.setattr(
+        flashinfer_all_reduce.envs,
+        "VLLM_ALLREDUCE_USE_FLASHINFER",
+        standalone_enabled,
+    )
+
+    max_token_num = flashinfer_all_reduce._promote_max_token_num_for_standalone(
+        max_token_num=128,
+        world_size=8,
+        hidden_dim=7168,
+        dtype=torch.bfloat16,
+        backend="mnnvl",
+        group=group,  # type: ignore[arg-type]
+    )
+
+    assert max_token_num == expected_max_token_num
+
+
+def test_enabled_flashinfer_dispatches_before_nccl(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    input_tensor = torch.empty(1)
+    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.pynccl_comm = Mock(world_size=8, disabled=False)
+    communicator.fi_ar_comm = fi_ar_comm
+    communicator.qr_comm = None
+    communicator.aiter_ar_comm = None
+    communicator.ca_comm = None
+    communicator.symm_mem_comm = None
+    nccl_selector = Mock(return_value=True)
+    monkeypatch.setattr(
+        cuda_communicator_module,
+        "should_nccl_symm_mem_allreduce",
+        nccl_selector,
+    )
+
+    assert communicator.all_reduce(input_tensor) is output
+    nccl_selector.assert_not_called()

vllm/distributed/device_communicators/all_reduce_utils.py

@@ -82,6 +82,12 @@
     },
 }
 
+# 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, 8, 2): 64,
+}
+
 # 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.

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,6 +276,16 @@ def _log_all_reduce_backend_selection(self) -> None:
         )
 
     def all_reduce(self, input_):
+        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_)
+        ):
+            out = fi_ar_comm.all_reduce(input_)
+            assert out is not None
+            return out
+
         # 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(
@@ -284,8 +294,6 @@ def all_reduce(self, 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,15 +303,6 @@ 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_)
-        ):
-            out = fi_ar_comm.all_reduce(input_)
-            assert out is not None
-            return out
         aiter_ar_comm = self.aiter_ar_comm
         if (
             aiter_ar_comm is not None

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,16 @@
 # The empirical value for small batch
 PDL_ADVANCE_LAUNCH_TOKENS = 16
 
+MiB = 1024 * 1024
+
+# Standalone FlashInfer all-reduce supports these dtypes. Fusion patterns have
+# their own, narrower dtype checks.
+FI_ALLREDUCE_SUPPORTED_DTYPES = (
+    torch.float16,
+    torch.bfloat16,
+    torch.float32,
+)
+
 fi_ar_available = False
 try:
     import flashinfer.comm as flashinfer_comm  # type: ignore[no-redef]
@@ -43,6 +56,44 @@
 _fi_ar_workspace_groups: dict[int, ProcessGroup] = {}
 
 
+def _get_tuned_standalone_max_size_mb(
+    world_size: int,
+    backend: str,
+    group: ProcessGroup,
+) -> float | None:
+    if backend != "mnnvl":
+        return None
+    capability = current_platform.get_device_capability()
+    if capability is None:
+        return None
+    return FI_MNNVL_ALLREDUCE_MAX_SIZE_MB.get(
+        (capability.to_int(), world_size, _node_count(group))
+    )
+
+
+def _promote_max_token_num_for_standalone(
+    max_token_num: int,
+    world_size: int,
+    hidden_dim: int,
+    dtype: torch.dtype,
+    backend: str,
+    group: ProcessGroup,
+) -> int:
+    """Make a first workspace allocation safe for standalone all-reduce."""
+    if not envs.VLLM_ALLREDUCE_USE_FLASHINFER:
+        return max_token_num
+    max_size_mb = _get_tuned_standalone_max_size_mb(
+        world_size,
+        backend,
+        group,
+    )
+    if max_size_mb is None:
+        return max_token_num
+    element_size = torch.empty((), dtype=dtype, device="cpu").element_size()
+    standalone_max_token_num = int(max_size_mb * MiB) // (hidden_dim * element_size)
+    return max(max_token_num, standalone_max_token_num)
+
+
 def _create_workspace(
     backend: str,
     world_size: int,
@@ -165,6 +216,15 @@ def get_fi_ar_workspace(
             "'trtllm' backend. Please use 'mnnvl' backend instead."
         )
 
+    max_token_num = _promote_max_token_num_for_standalone(
+        max_token_num,
+        world_size,
+        hidden_dim,
+        dtype,
+        backend,
+        group,
+    )
+
     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 +406,27 @@ 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_mb = _get_tuned_standalone_max_size_mb(
+            self.world_size,
+            backend,
+            self.group,
+        )
+        max_workspace_size_mb = tuned_max_size_mb or default_max_size_mb
+        assert max_workspace_size_mb is not None
+        self.max_workspace_size = int(max_workspace_size_mb * MiB)
+        self.max_size_is_exclusive = tuned_max_size_mb is not None
         self.max_num_tokens = 0
         self.disabled = False
 
@@ -394,6 +461,15 @@ 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 FI_ALLREDUCE_SUPPORTED_DTYPES:
+            return False
+
+        if input_tensor.nbytes > self.max_workspace_size or (
+            self.max_size_is_exclusive
+            and 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()