[Perf] Tune FlashInfer all-reduce selection on SM103 (#53318)
Signed-off-by: Summer Yang <girasoleyang@gmail.com>
Summer Yang committed
Aug 23, 2026 at 22:19 UTC
e6e1af4ca1909a00ffb4405783070fe42432d593
4 files changed
+164
-21
tests/distributed/test_comm_ops.py
+84
@@ -318,6 +318,90 @@ def test_cuda_communicator_checkpoints_flashinfer_workspaces(
318
workspace.checkpoint_restore.assert_called_once_with(group)
319
320
321
+@pytest.mark.parametrize(
322
+ ("backend", "capability", "world_size", "nodes", "expected"),
323
+ [
324
+ ("mnnvl", 103, 4, 1, 80 * flashinfer_all_reduce.MiB - 1),
325
+ ("mnnvl", 103, 8, 2, 64 * flashinfer_all_reduce.MiB - 1),
326
+ ("mnnvl", 103, 16, 4, 8 * flashinfer_all_reduce.MiB - 1),
327
+ ("mnnvl", 103, 2, 1, None),
328
+ ("mnnvl", 103, 12, 3, None),
329
+ ("mnnvl", 103, 8, 1, None),
330
+ ("trtllm", 103, 8, 2, None),
331
+ ("mnnvl", 90, 8, 2, None),
332
+ ],
333
+)
334
+def test_flashinfer_standalone_size_tuning(
335
+ monkeypatch: pytest.MonkeyPatch,
336
+ backend: str,
337
+ capability: int,
338
+ world_size: int,
339
+ nodes: int,
340
+ expected: int | None,
341
+) -> None:
342
+ monkeypatch.setattr(
343
+ flashinfer_all_reduce,
344
+ "current_platform",
345
+ Mock(get_device_capability=lambda: Mock(to_int=lambda: capability)),
346
+ )
347
+ monkeypatch.setattr(flashinfer_all_reduce, "_node_count", lambda _: nodes)
348
+
349
+ assert (
350
+ flashinfer_all_reduce._get_tuned_standalone_max_size(
351
+ world_size, backend, Mock()
352
+ )
353
+ == expected
354
+ )
355
+
356
+
357
+@pytest.mark.parametrize(("enabled", "expected"), [(True, 4681), (False, 128)])
358
+def test_flashinfer_standalone_workspace_size(
359
+ monkeypatch: pytest.MonkeyPatch, enabled: bool, expected: int
360
+) -> None:
361
+ create_workspace = Mock(return_value=Mock(backend="mnnvl"))
362
+ monkeypatch.setattr(
363
+ flashinfer_all_reduce.envs, "VLLM_ALLREDUCE_USE_FLASHINFER", enabled
364
+ )
365
+ monkeypatch.setattr(flashinfer_all_reduce, "_fi_ar_workspace", None)
366
+ monkeypatch.setattr(flashinfer_all_reduce, "_fi_ar_quant_workspace", None)
367
+ monkeypatch.setattr(
368
+ flashinfer_all_reduce,
369
+ "_resolve_fi_ar_backend",
370
+ Mock(return_value=("mnnvl", False)),
371
+ )
372
+ monkeypatch.setattr(flashinfer_all_reduce, "get_node_count", lambda: 2)
373
+ monkeypatch.setattr(
374
+ flashinfer_all_reduce,
375
+ "_get_tuned_standalone_max_size",
376
+ Mock(return_value=64 * flashinfer_all_reduce.MiB - 1),
377
+ )
378
+ monkeypatch.setattr(flashinfer_all_reduce, "_create_workspace", create_workspace)
379
+
380
+ flashinfer_all_reduce.get_fi_ar_workspace(8, 0, 128, 7168, torch.bfloat16, Mock())
381
+
382
+ assert create_workspace.call_args.args[3] == expected
383
+
384
+
385
+def test_flashinfer_all_reduce_precedes_nccl(monkeypatch: pytest.MonkeyPatch) -> None:
386
+ output = torch.empty(2)
387
+ fi_ar_comm = Mock(disabled=False)
388
+ fi_ar_comm.should_use_fi_ar.return_value = True
389
+ fi_ar_comm.all_reduce.return_value = output
390
+ communicator = CudaCommunicator.__new__(CudaCommunicator)
391
+ communicator.fi_ar_comm = fi_ar_comm
392
+ communicator.pynccl_comm = Mock(world_size=8)
393
+ communicator.qr_comm = None
394
+ nccl_selector = Mock(return_value=True)
395
+ monkeypatch.setattr(
396
+ "vllm.distributed.device_communicators.cuda_communicator."
397
+ "should_nccl_symm_mem_allreduce",
398
+ nccl_selector,
399
+ )
400
+
401
+ assert communicator.all_reduce(torch.empty(1)) is output
402
+ nccl_selector.assert_not_called()
403
+
404
+
405
def test_async_intermediate_tensors_lazy_wait() -> None:
406
work = _DummyWork()
407
post_calls = {"n": 0}
vllm/distributed/device_communicators/all_reduce_utils.py
+8
@@ -82,6 +82,14 @@ SYMM_MEM_ALL_REDUCE_MAX_SIZES = {
82
},
83
}
84
85
+# Per-rank input limit for standalone FlashInfer MNNVL all-reduce.
86
+# The key is (compute capability, world size, node count).
87
+FI_MNNVL_ALLREDUCE_MAX_SIZE_MB: dict[tuple[int, int, int], float] = {
88
+ (103, 4, 1): 80,
89
+ (103, 8, 2): 64,
90
+ (103, 16, 4): 8,
91
+}
92
+
93
# NCCL symmetric memory allreduce configuration based on H100 and GB200 benchmarks.
94
# PyNCCL-symm outperforms custom_AR for small and large tensor sizes,
95
# while custom_AR wins for mid-range sizes.
vllm/distributed/device_communicators/cuda_communicator.py
+16
-13
@@ -220,15 +220,17 @@ class CudaCommunicator(DeviceCommunicatorBase):
220
depends on the input tensor.
221
"""
222
all_potential_ar_backends = [
223
+ "FLASHINFER",
224
"NCCL_SYMM_MEM",
225
"QUICK_REDUCE",
225
- "FLASHINFER",
226
"AITER_CUSTOM",
227
"CUSTOM",
228
"SYMM_MEM",
229
"PYNCCL",
230
]
231
enabled_ar_backends: list[str] = []
232
+ if self.fi_ar_comm is not None and not self.fi_ar_comm.disabled:
233
+ enabled_ar_backends.append("FLASHINFER")
234
# Mirror the static preconditions of `should_nccl_symm_mem_allreduce`:
235
# VLLM_BATCH_INVARIANT off, NCCL symm mem enabled, world_size meets
236
# min_world_size, and world_size either has a tuned entry in
@@ -255,8 +257,6 @@ class CudaCommunicator(DeviceCommunicatorBase):
257
enabled_ar_backends.append("NCCL_SYMM_MEM")
258
if self.qr_comm is not None and not self.qr_comm.disabled:
259
enabled_ar_backends.append("QUICK_REDUCE")
258
- if self.fi_ar_comm is not None and not self.fi_ar_comm.disabled:
259
- enabled_ar_backends.append("FLASHINFER")
260
if self.aiter_ar_comm is not None and not self.aiter_ar_comm.disabled:
261
enabled_ar_backends.append("AITER_CUSTOM")
262
if self.ca_comm is not None and not self.ca_comm.disabled:
@@ -276,16 +276,23 @@ class CudaCommunicator(DeviceCommunicatorBase):
276
)
277
278
def all_reduce(self, input_):
279
+ fi_ar_comm = self.fi_ar_comm
280
+ use_fi_ar = (
281
+ fi_ar_comm is not None
282
+ and not fi_ar_comm.disabled
283
+ and fi_ar_comm.should_use_fi_ar(input_)
284
+ )
285
+
286
# since currently we perform copy input -> symm_input -> out-of-place AR
287
# return symm_output, we don't need to check if input is symmetric
281
- if self.pynccl_comm is not None and should_nccl_symm_mem_allreduce(
282
- self.pynccl_comm.world_size, input_
288
+ if (
289
+ self.pynccl_comm is not None
290
+ and not use_fi_ar
291
+ and should_nccl_symm_mem_allreduce(self.pynccl_comm.world_size, input_)
292
):
293
out = torch.ops.vllm.all_reduce_symmetric_with_copy(input_)
294
if out is not None:
295
return out
287
- # always try quick reduce first, then flashinfer, then the AITER or vLLM
288
- # custom allreduce, and then pynccl. (quick reduce just for ROCM MI3*)
296
qr_comm = self.qr_comm
297
if (
298
qr_comm is not None
@@ -295,12 +302,8 @@ class CudaCommunicator(DeviceCommunicatorBase):
302
out = qr_comm.quick_all_reduce(input_)
303
assert out is not None
304
return out
298
- fi_ar_comm = self.fi_ar_comm
299
- if (
300
- fi_ar_comm is not None
301
- and not fi_ar_comm.disabled
302
- and fi_ar_comm.should_use_fi_ar(input_)
303
- ):
305
+ if use_fi_ar:
306
+ assert fi_ar_comm is not None
307
out = fi_ar_comm.all_reduce(input_)
308
assert out is not None
309
return out
vllm/distributed/device_communicators/flashinfer_all_reduce.py
+56
-8
@@ -14,7 +14,10 @@ from torch.distributed import ProcessGroup
14
15
import vllm.envs as envs
16
from vllm.config.compilation import PassConfig
17
-from vllm.distributed.parallel_state import get_node_count
17
+from vllm.distributed.device_communicators.all_reduce_utils import (
18
+ FI_MNNVL_ALLREDUCE_MAX_SIZE_MB,
19
+)
20
+from vllm.distributed.parallel_state import _node_count, get_node_count
21
from vllm.logger import init_logger
22
from vllm.platforms import current_platform
23
@@ -23,6 +26,8 @@ logger = init_logger(__name__)
26
# The empirical value for small batch
27
PDL_ADVANCE_LAUNCH_TOKENS = 16
28
29
+MiB = 1024 * 1024
30
+
31
fi_ar_available = False
32
try:
33
import flashinfer.comm as flashinfer_comm # type: ignore[no-redef]
@@ -43,6 +48,23 @@ _fi_ar_quant_workspace = None
48
_fi_ar_workspace_groups: dict[int, ProcessGroup] = {}
49
50
51
+def _get_tuned_standalone_max_size(
52
+ world_size: int,
53
+ backend: str,
54
+ group: ProcessGroup,
55
+) -> int | None:
56
+ if backend != "mnnvl":
57
+ return None
58
+ capability = current_platform.get_device_capability()
59
+ if capability is None:
60
+ return None
61
+ max_size_mb = FI_MNNVL_ALLREDUCE_MAX_SIZE_MB.get(
62
+ (capability.to_int(), world_size, _node_count(group))
63
+ )
64
+ # Tuned cutoffs are exclusive; store the largest accepted size.
65
+ return None if max_size_mb is None else int(max_size_mb * MiB) - 1
66
+
67
+
68
def _create_workspace(
69
backend: str,
70
world_size: int,
@@ -165,6 +187,14 @@ def get_fi_ar_workspace(
187
"'trtllm' backend. Please use 'mnnvl' backend instead."
188
)
189
190
+ if (
191
+ envs.VLLM_ALLREDUCE_USE_FLASHINFER
192
+ and (max_size := _get_tuned_standalone_max_size(world_size, backend, group))
193
+ is not None
194
+ ):
195
+ element_size = torch.empty((), dtype=dtype, device="cpu").element_size()
196
+ max_token_num = max(max_token_num, max_size // (hidden_dim * element_size))
197
+
198
def _get_or_create(be: str):
199
# Reuse the quant workspace if it was already created with the same backend
200
if _fi_ar_quant_workspace is not None and _fi_ar_quant_workspace.backend == be:
@@ -346,20 +376,28 @@ class FlashInferAllReduce:
376
if self.world_size == 1:
377
return
378
349
- # Use the same threshold as the allreduce-rms fusion pass
350
- # TODO: tune the threshold
351
- MiB = 1024 * 1024
352
- max_workspace_size = PassConfig.default_fi_allreduce_fusion_max_size_mb().get(
353
- self.world_size, None
379
+ default_max_size_mb = PassConfig.default_fi_allreduce_fusion_max_size_mb().get(
380
+ self.world_size
381
)
355
- if not max_workspace_size:
382
+ if not default_max_size_mb:
383
logger.warning(
384
"FlashInfer All Reduce is disabled because it "
385
"is not supported for world_size=%d.",
386
self.world_size,
387
)
388
return
362
- self.max_workspace_size = max_workspace_size * MiB
389
+
390
+ backend, _ = _resolve_fi_ar_backend()
391
+ tuned_max_size = _get_tuned_standalone_max_size(
392
+ self.world_size,
393
+ backend,
394
+ self.group,
395
+ )
396
+ self.max_workspace_size = (
397
+ tuned_max_size
398
+ if tuned_max_size is not None
399
+ else int(default_max_size_mb * MiB)
400
+ )
401
self.max_num_tokens = 0
402
self.disabled = False
403
@@ -394,6 +432,16 @@ class FlashInferAllReduce:
432
if len(input_tensor.shape) != 2:
433
return False
434
435
+ if input_tensor.dtype not in (
436
+ torch.float16,
437
+ torch.bfloat16,
438
+ torch.float32,
439
+ ):
440
+ return False
441
+
442
+ if input_tensor.nbytes > self.max_workspace_size:
443
+ return False
444
+
445
num_tokens, hidden_dim = input_tensor.shape
446
if not self.max_num_tokens:
447
element_size = torch.tensor([], dtype=input_tensor.dtype).element_size()