A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 23:30 UTC
Fix multi-node allreduce fusion (#38136)
Signed-off-by: wzhao18 <wzhao18.sz@gmail.com> Co-authored-by: root <root@theia0053.lyris.clusters.nvidia.com>
Wei Zhao committed
Mar 26, 2026 at 16:24 UTC
0904b6550d8e2d11295762fabf3f0d7972bd0300
2 files changed
+63
-9
vllm/distributed/device_communicators/flashinfer_all_reduce.py
+61
-2
@@ -13,11 +13,13 @@ from torch.distributed import ProcessGroup
13
14
import vllm.envs as envs
15
from vllm.config.compilation import PassConfig
16
+from vllm.distributed.parallel_state import get_node_count
17
from vllm.logger import init_logger
18
from vllm.platforms import current_platform
19
20
logger = init_logger(__name__)
21
22
+
23
fi_ar_available = False
24
try:
25
import flashinfer.comm as flashinfer_comm # type: ignore[no-redef]
@@ -87,6 +89,27 @@ def _create_workspace(
89
return workspace
90
91
92
+def _resolve_fi_ar_backend() -> str:
93
+ backend = envs.VLLM_FLASHINFER_ALLREDUCE_BACKEND
94
+ if backend != "auto":
95
+ logger.info_once(f"Using flashinfer allreduce backend: {backend}")
96
+ return backend
97
+
98
+ if get_node_count() > 1: # noqa: SIM108
99
+ # Use mnnvl backend for multi-node setup since
100
+ # trtllm backend does not support multi-node allreduce
101
+ backend = "mnnvl"
102
+ else:
103
+ # Currently defaulting to trtllm backend for single-node
104
+ # setup since mnnvl has issues with cudagraph:
105
+ # https://github.com/vllm-project/vllm/issues/35772
106
+ # Should switch back to auto when the issue is resolved.
107
+ backend = "trtllm"
108
+
109
+ logger.info_once(f"Auto-selected flashinfer allreduce backend: {backend}")
110
+ return backend
111
+
112
+
113
def get_fi_ar_workspace(
114
world_size: int,
115
rank: int,
@@ -106,7 +129,13 @@ def get_fi_ar_workspace(
129
if _fi_ar_workspace is not None:
130
return _fi_ar_workspace
131
109
- backend = envs.VLLM_FLASHINFER_ALLREDUCE_BACKEND
132
+ backend = _resolve_fi_ar_backend()
133
+
134
+ if get_node_count() > 1 and backend == "trtllm":
135
+ raise ValueError(
136
+ "Flashinfer allreduce is not supported for multi-node allreduce with "
137
+ "'trtllm' backend. Please use 'mnnvl' backend instead."
138
+ )
139
140
# Reuse the quant workspace if it was already created with the same backend
141
if _fi_ar_quant_workspace is not None and _fi_ar_quant_workspace.backend == backend:
@@ -116,6 +145,17 @@ def get_fi_ar_workspace(
145
_fi_ar_workspace = _create_workspace(
146
backend, world_size, rank, max_token_num, hidden_dim, dtype, group
147
)
148
+ if _fi_ar_workspace is not None:
149
+ logger.info_once(
150
+ "Initialized FlashInfer Allreduce norm fusion workspace "
151
+ f"with backend={backend}"
152
+ )
153
+ else:
154
+ logger.warning_once(
155
+ "Failed to initialize FlashInfer Allreduce norm fusion workspace "
156
+ f"with backend={backend}"
157
+ )
158
+
159
return _fi_ar_workspace
160
161
@@ -131,12 +171,20 @@ def get_fi_ar_quant_workspace(
171
Return the allreduce workspace for quant patterns, initializing if needed.
172
173
Always uses trtllm backend as it is the only one supporting quantization
134
- fusion (FP8/FP4).
174
+ fusion (FP8/FP4). Returns None for multi-node setups since not supported
175
+ by trtllm backend.
176
"""
177
global _fi_ar_quant_workspace
178
if _fi_ar_quant_workspace is not None:
179
return _fi_ar_quant_workspace
180
181
+ if get_node_count() > 1:
182
+ logger.warning_once(
183
+ "Flashinfer allreduce quantization fusion is not supported for "
184
+ "multi-node allreduce. Disabling quant fusion."
185
+ )
186
+ return None
187
+
188
# Reuse the non-quant workspace if it was already created with trtllm
189
if _fi_ar_workspace is not None and _fi_ar_workspace.backend == "trtllm":
190
_fi_ar_quant_workspace = _fi_ar_workspace
@@ -145,6 +193,17 @@ def get_fi_ar_quant_workspace(
193
_fi_ar_quant_workspace = _create_workspace(
194
"trtllm", world_size, rank, max_token_num, hidden_dim, dtype, group
195
)
196
+ if _fi_ar_quant_workspace is not None:
197
+ logger.info_once(
198
+ "Initialized FlashInfer Allreduce norm quantization "
199
+ "fusion workspace with backend=trtllm"
200
+ )
201
+ else:
202
+ logger.warning_once(
203
+ "Failed to initialize FlashInfer Allreduce norm quantization "
204
+ "fusion workspace with backend=trtllm"
205
+ )
206
+
207
return _fi_ar_quant_workspace
208
209
vllm/envs.py
+2
-7
@@ -169,7 +169,7 @@ if TYPE_CHECKING:
169
VLLM_FLASHINFER_MOE_BACKEND: Literal["throughput", "latency", "masked_gemm"] = (
170
"latency"
171
)
172
- VLLM_FLASHINFER_ALLREDUCE_BACKEND: Literal["auto", "trtllm", "mnnvl"] = "trtllm"
172
+ VLLM_FLASHINFER_ALLREDUCE_BACKEND: Literal["auto", "trtllm", "mnnvl"] = "auto"
173
VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE: int = 394 * 1024 * 1024
174
VLLM_XGRAMMAR_CACHE_MB: int = 0
175
VLLM_MSGPACK_ZERO_COPY_THRESHOLD: int = 256
@@ -1305,14 +1305,9 @@ environment_variables: dict[str, Callable[[], Any]] = {
1305
["throughput", "latency", "masked_gemm"],
1306
),
1307
# Flashinfer fused allreduce backend.
1308
- # "auto" will default to "mnnvl", which performs mostly same/better than "trtllm".
1309
- # But "mnnvl" backend does not support fuse with quantization.
1310
- # TODO: Default is "trtllm" right now because "mnnvl" has issues with cudagraph:
1311
- # https://github.com/vllm-project/vllm/issues/35772
1312
- # Should switch back to "auto" if the issue is resolved.
1308
"VLLM_FLASHINFER_ALLREDUCE_BACKEND": env_with_choices(
1309
"VLLM_FLASHINFER_ALLREDUCE_BACKEND",
1315
- "trtllm",
1310
+ "auto",
1311
["auto", "trtllm", "mnnvl"],
1312
),
1313
# Control the workspace buffer size for the FlashInfer backend.