[Compilation]Fuse Transformers Residual Add + RMSNorm (#48757)
Signed-off-by: BadrBasowid <badr.basowid@gmail.com>
BadrBasowid committed
Jul 30, 2026 at 21:46 UTC
59e831c09a22bdd2732b86a158ff771584b14793
4 files changed
+298
-2
tests/compile/passes/test_rmsnorm_reshape_fusion.py
new
+105
@@ -0,0 +1,105 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+
7
+import vllm.ir.ops
8
+from tests.compile.backend import TestBackend
9
+from vllm.compilation.passes.fusion.add_rms_fusion import (
10
+ AddRMSNormFusionPass,
11
+ RMSNormReshapeFusionPass,
12
+)
13
+from vllm.compilation.passes.fx_utils import find_op_nodes, is_func
14
+from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
15
+from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
16
+from vllm.config import CompilationConfig, CompilationMode, VllmConfig
17
+from vllm.platforms import current_platform
18
+
19
+pytestmark = pytest.mark.skipif(
20
+ not current_platform.is_cuda_alike(), reason="Requires CUDA or ROCm"
21
+)
22
+
23
+
24
+class RMSNormModel(torch.nn.Module):
25
+ def __init__(self, hidden_size: int) -> None:
26
+ super().__init__()
27
+ self.weight = torch.nn.Parameter(torch.randn(hidden_size))
28
+
29
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
30
+ rms = vllm.ir.ops.rms_norm(x, self.weight, 1e-6)
31
+ return rms.reshape(-1, rms.shape[-1])
32
+
33
+
34
+class AddRMSNormModel(torch.nn.Module):
35
+ def __init__(self, hidden_size: int, residual_first: bool) -> None:
36
+ super().__init__()
37
+ self.weight = torch.nn.Parameter(torch.randn(hidden_size))
38
+ self.residual_first = residual_first
39
+
40
+ def forward(
41
+ self, x: torch.Tensor, residual: torch.Tensor
42
+ ) -> tuple[torch.Tensor, torch.Tensor]:
43
+ residual = residual + x if self.residual_first else x + residual
44
+ rms = vllm.ir.ops.rms_norm(residual, self.weight, 1e-6)
45
+ return rms.reshape(-1, rms.shape[-1]), residual
46
+
47
+ def ops_in_model_before(self):
48
+ return [torch.ops.aten.add, torch.ops.vllm_ir.rms_norm]
49
+
50
+ def ops_in_model_after(self):
51
+ return [torch.ops.vllm_ir.fused_add_rms_norm]
52
+
53
+
54
+def _run_fusion_test(model, config, passes, *inputs):
55
+ backend = TestBackend(NoOpEliminationPass(config), *passes, PostCleanupPass(config))
56
+ outputs_unfused = model(*inputs)
57
+ outputs_fused = torch.compile(model, backend=backend)(*inputs)
58
+ torch.testing.assert_close(outputs_fused, outputs_unfused)
59
+ return backend
60
+
61
+
62
+def _is_reshape(node):
63
+ return isinstance(node, torch.fx.Node) and is_func(
64
+ node, torch.ops.aten.reshape.default
65
+ )
66
+
67
+
68
+@pytest.fixture
69
+def vllm_config():
70
+ config = VllmConfig(
71
+ compilation_config=CompilationConfig(mode=CompilationMode.VLLM_COMPILE)
72
+ )
73
+ with vllm.config.set_current_vllm_config(config):
74
+ torch.set_default_device("cuda")
75
+ torch.set_default_dtype(torch.bfloat16)
76
+ torch.manual_seed(0)
77
+ yield config
78
+
79
+
80
+def test_rmsnorm_reshape_fusion(vllm_config):
81
+ fusion_pass = RMSNormReshapeFusionPass(vllm_config)
82
+ model = RMSNormModel(hidden_size=32)
83
+ x = torch.randn(2, 7, 32)
84
+ backend = _run_fusion_test(model, vllm_config, [fusion_pass], x)
85
+
86
+ assert fusion_pass.matched_count == 1
87
+ (rms_node,) = find_op_nodes(torch.ops.vllm_ir.rms_norm, backend.graph_post_pass)
88
+ assert _is_reshape(rms_node.args[0])
89
+
90
+
91
+@pytest.mark.parametrize("residual_first", [True, False])
92
+def test_add_rmsnorm_reshape_fusion(vllm_config, residual_first):
93
+ add_fusion = AddRMSNormFusionPass(vllm_config)
94
+ reshape_fusion = RMSNormReshapeFusionPass(vllm_config)
95
+ model = AddRMSNormModel(hidden_size=32, residual_first=residual_first)
96
+ x = torch.randn(2, 7, 32)
97
+ residual = torch.randn_like(x)
98
+ backend = _run_fusion_test(
99
+ model, vllm_config, [add_fusion, reshape_fusion], x, residual
100
+ )
101
+
102
+ assert add_fusion.matched_count == 1
103
+ assert reshape_fusion.matched_count == 1
104
+ backend.check_before_ops(model.ops_in_model_before())
105
+ backend.check_after_ops(model.ops_in_model_after())
vllm/compilation/passes/fusion/add_rms_fusion.py
new
+164
@@ -0,0 +1,164 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from itertools import product
5
+
6
+import torch
7
+
8
+import vllm.ir.ops
9
+from vllm.config import VllmConfig
10
+
11
+from ..vllm_inductor_pass import (
12
+ VllmFusionPatternMatcherPass,
13
+ VllmPatternReplacement,
14
+)
15
+
16
+
17
+class AddRMSNormPattern(VllmPatternReplacement):
18
+ def __init__(self, epsilon: float, residual_first: bool) -> None:
19
+ self.epsilon = epsilon
20
+ self.residual_first = residual_first
21
+
22
+ @property
23
+ def pattern(self):
24
+ def _pattern(
25
+ branch: torch.Tensor,
26
+ residual: torch.Tensor,
27
+ weight: torch.Tensor,
28
+ ) -> tuple[torch.Tensor, torch.Tensor]:
29
+ residual_out = (
30
+ residual + branch if self.residual_first else branch + residual
31
+ )
32
+ rms = vllm.ir.ops.rms_norm(residual_out, weight, self.epsilon)
33
+ return rms, residual_out
34
+
35
+ return _pattern
36
+
37
+ @property
38
+ def replacement(self):
39
+ def _replacement(
40
+ branch: torch.Tensor,
41
+ residual: torch.Tensor,
42
+ weight: torch.Tensor,
43
+ ) -> tuple[torch.Tensor, torch.Tensor]:
44
+ return vllm.ir.ops.fused_add_rms_norm(
45
+ branch, residual, weight, self.epsilon
46
+ )
47
+
48
+ return _replacement
49
+
50
+ def get_inputs(self) -> list[torch.Tensor]:
51
+ return [
52
+ self.empty_bf16(5, 16), # branch
53
+ self.empty_bf16(5, 16), # residual
54
+ self.empty_bf16(16), # weight
55
+ ]
56
+
57
+
58
+class RMSNormReshapePattern(VllmPatternReplacement):
59
+ """Move a prefix-flatten before RMSNorm."""
60
+
61
+ def __init__(self, epsilon: float) -> None:
62
+ self.epsilon = epsilon
63
+
64
+ @property
65
+ def pattern(self):
66
+ def _pattern(
67
+ input: torch.Tensor,
68
+ weight: torch.Tensor,
69
+ ) -> torch.Tensor:
70
+ rms = vllm.ir.ops.rms_norm(input, weight, self.epsilon)
71
+ return rms.reshape(-1, rms.shape[-1])
72
+
73
+ return _pattern
74
+
75
+ @property
76
+ def replacement(self):
77
+ def _replacement(
78
+ input: torch.Tensor,
79
+ weight: torch.Tensor,
80
+ ) -> torch.Tensor:
81
+ input = input.reshape(-1, input.shape[-1])
82
+ return vllm.ir.ops.rms_norm(input, weight, self.epsilon)
83
+
84
+ return _replacement
85
+
86
+ def get_inputs(self) -> list[torch.Tensor]:
87
+ return [
88
+ self.empty_bf16(1, 5, 16), # input
89
+ self.empty_bf16(16), # weight
90
+ ]
91
+
92
+
93
+class FusedAddRMSNormReshapePattern(VllmPatternReplacement):
94
+ """Move a prefix-flatten before fused add RMSNorm."""
95
+
96
+ def __init__(self, epsilon: float) -> None:
97
+ self.epsilon = epsilon
98
+
99
+ @property
100
+ def pattern(self):
101
+ def _pattern(
102
+ input: torch.Tensor,
103
+ residual: torch.Tensor,
104
+ weight: torch.Tensor,
105
+ ) -> tuple[torch.Tensor, torch.Tensor]:
106
+ rms, residual_out = vllm.ir.ops.fused_add_rms_norm(
107
+ input, residual, weight, self.epsilon
108
+ )
109
+ return rms.reshape(-1, rms.shape[-1]), residual_out
110
+
111
+ return _pattern
112
+
113
+ @property
114
+ def replacement(self):
115
+ def _replacement(
116
+ input: torch.Tensor,
117
+ residual: torch.Tensor,
118
+ weight: torch.Tensor,
119
+ ) -> tuple[torch.Tensor, torch.Tensor]:
120
+ original_shape = input.shape
121
+ hidden_size = input.shape[-1]
122
+ rms, residual_out = vllm.ir.ops.fused_add_rms_norm(
123
+ input.reshape(-1, hidden_size),
124
+ residual.reshape(-1, hidden_size),
125
+ weight,
126
+ self.epsilon,
127
+ )
128
+ return rms, residual_out.reshape(original_shape)
129
+
130
+ return _replacement
131
+
132
+ def get_inputs(self) -> list[torch.Tensor]:
133
+ return [
134
+ self.empty_bf16(1, 5, 16), # input
135
+ self.empty_bf16(1, 5, 16), # residual
136
+ self.empty_bf16(16), # weight
137
+ ]
138
+
139
+
140
+class AddRMSNormFusionPass(VllmFusionPatternMatcherPass):
141
+ """Fuse residual Add and RMSNorm emitted separately by the Transformers backend."""
142
+
143
+ def __init__(self, config: VllmConfig) -> None:
144
+ super().__init__(config, "add_rmsnorm_fusion_pass")
145
+
146
+ for epsilon, residual_first in product([1e-5, 1e-6], [True, False]):
147
+ self.register(AddRMSNormPattern(epsilon, residual_first))
148
+
149
+ self.dump_patterns(config, self.pm_pass)
150
+
151
+
152
+class RMSNormReshapeFusionPass(VllmFusionPatternMatcherPass):
153
+ """Move the Transformers backend's post-RMSNorm flatten before the norm
154
+ for downstream 2D fusions.
155
+ """
156
+
157
+ def __init__(self, config: VllmConfig) -> None:
158
+ super().__init__(config, "rmsnorm_reshape_fusion_pass")
159
+
160
+ for epsilon in [1e-5, 1e-6]:
161
+ self.register(FusedAddRMSNormReshapePattern(epsilon))
162
+ self.register(RMSNormReshapePattern(epsilon))
163
+
164
+ self.dump_patterns(config, self.pm_pass)
vllm/compilation/passes/pass_manager.py
+22
@@ -30,6 +30,10 @@ if rocm_aiter_ops.is_enabled():
30
)
31
32
if current_platform.is_cuda_alike() or current_platform.is_xpu():
33
+ from .fusion.add_rms_fusion import (
34
+ AddRMSNormFusionPass,
35
+ RMSNormReshapeFusionPass,
36
+ )
37
from .fusion.qk_norm_rope_fusion import QKNormRoPEFusionPass
38
from .fusion.sequence_parallelism import SequenceParallelismPass
39
from .utility.split_coalescing import SplitCoalescingPass
@@ -138,6 +142,16 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc]
142
143
def configure(self, config: VllmConfig) -> None:
144
self.pass_config = config.compilation_config.pass_config
145
+ model_config = config.model_config
146
+ enable_transformers_norm_canonicalization = (
147
+ (
148
+ self.pass_config.fuse_act_padding
149
+ or self.pass_config.fuse_allreduce_rms
150
+ or self.pass_config.fuse_norm_quant
151
+ )
152
+ and model_config is not None
153
+ and model_config.using_transformers_backend()
154
+ )
155
156
# Set the current vllm config to allow tracing CustomOp instances
157
with set_current_vllm_config(config, check_compile=False):
@@ -149,6 +163,9 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc]
163
if self.pass_config.fuse_gemm_comms:
164
self.passes += [AsyncTPPass(config)]
165
166
+ if enable_transformers_norm_canonicalization:
167
+ self.passes += [AddRMSNormFusionPass(config)]
168
+
169
if self.pass_config.fuse_act_padding and rocm_aiter_ops.is_enabled():
170
# Run the more specific RMSNorm+router-pad fusion before
171
# AR+RMS, since both consume fused_add_rms_norm.
@@ -160,6 +177,11 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc]
177
else:
178
self.passes += [AllReduceFusionPass(config)]
179
180
+ if enable_transformers_norm_canonicalization:
181
+ # Let AR+RMS match before moving output reshapes ahead of the
182
+ # remaining RMSNorms, exposing them to RMS+Quant fusion.
183
+ self.passes += [RMSNormReshapeFusionPass(config)]
184
+
185
if self.pass_config.fuse_norm_quant:
186
if rocm_aiter_ops.is_enabled():
187
self.passes += [
vllm/kernels/aiter_ops.py
+7
-2
@@ -75,7 +75,6 @@ def _rms_norm_fake(x: Tensor, weight: Tensor, variance_epsilon: float) -> Tensor
75
direct_register_aiter_op(
76
op_name="rms_norm", op_func=_rms_norm_impl, fake_impl=_rms_norm_fake
77
)
78
-
78
rms_add_no_var_16bit_only = (
79
lambda x, x_residual, weight, epsilon, variance_size=None: variance_size is None
80
and x.dtype in (torch.float16, torch.bfloat16)
@@ -112,6 +111,12 @@ def _rocm_aiter_rmsnorm2d_fwd_with_add_impl(
111
) -> tuple[torch.Tensor, torch.Tensor]:
112
from aiter import rmsnorm2d_fwd_with_add
113
114
+ x_shape = x.shape
115
+ residual_shape = residual.shape
116
+ hidden_size = x.shape[-1]
117
+ x = x.reshape(-1, hidden_size)
118
+ residual = residual.reshape(-1, hidden_size)
119
+
120
# TODO can out = x and residual_out = residual to save memory?
121
# Need to check if the kernel supports in-place residual output
122
# (if yes set mutates_args and inplace)
@@ -125,7 +130,7 @@ def _rocm_aiter_rmsnorm2d_fwd_with_add_impl(
130
weight,
131
variance_epsilon,
132
)
128
- return out, residual_out
133
+ return out.reshape(x_shape), residual_out.reshape(residual_shape)
134
135
136
def _rocm_aiter_rmsnorm2d_fwd_with_add_fake(