@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 71dfce6aa6
Fork 0

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 7, 17:45 UTC

[Kernel] Refactor FlashInfer allreduce for mnnvl backend (#34109)

Signed-off-by: hjjq <50634613+hjjq@users.noreply.github.com> Signed-off-by: wzhao18 <wzhao18.sz@gmail.com> Co-authored-by: wzhao18 <wzhao18.sz@gmail.com> Co-authored-by: Wei Zhao <51183510+wzhao18@users.noreply.github.com>

Hanjie Qiu committed Feb 25, 2026 at 19:17 UTC 71dfce6aa6cc14d016154b4e3fd8cc40c05415f9
7 files changed +592 -179
benchmarks/kernels/benchmark_device_communicators.py
+88 -25
@@ -30,6 +30,9 @@ import torch.distributed as dist
30 from torch.distributed import ProcessGroup
31
32 from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce
33 +from vllm.distributed.device_communicators.flashinfer_all_reduce import (
34 + FlashInferAllReduce,
35 +)
36 from vllm.distributed.device_communicators.pynccl import (
37 PyNcclCommunicator,
38 register_nccl_symmetric_ops,
@@ -44,7 +47,7 @@ from vllm.utils.argparse_utils import FlexibleArgumentParser
47 logger = init_logger(__name__)
48
49 # Default sequence lengths to benchmark
47 -DEFAULT_SEQUENCE_LENGTHS = [128, 512, 1024, 2048, 4096, 8192]
50 +DEFAULT_SEQUENCE_LENGTHS = [16, 64, 128, 512, 1024, 2048, 4096, 8192]
51
52 # Fixed hidden size and dtype for all benchmarks
53 HIDDEN_SIZE = 8192
@@ -81,6 +84,7 @@ class CommunicatorBenchmark:
84 self.symm_mem_comm = None
85 self.symm_mem_comm_multimem = None
86 self.symm_mem_comm_two_shot = None
87 + self.fi_ar_comm = None
88
89 self._init_communicators()
90
@@ -161,6 +165,22 @@ class CommunicatorBenchmark:
165 )
166 self.symm_mem_comm_two_shot = None
167
168 + try:
169 + self.fi_ar_comm = FlashInferAllReduce(
170 + group=self.cpu_group,
171 + device=self.device,
172 + )
173 + if not self.fi_ar_comm.disabled:
174 + logger.info("Rank %s: FlashInferAllReduce initialized", self.rank)
175 + else:
176 + logger.info("Rank %s: FlashInferAllReduce disabled", self.rank)
177 + self.fi_ar_comm = None
178 + except Exception as e:
179 + logger.warning(
180 + "Rank %s: Failed to initialize FlashInferAllReduce: %s", self.rank, e
181 + )
182 + self.fi_ar_comm = None
183 +
184 def benchmark_allreduce(
185 self, sequence_length: int, num_warmup: int, num_trials: int
186 ) -> dict[str, float]:
@@ -180,7 +200,8 @@ class CommunicatorBenchmark:
200 lambda t, c=comm: c.custom_all_reduce(t),
201 lambda t, c=comm: c.should_custom_ar(t),
202 comm.capture(),
183 - "1stage", # env variable value
203 + {"VLLM_CUSTOM_ALLREDUCE_ALGO": "1stage"},
204 + None, # no destroy function
205 )
206 )
207 # CustomAllreduce two-shot
@@ -190,7 +211,8 @@ class CommunicatorBenchmark:
211 lambda t, c=comm: c.custom_all_reduce(t),
212 lambda t, c=comm: c.should_custom_ar(t),
213 comm.capture(),
193 - "2stage", # env variable value
214 + {"VLLM_CUSTOM_ALLREDUCE_ALGO": "2stage"},
215 + None, # no destroy function
216 )
217 )
218
@@ -202,7 +224,8 @@ class CommunicatorBenchmark:
224 lambda t, c=comm: c.all_reduce(t),
225 lambda t: True, # Always available if initialized
226 nullcontext(),
205 - None, # no env variable needed
227 + {}, # no env variable needed
228 + None, # no destroy function
229 )
230 )
231 communicators.append(
@@ -211,7 +234,8 @@ class CommunicatorBenchmark:
234 lambda t: torch.ops.vllm.all_reduce_symmetric_with_copy(t),
235 lambda t: True, # Always available if initialized
236 nullcontext(),
214 - None, # no env variable needed
237 + {}, # no env variable needed
238 + None, # no destroy function
239 )
240 )
241
@@ -223,7 +247,8 @@ class CommunicatorBenchmark:
247 lambda t, c=comm: c.all_reduce(t),
248 lambda t, c=comm: c.should_use_symm_mem(t),
249 nullcontext(),
226 - None, # no env variable needed
250 + {}, # no env variable needed
251 + None, # no destroy function
252 )
253 )
254
@@ -235,29 +260,67 @@ class CommunicatorBenchmark:
260 lambda t, c=comm: c.all_reduce(t),
261 lambda t, c=comm: c.should_use_symm_mem(t),
262 nullcontext(),
238 - None, # no env variable needed
263 + {}, # no env variable needed
264 + None, # no destroy function needed
265 )
266 )
267
242 - # Benchmark each communicator
243 - for name, allreduce_fn, should_use_fn, context, env_var in communicators:
244 - # Set environment variable if needed
245 - if env_var is not None:
246 - os.environ["VLLM_CUSTOM_ALLREDUCE_ALGO"] = env_var
247 - else:
248 - # Clear the environment variable to avoid interference
249 - os.environ.pop("VLLM_CUSTOM_ALLREDUCE_ALGO", None)
250 -
251 - latency = self.benchmark_allreduce_single(
252 - sequence_length,
253 - allreduce_fn,
254 - should_use_fn,
255 - context,
256 - num_warmup,
257 - num_trials,
268 + if self.fi_ar_comm is not None:
269 + comm = self.fi_ar_comm
270 + communicators.append(
271 + (
272 + "flashinfer_trtllm",
273 + lambda t, c=comm: c.all_reduce(t),
274 + lambda t, c=comm: c.should_use_fi_ar(t),
275 + nullcontext(),
276 + {"VLLM_FLASHINFER_ALLREDUCE_BACKEND": "trtllm"},
277 + lambda c=comm: c.destroy(),
278 + )
279 )
259 - if latency is not None:
260 - results[name] = latency
280 + communicators.append(
281 + (
282 + "flashinfer_mnnvl",
283 + lambda t, c=comm: c.all_reduce(t),
284 + lambda t, c=comm: c.should_use_fi_ar(t),
285 + nullcontext(),
286 + {"VLLM_FLASHINFER_ALLREDUCE_BACKEND": "mnnvl"},
287 + lambda c=comm: c.destroy(),
288 + )
289 + )
290 +
291 + # Benchmark each communicator
292 + for (
293 + name,
294 + allreduce_fn,
295 + should_use_fn,
296 + context,
297 + env_dict,
298 + destroy_fn,
299 + ) in communicators:
300 + # Save original values and apply new environment variables
301 + saved_env = {key: os.environ.get(key) for key in env_dict}
302 + for key, value in env_dict.items():
303 + os.environ[key] = value
304 + try:
305 + latency = self.benchmark_allreduce_single(
306 + sequence_length,
307 + allreduce_fn,
308 + should_use_fn,
309 + context,
310 + num_warmup,
311 + num_trials,
312 + )
313 + if latency is not None:
314 + results[name] = latency
315 + finally:
316 + if destroy_fn is not None:
317 + destroy_fn()
318 + # Restore environment variables to their original state
319 + for key, original_value in saved_env.items():
320 + if original_value is None:
321 + os.environ.pop(key, None)
322 + else:
323 + os.environ[key] = original_value
324
325 return results
326
benchmarks/kernels/benchmark_fused_collective.py
+114 -96
@@ -5,8 +5,11 @@
5 Benchmark for FlashInfer fused collective operations vs standard operations.
6
7 This benchmark compares:
8 -1. FlashInfer's allreduce_fusion (fused allreduce + rmsnorm + optional quant)
9 -2. Standard tensor_model_parallel_all_reduce + separate rmsnorm/quant operations
8 +1. FlashInfer's allreduce_fusion with trtllm backend
9 + (fused allreduce + rmsnorm + optional FP8/FP4 quant)
10 +2. FlashInfer's allreduce_fusion with mnnvl backend
11 + (fused allreduce + rmsnorm only, no quantization support)
12 +3. Standard tensor_model_parallel_all_reduce + separate rmsnorm/quant operations
13
14 Usage with torchrun:
15 torchrun --nproc_per_node=2 benchmark_fused_collective.py
@@ -48,8 +51,12 @@ SCALED_FP4_QUANT_OP = torch.ops._C.scaled_fp4_quant
51 logger = init_logger(__name__)
52
53 # Try to import FlashInfer
54 +TorchDistBackend = None
55 try:
56 import flashinfer.comm as flashinfer_comm # type: ignore
57 + from flashinfer.comm.mnnvl import ( # type: ignore
58 + TorchDistBackend,
59 + )
60
61 if not (
62 hasattr(flashinfer_comm, "allreduce_fusion")
@@ -74,11 +81,15 @@ _FI_MAX_SIZES = {
81 8: 64 * MiB, # 64MB
82 }
83
77 -# Global workspace tensor for FlashInfer
78 -_FI_WORKSPACE = None
84 +# Global workspace tensors for FlashInfer (keyed by backend name)
85 +_FI_WORKSPACES: dict = {}
86 +
87 +# Backends to benchmark
88 +FLASHINFER_BACKENDS = ["trtllm", "mnnvl"]
89
90
91 def setup_flashinfer_workspace(
92 + backend: str,
93 world_size: int,
94 rank: int,
95 hidden_dim: int,
@@ -86,41 +97,54 @@ def setup_flashinfer_workspace(
97 dtype: torch.dtype,
98 ):
99 """Setup FlashInfer workspace for fused allreduce operations."""
89 - global _FI_WORKSPACE
100 + global FI_WORKSPACES
101
102 if flashinfer_comm is None:
92 - return None, None
103 + return None
104
105 if world_size not in _FI_MAX_SIZES:
106 logger.warning("FlashInfer not supported for world size %s", world_size)
96 - return None, None
107 + return None
108
109 try:
110 + kwargs = {}
111 + if TorchDistBackend is not None:
112 + kwargs["comm_backend"] = TorchDistBackend(group=dist.group.WORLD)
113 +
114 workspace = flashinfer_comm.create_allreduce_fusion_workspace(
100 - backend="trtllm",
115 + backend=backend,
116 world_size=world_size,
117 rank=rank,
118 max_token_num=max_token_num,
119 hidden_dim=hidden_dim,
120 dtype=dtype,
121 + **kwargs,
122 )
123
108 - _FI_WORKSPACE = workspace
124 + _FI_WORKSPACES[backend] = workspace
125 return workspace
126 except Exception as e:
111 - logger.error("Failed to setup FlashInfer workspace: %s", e)
127 + logger.error(
128 + "Failed to setup FlashInfer workspace (backend=%s): %s", backend, e
129 + )
130 return None
131
132
115 -def cleanup_flashinfer_workspace(workspace):
116 - """Cleanup FlashInfer workspace."""
117 - if flashinfer_comm is None or workspace is None:
133 +def cleanup_flashinfer_workspaces():
134 + """Cleanup all FlashInfer workspaces."""
135 + if flashinfer_comm is None:
136 return
137
120 - try:
121 - workspace.destroy()
122 - except Exception as e:
123 - logger.error("Failed to cleanup FlashInfer workspace: %s", e)
138 + for backend, workspace in _FI_WORKSPACES.items():
139 + try:
140 + workspace.destroy()
141 + except Exception as e:
142 + logger.error(
143 + "Failed to cleanup FlashInfer workspace (backend=%s): %s",
144 + backend,
145 + e,
146 + )
147 + _FI_WORKSPACES.clear()
148
149
150 class FlashInferFusedAllReduceParams:
@@ -134,7 +158,7 @@ class FlashInferFusedAllReduceParams:
158 self.fp32_acc = True
159 self.max_token_num = max_token_num
160
137 - def get_trtllm_fused_allreduce_kwargs(self):
161 + def get_flashinfer_fused_allreduce_kwargs(self):
162 return {
163 "launch_with_pdl": self.launch_with_pdl,
164 "fp32_acc": self.fp32_acc,
@@ -147,11 +171,12 @@ def flashinfer_fused_allreduce_rmsnorm(
171 rms_gamma: torch.Tensor,
172 rms_eps: float,
173 allreduce_params: "FlashInferFusedAllReduceParams",
174 + workspace: object,
175 use_oneshot: bool,
176 norm_out: torch.Tensor | None = None,
177 ):
178 """FlashInfer fused allreduce + rmsnorm operation."""
154 - if flashinfer_comm is None or _FI_WORKSPACE is None:
179 + if flashinfer_comm is None or workspace is None:
180 raise RuntimeError("FlashInfer not available or workspace not initialized")
181
182 if norm_out is None:
@@ -160,9 +185,13 @@ def flashinfer_fused_allreduce_rmsnorm(
185 else:
186 residual_out = input_tensor
187
188 + layout_code = None
189 + if workspace.backend == "trtllm":
190 + layout_code = flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4
191 +
192 flashinfer_comm.allreduce_fusion(
193 input=input_tensor,
165 - workspace=_FI_WORKSPACE,
194 + workspace=workspace,
195 pattern=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm,
196 residual_in=residual,
197 residual_out=residual_out,
@@ -171,10 +200,10 @@ def flashinfer_fused_allreduce_rmsnorm(
200 rms_eps=rms_eps,
201 quant_out=None,
202 scale_out=None,
174 - layout_code=flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4,
203 + layout_code=layout_code,
204 scale_factor=None,
205 use_oneshot=use_oneshot,
177 - **allreduce_params.get_trtllm_fused_allreduce_kwargs(),
206 + **allreduce_params.get_flashinfer_fused_allreduce_kwargs(),
207 )
208
209
@@ -185,12 +214,16 @@ def flashinfer_fused_allreduce_rmsnorm_fp8_quant(
214 rms_eps: float,
215 scale_factor: torch.Tensor,
216 allreduce_params: FlashInferFusedAllReduceParams,
217 + workspace: object,
218 use_oneshot: bool = True,
219 norm_out: torch.Tensor | None = None,
220 quant_out: torch.Tensor | None = None,
221 ):
192 - """FlashInfer fused allreduce + rmsnorm + FP8 quantization."""
193 - if flashinfer_comm is None or _FI_WORKSPACE is None:
222 + """FlashInfer fused allreduce + rmsnorm + FP8 quantization.
223 +
224 + Note: Only supported by the trtllm backend.
225 + """
226 + if flashinfer_comm is None or workspace is None:
227 raise RuntimeError("FlashInfer not available or workspace not initialized")
228
229 if norm_out is None:
@@ -201,7 +234,7 @@ def flashinfer_fused_allreduce_rmsnorm_fp8_quant(
234
235 flashinfer_comm.allreduce_fusion(
236 input=input_tensor,
204 - workspace=_FI_WORKSPACE,
237 + workspace=workspace,
238 pattern=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP8Quant,
239 residual_in=residual,
240 residual_out=residual_out,
@@ -213,7 +246,7 @@ def flashinfer_fused_allreduce_rmsnorm_fp8_quant(
246 layout_code=flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4,
247 scale_factor=scale_factor,
248 use_oneshot=use_oneshot,
216 - **allreduce_params.get_trtllm_fused_allreduce_kwargs(),
249 + **allreduce_params.get_flashinfer_fused_allreduce_kwargs(),
250 )
251
252
@@ -224,13 +257,17 @@ def flashinfer_fused_allreduce_rmsnorm_fp4_quant(
257 rms_eps: float,
258 input_global_scale: torch.Tensor,
259 allreduce_params: FlashInferFusedAllReduceParams,
260 + workspace: object,
261 quant_out: torch.Tensor,
262 use_oneshot: bool,
263 output_scale: torch.Tensor,
264 norm_out: torch.Tensor | None = None,
265 ):
232 - """FlashInfer fused allreduce + rmsnorm + FP4 quantization."""
233 - if flashinfer_comm is None or _FI_WORKSPACE is None:
266 + """FlashInfer fused allreduce + rmsnorm + FP4 quantization.
267 +
268 + Note: Only supported by the trtllm backend.
269 + """
270 + if flashinfer_comm is None or workspace is None:
271 raise RuntimeError("FlashInfer not available or workspace not initialized")
272
273 if norm_out is None:
@@ -241,7 +278,7 @@ def flashinfer_fused_allreduce_rmsnorm_fp4_quant(
278
279 flashinfer_comm.allreduce_fusion(
280 input=input_tensor,
244 - workspace=_FI_WORKSPACE,
281 + workspace=workspace,
282 pattern=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant,
283 residual_in=residual,
284 residual_out=residual_out,
@@ -253,7 +290,7 @@ def flashinfer_fused_allreduce_rmsnorm_fp4_quant(
290 layout_code=flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4,
291 scale_factor=input_global_scale,
292 use_oneshot=use_oneshot,
256 - **allreduce_params.get_trtllm_fused_allreduce_kwargs(),
293 + **allreduce_params.get_flashinfer_fused_allreduce_kwargs(),
294 )
295
296
@@ -386,13 +423,16 @@ def run_benchmarks(
423 dtype: torch.dtype,
424 use_residual: bool,
425 allreduce_params: FlashInferFusedAllReduceParams | None,
426 + workspaces: dict,
427 quant_modes: set[str],
428 no_oneshot: bool,
429 ):
430 """Run all benchmarks for given configuration.
431
432 Args:
395 - quant_mode: "none", "fp8_only", "fp4_only", or "all"
433 + allreduce_params: Shared parameters for FlashInfer fused allreduce.
434 + workspaces: Dict mapping backend name ("trtllm", "mnnvl") to workspace.
435 + quant_modes: Set of quantization modes: "none", "fp8", "fp4".
436 """
437 (
438 input_tensor,
@@ -454,10 +494,11 @@ def run_benchmarks(
494 logger.error("Standard AllReduce+RMSNorm Native Compiled failed: %s", e)
495 results["standard_allreduce_rmsnorm_native_compiled"] = float("inf")
496
457 - # FlashInfer Fused AllReduce + RMSNorm Oneshot/Twoshot
458 - if flashinfer_comm is not None and allreduce_params is not None:
497 + # FlashInfer Fused AllReduce + RMSNorm (all backends)
498 + for backend, workspace in workspaces.items():
499 for use_oneshot in use_oneshot_options:
500 suffix = "_oneshot" if use_oneshot else "_twoshot"
501 + key = f"flashinfer_{backend}_fused_allreduce_rmsnorm{suffix}"
502 try:
503 time_ms = benchmark_operation(
504 flashinfer_fused_allreduce_rmsnorm,
@@ -467,14 +508,17 @@ def run_benchmarks(
508 rms_gamma=rms_gamma,
509 rms_eps=rms_eps,
510 allreduce_params=allreduce_params,
511 + workspace=workspace,
512 use_oneshot=use_oneshot,
513 )
472 - results[f"flashinfer_fused_allreduce_rmsnorm{suffix}"] = time_ms
514 + results[key] = time_ms
515 except Exception as e:
474 - logger.error("FlashInfer Fused AllReduce+RMSNorm failed: %s", e)
475 - results[f"flashinfer_fused_allreduce_rmsnorm{suffix}"] = float(
476 - "inf"
516 + logger.error(
517 + "FlashInfer (%s) Fused AllReduce+RMSNorm failed: %s",
518 + backend,
519 + e,
520 )
521 + results[key] = float("inf")
522
523 if "fp8" in quant_modes:
524 # Standard AllReduce + RMSNorm + FP8 Quant
@@ -540,10 +584,12 @@ def run_benchmarks(
584 "inf"
585 )
586
543 - # FlashInfer Fused AllReduce + RMSNorm + FP8 Quant Oneshot
544 - if flashinfer_comm is not None and allreduce_params is not None:
587 + # FlashInfer Fused AllReduce + RMSNorm + FP8 Quant (trtllm only)
588 + if "trtllm" in workspaces:
589 + trtllm_ws = workspaces["trtllm"]
590 for use_oneshot in use_oneshot_options:
591 suffix = "_oneshot" if use_oneshot else "_twoshot"
592 + key = f"flashinfer_trtllm_fused_allreduce_rmsnorm_fp8_quant{suffix}"
593 try:
594 time_ms = benchmark_operation(
595 flashinfer_fused_allreduce_rmsnorm_fp8_quant,
@@ -555,19 +601,16 @@ def run_benchmarks(
601 scale_factor=scale_fp8,
602 quant_out=quant_out_fp8,
603 allreduce_params=allreduce_params,
604 + workspace=trtllm_ws,
605 use_oneshot=use_oneshot,
606 )
560 - results[f"flashinfer_fused_allreduce_rmsnorm_fp8_quant{suffix}"] = (
561 - time_ms
562 - )
607 + results[key] = time_ms
608 except Exception as e:
609 logger.error(
565 - "FlashInfer Fused AllReduce+RMSNorm+FP8 Oneshot failed: %s",
610 + "FlashInfer (trtllm) Fused AllReduce+RMSNorm+FP8 failed: %s",
611 e,
612 )
568 - results[f"flashinfer_fused_allreduce_rmsnorm_fp8_quant{suffix}"] = (
569 - float("inf")
570 - )
613 + results[key] = float("inf")
614
615 if "fp4" in quant_modes and current_platform.has_device_capability(100):
616 # Standard AllReduce + RMSNorm + FP4 Quant
@@ -627,10 +670,12 @@ def run_benchmarks(
670 "inf"
671 )
672
630 - # FlashInfer Fused AllReduce + RMSNorm + FP4 Quant Oneshot
631 - if flashinfer_comm is not None and allreduce_params is not None:
673 + # FlashInfer Fused AllReduce + RMSNorm + FP4 Quant (trtllm only)
674 + if "trtllm" in workspaces:
675 + trtllm_ws = workspaces["trtllm"]
676 for use_oneshot in use_oneshot_options:
677 suffix = "_oneshot" if use_oneshot else "_twoshot"
678 + key = f"flashinfer_trtllm_fused_allreduce_rmsnorm_fp4_quant{suffix}"
679 try:
680 time_ms = benchmark_operation(
681 flashinfer_fused_allreduce_rmsnorm_fp4_quant,
@@ -641,49 +686,18 @@ def run_benchmarks(
686 rms_eps=rms_eps,
687 input_global_scale=scale_fp4,
688 allreduce_params=allreduce_params,
689 + workspace=trtllm_ws,
690 quant_out=fp4_quant_out,
691 output_scale=fp4_output_scale,
692 use_oneshot=use_oneshot,
693 )
648 - results[f"flashinfer_fused_allreduce_rmsnorm_fp4_quant{suffix}"] = (
649 - time_ms
650 - )
694 + results[key] = time_ms
695 except Exception as e:
696 logger.error(
653 - "FlashInfer Fused AllReduce+RMSNorm+FP4 Oneshot failed: %s",
697 + "FlashInfer (trtllm) Fused AllReduce+RMSNorm+FP4 failed: %s",
698 e,
699 )
656 - results[f"flashinfer_fused_allreduce_rmsnorm_fp4_quant{suffix}"] = (
657 - float("inf")
658 - )
659 -
660 - # FlashInfer Fused AllReduce + RMSNorm + FP4 Quant Two-shot
661 - if flashinfer_comm is not None and allreduce_params is not None:
662 - try:
663 - time_ms = benchmark_operation(
664 - flashinfer_fused_allreduce_rmsnorm_fp4_quant,
665 - input_tensor,
666 - residual=residual,
667 - norm_out=norm_out,
668 - rms_gamma=rms_gamma,
669 - rms_eps=rms_eps,
670 - input_global_scale=scale_fp4,
671 - allreduce_params=allreduce_params,
672 - quant_out=fp4_quant_out,
673 - output_scale=fp4_output_scale,
674 - use_oneshot=False,
675 - )
676 - results["flashinfer_fused_allreduce_rmsnorm_fp4_quant_twoshot"] = (
677 - time_ms
678 - )
679 - except Exception as e:
680 - logger.error(
681 - "FlashInfer Fused AllReduce+RMSNorm+FP4 Two-shot failed: %s",
682 - e,
683 - )
684 - results["flashinfer_fused_allreduce_rmsnorm_fp4_quant_twoshot"] = float(
685 - "inf"
686 - )
700 + results[key] = float("inf")
701
702 return results
703
@@ -1021,8 +1035,7 @@ def main():
1035
1036 configs = list(itertools.product(args.num_tokens, dtypes, residual_options))
1037
1024 - # Setup FlashInfer workspace if available
1025 - workspace = None
1038 + # Setup FlashInfer workspaces for all backends
1039 allreduce_params = None
1040
1041 if flashinfer_comm is not None:
@@ -1037,15 +1050,17 @@ def main():
1050 args.hidden_dim * max_element_size
1051 )
1052
1040 - workspace = setup_flashinfer_workspace(
1041 - world_size,
1042 - rank,
1043 - args.hidden_dim,
1044 - max_num_token,
1045 - dtype=workspace_dtype,
1046 - )
1053 + for backend in FLASHINFER_BACKENDS:
1054 + setup_flashinfer_workspace(
1055 + backend=backend,
1056 + world_size=world_size,
1057 + rank=rank,
1058 + hidden_dim=args.hidden_dim,
1059 + max_token_num=max_num_token,
1060 + dtype=workspace_dtype,
1061 + )
1062
1048 - if workspace is not None:
1063 + if _FI_WORKSPACES:
1064 allreduce_params = FlashInferFusedAllReduceParams(
1065 max_token_num=max_num_token,
1066 )
@@ -1071,6 +1086,7 @@ def main():
1086 dtype,
1087 use_residual,
1088 allreduce_params,
1089 + workspaces=_FI_WORKSPACES,
1090 quant_modes=quant_modes,
1091 no_oneshot=args.no_oneshot,
1092 )
@@ -1109,11 +1125,13 @@ def main():
1125
1126 finally:
1127 # Cleanup
1112 - if workspace is not None:
1113 - cleanup_flashinfer_workspace(workspace)
1128 + cleanup_flashinfer_workspaces()
1129
1130 dist.barrier()
1131
1132
1133 if __name__ == "__main__":
1119 - main()
1134 + from vllm.config import VllmConfig, set_current_vllm_config
1135 +
1136 + with set_current_vllm_config(VllmConfig()):
1137 + main()
tests/compile/passes/distributed/test_fusion_all_reduce.py
+9 -1
@@ -142,7 +142,6 @@ class TestAllReduceFusedAddRMSNormStaticQuantFP4Model(torch.nn.Module):
142 *(scaled_fp4_quant(w, wg) for w, wg in zip(self.w, wgscale))
143 )
144 self.wq, self.wscale = list(wq_gen), list(wscale_gen)
145 - print(f"{self.wq=}, {self.wscale=}")
145
146 def forward(self, hidden_states):
147 # avoid having graph input be an arg to a pattern directly
@@ -199,6 +198,7 @@ class TestAllReduceFusedAddRMSNormStaticQuantFP4Model(torch.nn.Module):
198 @pytest.mark.parametrize("hidden_size", [64])
199 @pytest.mark.parametrize("dtype", [torch.bfloat16])
200 @pytest.mark.parametrize("enable_rms_norm_custom_op", [True, False])
201 +@pytest.mark.parametrize("flashinfer_allreduce_backend", ["trtllm", "mnnvl"])
202 @pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA")
203 @pytest.mark.skipif(
204 not find_spec("flashinfer")
@@ -215,6 +215,7 @@ def test_all_reduce_fusion_pass_replace(
215 dtype: torch.dtype,
216 enable_rms_norm_custom_op,
217 enable_quant_fp8_custom_op,
218 + flashinfer_allreduce_backend,
219 ):
220 num_processes = 2
221 if (
@@ -238,6 +239,7 @@ def test_all_reduce_fusion_pass_replace(
239 dtype,
240 enable_rms_norm_custom_op,
241 enable_quant_fp8_custom_op,
242 + flashinfer_allreduce_backend,
243 ),
244 nprocs=nprocs,
245 )
@@ -255,6 +257,7 @@ def all_reduce_fusion_pass_on_test_model(
257 dtype: torch.dtype,
258 enable_rms_norm_custom_op,
259 enable_quant_fp8_custom_op,
260 + flashinfer_allreduce_backend,
261 ):
262 set_random_seed(0)
263
@@ -270,6 +273,7 @@ def all_reduce_fusion_pass_on_test_model(
273 "WORLD_SIZE": str(world_size),
274 "MASTER_ADDR": "localhost",
275 "MASTER_PORT": "12345",
276 + "VLLM_FLASHINFER_ALLREDUCE_BACKEND": flashinfer_allreduce_backend,
277 }
278 )
279
@@ -317,6 +321,10 @@ def all_reduce_fusion_pass_on_test_model(
321 compiled_model = torch.compile(model, backend=backend)
322 compiled_model(hidden_states)
323
324 + results_unfused = model(hidden_states)
325 + results_fused = compiled_model(hidden_states)
326 + torch.testing.assert_close(results_unfused, results_fused, atol=1e-2, rtol=1e-2)
327 +
328 assert all_reduce_fusion_pass.matched_count == 4, (
329 f"{all_reduce_fusion_pass.matched_count=}"
330 )
vllm/compilation/passes/fusion/allreduce_rms_fusion.py
+88 -56
@@ -22,7 +22,9 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
22 kFp8StaticTensorSym,
23 )
24 from vllm.platforms import current_platform
25 -from vllm.utils.torch_utils import direct_register_custom_op
25 +from vllm.utils.torch_utils import (
26 + direct_register_custom_op,
27 +)
28
29 from ..inductor_pass import enable_fake_mode
30 from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
@@ -44,8 +46,6 @@ if find_spec("flashinfer"):
46 except ImportError:
47 pass
48
47 -logger = init_logger(__name__)
48 -
49 if hasattr(torch.ops._C, "scaled_fp4_quant"):
50 STATIC_FP4_QUANT_OP = torch.ops._C.scaled_fp4_quant.default
51
@@ -82,7 +82,16 @@ _FI_ALLREDUCE_ONE_SHOT_MAX_SIZES_MB: dict[int, dict[int, float]] = {
82
83
84 if flashinfer_comm is not None:
85 - _FI_WORKSPACE = None
85 + from vllm.distributed.device_communicators.flashinfer_all_reduce import (
86 + destroy_fi_ar_workspace,
87 + get_fi_ar_quant_workspace,
88 + get_fi_ar_workspace,
89 + initialize_fi_ar_quant_workspace,
90 + initialize_fi_ar_workspace,
91 + )
92 +
93 + ar_fusion_patterns = flashinfer_comm.AllReduceFusionPattern
94 +
95 MiB = 1024 * 1024
96
97 def call_trtllm_fused_allreduce_norm(
@@ -122,9 +131,19 @@ if flashinfer_comm is not None:
131 max_one_shot_size is None or current_tensor_size <= max_one_shot_size * MiB
132 )
133
125 - assert _FI_WORKSPACE is not None, (
126 - "Flashinfer must be enabled when using flashinfer"
134 + # Select workspace based on pattern: quant patterns use the
135 + # trtllm quant workspace, non-quant patterns use the primary workspace.
136 + if pattern_code in (
137 + ar_fusion_patterns.kARResidualRMSNormFP8Quant,
138 + ar_fusion_patterns.kARResidualRMSNormFP4Quant,
139 + ):
140 + workspace = get_fi_ar_quant_workspace()
141 + else:
142 + workspace = get_fi_ar_workspace()
143 + assert workspace is not None, (
144 + "Flashinfer workspace must be initialized when using flashinfer"
145 )
146 + assert flashinfer_comm is not None
147 if norm_out is None:
148 norm_out = allreduce_in
149 residual_out = residual
@@ -133,25 +152,30 @@ if flashinfer_comm is not None:
152 # as flashinfer does not support rms_norm
153 # and allreduce_out together
154 residual_out = allreduce_in
136 - # For the sizes that are smaller than the max size,
137 - # we only use flashinfer one shot allreduce
155 +
156 + layout_code = None
157 + # layout_code only supported by trtllm backend
158 + if workspace.backend == "trtllm":
159 + # in vllm we only support swizzled layout
160 + layout_code = flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4
161 +
162 flashinfer_comm.allreduce_fusion(
163 input=allreduce_in,
140 - workspace=_FI_WORKSPACE,
164 + workspace=workspace,
165 pattern=pattern_code,
142 - residual_in=residual,
166 + launch_with_pdl=launch_with_pdl,
167 + output=None,
168 residual_out=residual_out,
169 norm_out=norm_out,
170 + quant_out=quant_out,
171 + scale_out=scale_out,
172 + residual_in=residual,
173 rms_gamma=rms_gamma,
174 rms_eps=rms_eps,
147 - launch_with_pdl=launch_with_pdl,
175 + scale_factor=scale_factor,
176 + layout_code=layout_code,
177 use_oneshot=use_oneshot,
178 fp32_acc=fp32_acc,
150 - quant_out=quant_out,
151 - scale_out=scale_out,
152 - # in vllm we only support swizzled layout
153 - layout_code=flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4,
154 - scale_factor=scale_factor,
179 )
180
181 def call_trtllm_fused_allreduce_norm_fake(
@@ -729,29 +753,36 @@ class AllReduceFusionPass(VllmPatternMatcherPass):
753 scope="global",
754 )
755
732 - try:
733 - self.workspace = flashinfer_comm.create_allreduce_fusion_workspace(
734 - backend="trtllm",
735 - world_size=self.tp_size,
736 - rank=rank,
737 - max_token_num=self.max_token_num,
738 - hidden_dim=self.hidden_dim,
739 - dtype=self.model_dtype,
740 - )
741 - except RuntimeError as e:
742 - if "multicast" not in str(e).lower():
743 - raise
744 - logger.warning_once(
745 - "AllReduce fusion pass is disabled: flashinfer workspace "
746 - "creation failed: %s. This is expected on GPUs without "
747 - "NVSwitch (e.g., NVLink bridge-only or PCIe topologies). "
748 - "Falling back to non-fused allreduce.",
749 - str(e),
750 - )
751 - return
756 + for workspace_init_fn in [
757 + initialize_fi_ar_workspace,
758 + initialize_fi_ar_quant_workspace,
759 + ]:
760 + try:
761 + workspace_init_fn(
762 + world_size=self.tp_size,
763 + rank=rank,
764 + max_token_num=self.max_token_num,
765 + hidden_dim=self.hidden_dim,
766 + dtype=self.model_dtype,
767 + group=self.group,
768 + )
769 + except Exception as e:
770 + if "multicast" in str(e).lower():
771 + logger.warning(
772 + "AllReduce fusion pass is disabled: flashinfer workspace "
773 + "creation failed: %s. This is expected on GPUs without "
774 + "NVSwitch (e.g., NVLink bridge-only or PCIe topologies). "
775 + "Falling back to non-fused allreduce.",
776 + str(e),
777 + )
778 + else:
779 + logger.warning(
780 + "Failed to initialize FlashInfer All Reduce workspace: %s. "
781 + "AllReduce fusion pass will be disabled.",
782 + e,
783 + )
784 + return
785
753 - global _FI_WORKSPACE
754 - _FI_WORKSPACE = self.workspace
786 self.allreduce_params = FlashInferFusedAllReduceParams(
787 world_size=self.tp_size,
788 max_token_num=self.max_token_num,
@@ -762,32 +793,34 @@ class AllReduceFusionPass(VllmPatternMatcherPass):
793
794 @enable_fake_mode
795 def register_patterns(self) -> None:
796 + supports_quantization = get_fi_ar_quant_workspace() is not None
797 for epsilon in [1e-5, 1e-6]:
766 - AllReduceFusedRMSNormStaticQuantFP8Pattern(
767 - epsilon,
768 - self.model_dtype,
769 - self.device,
770 - self.allreduce_params,
771 - ).register(self.patterns)
772 - AllReduceFusedAddRMSNormStaticQuantFP8Pattern(
773 - epsilon,
774 - self.model_dtype,
775 - self.device,
776 - self.allreduce_params,
777 - ).register(self.patterns)
778 - if current_platform.has_device_capability(100):
779 - AllReduceFusedRMSNormStaticQuantNVFP4Pattern(
798 + if supports_quantization:
799 + AllReduceFusedRMSNormStaticQuantFP8Pattern(
800 epsilon,
801 self.model_dtype,
802 self.device,
803 self.allreduce_params,
804 ).register(self.patterns)
785 - AllReduceFusedAddRMSNormStaticQuantNVFP4Pattern(
805 + AllReduceFusedAddRMSNormStaticQuantFP8Pattern(
806 epsilon,
807 self.model_dtype,
808 self.device,
809 self.allreduce_params,
810 ).register(self.patterns)
811 + if current_platform.has_device_capability(100):
812 + AllReduceFusedRMSNormStaticQuantNVFP4Pattern(
813 + epsilon,
814 + self.model_dtype,
815 + self.device,
816 + self.allreduce_params,
817 + ).register(self.patterns)
818 + AllReduceFusedAddRMSNormStaticQuantNVFP4Pattern(
819 + epsilon,
820 + self.model_dtype,
821 + self.device,
822 + self.allreduce_params,
823 + ).register(self.patterns)
824 AllReduceRMSNormPattern(
825 epsilon,
826 self.model_dtype,
@@ -825,6 +858,5 @@ class AllReduceFusionPass(VllmPatternMatcherPass):
858 def __del__(self) -> None:
859 if getattr(self, "disabled", True):
860 return
828 - if getattr(self, "workspace", None) is not None:
829 - with contextlib.suppress(Exception):
830 - self.workspace.destroy()
861 + with contextlib.suppress(Exception):
862 + destroy_fi_ar_workspace()
vllm/distributed/device_communicators/cuda_communicator.py
+27 -1
@@ -34,19 +34,25 @@ class CudaCommunicator(DeviceCommunicatorBase):
34 # custom allreduce or torch symm mem can be used only by tp
35 use_custom_allreduce = False
36 use_torch_symm_mem = False
37 + use_flashinfer_allreduce = False
38 else:
39 from vllm.distributed.parallel_state import _ENABLE_CUSTOM_ALL_REDUCE
40
41 use_custom_allreduce = _ENABLE_CUSTOM_ALL_REDUCE
42 use_torch_symm_mem = envs.VLLM_ALLREDUCE_USE_SYMM_MEM
43 + use_flashinfer_allreduce = envs.VLLM_ALLREDUCE_USE_FLASHINFER
44
45 self.use_custom_allreduce = use_custom_allreduce
46 self.use_torch_symm_mem = use_torch_symm_mem
47 + self.use_flashinfer_allreduce = use_flashinfer_allreduce
48
49 # lazy import to avoid documentation build error
50 from vllm.distributed.device_communicators.custom_all_reduce import (
51 CustomAllreduce,
52 )
53 + from vllm.distributed.device_communicators.flashinfer_all_reduce import (
54 + FlashInferAllReduce,
55 + )
56 from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
57 from vllm.distributed.device_communicators.quick_all_reduce import (
58 QuickAllReduce,
@@ -65,12 +71,20 @@ class CudaCommunicator(DeviceCommunicatorBase):
71 self.ca_comm: CustomAllreduce | None = None
72 self.qr_comm: QuickAllReduce | None = None
73 self.symm_mem_comm: SymmMemCommunicator | None = None
74 + self.fi_ar_comm: FlashInferAllReduce | None = None
75 +
76 if use_torch_symm_mem and current_platform.is_cuda():
77 self.symm_mem_comm = SymmMemCommunicator(
78 group=self.cpu_group,
79 device=self.device,
80 )
81
82 + if self.use_flashinfer_allreduce and self.world_size > 1:
83 + self.fi_ar_comm = FlashInferAllReduce(
84 + group=self.cpu_group,
85 + device=self.device,
86 + )
87 +
88 if use_custom_allreduce and self.world_size > 1:
89 # Initialize a custom fast all-reduce implementation.
90 self.ca_comm = CustomAllreduce(
@@ -136,7 +150,7 @@ class CudaCommunicator(DeviceCommunicatorBase):
150 out = torch.ops.vllm.all_reduce_symmetric_with_copy(input_)
151 if out is not None:
152 return out
139 - # always try quick reduce first, then custom allreduce,
153 + # always try quick reduce first, then flashinfer, then custom allreduce,
154 # and then pynccl. (quick reduce just for ROCM MI3*)
155 qr_comm = self.qr_comm
156 if (
@@ -147,6 +161,15 @@ class CudaCommunicator(DeviceCommunicatorBase):
161 out = qr_comm.quick_all_reduce(input_)
162 assert out is not None
163 return out
164 + fi_ar_comm = self.fi_ar_comm
165 + if (
166 + fi_ar_comm is not None
167 + and not fi_ar_comm.disabled
168 + and fi_ar_comm.should_use_fi_ar(input_)
169 + ):
170 + out = fi_ar_comm.all_reduce(input_)
171 + assert out is not None
172 + return out
173 ca_comm = self.ca_comm
174 if (
175 ca_comm is not None
@@ -270,6 +293,9 @@ class CudaCommunicator(DeviceCommunicatorBase):
293 self.pynccl_comm = None
294 if self.ca_comm is not None:
295 self.ca_comm = None
296 + if self.fi_ar_comm is not None:
297 + self.fi_ar_comm.destroy()
298 + self.fi_ar_comm = None
299 if self.all2all_manager is not None:
300 self.all2all_manager.destroy()
301 self.all2all_manager = None
vllm/distributed/device_communicators/flashinfer_all_reduce.py new
+252
@@ -0,0 +1,252 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +
5 +import torch
6 +import torch.distributed as dist
7 +from torch.distributed import ProcessGroup
8 +
9 +import vllm.envs as envs
10 +from vllm.config.compilation import PassConfig
11 +from vllm.logger import init_logger
12 +from vllm.platforms import current_platform
13 +
14 +logger = init_logger(__name__)
15 +
16 +fi_ar_available = False
17 +try:
18 + import flashinfer.comm as flashinfer_comm # type: ignore[no-redef]
19 + from flashinfer.comm.mnnvl import (
20 + TorchDistBackend, # type: ignore[import-not-found, no-redef]
21 + )
22 +
23 + fi_ar_available = hasattr(flashinfer_comm, "allreduce_fusion")
24 +except ImportError:
25 + pass
26 +
27 +# Global workspace for standalone allreduce and non-quant ar+rms fusion
28 +_fi_ar_workspace = None
29 +# Extra workspace for quant fusion patterns (only supported by trtllm backend)
30 +# Only created if primary workspace is not already trtllm
31 +_fi_ar_quant_workspace = None
32 +
33 +
34 +def get_fi_ar_workspace():
35 + return _fi_ar_workspace
36 +
37 +
38 +def get_fi_ar_quant_workspace():
39 + return _fi_ar_quant_workspace
40 +
41 +
42 +def initialize_fi_ar_workspace(
43 + world_size: int,
44 + rank: int,
45 + max_token_num: int,
46 + hidden_dim: int,
47 + dtype: torch.dtype,
48 + group: ProcessGroup,
49 +) -> None:
50 + """
51 + Initialize the workspace if not already initialized.
52 +
53 + Currently, this function is called by either the AllReduceFusionPass
54 + or the FlashInferAllReduce backend for standalone allreduce.
55 + If the fusion pass is enabled via
56 + --compilation-config.pass_config.fuse_allreduce_rms=true,
57 + it will create the workspace first, and the standalone backend
58 + will reuse the workspace. Otherwise, the standalone backend will
59 + create the workspace.
60 + """
61 + global _fi_ar_workspace
62 + if _fi_ar_workspace is not None:
63 + return
64 +
65 + backend = envs.VLLM_FLASHINFER_ALLREDUCE_BACKEND
66 + comm_backend = TorchDistBackend(group=group)
67 + _fi_ar_workspace = flashinfer_comm.create_allreduce_fusion_workspace(
68 + backend=backend,
69 + world_size=world_size,
70 + rank=rank,
71 + max_token_num=max_token_num,
72 + hidden_dim=hidden_dim,
73 + dtype=dtype,
74 + comm_backend=comm_backend,
75 + )
76 + assert _fi_ar_workspace is not None
77 + logger.debug(
78 + "Initialized FlashInfer All Reduce workspace: backend=%s, "
79 + "world_size=%d, rank=%d, max_token_num=%d, hidden_dim=%d, dtype=%s",
80 + backend,
81 + world_size,
82 + rank,
83 + max_token_num,
84 + hidden_dim,
85 + dtype,
86 + )
87 +
88 +
89 +def initialize_fi_ar_quant_workspace(
90 + world_size: int,
91 + rank: int,
92 + max_token_num: int,
93 + hidden_dim: int,
94 + dtype: torch.dtype,
95 + group: ProcessGroup,
96 +) -> None:
97 + """
98 + Initialize the workspace used by quantization fusion patterns.
99 +
100 + Currently this always creates a workspace for trtllm backend as only it
101 + supports quantization fusion (FP8/FP4). If the primary workspace
102 + is already trtllm, the quant workspace aliases to it.
103 + """
104 + global _fi_ar_quant_workspace
105 + if _fi_ar_quant_workspace is not None:
106 + return
107 +
108 + # If primary workspace is already trtllm, reuse it
109 + if _fi_ar_workspace is not None and _fi_ar_workspace.backend == "trtllm":
110 + _fi_ar_quant_workspace = _fi_ar_workspace
111 + return
112 +
113 + comm_backend = TorchDistBackend(group=group)
114 + _fi_ar_quant_workspace = flashinfer_comm.create_allreduce_fusion_workspace(
115 + backend="trtllm",
116 + world_size=world_size,
117 + rank=rank,
118 + max_token_num=max_token_num,
119 + hidden_dim=hidden_dim,
120 + dtype=dtype,
121 + comm_backend=comm_backend,
122 + )
123 + assert _fi_ar_quant_workspace is not None
124 + logger.debug(
125 + "Initialized FlashInfer All Reduce workspace: backend=trtllm, "
126 + "world_size=%d, rank=%d, max_token_num=%d, hidden_dim=%d, dtype=%s",
127 + world_size,
128 + rank,
129 + max_token_num,
130 + hidden_dim,
131 + dtype,
132 + )
133 +
134 +
135 +def destroy_fi_ar_workspace():
136 + global _fi_ar_workspace
137 + global _fi_ar_quant_workspace
138 + if (
139 + _fi_ar_quant_workspace is not None
140 + and _fi_ar_quant_workspace is not _fi_ar_workspace
141 + ):
142 + _fi_ar_quant_workspace.destroy()
143 + _fi_ar_quant_workspace = None
144 + if _fi_ar_workspace is not None:
145 + _fi_ar_workspace.destroy()
146 + _fi_ar_workspace = None
147 +
148 +
149 +class FlashInferAllReduce:
150 + def __init__(
151 + self,
152 + group: ProcessGroup,
153 + device: int | str | torch.device,
154 + ):
155 + self.disabled = True
156 +
157 + if not fi_ar_available:
158 + logger.info(
159 + "FlashInfer All Reduce is disabled because flashinfer is not available"
160 + )
161 + return
162 +
163 + if not current_platform.is_cuda():
164 + logger.info(
165 + "FlashInfer All Reduce is disabled because it requires CUDA platform"
166 + )
167 + return
168 +
169 + self.group = group
170 + self.world_size = dist.get_world_size(self.group)
171 + self.rank = dist.get_rank(self.group)
172 + self.device = device
173 + if self.world_size == 1:
174 + return
175 +
176 + # Use the same threshold as the allreduce-rms fusion pass
177 + # TODO: tune the threshold
178 + MiB = 1024 * 1024
179 + max_workspace_size = PassConfig.default_fi_allreduce_fusion_max_size_mb().get(
180 + self.world_size, None
181 + )
182 + if not max_workspace_size:
183 + logger.warning(
184 + "FlashInfer All Reduce is disabled because it "
185 + "is not supported for world_size=%d.",
186 + self.world_size,
187 + )
188 + return
189 + self.max_workspace_size = max_workspace_size * MiB
190 + self.max_num_tokens = 0
191 + self.disabled = False
192 +
193 + def _ensure_workspace(self, hidden_dim: int, dtype: torch.dtype) -> bool:
194 + """Ensure the all reduce workspace is initialized."""
195 + if get_fi_ar_workspace() is not None:
196 + return True
197 + if self.max_num_tokens == 0:
198 + element_size = torch.tensor([], dtype=dtype, device="cpu").element_size()
199 + self.max_num_tokens = self.max_workspace_size // (hidden_dim * element_size)
200 + try:
201 + initialize_fi_ar_workspace(
202 + world_size=self.world_size,
203 + rank=self.rank,
204 + max_token_num=self.max_num_tokens,
205 + hidden_dim=hidden_dim,
206 + dtype=dtype,
207 + group=self.group,
208 + )
209 + return True
210 + except Exception as e:
211 + logger.warning(
212 + "Failed to initialize FlashInfer All Reduce workspace: %s. "
213 + "FlashInfer All Reduce will be disabled.",
214 + e,
215 + )
216 + self.disabled = True
217 + return False
218 +
219 + def should_use_fi_ar(self, input_tensor: torch.Tensor) -> bool:
220 + if self.disabled:
221 + return False
222 +
223 + if not input_tensor.is_cuda:
224 + return False
225 +
226 + if not input_tensor.is_contiguous():
227 + return False
228 +
229 + if len(input_tensor.shape) != 2:
230 + return False
231 +
232 + num_tokens, hidden_dim = input_tensor.shape
233 + if not self.max_num_tokens:
234 + element_size = torch.tensor([], dtype=input_tensor.dtype).element_size()
235 + self.max_num_tokens = self.max_workspace_size // (hidden_dim * element_size)
236 +
237 + if num_tokens > self.max_num_tokens:
238 + return False
239 +
240 + return self._ensure_workspace(hidden_dim, input_tensor.dtype)
241 +
242 + def all_reduce(self, input_tensor: torch.Tensor) -> torch.Tensor:
243 + workspace = get_fi_ar_workspace()
244 + return flashinfer_comm.allreduce_fusion(
245 + input=input_tensor,
246 + workspace=workspace,
247 + pattern=flashinfer_comm.AllReduceFusionPattern.kAllReduce,
248 + )
249 +
250 + def destroy(self):
251 + if not self.disabled:
252 + destroy_fi_ar_workspace()
vllm/envs.py
+14
@@ -168,6 +168,7 @@ if TYPE_CHECKING:
168 VLLM_FLASHINFER_MOE_BACKEND: Literal["throughput", "latency", "masked_gemm"] = (
169 "latency"
170 )
171 + VLLM_FLASHINFER_ALLREDUCE_BACKEND: Literal["auto", "trtllm", "mnnvl"] = "auto"
172 VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE: int = 394 * 1024 * 1024
173 VLLM_XGRAMMAR_CACHE_MB: int = 0
174 VLLM_MSGPACK_ZERO_COPY_THRESHOLD: int = 256
@@ -206,6 +207,7 @@ if TYPE_CHECKING:
207 VLLM_ROCM_FP8_MFMA_PAGE_ATTN: bool = False
208 VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8_CUTLASS: bool = False
209 VLLM_ALLREDUCE_USE_SYMM_MEM: bool = True
210 + VLLM_ALLREDUCE_USE_FLASHINFER: bool = False
211 VLLM_TUNED_CONFIG_FOLDER: str | None = None
212 VLLM_GPT_OSS_SYSTEM_TOOL_MCP_LABELS: set[str] = set()
213 VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT: bool = False
@@ -1290,6 +1292,14 @@ environment_variables: dict[str, Callable[[], Any]] = {
1292 "latency",
1293 ["throughput", "latency", "masked_gemm"],
1294 ),
1295 + # Flashinfer fused allreduce backend.
1296 + # "auto" will default to "mnnvl", which performs mostly same/better than "trtllm".
1297 + # But "mnnvl" backend does not support fuse with quantization.
1298 + "VLLM_FLASHINFER_ALLREDUCE_BACKEND": env_with_choices(
1299 + "VLLM_FLASHINFER_ALLREDUCE_BACKEND",
1300 + "auto",
1301 + ["auto", "trtllm", "mnnvl"],
1302 + ),
1303 # Control the workspace buffer size for the FlashInfer backend.
1304 "VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE": lambda: int(
1305 os.getenv("VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE", str(394 * 1024 * 1024))
@@ -1448,6 +1458,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
1458 "VLLM_ALLREDUCE_USE_SYMM_MEM": lambda: bool(
1459 int(os.getenv("VLLM_ALLREDUCE_USE_SYMM_MEM", "1"))
1460 ),
1461 + # Whether to use FlashInfer allreduce
1462 + "VLLM_ALLREDUCE_USE_FLASHINFER": lambda: bool(
1463 + int(os.getenv("VLLM_ALLREDUCE_USE_FLASHINFER", "0"))
1464 + ),
1465 # Experimental: use this to enable MCP tool calling for non harmony models
1466 "VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT": lambda: bool(
1467 int(os.getenv("VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT", "0"))