[Perf] Tune FlashInfer all-reduce selection on SM103

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}

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.

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

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()