[Distributed] Default FlashInfer allreduce to mnnvl on single node (#47219)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Woosuk Kwon committed
Jun 30, 2026 at 18:35 UTC
a264e419751a9248865aefb1547e0109e246934a
1 file changed
+38
-22
vllm/distributed/device_communicators/flashinfer_all_reduce.py
+38
-22
@@ -92,25 +92,32 @@ def _create_workspace(
92
return workspace
93
94
95
-def _resolve_fi_ar_backend() -> str:
95
+def _resolve_fi_ar_backend() -> tuple[str, bool]:
96
+ """Resolve the flashinfer allreduce backend for the current setup.
97
+
98
+ Returns:
99
+ A ``(backend, allow_trtllm_fallback)`` tuple. ``allow_trtllm_fallback``
100
+ is True only when ``auto`` selects mnnvl for a single node, so that
101
+ workspace creation can fall back to trtllm on single-node topologies
102
+ without NVSwitch multicast support (where mnnvl is unavailable).
103
+ """
104
backend = envs.VLLM_FLASHINFER_ALLREDUCE_BACKEND
105
if backend != "auto":
106
logger.info_once(f"Using flashinfer allreduce backend: {backend}")
99
- return backend
100
-
101
- if get_node_count() > 1: # noqa: SIM108
102
- # Use mnnvl backend for multi-node setup since
103
- # trtllm backend does not support multi-node allreduce
104
- backend = "mnnvl"
105
- else:
106
- # Currently defaulting to trtllm backend for single-node
107
- # setup since mnnvl has issues with cudagraph:
108
- # https://github.com/vllm-project/vllm/issues/35772
109
- # Should switch back to auto when the issue is resolved.
110
- backend = "trtllm"
107
+ return backend, False
108
+
109
+ # Default to mnnvl for both single- and multi-node setups. The mnnvl
110
+ # cudagraph hang that previously forced single-node to trtllm
111
+ # (https://github.com/vllm-project/vllm/issues/35772) was fixed upstream in
112
+ # FlashInfer (>= 0.6.12, vLLM pins 0.6.13), so mnnvl is safe here. trtllm
113
+ # does not support multi-node allreduce, so mnnvl is required there anyway.
114
+ # mnnvl needs NVSwitch multicast; on single-node topologies without it,
115
+ # fall back to trtllm so fused allreduce stays enabled.
116
+ backend = "mnnvl"
117
+ allow_trtllm_fallback = get_node_count() == 1
118
119
logger.info_once(f"Auto-selected flashinfer allreduce backend: {backend}")
113
- return backend
120
+ return backend, allow_trtllm_fallback
121
122
123
def get_fi_ar_workspace(
@@ -132,7 +139,7 @@ def get_fi_ar_workspace(
139
if _fi_ar_workspace is not None:
140
return _fi_ar_workspace
141
135
- backend = _resolve_fi_ar_backend()
142
+ backend, allow_trtllm_fallback = _resolve_fi_ar_backend()
143
144
if get_node_count() > 1 and backend == "trtllm":
145
raise ValueError(
@@ -140,14 +147,23 @@ def get_fi_ar_workspace(
147
"'trtllm' backend. Please use 'mnnvl' backend instead."
148
)
149
143
- # Reuse the quant workspace if it was already created with the same backend
144
- if _fi_ar_quant_workspace is not None and _fi_ar_quant_workspace.backend == backend:
145
- _fi_ar_workspace = _fi_ar_quant_workspace
146
- return _fi_ar_workspace
150
+ def _get_or_create(be: str):
151
+ # Reuse the quant workspace if it was already created with the same backend
152
+ if _fi_ar_quant_workspace is not None and _fi_ar_quant_workspace.backend == be:
153
+ return _fi_ar_quant_workspace
154
+ return _create_workspace(
155
+ be, world_size, rank, max_token_num, hidden_dim, dtype, group
156
+ )
157
+
158
+ _fi_ar_workspace = _get_or_create(backend)
159
+ if _fi_ar_workspace is None and allow_trtllm_fallback and backend != "trtllm":
160
+ logger.warning_once(
161
+ "FlashInfer mnnvl allreduce workspace unavailable (likely no NVSwitch "
162
+ "multicast support); falling back to trtllm backend for single node."
163
+ )
164
+ backend = "trtllm"
165
+ _fi_ar_workspace = _get_or_create(backend)
166
148
- _fi_ar_workspace = _create_workspace(
149
- backend, world_size, rank, max_token_num, hidden_dim, dtype, group
150
- )
167
if _fi_ar_workspace is not None:
168
logger.info_once(
169
"Initialized FlashInfer Allreduce norm fusion workspace "