[Distributed][Perf] Enable FlashInfer MNNVL allreduce RMS quant fusion (#48064)
Signed-off-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>
Mohammad Miadh Angkad committed
Jul 13, 2026 at 15:02 UTC
75fe92a3162a68d74581ec324f04684a752e3ad2
2 files changed
+74
-22
vllm/compilation/passes/fusion/allreduce_rms_fusion.py
+33
-9
@@ -50,6 +50,20 @@ _IR_RMS_NORM_OP = torch.ops.vllm_ir.rms_norm.default
50
_IR_FUSED_ADD_RMS_NORM_OP = torch.ops.vllm_ir.fused_add_rms_norm.default
51
52
53
+def _view_nvfp4_scale_out_for_flashinfer(
54
+ scale_out: torch.Tensor,
55
+) -> torch.Tensor:
56
+ """View vLLM's packed NVFP4 scale buffer as FP8 for FlashInfer."""
57
+ return torch.ops.aten.view.dtype(scale_out, FP8_DTYPE)
58
+
59
+
60
+def _view_flashinfer_nvfp4_scale_out_as_int32(
61
+ scale_out: torch.Tensor,
62
+) -> torch.Tensor:
63
+ """View FlashInfer's NVFP4 scale buffer back as vLLM's int32 format."""
64
+ return torch.ops.aten.view.dtype(scale_out, torch.int32)
65
+
66
+
67
def _norm_input_weight_dtype_match(match: pm.Match) -> bool:
68
"""Prevent fusion when the norm input and weight dtypes differ (e.g. a Gemma
69
fp32 weight.float()+1 gamma), covering rms_norm and fused_add_rms_norm."""
@@ -196,8 +210,8 @@ if flashinfer_comm is not None:
210
curr_device = current_platform.get_device_capability()
211
device_capability = curr_device.to_int() if curr_device is not None else None
212
199
- # Select workspace based on pattern: quant patterns use the
200
- # trtllm quant workspace, non-quant patterns use the primary workspace.
213
+ # Select workspace based on pattern: quant patterns use the quant
214
+ # workspace, non-quant patterns use the primary workspace.
215
is_quant_pattern = pattern_code in (
216
ar_fusion_patterns.kARResidualRMSNormFP8Quant,
217
ar_fusion_patterns.kARResidualRMSNormFP4Quant,
@@ -233,9 +247,9 @@ if flashinfer_comm is not None:
247
residual_out = allreduce_in
248
249
layout_code = None
236
- # layout_code only supported by trtllm backend
237
- if workspace.backend == "trtllm":
238
- # in vllm we only support swizzled layout
250
+ # vLLM quant patterns use swizzled scale-factor layout. Non-quant
251
+ # patterns ignore layout_code.
252
+ if workspace.backend in ("trtllm", "mnnvl"):
253
layout_code = flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4
254
255
flashinfer_comm.allreduce_fusion(
@@ -822,6 +836,7 @@ class AllReduceFusedRMSNormStaticQuantNVFP4Pattern(BasePattern):
836
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
837
residual = torch.zeros_like(input)
838
result_rms = torch.empty_like(input)
839
+ output_scale_fp8 = _view_nvfp4_scale_out_for_flashinfer(output_scale)
840
assert flashinfer_comm is not None, "FlashInfer must be enabled"
841
allreduce = auto_functionalized(
842
flashinfer_trtllm_fused_allreduce_norm,
@@ -829,7 +844,7 @@ class AllReduceFusedRMSNormStaticQuantNVFP4Pattern(BasePattern):
844
residual=residual,
845
norm_out=result_rms,
846
quant_out=quant_result,
832
- scale_out=output_scale,
847
+ scale_out=output_scale_fp8,
848
rms_gamma=weight,
849
rms_eps=self.epsilon,
850
# We don't use norm_out afterwards
@@ -841,7 +856,11 @@ class AllReduceFusedRMSNormStaticQuantNVFP4Pattern(BasePattern):
856
)
857
858
# quant_out, allreduce_output, output_scale
844
- return allreduce[4], allreduce[1], allreduce[5]
859
+ return (
860
+ allreduce[4],
861
+ allreduce[1],
862
+ _view_flashinfer_nvfp4_scale_out_as_int32(allreduce[5]),
863
+ )
864
865
pm.register_replacement(
866
pattern,
@@ -925,6 +944,7 @@ class AllReduceFusedAddRMSNormStaticQuantNVFP4Pattern(BasePattern):
944
weight: torch.Tensor,
945
input_global_scale: torch.Tensor,
946
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
947
+ output_scale_fp8 = _view_nvfp4_scale_out_for_flashinfer(output_scale)
948
assert flashinfer_comm is not None, "FlashInfer must be enabled"
949
allreduce = auto_functionalized(
950
flashinfer_trtllm_fused_allreduce_norm,
@@ -932,7 +952,7 @@ class AllReduceFusedAddRMSNormStaticQuantNVFP4Pattern(BasePattern):
952
residual=residual,
953
norm_out=None,
954
quant_out=quant_result,
935
- scale_out=output_scale,
955
+ scale_out=output_scale_fp8,
956
rms_gamma=weight,
957
rms_eps=self.epsilon,
958
# We don't use norm_out afterwards
@@ -943,7 +963,11 @@ class AllReduceFusedAddRMSNormStaticQuantNVFP4Pattern(BasePattern):
963
**self.allreduce_params.get_trtllm_fused_allreduce_kwargs(),
964
)
965
# quant_out, rms_norm_residual, output_scale
946
- return allreduce[4], allreduce[2], allreduce[5]
966
+ return (
967
+ allreduce[4],
968
+ allreduce[2],
969
+ _view_flashinfer_nvfp4_scale_out_as_int32(allreduce[5]),
970
+ )
971
972
pm.register_replacement(
973
pattern,
vllm/distributed/device_communicators/flashinfer_all_reduce.py
+41
-13
@@ -35,7 +35,9 @@ except ImportError:
35
36
# Workspace for standalone allreduce and non-quant ar+rms fusion
37
_fi_ar_workspace = None
38
-# Extra workspace for quant fusion patterns (only supported by trtllm backend)
38
+# Extra workspace for quant fusion patterns. This may use either the primary
39
+# allreduce backend or a fallback backend when the primary workspace is not
40
+# available on the current topology.
41
_fi_ar_quant_workspace = None
42
43
@@ -189,38 +191,64 @@ def get_fi_ar_quant_workspace(
191
"""
192
Return the allreduce workspace for quant patterns, initializing if needed.
193
192
- Always uses trtllm backend as it is the only one supporting quantization
193
- fusion (FP8/FP4). Returns None for multi-node setups since not supported
194
- by trtllm backend.
194
+ Backend is controlled by VLLM_FLASHINFER_ALLREDUCE_BACKEND env var, matching
195
+ non-quant fusion. With ``auto`` this prefers mnnvl and falls back to trtllm
196
+ only on single-node topologies where mnnvl multicast is unavailable.
197
"""
198
global _fi_ar_quant_workspace
199
if _fi_ar_quant_workspace is not None:
200
return _fi_ar_quant_workspace
201
200
- if get_node_count() > 1:
201
- logger.warning_once(
202
+ backend, allow_trtllm_fallback = _resolve_fi_ar_backend()
203
+
204
+ if get_node_count() > 1 and backend == "trtllm":
205
+ raise ValueError(
206
"Flashinfer allreduce quantization fusion is not supported for "
203
- "multi-node allreduce. Disabling quant fusion."
207
+ "multi-node allreduce with 'trtllm' backend. Please use 'mnnvl' "
208
+ "backend instead."
209
)
205
- return None
210
207
- # Reuse the non-quant workspace if it was already created with trtllm
208
- if _fi_ar_workspace is not None and _fi_ar_workspace.backend == "trtllm":
211
+ # Reuse the non-quant workspace if it was already created with the same
212
+ # backend.
213
+ if _fi_ar_workspace is not None and _fi_ar_workspace.backend == backend:
214
+ _fi_ar_quant_workspace = _fi_ar_workspace
215
+ return _fi_ar_quant_workspace
216
+
217
+ if (
218
+ _fi_ar_workspace is not None
219
+ and _fi_ar_workspace.backend == "trtllm"
220
+ and allow_trtllm_fallback
221
+ and backend != "trtllm"
222
+ ):
223
_fi_ar_quant_workspace = _fi_ar_workspace
224
return _fi_ar_quant_workspace
225
226
_fi_ar_quant_workspace = _create_workspace(
213
- "trtllm", world_size, rank, max_token_num, hidden_dim, dtype, group
227
+ backend, world_size, rank, max_token_num, hidden_dim, dtype, group
228
)
229
+ if _fi_ar_quant_workspace is None and allow_trtllm_fallback and backend != "trtllm":
230
+ logger.warning_once(
231
+ "FlashInfer mnnvl allreduce quantization fusion workspace unavailable "
232
+ "(likely no NVSwitch multicast support); falling back to trtllm "
233
+ "backend for single node."
234
+ )
235
+ backend = "trtllm"
236
+ if _fi_ar_workspace is not None and _fi_ar_workspace.backend == backend:
237
+ _fi_ar_quant_workspace = _fi_ar_workspace
238
+ else:
239
+ _fi_ar_quant_workspace = _create_workspace(
240
+ backend, world_size, rank, max_token_num, hidden_dim, dtype, group
241
+ )
242
+
243
if _fi_ar_quant_workspace is not None:
244
logger.info_once(
245
"Initialized FlashInfer Allreduce norm quantization "
218
- "fusion workspace with backend=trtllm"
246
+ f"fusion workspace with backend={backend}"
247
)
248
else:
249
logger.warning_once(
250
"Failed to initialize FlashInfer Allreduce norm quantization "
223
- "fusion workspace with backend=trtllm"
251
+ f"fusion workspace with backend={backend}"
252
)
253
254
return _fi_ar_quant_workspace