@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / d17ecc6b19
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 13, 13:15 UTC

[PERF] Allreduce fusion. Support torch native matching. Tuning of the thresholds (#24248)

Signed-off-by: Luka Govedič <lgovedic@redhat.com> Signed-off-by: Luka Govedič <ProExpertProg@users.noreply.github.com> Signed-off-by: ilmarkov <markovilya197@gmail.com> Co-authored-by: Luka Govedič <lgovedic@redhat.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>

Ilya Markov committed Nov 11, 2025 at 00:33 UTC d17ecc6b19b597615893be6c0eb61c9b4a9c9455
6 files changed +1284 -83
.buildkite/test-pipeline.yaml
+2 -2
@@ -463,8 +463,8 @@ steps:
463 - pytest -v -s compile/test_multimodal_compile.py
464 - pytest -v -s compile/piecewise/
465
466 -- label: PyTorch Fullgraph Test # 22min
467 - timeout_in_minutes: 35
466 +- label: PyTorch Fullgraph Test # 27min
467 + timeout_in_minutes: 40
468 mirror_hardwares: [amdexperimental]
469 torch_nightly: true
470 source_file_dependencies:
benchmarks/kernels/benchmark_fused_collective.py new
+1129
@@ -0,0 +1,1129 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +"""
5 +Benchmark for FlashInfer fused collective operations vs standard operations.
6 +
7 +This benchmark compares:
8 +1. FlashInfer's trtllm_allreduce_fusion (fused allreduce + rmsnorm + optional quant)
9 +2. Standard tensor_model_parallel_all_reduce + separate rmsnorm/quant operations
10 +
11 +Usage with torchrun:
12 + torchrun --nproc_per_node=2 benchmark_fused_collective.py
13 +
14 +"""
15 +
16 +import argparse
17 +import itertools
18 +import os
19 +import time
20 +
21 +import pandas as pd
22 +import torch # type: ignore
23 +import torch.distributed as dist # type: ignore
24 +
25 +from vllm.config.vllm import CompilationConfig, VllmConfig, set_current_vllm_config
26 +from vllm.distributed import (
27 + get_tp_group,
28 + tensor_model_parallel_all_reduce,
29 +)
30 +from vllm.distributed.parallel_state import (
31 + graph_capture,
32 + init_distributed_environment,
33 + initialize_model_parallel,
34 +)
35 +from vllm.logger import init_logger
36 +from vllm.model_executor.layers.layernorm import RMSNorm # noqa
37 +from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8 # noqa
38 +from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape # noqa
39 +from vllm.platforms import current_platform # noqa
40 +
41 +RMS_NORM_OP = torch.ops._C.rms_norm
42 +FUSED_ADD_RMS_NORM_OP = torch.ops._C.fused_add_rms_norm
43 +RMS_NORM_STATIC_FP8_QUANT_OP = torch.ops._C.rms_norm_static_fp8_quant
44 +FUSED_ADD_RMS_NORM_STATIC_FP8_QUANT_OP = (
45 + torch.ops._C.fused_add_rms_norm_static_fp8_quant
46 +)
47 +SCALED_FP4_QUANT_OP = torch.ops._C.scaled_fp4_quant
48 +
49 +logger = init_logger(__name__)
50 +
51 +# Try to import FlashInfer
52 +try:
53 + import flashinfer.comm as flashinfer_comm # type: ignore
54 +
55 + if not hasattr(flashinfer_comm, "trtllm_allreduce_fusion"):
56 + flashinfer_comm = None
57 + logger.warning(
58 + "FlashInfer comm module found but missing trtllm_allreduce_fusion"
59 + )
60 +except ImportError:
61 + flashinfer_comm = None
62 + logger.warning("FlashInfer not found, only benchmarking standard operations")
63 +
64 +# Constants
65 +FP8_DTYPE = current_platform.fp8_dtype()
66 +MiB = 1024 * 1024
67 +
68 +# FlashInfer max sizes per world size
69 +# Enable 64MB for 2, 4, 8 world sizes to verify large input sizes
70 +# use --disable-oneshot to disable oneshot mode for very large input sizes
71 +_FI_MAX_SIZES = {
72 + 2: 64 * MiB, # 64MB
73 + 4: 64 * MiB, # 64MB
74 + 8: 64 * MiB, # 64MB
75 +}
76 +
77 +# Global workspace tensor for FlashInfer
78 +_FI_WORKSPACE_TENSOR = None
79 +
80 +
81 +def setup_flashinfer_workspace(
82 + world_size: int,
83 + rank: int,
84 + hidden_dim: int,
85 + max_token_num: int,
86 + use_fp32_lamport: bool = False,
87 +):
88 + """Setup FlashInfer workspace for fused allreduce operations."""
89 + global _FI_WORKSPACE_TENSOR
90 +
91 + if flashinfer_comm is None:
92 + return None, None
93 +
94 + if world_size not in _FI_MAX_SIZES:
95 + logger.warning("FlashInfer not supported for world size %s", world_size)
96 + return None, None
97 +
98 + try:
99 + # Create IPC workspace
100 + ipc_handles, workspace_tensor = (
101 + flashinfer_comm.trtllm_create_ipc_workspace_for_all_reduce_fusion(
102 + tp_rank=rank,
103 + tp_size=world_size,
104 + max_token_num=max_token_num,
105 + hidden_dim=hidden_dim,
106 + group=get_tp_group().device_group,
107 + use_fp32_lamport=use_fp32_lamport,
108 + )
109 + )
110 +
111 + _FI_WORKSPACE_TENSOR = workspace_tensor
112 + return ipc_handles, workspace_tensor
113 + except Exception as e:
114 + logger.error("Failed to setup FlashInfer workspace: %s", e)
115 + return None, None
116 +
117 +
118 +def cleanup_flashinfer_workspace(ipc_handles):
119 + """Cleanup FlashInfer workspace."""
120 + if flashinfer_comm is None or ipc_handles is None:
121 + return
122 +
123 + try:
124 + group = get_tp_group().device_group
125 + flashinfer_comm.trtllm_destroy_ipc_workspace_for_all_reduce(ipc_handles, group)
126 + except Exception as e:
127 + logger.error("Failed to cleanup FlashInfer workspace: %s", e)
128 +
129 +
130 +class FlashInferFusedAllReduceParams:
131 + """Parameters for FlashInfer fused allreduce operations."""
132 +
133 + def __init__(
134 + self,
135 + rank: int,
136 + world_size: int,
137 + use_fp32_lamport: bool = False,
138 + max_token_num: int = 1024,
139 + ):
140 + self.rank = rank
141 + self.world_size = world_size
142 + self.use_fp32_lamport = use_fp32_lamport
143 + self.trigger_completion_at_end = True
144 + self.launch_with_pdl = True
145 + self.fp32_acc = True
146 + self.max_token_num = max_token_num
147 +
148 + def get_trtllm_fused_allreduce_kwargs(self):
149 + return {
150 + "world_rank": self.rank,
151 + "world_size": self.world_size,
152 + "launch_with_pdl": self.launch_with_pdl,
153 + "trigger_completion_at_end": self.trigger_completion_at_end,
154 + "fp32_acc": self.fp32_acc,
155 + }
156 +
157 +
158 +def flashinfer_fused_allreduce_rmsnorm(
159 + input_tensor: torch.Tensor,
160 + residual: torch.Tensor | None,
161 + rms_gamma: torch.Tensor,
162 + rms_eps: float,
163 + allreduce_params: "FlashInferFusedAllReduceParams",
164 + use_oneshot: bool,
165 + norm_out: torch.Tensor | None = None,
166 +):
167 + """FlashInfer fused allreduce + rmsnorm operation."""
168 + if flashinfer_comm is None or _FI_WORKSPACE_TENSOR is None:
169 + raise RuntimeError("FlashInfer not available or workspace not initialized")
170 +
171 + if norm_out is None:
172 + norm_out = input_tensor
173 + residual_out = residual
174 + else:
175 + residual_out = input_tensor
176 +
177 + flashinfer_comm.trtllm_allreduce_fusion(
178 + allreduce_in=input_tensor,
179 + token_num=input_tensor.shape[0],
180 + residual_in=residual,
181 + residual_out=residual_out,
182 + norm_out=norm_out,
183 + rms_gamma=rms_gamma,
184 + rms_eps=rms_eps,
185 + hidden_dim=input_tensor.shape[-1],
186 + workspace_ptrs=_FI_WORKSPACE_TENSOR,
187 + pattern_code=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm,
188 + allreduce_out=None,
189 + quant_out=None,
190 + scale_out=None,
191 + layout_code=flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4,
192 + scale_factor=None,
193 + use_oneshot=use_oneshot,
194 + **allreduce_params.get_trtllm_fused_allreduce_kwargs(),
195 + )
196 +
197 +
198 +def flashinfer_fused_allreduce_rmsnorm_fp8_quant(
199 + input_tensor: torch.Tensor,
200 + residual: torch.Tensor | None,
201 + rms_gamma: torch.Tensor,
202 + rms_eps: float,
203 + scale_factor: torch.Tensor,
204 + allreduce_params: FlashInferFusedAllReduceParams,
205 + use_oneshot: bool = True,
206 + norm_out: torch.Tensor | None = None,
207 + quant_out: torch.Tensor | None = None,
208 +):
209 + """FlashInfer fused allreduce + rmsnorm + FP8 quantization."""
210 + if flashinfer_comm is None or _FI_WORKSPACE_TENSOR is None:
211 + raise RuntimeError("FlashInfer not available or workspace not initialized")
212 +
213 + if norm_out is None:
214 + norm_out = input_tensor
215 + residual_out = residual
216 + else:
217 + residual_out = input_tensor
218 +
219 + flashinfer_comm.trtllm_allreduce_fusion(
220 + allreduce_in=input_tensor,
221 + token_num=input_tensor.shape[0],
222 + residual_in=residual,
223 + residual_out=residual_out,
224 + norm_out=norm_out,
225 + rms_gamma=rms_gamma,
226 + rms_eps=rms_eps,
227 + hidden_dim=input_tensor.shape[-1],
228 + workspace_ptrs=_FI_WORKSPACE_TENSOR,
229 + pattern_code=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP8Quant,
230 + allreduce_out=None,
231 + quant_out=quant_out,
232 + scale_out=None,
233 + layout_code=flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4,
234 + scale_factor=scale_factor,
235 + use_oneshot=use_oneshot,
236 + **allreduce_params.get_trtllm_fused_allreduce_kwargs(),
237 + )
238 +
239 +
240 +def flashinfer_fused_allreduce_rmsnorm_fp4_quant(
241 + input_tensor: torch.Tensor,
242 + residual: torch.Tensor | None,
243 + rms_gamma: torch.Tensor,
244 + rms_eps: float,
245 + input_global_scale: torch.Tensor,
246 + allreduce_params: FlashInferFusedAllReduceParams,
247 + quant_out: torch.Tensor,
248 + use_oneshot: bool,
249 + output_scale: torch.Tensor,
250 + norm_out: torch.Tensor | None = None,
251 +):
252 + """FlashInfer fused allreduce + rmsnorm + FP4 quantization."""
253 + if flashinfer_comm is None or _FI_WORKSPACE_TENSOR is None:
254 + raise RuntimeError("FlashInfer not available or workspace not initialized")
255 +
256 + if norm_out is None:
257 + norm_out = input_tensor
258 + residual_out = residual
259 + else:
260 + residual_out = input_tensor
261 +
262 + flashinfer_comm.trtllm_allreduce_fusion(
263 + allreduce_in=input_tensor,
264 + token_num=input_tensor.shape[0],
265 + residual_in=residual,
266 + residual_out=residual_out,
267 + norm_out=norm_out,
268 + rms_gamma=rms_gamma,
269 + rms_eps=rms_eps,
270 + hidden_dim=input_tensor.shape[-1],
271 + workspace_ptrs=_FI_WORKSPACE_TENSOR,
272 + pattern_code=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant,
273 + allreduce_out=None,
274 + quant_out=quant_out,
275 + scale_out=output_scale,
276 + layout_code=flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4,
277 + scale_factor=input_global_scale,
278 + use_oneshot=use_oneshot,
279 + **allreduce_params.get_trtllm_fused_allreduce_kwargs(),
280 + )
281 +
282 +
283 +class VllmFusedAllreduce:
284 + def __init__(self, hidden_dim, dtype):
285 + self.rms_eps = 1e-6
286 + self.rms_norm = RMSNorm(hidden_dim, eps=self.rms_eps, dtype=dtype)
287 + self.fp8_quant = QuantFP8(
288 + static=True,
289 + group_shape=GroupShape.PER_TENSOR,
290 + )
291 +
292 + def allreduce_rmsnorm(
293 + self, input_tensor: torch.Tensor, residual: torch.Tensor | None
294 + ):
295 + allreduce_out = tensor_model_parallel_all_reduce(input_tensor)
296 + return self.rms_norm(allreduce_out, residual)
297 +
298 + def allreduce_rmsnorm_fp8_quant(
299 + self,
300 + input_tensor: torch.Tensor,
301 + residual: torch.Tensor | None,
302 + scale_factor: torch.Tensor,
303 + ):
304 + allreduce_out = tensor_model_parallel_all_reduce(input_tensor)
305 + rms_out = self.rms_norm(allreduce_out, residual)
306 + if residual is None:
307 + quant_out = self.fp8_quant(rms_out, scale_factor)
308 + return quant_out
309 + else:
310 + rms_out, residual_out = rms_out
311 + quant_out = self.fp8_quant(rms_out, scale_factor)
312 + return quant_out, residual_out
313 +
314 + def allreduce_rmsnorm_fp4_quant(
315 + self,
316 + input_tensor: torch.Tensor,
317 + residual: torch.Tensor | None,
318 + input_global_scale: torch.Tensor,
319 + quant_out: torch.Tensor,
320 + output_scale: torch.Tensor,
321 + ):
322 + allreduce_out = tensor_model_parallel_all_reduce(input_tensor)
323 + rms_out = self.rms_norm(allreduce_out, residual)
324 + if residual is None:
325 + SCALED_FP4_QUANT_OP(quant_out, rms_out, output_scale, input_global_scale)
326 + return quant_out, output_scale
327 + else:
328 + rms_out, residual_out = rms_out
329 + SCALED_FP4_QUANT_OP(quant_out, rms_out, output_scale, input_global_scale)
330 + return quant_out, residual_out, output_scale
331 +
332 +
333 +def create_test_tensors(
334 + num_tokens: int, hidden_dim: int, dtype: torch.dtype, use_residual: bool = True
335 +):
336 + """Create test tensors for benchmarking."""
337 + input_tensor = torch.randn(num_tokens, hidden_dim, dtype=dtype)
338 + residual = (
339 + torch.randn_like(input_tensor)
340 + if use_residual
341 + else torch.zeros_like(input_tensor)
342 + )
343 + rms_gamma = torch.ones(hidden_dim, dtype=dtype)
344 + norm_out = None if use_residual else torch.empty_like(input_tensor)
345 +
346 + # Quantization scales
347 + scale_fp8 = torch.tensor(1.0, dtype=torch.float32)
348 + scale_fp4 = torch.tensor(1.0, dtype=torch.float32)
349 + quant_out_fp8 = torch.empty_like(input_tensor, dtype=FP8_DTYPE)
350 + # Pre-allocate FP4 output tensors (to avoid allocation overhead in benchmarks)
351 + fp4_quant_out = torch.empty((num_tokens, hidden_dim // 2), dtype=torch.uint8)
352 + fp4_output_scale = torch.empty((128, 4), dtype=torch.int32)
353 +
354 + return (
355 + input_tensor,
356 + norm_out,
357 + residual,
358 + rms_gamma,
359 + scale_fp8,
360 + quant_out_fp8,
361 + scale_fp4,
362 + fp4_quant_out,
363 + fp4_output_scale,
364 + )
365 +
366 +
367 +def benchmark_operation(
368 + operation_func, *args, warmup: int = 5, trials: int = 20, **kwargs
369 +):
370 + """Benchmark a single operation using CUDA graphs."""
371 + # Warmup before graph capture
372 + for _ in range(warmup):
373 + operation_func(*args, **kwargs)
374 + torch.cuda.synchronize()
375 +
376 + # Create CUDA graph
377 + graph = torch.cuda.CUDAGraph()
378 + num_op_per_cudagraph = 10
379 +
380 + # Use vLLM's graph_capture to make tensor_model_parallel_all_reduce graph-safe
381 + device = torch.device(f"cuda:{torch.cuda.current_device()}")
382 + with graph_capture(device=device), torch.cuda.graph(graph):
383 + for _ in range(num_op_per_cudagraph):
384 + operation_func(*args, **kwargs)
385 +
386 + # Graph warmup
387 + torch.cuda.synchronize()
388 + for _ in range(warmup):
389 + graph.replay()
390 +
391 + # Benchmark with CUDA graph
392 + torch.cuda.synchronize()
393 + start_time = time.perf_counter()
394 +
395 + for _ in range(trials // num_op_per_cudagraph):
396 + # operation_func(*args, **kwargs)
397 + graph.replay()
398 +
399 + torch.cuda.synchronize()
400 + end_time = time.perf_counter()
401 +
402 + avg_time_ms = ((end_time - start_time) / trials) * 1000
403 + return avg_time_ms
404 +
405 +
406 +def run_benchmarks(
407 + num_tokens: int,
408 + hidden_dim: int,
409 + dtype: torch.dtype,
410 + use_residual: bool,
411 + allreduce_params: FlashInferFusedAllReduceParams | None,
412 + quant_modes: set[str],
413 + no_oneshot: bool,
414 +):
415 + """Run all benchmarks for given configuration.
416 +
417 + Args:
418 + quant_mode: "none", "fp8_only", "fp4_only", or "all"
419 + """
420 + (
421 + input_tensor,
422 + norm_out,
423 + residual,
424 + rms_gamma,
425 + scale_fp8,
426 + quant_out_fp8,
427 + scale_fp4,
428 + fp4_quant_out,
429 + fp4_output_scale,
430 + ) = create_test_tensors(num_tokens, hidden_dim, dtype, use_residual)
431 +
432 + rms_eps = 1e-6
433 + results = {}
434 + vllm_fused_allreduce = VllmFusedAllreduce(hidden_dim, dtype)
435 + use_oneshot_options = [False] if no_oneshot else [True, False]
436 +
437 + # Create RMSNorm and QuantFP8 layers once for native benchmarks
438 +
439 + if "none" in quant_modes:
440 + # Standard AllReduce + RMSNorm
441 + for custom_op in ["-rms_norm", "+rms_norm"]:
442 + with set_current_vllm_config(
443 + VllmConfig(compilation_config=CompilationConfig(custom_ops=[custom_op]))
444 + ):
445 + try:
446 + suffix = (
447 + "_custom_rms_norm" if "+" in custom_op else "_native_rms_norm"
448 + )
449 + time_ms = benchmark_operation(
450 + vllm_fused_allreduce.allreduce_rmsnorm,
451 + input_tensor,
452 + residual=residual,
453 + )
454 + results[f"standard_allreduce_{suffix}"] = time_ms
455 + except Exception as e:
456 + logger.error("Standard AllReduce+RMSNorm failed: %s", e)
457 + results[f"standard_allreduce_{suffix}"] = float("inf")
458 +
459 + # Standard AllReduce + RMSNorm Native Compiled
460 + with set_current_vllm_config(
461 + VllmConfig(compilation_config=CompilationConfig(custom_ops=["-rms_norm"]))
462 + ):
463 + try:
464 + standard_allreduce_rmsnorm_native_compiled = torch.compile(
465 + vllm_fused_allreduce.allreduce_rmsnorm,
466 + fullgraph=True,
467 + dynamic=False,
468 + )
469 + time_ms = benchmark_operation(
470 + standard_allreduce_rmsnorm_native_compiled,
471 + input_tensor,
472 + residual=residual,
473 + )
474 + results["standard_allreduce_rmsnorm_native_compiled"] = time_ms
475 + except Exception as e:
476 + logger.error("Standard AllReduce+RMSNorm Native Compiled failed: %s", e)
477 + results["standard_allreduce_rmsnorm_native_compiled"] = float("inf")
478 +
479 + # FlashInfer Fused AllReduce + RMSNorm Oneshot/Twoshot
480 + if flashinfer_comm is not None and allreduce_params is not None:
481 + for use_oneshot in use_oneshot_options:
482 + suffix = "_oneshot" if use_oneshot else "_twoshot"
483 + try:
484 + time_ms = benchmark_operation(
485 + flashinfer_fused_allreduce_rmsnorm,
486 + input_tensor,
487 + residual=residual,
488 + norm_out=norm_out,
489 + rms_gamma=rms_gamma,
490 + rms_eps=rms_eps,
491 + allreduce_params=allreduce_params,
492 + use_oneshot=use_oneshot,
493 + )
494 + results[f"flashinfer_fused_allreduce_rmsnorm{suffix}"] = time_ms
495 + except Exception as e:
496 + logger.error("FlashInfer Fused AllReduce+RMSNorm failed: %s", e)
497 + results[f"flashinfer_fused_allreduce_rmsnorm{suffix}"] = float(
498 + "inf"
499 + )
500 +
501 + if "fp8" in quant_modes:
502 + # Standard AllReduce + RMSNorm + FP8 Quant
503 + for rms_norm_custom_op in ["-rms_norm", "+rms_norm"]:
504 + suffix = (
505 + "_custom_rms_norm" if "+" in rms_norm_custom_op else "_native_rms_norm"
506 + )
507 + for quant_fp8_custom_op in ["-quant_fp8", "+quant_fp8"]:
508 + suffix += (
509 + "_custom_quant_fp8"
510 + if "+" in quant_fp8_custom_op
511 + else "_native_quant_fp8"
512 + )
513 + with set_current_vllm_config(
514 + VllmConfig(
515 + compilation_config=CompilationConfig(
516 + custom_ops=[rms_norm_custom_op, quant_fp8_custom_op]
517 + )
518 + )
519 + ):
520 + try:
521 + time_ms = benchmark_operation(
522 + vllm_fused_allreduce.allreduce_rmsnorm_fp8_quant,
523 + input_tensor,
524 + residual=residual,
525 + scale_factor=scale_fp8,
526 + )
527 + results[f"standard_allreduce{suffix}"] = time_ms
528 + except Exception as e:
529 + logger.error("Standard AllReduce+RMSNorm+FP8 failed: %s", e)
530 + results[f"standard_allreduce{suffix}"] = float("inf")
531 +
532 + # Standard AllReduce + RMSNorm + FP8 Quant Native Compiled
533 + with set_current_vllm_config(
534 + VllmConfig(
535 + compilation_config=CompilationConfig(
536 + custom_ops=["-rms_norm", "-quant_fp8"]
537 + )
538 + )
539 + ):
540 + try:
541 + standard_allreduce_rmsnorm_fp8_quant_native_compiled = torch.compile(
542 + vllm_fused_allreduce.allreduce_rmsnorm_fp8_quant,
543 + fullgraph=True,
544 + dynamic=False,
545 + )
546 + time_ms = benchmark_operation(
547 + standard_allreduce_rmsnorm_fp8_quant_native_compiled,
548 + input_tensor,
549 + residual=residual,
550 + scale_factor=scale_fp8,
551 + )
552 + results["standard_allreduce_rmsnorm_fp8_quant_native_compiled"] = (
553 + time_ms
554 + )
555 + except Exception as e:
556 + logger.error(
557 + "Standard AllReduce+RMSNorm+FP8 Native Compiled failed: %s", e
558 + )
559 + results["standard_allreduce_rmsnorm_fp8_quant_native_compiled"] = float(
560 + "inf"
561 + )
562 +
563 + # FlashInfer Fused AllReduce + RMSNorm + FP8 Quant Oneshot
564 + if flashinfer_comm is not None and allreduce_params is not None:
565 + for use_oneshot in use_oneshot_options:
566 + suffix = "_oneshot" if use_oneshot else "_twoshot"
567 + try:
568 + time_ms = benchmark_operation(
569 + flashinfer_fused_allreduce_rmsnorm_fp8_quant,
570 + input_tensor,
571 + norm_out=norm_out,
572 + residual=residual,
573 + rms_gamma=rms_gamma,
574 + rms_eps=rms_eps,
575 + scale_factor=scale_fp8,
576 + quant_out=quant_out_fp8,
577 + allreduce_params=allreduce_params,
578 + use_oneshot=use_oneshot,
579 + )
580 + results[f"flashinfer_fused_allreduce_rmsnorm_fp8_quant{suffix}"] = (
581 + time_ms
582 + )
583 + except Exception as e:
584 + logger.error(
585 + "FlashInfer Fused AllReduce+RMSNorm+FP8 Oneshot failed: %s",
586 + e,
587 + )
588 + results[f"flashinfer_fused_allreduce_rmsnorm_fp8_quant{suffix}"] = (
589 + float("inf")
590 + )
591 +
592 + if "fp4" in quant_modes and current_platform.has_device_capability(100):
593 + # Standard AllReduce + RMSNorm + FP4 Quant
594 + for rms_norm_custom_op in ["-rms_norm", "+rms_norm"]:
595 + suffix = (
596 + "_custom_rms_norm" if "+" in rms_norm_custom_op else "_native_rms_norm"
597 + )
598 + with set_current_vllm_config(
599 + VllmConfig(
600 + compilation_config=CompilationConfig(
601 + custom_ops=[rms_norm_custom_op]
602 + )
603 + )
604 + ):
605 + try:
606 + time_ms = benchmark_operation(
607 + vllm_fused_allreduce.allreduce_rmsnorm_fp4_quant,
608 + input_tensor,
609 + residual=residual,
610 + input_global_scale=scale_fp4,
611 + quant_out=fp4_quant_out,
612 + output_scale=fp4_output_scale,
613 + )
614 + results[f"standard_allreduce_{suffix}_fp4_quant"] = time_ms
615 + except Exception as e:
616 + logger.error("Standard AllReduce+RMSNorm+FP4 failed: %s", e)
617 + results[f"standard_allreduce_{suffix}_fp4_quant"] = float("inf")
618 +
619 + # Standard AllReduce + RMSNorm + FP4 Quant Native Compiled
620 + with set_current_vllm_config(
621 + VllmConfig(compilation_config=CompilationConfig(custom_ops=["-rms_norm"]))
622 + ):
623 + try:
624 + standard_allreduce_rmsnorm_fp4_quant_native_compiled = torch.compile(
625 + vllm_fused_allreduce.allreduce_rmsnorm_fp4_quant,
626 + fullgraph=True,
627 + dynamic=False,
628 + )
629 + time_ms = benchmark_operation(
630 + standard_allreduce_rmsnorm_fp4_quant_native_compiled,
631 + input_tensor,
632 + residual=residual,
633 + quant_out=fp4_quant_out,
634 + input_global_scale=scale_fp4,
635 + output_scale=fp4_output_scale,
636 + )
637 + results["standard_allreduce_rmsnorm_fp4_quant_native_compiled"] = (
638 + time_ms
639 + )
640 + except Exception as e:
641 + logger.error(
642 + "Standard AllReduce+RMSNorm+FP4 Native Compiled failed: %s", e
643 + )
644 + results["standard_allreduce_rmsnorm_fp4_quant_native_compiled"] = float(
645 + "inf"
646 + )
647 +
648 + # FlashInfer Fused AllReduce + RMSNorm + FP4 Quant Oneshot
649 + if flashinfer_comm is not None and allreduce_params is not None:
650 + for use_oneshot in use_oneshot_options:
651 + suffix = "_oneshot" if use_oneshot else "_twoshot"
652 + try:
653 + time_ms = benchmark_operation(
654 + flashinfer_fused_allreduce_rmsnorm_fp4_quant,
655 + input_tensor,
656 + residual=residual,
657 + norm_out=norm_out,
658 + rms_gamma=rms_gamma,
659 + rms_eps=rms_eps,
660 + input_global_scale=scale_fp4,
661 + allreduce_params=allreduce_params,
662 + quant_out=fp4_quant_out,
663 + output_scale=fp4_output_scale,
664 + use_oneshot=use_oneshot,
665 + )
666 + results[f"flashinfer_fused_allreduce_rmsnorm_fp4_quant{suffix}"] = (
667 + time_ms
668 + )
669 + except Exception as e:
670 + logger.error(
671 + "FlashInfer Fused AllReduce+RMSNorm+FP4 Oneshot failed: %s",
672 + e,
673 + )
674 + results[f"flashinfer_fused_allreduce_rmsnorm_fp4_quant{suffix}"] = (
675 + float("inf")
676 + )
677 +
678 + # FlashInfer Fused AllReduce + RMSNorm + FP4 Quant Two-shot
679 + if flashinfer_comm is not None and allreduce_params is not None:
680 + try:
681 + time_ms = benchmark_operation(
682 + flashinfer_fused_allreduce_rmsnorm_fp4_quant,
683 + input_tensor,
684 + residual=residual,
685 + norm_out=norm_out,
686 + rms_gamma=rms_gamma,
687 + rms_eps=rms_eps,
688 + input_global_scale=scale_fp4,
689 + allreduce_params=allreduce_params,
690 + quant_out=fp4_quant_out,
691 + output_scale=fp4_output_scale,
692 + use_oneshot=False,
693 + )
694 + results["flashinfer_fused_allreduce_rmsnorm_fp4_quant_twoshot"] = (
695 + time_ms
696 + )
697 + except Exception as e:
698 + logger.error(
699 + "FlashInfer Fused AllReduce+RMSNorm+FP4 Two-shot failed: %s",
700 + e,
701 + )
702 + results["flashinfer_fused_allreduce_rmsnorm_fp4_quant_twoshot"] = float(
703 + "inf"
704 + )
705 +
706 + return results
707 +
708 +
709 +def prepare_results_with_speedups(results_dict):
710 + """Prepare results with speedup calculations based on dynamic baseline selection."""
711 + prepared_results = []
712 +
713 + # Determine the fastest baseline for each operation type
714 + def get_fastest_baseline(op_name, results_dict):
715 + """Get the fastest baseline between standard and native_compiled versions."""
716 + if "fp8_quant" in op_name:
717 + candidates = [
718 + "standard_allreduce_rmsnorm_fp8_quant",
719 + "standard_allreduce_rmsnorm_fp8_quant_native_compiled",
720 + ]
721 + elif "fp4_quant" in op_name:
722 + candidates = [
723 + "standard_allreduce_rmsnorm_fp4_quant",
724 + "standard_allreduce_rmsnorm_fp4_quant_native_compiled",
725 + ]
726 + else:
727 + candidates = [
728 + "standard_allreduce_rmsnorm",
729 + "standard_allreduce_rmsnorm_native_compiled",
730 + ]
731 +
732 + # Find the fastest among available candidates
733 + fastest_time = float("inf")
734 + fastest_baseline = None
735 +
736 + for candidate in candidates:
737 + if (
738 + candidate in results_dict
739 + and results_dict[candidate] != float("inf")
740 + and results_dict[candidate] < fastest_time
741 + ):
742 + fastest_time = results_dict[candidate]
743 + fastest_baseline = candidate
744 +
745 + return fastest_baseline
746 +
747 + # Create dynamic baseline mapping
748 + dynamic_baseline_mapping = {}
749 + for op_name in results_dict:
750 + if (
751 + op_name.startswith("flashinfer_")
752 + or op_name.startswith("standard_")
753 + and not op_name.endswith("_native_compiled")
754 + ):
755 + dynamic_baseline_mapping[op_name] = get_fastest_baseline(
756 + op_name, results_dict
757 + )
758 +
759 + for op_name, time_ms in results_dict.items():
760 + if time_ms == float("inf"):
761 + speedup_str = "FAILED"
762 + time_str = "FAILED"
763 + else:
764 + time_str = f"{time_ms:.3f}"
765 + # Find the appropriate baseline for this operation
766 + baseline_op = dynamic_baseline_mapping.get(op_name)
767 + if baseline_op and baseline_op in results_dict:
768 + baseline_time = results_dict[baseline_op]
769 + if baseline_time != float("inf") and baseline_time > 0:
770 + speedup = baseline_time / time_ms
771 + speedup_str = f"{speedup:.2f}x"
772 + else:
773 + speedup_str = "N/A"
774 + else:
775 + # For baseline operations, determine if this is the fastest baseline
776 + if op_name.endswith("_native_compiled") or (
777 + op_name.startswith("standard_")
778 + and not op_name.endswith("_native_compiled")
779 + ):
780 + fastest_baseline = get_fastest_baseline(op_name, results_dict)
781 + if fastest_baseline == op_name:
782 + speedup_str = "baseline"
783 + else:
784 + if fastest_baseline and fastest_baseline in results_dict:
785 + baseline_time = results_dict[fastest_baseline]
786 + if baseline_time != float("inf") and baseline_time > 0:
787 + speedup = baseline_time / time_ms
788 + speedup_str = f"{speedup:.2f}x"
789 + else:
790 + speedup_str = "N/A"
791 + else:
792 + speedup_str = "N/A"
793 + else:
794 + speedup_str = "N/A"
795 +
796 + prepared_results.append(
797 + {
798 + "operation": op_name,
799 + "time_ms": time_ms,
800 + "time_str": time_str,
801 + "speedup_str": speedup_str,
802 + }
803 + )
804 +
805 + return prepared_results
806 +
807 +
808 +def print_results(
809 + results_dict,
810 + num_tokens,
811 + hidden_dim,
812 + dtype,
813 + use_residual,
814 + quant_modes,
815 + input_size_mb,
816 +):
817 + """Print benchmark results in a formatted table."""
818 + print(f"\n{'=' * 80}")
819 + print(
820 + f"Results: num_tokens={num_tokens}, hidden_dim={hidden_dim} "
821 + f"(input size: {input_size_mb:.2f} MB)"
822 + )
823 + print(
824 + f"dtype={dtype}, residual={'yes' if use_residual else 'no'}, "
825 + f"quant_modes={','.join(sorted(list(quant_modes)))}"
826 + )
827 + print(f"{'=' * 80}")
828 + print(f"{'Operation':<50} {'Time (ms)':<12} {'Speedup':<10}")
829 + print(f"{'-' * 80}")
830 +
831 + # Prepare results with speedup calculations
832 + prepared_results = prepare_results_with_speedups(results_dict)
833 +
834 + for result in prepared_results:
835 + if result["time_ms"] == float("inf"):
836 + time_display = result["time_str"]
837 + else:
838 + time_display = f"{result['time_ms']:.3f}"
839 +
840 + print(
841 + f"{result['operation']:<50} {time_display:<12} {result['speedup_str']:<10}"
842 + )
843 +
844 +
845 +def format_results_markdown(
846 + all_results: list[dict], world_size: int, args: argparse.Namespace
847 +) -> str:
848 + """Format all benchmark results as markdown."""
849 + lines: list[str] = []
850 + lines.append("# FlashInfer Fused Collective Operations Benchmark Results")
851 + lines.append("")
852 + lines.append(f"**World Size:** {world_size} ")
853 + lines.append(f"**Hidden Dimension:** {args.hidden_dim} ")
854 + lines.append(f"**Warmup Iterations:** {args.warmup} ")
855 + lines.append(f"**Benchmark Trials:** {args.trials} ")
856 + modes = ",".join(all_results[0]["quant_modes"]) if all_results else "N/A"
857 + lines.append(f"**Quantization Modes:** {modes} ")
858 + lines.append("")
859 + lines.append("---")
860 + lines.append("")
861 +
862 + for entry in all_results:
863 + num_tokens = entry["num_tokens"]
864 + dtype = entry["dtype"]
865 + use_residual = entry["use_residual"]
866 + results_dict = entry["results"]
867 + input_size_mb = entry["input_size_mb"]
868 + residual_str = "with residual" if use_residual else "no residual"
869 +
870 + lines.append(
871 + f"## Configuration: num_tokens={num_tokens}, dtype={dtype}, {residual_str}"
872 + )
873 + lines.append(f"**Input Size:** {input_size_mb:.2f} MB")
874 + lines.append("")
875 +
876 + prepared = prepare_results_with_speedups(results_dict)
877 + # Build DataFrame for markdown export
878 + rows = [
879 + {
880 + "Operation": r["operation"].replace("_", " ").title(),
881 + "Time (ms)": r["time_str"],
882 + "Speedup": r["speedup_str"],
883 + }
884 + for r in prepared
885 + ]
886 + df = pd.DataFrame(rows)
887 + if df.empty:
888 + lines.append("No results.")
889 + else:
890 + lines.append(df.to_markdown(index=False))
891 + lines.append("")
892 +
893 + return "\n".join(lines)
894 +
895 +
896 +def save_results_to_file(
897 + all_results: list[dict], world_size: int, args: argparse.Namespace, rank: int
898 +):
899 + """Save benchmark results to markdown file (only on rank 0)."""
900 + if rank != 0:
901 + return
902 +
903 + if not all_results:
904 + logger.warning("No results to save")
905 + return
906 +
907 + output_path = args.output_file
908 +
909 + try:
910 + markdown_content = format_results_markdown(all_results, world_size, args)
911 +
912 + with open(output_path, "a") as f:
913 + f.write(markdown_content)
914 +
915 + except Exception as e:
916 + logger.error("Failed to save results to file: %s", e)
917 +
918 +
919 +def main():
920 + parser = argparse.ArgumentParser(
921 + description="Benchmark fused collective operations"
922 + )
923 + parser.add_argument(
924 + "--num-tokens",
925 + type=int,
926 + nargs="+",
927 + default=[128, 512, 1024, 2048],
928 + help="Numbers of tokens to test",
929 + )
930 + parser.add_argument(
931 + "--hidden-dim", type=int, default=8192, help="Hidden dimension size"
932 + )
933 + parser.add_argument(
934 + "--dtypes",
935 + type=str,
936 + nargs="+",
937 + default=["bfloat16"],
938 + choices=["float16", "bfloat16", "float32"],
939 + help="Data types to test",
940 + )
941 + parser.add_argument(
942 + "--no-residual",
943 + action="store_true",
944 + help="Skip residual connection tests",
945 + )
946 +
947 + parser.add_argument(
948 + "--quant-modes",
949 + type=str,
950 + default="none,fp8,fp4",
951 + help=(
952 + "Comma-separated quantization modes to run: none, fp8, fp4. "
953 + "Default: none,fp8,fp4"
954 + ),
955 + )
956 +
957 + parser.add_argument(
958 + "--warmup", type=int, default=5, help="Number of warmup iterations"
959 + )
960 + parser.add_argument(
961 + "--trials", type=int, default=20, help="Number of benchmark trials"
962 + )
963 + parser.add_argument(
964 + "--output-file",
965 + type=str,
966 + help="""Output file path for markdown results
967 + (default: benchmark_results_<timestamp>.md)
968 + """,
969 + )
970 +
971 + parser.add_argument(
972 + "--no-oneshot",
973 + action="store_true",
974 + help="Skip oneshot benchmarks",
975 + )
976 +
977 + args = parser.parse_args()
978 +
979 + # Check if running with torchrun (required for collective operations)
980 + if "RANK" not in os.environ or "WORLD_SIZE" not in os.environ:
981 + raise RuntimeError(
982 + "Must run with torchrun for distributed benchmarking. "
983 + "Example: torchrun --nproc_per_node=2 benchmark_fused_collective.py"
984 + )
985 +
986 + # Initialize distributed environment
987 + rank = int(os.environ["RANK"])
988 + world_size = int(os.environ["WORLD_SIZE"])
989 +
990 + device = torch.device(f"cuda:{rank}")
991 + torch.cuda.set_device(device)
992 + torch.set_default_device(device)
993 +
994 + init_distributed_environment()
995 + initialize_model_parallel(tensor_model_parallel_size=world_size)
996 +
997 + # Validate world size (must be > 1 for collective operations)
998 + if world_size <= 1:
999 + raise ValueError(
1000 + "World size must be > 1 for collective operations benchmarking. "
1001 + f"Current world size: {world_size}. Use torchrun with --nproc_per_node > 1."
1002 + )
1003 +
1004 + # Parse quantization modes
1005 + valid_quant_modes = {"none", "fp8", "fp4"}
1006 + raw_modes = [
1007 + m.strip().lower() for m in (args.quant_modes or "").split(",") if m.strip()
1008 + ]
1009 + quant_modes = set(raw_modes) if raw_modes else {"none", "fp8", "fp4"}
1010 + invalid = sorted(list(quant_modes - valid_quant_modes))
1011 + if invalid:
1012 + raise ValueError(
1013 + f"Invalid --quant-modes entries: {','.join(invalid)}. "
1014 + f"Valid options are: {','.join(sorted(valid_quant_modes))}."
1015 + )
1016 +
1017 + if rank == 0:
1018 + logger.info("Running benchmark with world_size=%s, rank=%s", world_size, rank)
1019 + logger.info("Quantization modes: %s", ",".join(sorted(list(quant_modes))))
1020 + if flashinfer_comm is not None:
1021 + logger.info(
1022 + "FlashInfer available - will benchmark fused operations",
1023 + )
1024 + else:
1025 + logger.info(
1026 + "FlashInfer not available - only benchmarking standard operations"
1027 + )
1028 +
1029 + # Convert dtype strings to torch dtypes
1030 + dtype_map = {
1031 + "float16": torch.float16,
1032 + "bfloat16": torch.bfloat16,
1033 + "float32": torch.float32,
1034 + }
1035 + dtypes = [dtype_map[dt] for dt in args.dtypes]
1036 +
1037 + # Test configurations
1038 + residual_options = [True] if not args.no_residual else [False]
1039 +
1040 + configs = list(itertools.product(args.num_tokens, dtypes, residual_options))
1041 +
1042 + # Setup FlashInfer workspace if available
1043 + ipc_handles = None
1044 + allreduce_params = None
1045 +
1046 + if flashinfer_comm is not None:
1047 + # Use the largest hidden dimension for workspace setup
1048 + max_num_token = _FI_MAX_SIZES.get(world_size) // (
1049 + args.hidden_dim * world_size * 2
1050 + )
1051 +
1052 + ipc_handles, workspace_tensor = setup_flashinfer_workspace(
1053 + world_size, rank, args.hidden_dim, max_num_token
1054 + )
1055 +
1056 + if workspace_tensor is not None:
1057 + allreduce_params = FlashInferFusedAllReduceParams(
1058 + rank=rank,
1059 + world_size=world_size,
1060 + max_token_num=max_num_token,
1061 + )
1062 +
1063 + # Collect all results for markdown export
1064 + all_results = []
1065 +
1066 + try:
1067 + # Run benchmarks
1068 + for num_tokens, dtype, use_residual in configs:
1069 + if rank == 0:
1070 + logger.info(
1071 + "\nTesting: num_tokens=%s, hidden_dim=%s, dtype=%s, residual=%s",
1072 + num_tokens,
1073 + args.hidden_dim,
1074 + dtype,
1075 + use_residual,
1076 + )
1077 +
1078 + results = run_benchmarks(
1079 + num_tokens,
1080 + args.hidden_dim,
1081 + dtype,
1082 + use_residual,
1083 + allreduce_params,
1084 + quant_modes=quant_modes,
1085 + no_oneshot=args.no_oneshot,
1086 + )
1087 +
1088 + # Store results for markdown export
1089 + if rank == 0:
1090 + # Calculate input size in MB
1091 + input_size_mb = (
1092 + num_tokens * args.hidden_dim * torch.finfo(dtype).bits
1093 + ) / (8 * 1024 * 1024)
1094 + all_results.append(
1095 + {
1096 + "num_tokens": num_tokens,
1097 + "hidden_dim": args.hidden_dim,
1098 + "dtype": str(dtype).replace("torch.", ""),
1099 + "use_residual": use_residual,
1100 + "quant_modes": sorted(list(quant_modes)),
1101 + "input_size_mb": input_size_mb,
1102 + "results": results,
1103 + }
1104 + )
1105 +
1106 + print_results(
1107 + results,
1108 + num_tokens,
1109 + args.hidden_dim,
1110 + dtype,
1111 + use_residual,
1112 + quant_modes,
1113 + input_size_mb,
1114 + )
1115 +
1116 + # Save results to markdown file
1117 + if args.output_file and rank == 0:
1118 + save_results_to_file(all_results, world_size, args, rank)
1119 +
1120 + finally:
1121 + # Cleanup
1122 + if ipc_handles is not None:
1123 + cleanup_flashinfer_workspace(ipc_handles)
1124 +
1125 + dist.barrier()
1126 +
1127 +
1128 +if __name__ == "__main__":
1129 + main()
tests/compile/test_fusions_e2e.py
+7
@@ -71,6 +71,13 @@ if current_platform.is_cuda():
71 attention_fusions=0,
72 allreduce_fusions=65,
73 ),
74 + ModelBackendTestCase(
75 + model_name="Qwen/Qwen3-30B-A3B",
76 + model_kwargs=dict(max_model_len=1024),
77 + backend=_Backend.TRITON_ATTN,
78 + attention_fusions=0,
79 + allreduce_fusions=97,
80 + ),
81 ]
82
83 elif current_platform.is_rocm():
vllm/compilation/collective_fusion.py
+72 -60
@@ -9,7 +9,6 @@ from torch._higher_order_ops.auto_functionalize import auto_functionalized
9 from torch._inductor.pattern_matcher import PatternMatcherPass
10 from torch.distributed._symmetric_memory import enable_symm_mem_for_group
11
12 -import vllm.envs as envs
12 from vllm.config import VllmConfig
13 from vllm.distributed import get_tp_group, tensor_model_parallel_all_reduce
14 from vllm.distributed.parallel_state import (
@@ -450,34 +449,41 @@ class AsyncTPPass(VllmPatternMatcherPass):
449 logger.debug("Replaced %s patterns", self.matched_count)
450
451
452 +# Max size of the input tensor per world size per device capability
453 +# to use flashinfer fused allreduce
454 +FI_ALLREDUCE_FUSION_MAX_SIZE_MB: dict[int, dict[int, float]] = {
455 + 90: {
456 + 2: 64, # 64MB
457 + 4: 2, # 2MB
458 + 8: 0.5, # 0.5MB
459 + },
460 + 100: {
461 + 2: 64, # 64MB
462 + 4: 32, # 32MB
463 + 8: 1, # 1MB
464 + },
465 +}
466 +
467 +# Max size of the input tensor per world size per device capability
468 +# to use flashinfer one shot fused allreduce
469 +# OneShot max size is at most 64MB / world size (FlashInfer restriction)
470 +_FI_ALLREDUCE_ONE_SHOT_MAX_SIZES_MB: dict[int, dict[int, float]] = {
471 + 90: {
472 + 2: 32, # 32MB
473 + 4: 2, # 2MB
474 + 8: 0.5, # 0.5MB
475 + },
476 + 100: {
477 + 2: 32, # 32MB
478 + 4: 4, # 4MB
479 + 8: 1, # 1MB
480 + },
481 +}
482 +
483 +
484 if flashinfer_comm is not None:
485 _FI_WORKSPACE_TENSOR = None
455 -
486 MiB = 1024 * 1024
457 - # Max size of the input tensor per world size
458 - # to use flashinfer fused allreduce
459 - _FI_MAX_SIZES = {
460 - 2: 64 * MiB, # 64MB
461 - 4: MiB, # 1MB
462 - 6: MiB // 2, # 512KB
463 - 8: MiB // 2, # 512KB
464 - }
465 -
466 - try:
467 - _FI_MAX_SIZES.update(
468 - {
469 - int(k): int(float(v) * MiB)
470 - for k, v in envs.VLLM_FLASHINFER_ALLREDUCE_FUSION_THRESHOLDS_MB.items()
471 - }
472 - )
473 - except Exception as e:
474 - raise ValueError(
475 - "Failed to parse VLLM_FLASHINFER_ALLREDUCE_FUSION_THRESHOLDS_MB: " + str(e)
476 - ) from e
477 -
478 - # opt for a more conservative default value
479 - # when world size is not in _FI_MAX_SIZES
480 - _DEFAULT_FI_MAX_SIZE = MiB // 2
487
488 def call_trtllm_fused_allreduce_norm(
489 allreduce_in: torch.Tensor,
@@ -491,7 +497,6 @@ if flashinfer_comm is not None:
497 fp32_acc: bool,
498 max_token_num: int,
499 pattern_code: int,
494 - fuse_rms_quant: bool,
500 norm_out: torch.Tensor | None = None,
501 quant_out: torch.Tensor | None = None,
502 scale_out: torch.Tensor | None = None,
@@ -500,12 +505,20 @@ if flashinfer_comm is not None:
505 num_tokens, hidden_size = allreduce_in.shape
506 element_size = allreduce_in.element_size()
507 current_tensor_size = num_tokens * hidden_size * element_size
503 - max_fusion_size = max_token_num * hidden_size * element_size
504 - use_flashinfer = current_tensor_size <= min(
505 - _FI_MAX_SIZES.get(world_size, _DEFAULT_FI_MAX_SIZE),
506 - max_fusion_size,
507 - )
508 - if use_flashinfer:
508 +
509 + if num_tokens <= max_token_num:
510 + device_capability = current_platform.get_device_capability().to_int()
511 + # Get one shot input size limit for the current world size
512 + # for the current device capability
513 + max_one_shot_size_mb = _FI_ALLREDUCE_ONE_SHOT_MAX_SIZES_MB.get(
514 + device_capability, {}
515 + ).get(world_size, None)
516 + # Use one shot if no max size for one shot is specified
517 + use_oneshot = (
518 + max_one_shot_size_mb is None
519 + or current_tensor_size <= max_one_shot_size_mb * MiB
520 + )
521 +
522 assert _FI_WORKSPACE_TENSOR is not None, (
523 "Flashinfer must be enabled when using flashinfer"
524 )
@@ -532,7 +545,7 @@ if flashinfer_comm is not None:
545 hidden_dim=allreduce_in.shape[-1],
546 workspace_ptrs=_FI_WORKSPACE_TENSOR,
547 launch_with_pdl=launch_with_pdl,
535 - use_oneshot=True,
548 + use_oneshot=use_oneshot,
549 trigger_completion_at_end=trigger_completion_at_end,
550 fp32_acc=fp32_acc,
551 pattern_code=pattern_code,
@@ -545,7 +558,7 @@ if flashinfer_comm is not None:
558 )
559 else:
560 allreduce_out = tensor_model_parallel_all_reduce(allreduce_in)
548 - if scale_factor is not None and scale_out is None and fuse_rms_quant:
561 + if scale_factor is not None and scale_out is None:
562 # Do fused rms norm static fp8 quant fused op
563 if norm_out is None:
564 torch.ops._C.fused_add_rms_norm_static_fp8_quant(
@@ -568,15 +581,10 @@ if flashinfer_comm is not None:
581 norm_out = allreduce_out
582 else:
583 torch.ops._C.rms_norm(norm_out, allreduce_out, rms_gamma, rms_eps)
571 - if scale_factor is not None:
572 - if scale_out is not None:
573 - torch.ops._C.scaled_fp4_quant(
574 - quant_out, norm_out, scale_out, scale_factor
575 - )
576 - else:
577 - torch.ops._C.static_scaled_fp8_quant(
578 - quant_out, norm_out, scale_factor
579 - )
584 + if scale_factor is not None and scale_out is not None:
585 + torch.ops._C.scaled_fp4_quant(
586 + quant_out, norm_out, scale_out, scale_factor
587 + )
588 if scale_factor is None or norm_out is not None:
589 # we need to return allreduce output
590 # in cases of non quant fused AR + RMS norm
@@ -595,7 +603,6 @@ if flashinfer_comm is not None:
603 fp32_acc: bool,
604 max_token_num: int,
605 pattern_code: int,
598 - fuse_rms_quant: bool,
606 norm_out: torch.Tensor | None = None,
607 quant_out: torch.Tensor | None = None,
608 scale_out: torch.Tensor | None = None,
@@ -629,7 +636,6 @@ class FlashInferFusedAllReduceParams:
636 world_size: int,
637 use_fp32_lamport: bool = False,
638 max_token_num: int = 1024,
632 - fuse_rms_quant: bool = False,
639 ):
640 self.rank = rank
641 self.world_size = world_size
@@ -637,9 +643,7 @@ class FlashInferFusedAllReduceParams:
643 self.trigger_completion_at_end = True
644 self.launch_with_pdl = True
645 self.fp32_acc = True
640 - self.use_oneshot = False
646 self.max_token_num = max_token_num
642 - self.fuse_rms_quant = fuse_rms_quant
647
648 def get_trtllm_fused_allreduce_kwargs(self):
649 return {
@@ -649,7 +653,6 @@ class FlashInferFusedAllReduceParams:
653 "trigger_completion_at_end": self.trigger_completion_at_end,
654 "fp32_acc": self.fp32_acc,
655 "max_token_num": self.max_token_num,
652 - "fuse_rms_quant": self.fuse_rms_quant,
656 }
657
658
@@ -1119,23 +1122,35 @@ class AllReduceFusionPass(VllmPatternMatcherPass):
1122 "skipping allreduce fusion pass"
1123 )
1124 return
1122 - # Check if the world size is supported
1123 - if self.tp_size not in _FI_MAX_SIZES:
1125 + max_size = config.compilation_config.pass_config.flashinfer_max_size(
1126 + self.tp_size
1127 + )
1128 + if max_size is None:
1129 + # Flashinfer doesn't support current world size
1130 logger.warning(
1131 "Flashinfer allreduce fusion is not supported for world size %s",
1132 self.tp_size,
1133 )
1134 return
1129 - max_num_token = min(
1130 - _FI_MAX_SIZES.get(self.tp_size, _DEFAULT_FI_MAX_SIZE)
1131 - // (self.hidden_dim * self.tp_size * (4 if use_fp32_lamport else 2)),
1132 - config.compilation_config.pass_config.fi_allreduce_fusion_max_token_num,
1135 + element_size = 4 if use_fp32_lamport else 2
1136 + self.max_token_num = max_size // (self.hidden_dim * element_size)
1137 + # take the min to save workspace size and we'll never use more
1138 + # than max_num_batched_tokens anyways
1139 + self.max_token_num = min(
1140 + self.max_token_num, config.scheduler_config.max_num_batched_tokens
1141 + )
1142 + logger.debug_once(
1143 + f"Flashinfer max size: {max_size // (1024 * 1024)} MB,"
1144 + "Maximal number of tokens used by "
1145 + f"Flashinfer Allreduce Fusion: {self.max_token_num}",
1146 + scope="global",
1147 )
1148 +
1149 self.ipc_handles, workspace_tensor = (
1150 flashinfer_comm.trtllm_create_ipc_workspace_for_all_reduce_fusion(
1151 tp_rank=rank,
1152 tp_size=self.tp_size,
1138 - max_token_num=max_num_token,
1153 + max_token_num=self.max_token_num,
1154 hidden_dim=self.hidden_dim,
1155 group=self.group,
1156 use_fp32_lamport=use_fp32_lamport,
@@ -1148,10 +1163,7 @@ class AllReduceFusionPass(VllmPatternMatcherPass):
1163 rank=rank,
1164 world_size=self.tp_size,
1165 use_fp32_lamport=use_fp32_lamport,
1151 - max_token_num=max_num_token,
1152 - # fuse rms norm static fp8 quant fused op
1153 - # in fallback path, when we don't use flashinfer
1154 - fuse_rms_quant=config.compilation_config.pass_config.enable_fusion,
1166 + max_token_num=self.max_token_num,
1167 )
1168
1169 self.register_patterns()
vllm/config/compilation.py
+48 -2
@@ -111,11 +111,52 @@ class PassConfig:
111 """Whether to enable async TP."""
112 enable_fi_allreduce_fusion: bool = False
113 """Whether to enable flashinfer allreduce fusion."""
114 - fi_allreduce_fusion_max_token_num: int = 16384
115 - """Max number of tokens to used in flashinfer allreduce fusion."""
114 + fi_allreduce_fusion_max_size_mb: float | None = None
115 + """The threshold of the communicated tensor sizes under which
116 + vllm should use flashinfer fused allreduce. Specified as a
117 + float in MB.
118 + Unspecified will fallback to default values
119 + which are compute capability and world size dependent.
120 + FI_ALLREDUCE_FUSION_MAX_SIZE_MB = {
121 + 90: {
122 + 2: 64, # 64MB
123 + 4: 2, # 2MB
124 + 8: 1, # 1MB
125 + },
126 + 100: {
127 + 2: 64, # 64MB
128 + 4: 32, # 32MB
129 + 8: 1, # 1MB
130 + },
131 + }, where key is the device capability"""
132
133 # TODO(luka) better pass enabling system.
134
135 + def flashinfer_max_size(self, world_size: int) -> int | None:
136 + """
137 + Returns the max communication size in bytes for flashinfer
138 + allreduce fusion for the given world size. Returns None if world size
139 + is not supported by configs as it's not supported by flashinfer.
140 + """
141 +
142 + MiB = 1024 * 1024
143 + max_size_mb = self.fi_allreduce_fusion_max_size_mb
144 + if max_size_mb is None:
145 + max_size_mb = self.default_fi_allreduce_fusion_max_size_mb().get(world_size)
146 +
147 + return int(max_size_mb * MiB) if max_size_mb is not None else None
148 +
149 + @staticmethod
150 + def default_fi_allreduce_fusion_max_size_mb() -> dict[int, float]:
151 + from vllm.compilation.collective_fusion import FI_ALLREDUCE_FUSION_MAX_SIZE_MB
152 + from vllm.platforms import current_platform
153 +
154 + if not current_platform.is_cuda():
155 + return {}
156 + return FI_ALLREDUCE_FUSION_MAX_SIZE_MB.get(
157 + current_platform.get_device_capability().to_int(), {}
158 + )
159 +
160 def uuid(self):
161 """
162 Produces a hash unique to the pass configuration.
@@ -136,6 +177,11 @@ class PassConfig:
177 "Fusion enabled but reshape elimination disabled. "
178 "Attention + quant (fp8) fusion might not work"
179 )
180 + if self.enable_fi_allreduce_fusion:
181 + logger.warning_once(
182 + "Fusion enabled but reshape elimination disabled. "
183 + "Allreduce + rms norm + quant (fp8) fusion might not work"
184 + )
185
186
187 @config
vllm/model_executor/layers/fused_moe/layer.py
+26 -19
@@ -2356,6 +2356,16 @@ class FusedMoE(CustomOp):
2356 value=0.0,
2357 )
2358
2359 + def reduce_output(states: torch.Tensor) -> torch.Tensor:
2360 + if (
2361 + not self.is_sequence_parallel
2362 + and not self.use_dp_chunking
2363 + and self.reduce_results
2364 + and (self.tp_size > 1 or self.ep_size > 1)
2365 + ):
2366 + states = self.maybe_all_reduce_tensor_model_parallel(states)
2367 + return states
2368 +
2369 if self.shared_experts is None:
2370 if current_platform.is_tpu():
2371 # TODO: Once the OOM issue for the TPU backend is resolved, we
@@ -2366,7 +2376,14 @@ class FusedMoE(CustomOp):
2376 fused_output = torch.ops.vllm.moe_forward(
2377 hidden_states, router_logits, self.layer_name
2378 )
2369 - return fused_output[..., :og_hidden_states]
2379 + if self.zero_expert_num is not None and self.zero_expert_num > 0:
2380 + assert isinstance(fused_output, tuple)
2381 + fused_output, zero_expert_result = fused_output
2382 + return (reduce_output(fused_output) + zero_expert_result)[
2383 + ..., :og_hidden_states
2384 + ]
2385 + else:
2386 + return reduce_output(fused_output)[..., :og_hidden_states]
2387 else:
2388 if current_platform.is_tpu():
2389 # TODO: Once the OOM issue for the TPU backend is resolved, we
@@ -2379,8 +2396,8 @@ class FusedMoE(CustomOp):
2396 hidden_states, router_logits, self.layer_name
2397 )
2398 return (
2382 - shared_output[..., :og_hidden_states],
2383 - fused_output[..., :og_hidden_states],
2399 + reduce_output(shared_output)[..., :og_hidden_states],
2400 + reduce_output(fused_output)[..., :og_hidden_states],
2401 )
2402
2403 def forward_cuda(
@@ -2667,31 +2684,21 @@ class FusedMoE(CustomOp):
2684 assert isinstance(final_hidden_states, tuple)
2685 final_hidden_states, zero_expert_result = final_hidden_states
2686
2670 - def reduce_output(
2671 - states: torch.Tensor, do_combine: bool = True
2672 - ) -> torch.Tensor:
2673 - if do_naive_dispatch_combine and do_combine:
2687 + def combine_output(states: torch.Tensor) -> torch.Tensor:
2688 + if do_naive_dispatch_combine:
2689 states = get_ep_group().combine(states, self.is_sequence_parallel)
2675 -
2676 - if (
2677 - not self.is_sequence_parallel
2678 - and self.reduce_results
2679 - and (self.tp_size > 1 or self.ep_size > 1)
2680 - ):
2681 - states = self.maybe_all_reduce_tensor_model_parallel(states)
2682 -
2690 return states
2691
2692 if self.shared_experts is not None:
2693 return (
2687 - reduce_output(final_hidden_states[0], do_combine=False),
2688 - reduce_output(final_hidden_states[1]),
2694 + final_hidden_states[0],
2695 + combine_output(final_hidden_states[1]),
2696 )
2697 elif self.zero_expert_num is not None and self.zero_expert_num > 0:
2698 assert isinstance(final_hidden_states, torch.Tensor)
2692 - return reduce_output(final_hidden_states) + zero_expert_result
2699 + return (combine_output(final_hidden_states), zero_expert_result)
2700 else:
2694 - return reduce_output(final_hidden_states)
2701 + return combine_output(final_hidden_states)
2702
2703 @classmethod
2704 def make_expert_params_mapping(