@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 6e672daf62
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 14, 22:35 UTC

Add FlashInfer allreduce RMSNorm Quant fusion (#21069)

Signed-off-by: ilmarkov <imarkov@redhat.com> Signed-off-by: ilmarkov <markovilya197@gmail.com> Co-authored-by: ilmarkov <imarkov@redhat.com>

Ilya Markov committed Jul 31, 2025 at 22:58 UTC 6e672daf62e7b03ff1dcf74e4206dad07d39d4ec
5 files changed +606 -68
.buildkite/test-pipeline.yaml
+1
@@ -353,6 +353,7 @@ steps:
353 - pytest -v -s compile/test_silu_mul_quant_fusion.py
354 - pytest -v -s compile/test_sequence_parallelism.py
355 - pytest -v -s compile/test_async_tp.py
356 + - pytest -v -s compile/test_fusion_all_reduce.py
357
358 - label: PyTorch Fullgraph Smoke Test # 9min
359 mirror_hardwares: [amdexperimental]
tests/compile/test_fusion_all_reduce.py
+104 -22
@@ -7,22 +7,26 @@ import torch
7
8 import vllm.envs as envs
9 from vllm.compilation.collective_fusion import AllReduceFusionPass
10 +from vllm.compilation.fix_functionalization import FixFunctionalizationPass
11 +from vllm.compilation.noop_elimination import NoOpEliminationPass
12 from vllm.config import (CompilationConfig, CompilationLevel, DeviceConfig,
13 ModelConfig, PassConfig, VllmConfig)
14 from vllm.distributed import tensor_model_parallel_all_reduce
15 from vllm.distributed.parallel_state import (init_distributed_environment,
16 initialize_model_parallel)
17 from vllm.model_executor.layers.layernorm import RMSNorm
18 +from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
19 + GroupShape, QuantFP8)
20 from vllm.platforms import current_platform
21 from vllm.utils import update_environment_variables
22
19 -from ..utils import multi_gpu_test
23 +from ..utils import has_module_attribute, multi_gpu_test
24 from .backend import TestBackend
25
26
27 class TestAllReduceRMSNormModel(torch.nn.Module):
28
25 - def __init__(self, hidden_size=16, eps=1e-6):
29 + def __init__(self, hidden_size=16, token_num=16, eps=1e-6):
30 super().__init__()
31 self.hidden_size = hidden_size
32 self.eps = eps
@@ -43,7 +47,7 @@ class TestAllReduceRMSNormModel(torch.nn.Module):
47
48 class TestAllReduceFusedAddRMSNormModel(torch.nn.Module):
49
46 - def __init__(self, hidden_size=16, eps=1e-6):
50 + def __init__(self, hidden_size=16, token_num=16, eps=1e-6):
51 super().__init__()
52 self.hidden_size = hidden_size
53 self.eps = eps
@@ -62,24 +66,101 @@ class TestAllReduceFusedAddRMSNormModel(torch.nn.Module):
66 return [torch.ops.vllm.flashinfer_trtllm_fused_allreduce_norm.default]
67
68
69 +class TestAllReduceFusedAddRMSNormStaticQuantFP8Model(torch.nn.Module):
70 +
71 + def __init__(self, hidden_size=16, token_num=16, eps=1e-6):
72 + super().__init__()
73 + self.hidden_size = hidden_size
74 + self.eps = eps
75 + self.norm = RMSNorm(hidden_size, eps)
76 + self.quant_fp8 = QuantFP8(static=True,
77 + group_shape=GroupShape.PER_TENSOR)
78 + self.scale = torch.rand(1, dtype=torch.float32)
79 + self.output = torch.empty((token_num, hidden_size),
80 + dtype=torch.float32)
81 +
82 + def forward(self, hidden_states, residual):
83 + view = hidden_states.reshape(-1, self.hidden_size)
84 + all_reduce = tensor_model_parallel_all_reduce(view)
85 + norm_output, residual_output = self.norm(all_reduce, residual)
86 + torch.ops._C.static_scaled_fp8_quant(self.output,
87 + norm_output.contiguous(),
88 + self.scale)
89 + return self.output, residual_output
90 +
91 + def ops_in_model_after(self):
92 + return [torch.ops.vllm.flashinfer_trtllm_fused_allreduce_norm.default]
93 +
94 + def ops_in_model_before(self):
95 + return [
96 + torch.ops.vllm.all_reduce.default,
97 + torch.ops._C.static_scaled_fp8_quant.default
98 + ]
99 +
100 +
101 +class TestAllReduceFusedAddRMSNormStaticQuantFP4Model(torch.nn.Module):
102 +
103 + def __init__(self, hidden_size=16, token_num=16, eps=1e-6):
104 + super().__init__()
105 + self.hidden_size = hidden_size
106 + self.eps = eps
107 + self.norm = RMSNorm(hidden_size, eps)
108 + self.scale = torch.rand(1, dtype=torch.float32)
109 + self.output = torch.empty((token_num, hidden_size),
110 + dtype=torch.float32)
111 +
112 + round_up = lambda x, y: (x + y - 1) // y * y
113 + rounded_m = round_up(token_num, 128)
114 + scale_n = hidden_size // 16
115 + rounded_n = round_up(scale_n, 4)
116 + self.output_scale = torch.empty((rounded_m, rounded_n // 4),
117 + dtype=torch.int32)
118 +
119 + def forward(self, hidden_states, residual):
120 + view = hidden_states.reshape(-1, self.hidden_size)
121 + all_reduce = tensor_model_parallel_all_reduce(view)
122 + norm_output, residual_output = self.norm(all_reduce, residual)
123 + norm_output = norm_output.reshape(-1, norm_output.shape[-1])
124 + torch.ops._C.scaled_fp4_quant(self.output, norm_output,
125 + self.output_scale, self.scale)
126 + return self.output, residual_output, self.output_scale
127 +
128 + def ops_in_model_after(self):
129 + return [torch.ops.vllm.flashinfer_trtllm_fused_allreduce_norm.default]
130 +
131 + def ops_in_model_before(self):
132 + return [
133 + torch.ops.vllm.all_reduce.default,
134 + torch.ops._C.scaled_fp4_quant.default
135 + ]
136 +
137 +
138 @multi_gpu_test(num_gpus=2)
66 -@pytest.mark.parametrize(
67 - "test_model",
68 - [TestAllReduceRMSNormModel, TestAllReduceFusedAddRMSNormModel])
139 +@pytest.mark.parametrize("test_model", [
140 + TestAllReduceRMSNormModel,
141 + TestAllReduceFusedAddRMSNormModel,
142 + TestAllReduceFusedAddRMSNormStaticQuantFP8Model,
143 + TestAllReduceFusedAddRMSNormStaticQuantFP4Model,
144 +])
145 @pytest.mark.parametrize("batch_size", [8])
146 @pytest.mark.parametrize("seq_len", [8])
71 -@pytest.mark.parametrize("hidden_size", [4096])
147 +@pytest.mark.parametrize("hidden_size", [16])
148 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
149 @pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"],
150 reason="Only test on CUDA")
75 -@pytest.mark.skipif(not find_spec("flashinfer"),
76 - reason="flashinfer is not installed")
77 -@pytest.mark.skipif(not current_platform.is_device_capability(100),
78 - reason="Only test on SM100")
151 +@pytest.mark.skipif(
152 + not find_spec("flashinfer")
153 + or not has_module_attribute("flashinfer.comm", "trtllm_allreduce_fusion"),
154 + reason="flashinfer is not found or flashinfer "
155 + "is not compiled with trtllm_allreduce_fusion")
156 def test_all_reduce_fusion_pass_replace(test_model: torch.nn.Module,
157 batch_size: int, seq_len: int,
158 hidden_size: int, dtype: torch.dtype):
159 num_processes = 2
160 + if (test_model == TestAllReduceFusedAddRMSNormStaticQuantFP4Model
161 + and not current_platform.has_device_capability(100)):
162 + pytest.skip("Skip as nvfp4 is only supported on "
163 + "devices with compute capability 10.0 (Blackwell)")
164
165 def run_torch_spawn(fn, nprocs):
166 torch.multiprocessing.spawn(fn,
@@ -113,12 +194,11 @@ def all_reduce_fusion_pass_on_test_model(local_rank: int, world_size: int,
194 init_distributed_environment()
195 initialize_model_parallel(tensor_model_parallel_size=world_size)
196
116 - vllm_config = VllmConfig(
117 - compilation_config=CompilationConfig(level=CompilationLevel.PIECEWISE,
118 - custom_ops=["+rms_norm"],
119 - compile_sizes=[2, 4, 8]))
197 + vllm_config = VllmConfig(compilation_config=CompilationConfig(
198 + level=CompilationLevel.PIECEWISE,
199 + custom_ops=["+rms_norm", "+quant_fp8"]))
200 vllm_config.compilation_config.pass_config = PassConfig(
121 - enable_fi_allreduce_fusion=True)
201 + enable_fi_allreduce_fusion=True, enable_noop=True)
202 vllm_config.device_config = DeviceConfig(device=torch.device("cuda"))
203
204 # this is a fake model name to construct the model config
@@ -130,14 +210,16 @@ def all_reduce_fusion_pass_on_test_model(local_rank: int, world_size: int,
210 seed=42)
211
212 all_reduce_fusion_pass = AllReduceFusionPass(vllm_config)
133 - backend = TestBackend(all_reduce_fusion_pass)
213 + noop_pass = NoOpEliminationPass(vllm_config)
214 + func_pass = FixFunctionalizationPass(vllm_config)
215 +
216 + backend = TestBackend(all_reduce_fusion_pass, noop_pass, func_pass)
217
135 - model = test_model_cls(hidden_size)
218 + token_num = batch_size * seq_len
219 + model = test_model_cls(hidden_size, token_num)
220
137 - hidden_states = torch.randn((batch_size * seq_len, hidden_size),
138 - requires_grad=False)
139 - residual = torch.randn((batch_size * seq_len, hidden_size),
140 - requires_grad=False)
221 + hidden_states = torch.randn((token_num, hidden_size), requires_grad=False)
222 + residual = torch.randn((token_num, hidden_size), requires_grad=False)
223
224 compiled_model = torch.compile(model, backend=backend)
225 compiled_model(hidden_states, residual)
tests/utils.py
+12
@@ -4,6 +4,7 @@
4 import asyncio
5 import copy
6 import functools
7 +import importlib
8 import os
9 import signal
10 import subprocess
@@ -974,3 +975,14 @@ def get_client_text_logprob_generations(
975 return [(text_generations, text,
976 (None if x.logprobs is None else x.logprobs.top_logprobs))
977 for completion in completions for x in completion.choices]
978 +
979 +
980 +def has_module_attribute(module_name, attribute_name):
981 + """
982 + Helper function to check if a module has a specific attribute.
983 + """
984 + try:
985 + module = importlib.import_module(module_name)
986 + return hasattr(module, attribute_name)
987 + except ImportError:
988 + return False
vllm/compilation/collective_fusion.py
+488 -45
@@ -37,6 +37,8 @@ logger = init_logger(__name__)
37 ALLREDUCE_OP = torch.ops.vllm.all_reduce.default
38 RMS_OP = torch.ops._C.rms_norm.default
39 RMS_ADD_OP = torch.ops._C.fused_add_rms_norm.default
40 +STATIC_FP8_QUANT_OP = torch.ops._C.static_scaled_fp8_quant.default
41 +STATIC_FP4_QUANT_OP = torch.ops._C.scaled_fp4_quant.default
42
43
44 class BasePattern:
@@ -394,7 +396,7 @@ if flashinfer_comm is not None:
396 # Max size of the input tensor per world size
397 # to use flashinfer fused allreduce
398 _FI_MAX_SIZES = {
397 - 2: MiB, # 1MB
399 + 2: 64 * MiB, # 64MB
400 4: MiB, # 1MB
401 6: MiB // 2, # 512KB
402 8: MiB // 2, # 512KB
@@ -414,9 +416,13 @@ if flashinfer_comm is not None:
416 trigger_completion_at_end: bool,
417 fp32_acc: bool,
418 max_token_num: int,
419 + pattern_code: int,
420 + fuse_rms_quant: bool,
421 norm_out: Optional[torch.Tensor] = None,
422 + quant_out: Optional[torch.Tensor] = None,
423 + scale_out: Optional[torch.Tensor] = None,
424 + scale_factor: Optional[torch.Tensor] = None,
425 ) -> None:
419 -
426 num_tokens, hidden_size = allreduce_in.shape
427 element_size = allreduce_in.element_size()
428 current_tensor_size = num_tokens * hidden_size * element_size
@@ -425,7 +431,6 @@ if flashinfer_comm is not None:
431 _FI_MAX_SIZES.get(world_size, _DEFAULT_FI_MAX_SIZE),
432 max_fusion_size,
433 )
428 -
434 if use_flashinfer:
435 assert (_FI_WORKSPACE_TENSOR is not None
436 ), "Flashinfer must be enabled when using flashinfer"
@@ -455,37 +460,65 @@ if flashinfer_comm is not None:
460 use_oneshot=True,
461 trigger_completion_at_end=trigger_completion_at_end,
462 fp32_acc=fp32_acc,
458 - pattern_code=flashinfer_comm.AllReduceFusionPattern.
459 - kARResidualRMSNorm,
463 + pattern_code=pattern_code,
464 allreduce_out=None,
461 - quant_out=None,
462 - scale_out=None,
463 - layout_code=None,
464 - scale_factor=None,
465 + quant_out=quant_out,
466 + scale_out=scale_out,
467 + # in vllm we only support swizzled layout
468 + layout_code=flashinfer_comm.FP4QuantizationSFLayout.SWIZZLED,
469 + scale_factor=scale_factor,
470 )
471 else:
472 allreduce_out = tensor_model_parallel_all_reduce(allreduce_in)
468 - if norm_out is None:
469 - torch.ops._C.fused_add_rms_norm(allreduce_out, residual,
470 - rms_gamma, rms_eps)
473 + if (scale_factor is not None and scale_out is None
474 + and fuse_rms_quant):
475 + # Do fused rms norm static fp8 quant fused op
476 + if norm_out is None:
477 + torch.ops._C.fused_add_rms_norm_static_fp8_quant(
478 + quant_out, allreduce_out, residual, rms_gamma,
479 + scale_factor, rms_eps)
480 + else:
481 + torch.ops._C.rms_norm_static_fp8_quant(
482 + quant_out, allreduce_out, rms_gamma, scale_factor,
483 + rms_eps)
484 else:
472 - torch.ops._C.rms_norm(norm_out, allreduce_out, rms_gamma,
473 - rms_eps)
474 - allreduce_in.copy_(allreduce_out)
485 + if norm_out is None:
486 + torch.ops._C.fused_add_rms_norm(allreduce_out, residual,
487 + rms_gamma, rms_eps)
488 + norm_out = allreduce_out
489 + else:
490 + torch.ops._C.rms_norm(norm_out, allreduce_out, rms_gamma,
491 + rms_eps)
492 + if scale_factor is not None:
493 + if scale_out is not None:
494 + torch.ops._C.scaled_fp4_quant(quant_out, norm_out,
495 + scale_out, scale_factor)
496 + else:
497 + torch.ops._C.static_scaled_fp8_quant(
498 + quant_out, norm_out, scale_factor)
499 + if scale_factor is None or norm_out is not None:
500 + # we need to return allreduce outpput
501 + # in cases of non quant fused AR + RMS norm
502 + # and fused AR + RMS norm + quant without fused add
503 + allreduce_in.copy_(allreduce_out)
504
505 def call_trtllm_fused_allreduce_norm_fake(
477 - allreduce_in: torch.Tensor,
478 - residual: torch.Tensor,
479 - rms_gamma: torch.Tensor,
480 - rms_eps: float,
481 - world_rank: int,
482 - world_size: int,
483 - launch_with_pdl: bool,
484 - trigger_completion_at_end: bool,
485 - fp32_acc: bool,
486 - max_token_num: int,
487 - norm_out: Optional[torch.Tensor] = None,
488 - ) -> None:
506 + allreduce_in: torch.Tensor,
507 + residual: torch.Tensor,
508 + rms_gamma: torch.Tensor,
509 + rms_eps: float,
510 + world_rank: int,
511 + world_size: int,
512 + launch_with_pdl: bool,
513 + trigger_completion_at_end: bool,
514 + fp32_acc: bool,
515 + max_token_num: int,
516 + pattern_code: int,
517 + fuse_rms_quant: bool,
518 + norm_out: Optional[torch.Tensor] = None,
519 + quant_out: Optional[torch.Tensor] = None,
520 + scale_out: Optional[torch.Tensor] = None,
521 + scale_factor: Optional[torch.Tensor] = None) -> None:
522 pass
523
524 direct_register_custom_op(
@@ -495,6 +528,8 @@ if flashinfer_comm is not None:
528 "allreduce_in",
529 "residual",
530 "norm_out",
531 + "quant_out",
532 + "scale_out",
533 ],
534 fake_impl=call_trtllm_fused_allreduce_norm_fake,
535 dispatch_key=current_platform.dispatch_key,
@@ -512,6 +547,7 @@ class FlashInferFusedAllReduceParams:
547 world_size: int,
548 use_fp32_lamport: bool = False,
549 max_token_num: int = 1024,
550 + fuse_rms_quant: bool = False,
551 ):
552 self.rank = rank
553 self.world_size = world_size
@@ -521,6 +557,7 @@ class FlashInferFusedAllReduceParams:
557 self.fp32_acc = True
558 self.use_oneshot = False
559 self.max_token_num = max_token_num
560 + self.fuse_rms_quant = fuse_rms_quant
561
562 def get_trtllm_fused_allreduce_kwargs(self):
563 return {
@@ -530,10 +567,16 @@ class FlashInferFusedAllReduceParams:
567 "trigger_completion_at_end": self.trigger_completion_at_end,
568 "fp32_acc": self.fp32_acc,
569 "max_token_num": self.max_token_num,
570 + "fuse_rms_quant": self.fuse_rms_quant,
571 }
572
573
536 -class AllReduceRMSNORMPattern(BasePattern):
574 +class AllReduceRMSNormPattern(BasePattern):
575 + """
576 + This pattern replaces the allreduce + rms norm (without residual)
577 + with fused flashinfer implementation.
578 + Applies to allreduce + rmsnorm before attn in the first Transformer block.
579 + """
580
581 def __init__(
582 self,
@@ -559,29 +602,34 @@ class AllReduceRMSNORMPattern(BasePattern):
602
603 def pattern(input: torch.Tensor, rms_result: torch.Tensor,
604 weight: torch.Tensor):
562 - all_reduce_output = tensor_model_parallel_all_reduce(input)
605 + allreduce_output = tensor_model_parallel_all_reduce(input)
606 rms = auto_functionalized(
607 RMS_OP,
608 result=rms_result,
566 - input=all_reduce_output,
609 + input=allreduce_output,
610 weight=weight,
611 epsilon=self.epsilon,
612 )
570 - return rms[1], all_reduce_output
613 + # rms_result, allreduce_output
614 + return rms[1], allreduce_output
615
616 def replacement(input: torch.Tensor, rms_result: torch.Tensor,
617 weight: torch.Tensor):
618 residual = torch.zeros_like(input)
619 allreduce = auto_functionalized(
576 - torch.ops.vllm.flashinfer_trtllm_fused_allreduce_norm.default,
620 + flashinfer_trtllm_fused_allreduce_norm,
621 allreduce_in=input,
622 residual=residual,
623 norm_out=rms_result,
624 + quant_out=None,
625 + scale_out=None,
626 rms_gamma=weight,
627 rms_eps=self.epsilon,
628 + pattern_code=flashinfer_comm.AllReduceFusionPattern.
629 + kARResidualRMSNorm,
630 **self.allreduce_params.get_trtllm_fused_allreduce_kwargs(),
631 )
584 -
632 + # rms_result, allreduce_in
633 return allreduce[3], allreduce[1]
634
635 pm.register_replacement(pattern, replacement, self.get_inputs(),
@@ -589,6 +637,11 @@ class AllReduceRMSNORMPattern(BasePattern):
637
638
639 class AllReduceFusedAddRMSNormPattern(BasePattern):
640 + """
641 + This pattern replaces the allreduce + rms norm (with residual)
642 + with fused flashinfer implementation.
643 + Applies to o_proj + rmsnorm after attn and mlp + rmsnorm before attn.
644 + """
645
646 def __init__(
647 self,
@@ -615,33 +668,390 @@ class AllReduceFusedAddRMSNormPattern(BasePattern):
668
669 def pattern(residual: torch.Tensor, input: torch.Tensor,
670 weight: torch.Tensor):
618 - all_reduce_output = tensor_model_parallel_all_reduce(input)
671 + allreduce_output = tensor_model_parallel_all_reduce(input)
672 rms = auto_functionalized(
673 RMS_ADD_OP,
621 - input=all_reduce_output,
674 + input=allreduce_output,
675 residual=residual,
676 weight=weight,
677 epsilon=self.epsilon,
678 )
679 + # input, residual
680 return rms[1], rms[2]
681
682 def replacement(residual: torch.Tensor, input: torch.Tensor,
683 weight: torch.Tensor):
684 allreduce = auto_functionalized(
631 - torch.ops.vllm.flashinfer_trtllm_fused_allreduce_norm.default,
685 + flashinfer_trtllm_fused_allreduce_norm,
686 allreduce_in=input,
687 residual=residual,
688 + norm_out=None,
689 + quant_out=None,
690 + scale_out=None,
691 rms_gamma=weight,
692 rms_eps=self.epsilon,
636 - norm_out=None,
693 + pattern_code=flashinfer_comm.AllReduceFusionPattern.
694 + kARResidualRMSNorm,
695 **self.allreduce_params.get_trtllm_fused_allreduce_kwargs(),
696 )
697 + # allreduce_in, residual
698 return allreduce[1], allreduce[2]
699
700 pm.register_replacement(pattern, replacement, self.get_inputs(),
701 pm.fwd_only, pm_pass)
702
703
704 +class AllReduceFusedRMSNormStaticQuantFP8Pattern(BasePattern):
705 + """
706 + This pattern replaces the allreduce + rms norm (without residual)
707 + + static fp8 quant with fused flashinfer implementation.
708 + Applies to allreduce + rmsnorm + quant before attn
709 + in the first Transformer block.
710 + """
711 +
712 + def __init__(self, epsilon: float, dtype: torch.dtype, device: str,
713 + allreduce_params: FlashInferFusedAllReduceParams):
714 + super().__init__(dtype, device)
715 + self.epsilon = epsilon
716 + self.allreduce_params = allreduce_params
717 + self.quant_dtype = torch.float8_e4m3fn
718 +
719 + def register(self, pm_pass: PatternMatcherPass):
720 +
721 + def get_inputs():
722 + input = torch.zeros([1, 8, 4],
723 + device=self.device,
724 + dtype=self.dtype)
725 + rmsnorm_result = torch.empty([1, 8, 4],
726 + device=self.device,
727 + dtype=self.dtype)
728 + quant_result = torch.empty([1, 8, 4],
729 + device=self.device,
730 + dtype=self.quant_dtype)
731 + weight = torch.empty([4], device=self.device, dtype=self.dtype)
732 + scale = torch.tensor(1.0, device=self.device, dtype=torch.float32)
733 + return [input, rmsnorm_result, quant_result, weight, scale]
734 +
735 + def pattern(
736 + input: torch.Tensor,
737 + rmsnorm_result: torch.Tensor,
738 + quant_result: torch.Tensor,
739 + weight: torch.Tensor,
740 + scale: torch.Tensor,
741 + ):
742 + all_reduce = tensor_model_parallel_all_reduce(input)
743 + rmsnorm_out_tuple = auto_functionalized(RMS_OP,
744 + result=rmsnorm_result,
745 + input=all_reduce,
746 + weight=weight,
747 + epsilon=self.epsilon)
748 +
749 + quant_out_tuple = auto_functionalized(STATIC_FP8_QUANT_OP,
750 + result=quant_result,
751 + input=rmsnorm_out_tuple[1],
752 + scale=scale)
753 +
754 + # quant_out, allreduce_output
755 + return quant_out_tuple[1], all_reduce
756 +
757 + def replacement(input: torch.Tensor, result_rms: torch.Tensor,
758 + quant_result: torch.Tensor, weight: torch.Tensor,
759 + scale: torch.Tensor):
760 + residual = torch.zeros_like(input)
761 + allreduce = auto_functionalized(
762 + flashinfer_trtllm_fused_allreduce_norm,
763 + allreduce_in=input,
764 + residual=residual,
765 + norm_out=result_rms,
766 + quant_out=quant_result,
767 + scale_out=None,
768 + rms_gamma=weight,
769 + rms_eps=self.epsilon,
770 + pattern_code=flashinfer_comm.AllReduceFusionPattern.
771 + kARResidualRMSNormFP8Quant, # we don't use norm_out afterwards
772 + scale_factor=scale,
773 + **self.allreduce_params.get_trtllm_fused_allreduce_kwargs(),
774 + )
775 +
776 + # quant_out, allreduce_output
777 + return allreduce[4], allreduce[1]
778 +
779 + pm.register_replacement(pattern, replacement, get_inputs(),
780 + pm.fwd_only, pm_pass)
781 +
782 +
783 +class AllReduceFusedAddRMSNormStaticQuantFP8Pattern(BasePattern):
784 + """
785 + This pattern replaces the allreduce + rms norm (with residual)
786 + + static fp8 quant with fused flashinfer implementation.
787 + Applies to o_proj + rmsnorm after attn + quant and
788 + mlp + rmsnorm + quant before attn.
789 + """
790 +
791 + def __init__(self, epsilon: float, dtype: torch.dtype, device: str,
792 + allreduce_params: FlashInferFusedAllReduceParams):
793 + super().__init__(dtype, device)
794 + self.epsilon = epsilon
795 + self.allreduce_params = allreduce_params
796 + self.quant_dtype = torch.float8_e4m3fn
797 +
798 + def register(self, pm_pass: PatternMatcherPass):
799 +
800 + def get_inputs():
801 + input = torch.empty([4, 4], device=self.device, dtype=self.dtype)
802 +
803 + residual = torch.empty([4, 4],
804 + device=self.device,
805 + dtype=self.dtype)
806 + weight = torch.empty([4, 4], device=self.device, dtype=self.dtype)
807 + quant_result = torch.empty([4, 4],
808 + device=self.device,
809 + dtype=self.quant_dtype)
810 + scale = torch.empty([1, 1],
811 + device=self.device,
812 + dtype=torch.float32)
813 +
814 + return [
815 + quant_result,
816 + residual,
817 + input,
818 + weight,
819 + scale,
820 + ]
821 +
822 + def pattern(
823 + quant_result: torch.Tensor,
824 + residual: torch.Tensor,
825 + input: torch.Tensor,
826 + weight: torch.Tensor,
827 + scale: torch.Tensor,
828 + ):
829 + allreduce_output = tensor_model_parallel_all_reduce(input)
830 +
831 + fused_add_rmsnorm_out_tuple = \
832 + auto_functionalized(
833 + RMS_ADD_OP,
834 + input=allreduce_output,
835 + residual=residual,
836 + weight=weight,
837 + epsilon=self.epsilon)
838 + quant_out_tuple = auto_functionalized(
839 + STATIC_FP8_QUANT_OP,
840 + result=quant_result,
841 + input=fused_add_rmsnorm_out_tuple[1],
842 + scale=scale)
843 +
844 + # quant_out, allreduce_output
845 + return quant_out_tuple[1], fused_add_rmsnorm_out_tuple[2]
846 +
847 + def replacement(quant_result: torch.Tensor, residual: torch.Tensor,
848 + input: torch.Tensor, weight: torch.Tensor,
849 + scale: torch.Tensor):
850 + allreduce = auto_functionalized(
851 + flashinfer_trtllm_fused_allreduce_norm,
852 + allreduce_in=input,
853 + residual=residual,
854 + norm_out=None,
855 + quant_out=quant_result,
856 + scale_out=None,
857 + rms_gamma=weight,
858 + rms_eps=self.epsilon,
859 + pattern_code=flashinfer_comm.AllReduceFusionPattern.
860 + kARResidualRMSNormFP8Quant, # we don't use norm_out afterwards
861 + scale_factor=scale,
862 + **self.allreduce_params.get_trtllm_fused_allreduce_kwargs(),
863 + )
864 + # # quant_out, rms_norm_residual
865 + return allreduce[4], allreduce[2]
866 +
867 + pm.register_replacement(pattern, replacement, get_inputs(),
868 + pm.fwd_only, pm_pass)
869 +
870 +
871 +class AllReduceFusedRMSNormStaticQuantNVFP4Pattern(BasePattern):
872 + """
873 + This pattern replaces the allreduce + rms norm (without residual)
874 + + static nvfp4 quant with fused flashinfer implementation.
875 + Applies to allreduce + rmsnorm + quant before attn
876 + in the first Transformer block.
877 + """
878 +
879 + def __init__(self, epsilon: float, dtype: torch.dtype, device: str,
880 + allreduce_params: FlashInferFusedAllReduceParams):
881 + super().__init__(dtype, device)
882 + self.epsilon = epsilon
883 + self.allreduce_params = allreduce_params
884 +
885 + def register(self, pm_pass: PatternMatcherPass):
886 +
887 + def get_inputs():
888 + input = torch.empty([1, 16, 16],
889 + device=self.device,
890 + dtype=self.dtype)
891 +
892 + rmsnorm_result = torch.empty([1, 16, 16],
893 + device=self.device,
894 + dtype=self.dtype)
895 + quant_result = torch.empty((16, 8),
896 + device=self.device,
897 + dtype=torch.uint8)
898 + input_global_scale = torch.empty([1, 1],
899 + device=self.device,
900 + dtype=torch.float32)
901 + weight = torch.empty([16], device=self.device, dtype=self.dtype)
902 + output_scale = torch.empty([128, 4],
903 + device=self.device,
904 + dtype=torch.int32)
905 +
906 + return [
907 + input, rmsnorm_result, quant_result, weight,
908 + input_global_scale, output_scale
909 + ]
910 +
911 + def pattern(
912 + input: torch.Tensor,
913 + rmsnorm_result: torch.Tensor,
914 + quant_result: torch.Tensor,
915 + weight: torch.Tensor,
916 + input_global_scale: torch.Tensor,
917 + output_scale: torch.Tensor,
918 + ):
919 + all_reduce = tensor_model_parallel_all_reduce(input)
920 + rmsnorm_out_tuple = auto_functionalized(RMS_OP,
921 + result=rmsnorm_result,
922 + input=all_reduce,
923 + weight=weight,
924 + epsilon=self.epsilon)
925 +
926 + quant_out_tuple = auto_functionalized(
927 + STATIC_FP4_QUANT_OP,
928 + output=quant_result,
929 + input=rmsnorm_out_tuple[1],
930 + output_scale=output_scale,
931 + input_scale=input_global_scale)
932 +
933 + # quant_out, allreduce_output, output_scale
934 + return quant_out_tuple[1], all_reduce, quant_out_tuple[2]
935 +
936 + def replacement(input: torch.Tensor, result_rms: torch.Tensor,
937 + quant_result: torch.Tensor, weight: torch.Tensor,
938 + input_global_scale: torch.Tensor,
939 + output_scale: torch.Tensor):
940 + residual = torch.zeros_like(input)
941 + allreduce = auto_functionalized(
942 + flashinfer_trtllm_fused_allreduce_norm,
943 + allreduce_in=input,
944 + residual=residual,
945 + norm_out=result_rms,
946 + quant_out=quant_result,
947 + scale_out=output_scale,
948 + rms_gamma=weight,
949 + rms_eps=self.epsilon,
950 + pattern_code=flashinfer_comm.AllReduceFusionPattern.
951 + kARResidualRMSNormFP4Quant, # we don't use norm_out afterwards
952 + scale_factor=input_global_scale,
953 + **self.allreduce_params.get_trtllm_fused_allreduce_kwargs(),
954 + )
955 +
956 + # quant_out, allreduce_output, output_scale
957 + return allreduce[4], allreduce[1], allreduce[5]
958 +
959 + pm.register_replacement(pattern, replacement, get_inputs(),
960 + pm.fwd_only, pm_pass)
961 +
962 +
963 +class AllReduceFusedAddRMSNormStaticQuantNVFP4Pattern(BasePattern):
964 + """
965 + This pattern replaces the allreduce + rms norm (with residual)
966 + + static nvfp4 quant with fused flashinfer implementation.
967 + Applies to o_proj + rmsnorm after attn + quant and
968 + mlp + rmsnorm + quant before attn.
969 + """
970 +
971 + def __init__(self, epsilon: float, dtype: torch.dtype, device: str,
972 + allreduce_params: FlashInferFusedAllReduceParams):
973 + super().__init__(dtype, device)
974 + self.epsilon = epsilon
975 + self.allreduce_params = allreduce_params
976 +
977 + def register(self, pm_pass: PatternMatcherPass):
978 +
979 + def get_inputs():
980 + input = torch.empty([16, 16], device=self.device, dtype=self.dtype)
981 +
982 + residual = torch.empty([16, 16],
983 + device=self.device,
984 + dtype=self.dtype)
985 + weight = torch.empty([16, 16],
986 + device=self.device,
987 + dtype=self.dtype)
988 + quant_result = torch.empty((16, 8),
989 + device=self.device,
990 + dtype=torch.uint8)
991 + input_global_scale = torch.empty([1, 1],
992 + device=self.device,
993 + dtype=torch.float32)
994 + output_scale = torch.empty([128, 4],
995 + device=self.device,
996 + dtype=torch.int32)
997 +
998 + return [
999 + quant_result,
1000 + residual,
1001 + input,
1002 + output_scale,
1003 + weight,
1004 + input_global_scale,
1005 + ]
1006 +
1007 + def pattern(quant_result: torch.Tensor, residual: torch.Tensor,
1008 + input: torch.Tensor, output_scale: torch.Tensor,
1009 + weight: torch.Tensor, input_global_scale: torch.Tensor):
1010 + allreduce_output = tensor_model_parallel_all_reduce(input)
1011 +
1012 + fused_add_rmsnorm_out_tuple = \
1013 + auto_functionalized(
1014 + RMS_ADD_OP,
1015 + input=allreduce_output,
1016 + residual=residual,
1017 + weight=weight,
1018 + epsilon=self.epsilon)
1019 + quant_out_tuple = auto_functionalized(
1020 + STATIC_FP4_QUANT_OP,
1021 + output=quant_result,
1022 + input=fused_add_rmsnorm_out_tuple[1],
1023 + output_scale=output_scale,
1024 + input_scale=input_global_scale)
1025 +
1026 + # quant_out, allreduce_output, output_scale
1027 + return quant_out_tuple[1], fused_add_rmsnorm_out_tuple[
1028 + 2], quant_out_tuple[2]
1029 +
1030 + def replacement(quant_result: torch.Tensor, residual: torch.Tensor,
1031 + input: torch.Tensor, output_scale: torch.Tensor,
1032 + weight: torch.Tensor,
1033 + input_global_scale: torch.Tensor):
1034 + allreduce = auto_functionalized(
1035 + flashinfer_trtllm_fused_allreduce_norm,
1036 + allreduce_in=input,
1037 + residual=residual,
1038 + norm_out=None,
1039 + quant_out=quant_result,
1040 + scale_out=output_scale,
1041 + rms_gamma=weight,
1042 + rms_eps=self.epsilon,
1043 + pattern_code=flashinfer_comm.AllReduceFusionPattern.
1044 + kARResidualRMSNormFP4Quant, # we don't use norm_out afterwards
1045 + scale_factor=input_global_scale,
1046 + **self.allreduce_params.get_trtllm_fused_allreduce_kwargs(),
1047 + )
1048 + # quant_out, rms_norm_residual, output_scale
1049 + return allreduce[4], allreduce[2], allreduce[5]
1050 +
1051 + pm.register_replacement(pattern, replacement, get_inputs(),
1052 + pm.fwd_only, pm_pass)
1053 +
1054 +
1055 class AllReduceFusionPass(VllmInductorPass):
1056
1057 def __init__(self, config: VllmConfig):
@@ -671,13 +1081,16 @@ class AllReduceFusionPass(VllmInductorPass):
1081 self.tp_size,
1082 )
1083 return
674 -
1084 + max_num_token = min(
1085 + _FI_MAX_SIZES.get(self.tp_size, _DEFAULT_FI_MAX_SIZE) //
1086 + (self.hidden_dim * self.tp_size * (4 if use_fp32_lamport else 2)),
1087 + config.compilation_config.pass_config.
1088 + fi_allreduce_fusion_max_token_num)
1089 self.ipc_handles, workspace_tensor = (
1090 flashinfer_comm.trtllm_create_ipc_workspace_for_all_reduce_fusion(
1091 tp_rank=rank,
1092 tp_size=self.tp_size,
679 - max_token_num=config.compilation_config.pass_config.
680 - fi_allreduce_fusion_max_token_num,
1093 + max_token_num=max_num_token,
1094 hidden_dim=self.hidden_dim,
1095 group=self.group,
1096 use_fp32_lamport=use_fp32_lamport,
@@ -689,12 +1102,38 @@ class AllReduceFusionPass(VllmInductorPass):
1102 rank=rank,
1103 world_size=self.tp_size,
1104 use_fp32_lamport=use_fp32_lamport,
692 - max_token_num=config.compilation_config.pass_config.
693 - fi_allreduce_fusion_max_token_num,
694 - )
1105 + max_token_num=max_num_token,
1106 + # fuse rms norm static fp8 quant fused op
1107 + # in fallback path, when we don't use flashinfer
1108 + fuse_rms_quant=config.compilation_config.pass_config.enable_fusion)
1109
1110 for epsilon in [1e-5, 1e-6]:
697 - AllReduceRMSNORMPattern(
1111 + AllReduceFusedRMSNormStaticQuantFP8Pattern(
1112 + epsilon,
1113 + self.model_dtype,
1114 + self.device,
1115 + self.allreduce_params,
1116 + ).register(self.patterns)
1117 + AllReduceFusedAddRMSNormStaticQuantFP8Pattern(
1118 + epsilon,
1119 + self.model_dtype,
1120 + self.device,
1121 + self.allreduce_params,
1122 + ).register(self.patterns)
1123 + if current_platform.has_device_capability(100):
1124 + AllReduceFusedRMSNormStaticQuantNVFP4Pattern(
1125 + epsilon,
1126 + self.model_dtype,
1127 + self.device,
1128 + self.allreduce_params,
1129 + ).register(self.patterns)
1130 + AllReduceFusedAddRMSNormStaticQuantNVFP4Pattern(
1131 + epsilon,
1132 + self.model_dtype,
1133 + self.device,
1134 + self.allreduce_params,
1135 + ).register(self.patterns)
1136 + AllReduceRMSNormPattern(
1137 epsilon,
1138 self.model_dtype,
1139 self.device,
@@ -707,6 +1146,10 @@ class AllReduceFusionPass(VllmInductorPass):
1146 self.allreduce_params,
1147 ).register(self.patterns)
1148
1149 + # WARNING: This is a hack to clear the pattern matcher cache
1150 + # and allow multiple values of epsilon.
1151 + torch._inductor.pattern_matcher._seen_patterns.clear()
1152 +
1153 self.disabled = False
1154
1155 def __call__(self, graph: fx.Graph):
@@ -723,5 +1166,5 @@ class AllReduceFusionPass(VllmInductorPass):
1166 if self.disabled:
1167 return
1168 if flashinfer_comm is not None:
726 - flashinfer_comm.trtllm_destroy_ipc_workspace(
1169 + flashinfer_comm.trtllm_destroy_ipc_workspace_for_all_reduce(
1170 self.ipc_handles, self.group)
vllm/config.py
+1 -1
@@ -4051,7 +4051,7 @@ class PassConfig:
4051 """Whether to enable async TP."""
4052 enable_fi_allreduce_fusion: bool = False
4053 """Whether to enable flashinfer allreduce fusion."""
4054 - fi_allreduce_fusion_max_token_num: int = 1024
4054 + fi_allreduce_fusion_max_token_num: int = 16384
4055 """Max number of tokens to used in flashinfer allreduce fusion."""
4056
4057 # TODO(luka) better pass enabling system.