Fix Transformers modelling backend `RMSNormFuser.fuse` performance (#52766)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Aug 20, 2026 at 16:21 UTC
2dd17225c6a579c8ed0041943c636d74684fcbb5
3 files changed
+157
-10
tests/models/transformers/fusers/test_rms_norm.py
+90
@@ -60,6 +60,54 @@ class LayerNorm(RMSNorm):
60
return self.weight * self._rms(x)
61
62
63
+class NamedEpsRMSNorm(nn.Module):
64
+ """Holds eps under `attr`, with no `variance_epsilon` to find instead."""
65
+
66
+ attr = "eps"
67
+
68
+ def __init__(self, hidden: int = 16, eps: float = 1e-6):
69
+ super().__init__()
70
+ self.weight = nn.Parameter(torch.ones(hidden))
71
+ setattr(self, self.attr, eps)
72
+
73
+ def forward(self, x):
74
+ eps = getattr(self, self.attr)
75
+ return self.weight * (x * torch.rsqrt(x.pow(2).mean(-1, True) + eps))
76
+
77
+
78
+class ToleranceRMSNorm(NamedEpsRMSNorm):
79
+ """Names eps something other than `eps`, so only its value identifies it."""
80
+
81
+ attr = "tolerance"
82
+
83
+
84
+class LiteralEpsRMSNorm(RMSNorm):
85
+ """eps is a literal in the source, and an unrelated attribute happens to
86
+ hold the same value: matching on the value alone would bind to it."""
87
+
88
+ def __init__(self, hidden: int = 16, eps: float = 1e-5):
89
+ super().__init__(hidden, eps)
90
+ self.epsilon_lookalike = 1e-4
91
+
92
+ def forward(self, x):
93
+ return self.weight * (x * torch.rsqrt(x.pow(2).mean(-1, True) + 1e-4))
94
+
95
+
96
+class AmbiguousEpsRMSNorm(nn.Module):
97
+ """Two attributes hold the eps value, and the forward reads the second, so
98
+ their order cannot pick it."""
99
+
100
+ def __init__(self, hidden: int = 16, eps: float = 1e-6):
101
+ super().__init__()
102
+ self.weight = nn.Parameter(torch.ones(hidden))
103
+ self.decoy_eps = eps
104
+ self.variance_epsilon = eps
105
+
106
+ def forward(self, x):
107
+ var = x.pow(2).mean(-1, True)
108
+ return self.weight * (x * torch.rsqrt(var + self.variance_epsilon))
109
+
110
+
111
class NotAnRMSNorm(RMSNorm):
112
"""Mean-subtracting LayerNorm-like math -> not an RMSNorm."""
113
@@ -207,6 +255,48 @@ def test_eps_is_derived_per_instance(default_vllm_config):
255
assert built.variance_epsilon == eps
256
257
258
+@pytest.mark.parametrize("cls", [NamedEpsRMSNorm, ToleranceRMSNorm])
259
+def test_eps_attr_is_found_by_value_not_name(cls, default_vllm_config):
260
+ """The eps attribute is identified by holding the traced value, so a norm
261
+ stays per-instance correct whatever it names it."""
262
+ with torch.device("meta"):
263
+ for eps in (1e-5, 1e-6):
264
+ module = cls(16, eps=eps)
265
+ fuser = get_fuser(module)
266
+ assert fuser.eps_attr == cls.attr
267
+ built = fuser.fuse(module, "norm", default_vllm_config)
268
+ assert built.variance_epsilon == eps
269
+
270
+
271
+def test_literal_eps_is_not_mistaken_for_an_attribute(default_vllm_config, caplog):
272
+ """A literal eps is recognised as coming from no attribute, even when one
273
+ holds the same value, and is taken from the traced source instead."""
274
+ logger = "vllm.model_executor.models.transformers.fusers.rms_norm"
275
+ with caplog.at_level("DEBUG", logger=logger), torch.device("meta"):
276
+ module = LiteralEpsRMSNorm()
277
+ fuser = get_fuser(module)
278
+ built = fuser.fuse(module, "norm", default_vllm_config)
279
+ assert fuser.eps_attr is None
280
+ assert built.variance_epsilon == 1e-4
281
+ assert "does not hold its eps" in caplog.text
282
+
283
+
284
+def test_ambiguous_eps_attrs_are_disambiguated(default_vllm_config):
285
+ """When several attributes hold the eps value, the one the forward actually
286
+ reads is identified, and they are all left as they were found."""
287
+ with torch.device("meta"):
288
+ module = AmbiguousEpsRMSNorm(16, eps=1e-6)
289
+ before = dict(vars(module))
290
+ fuser = get_fuser(module)
291
+ assert fuser.eps_attr == "variance_epsilon"
292
+ assert vars(module) == before
293
+
294
+ other = AmbiguousEpsRMSNorm(16, eps=1e-5)
295
+ other.decoy_eps = 1e-3
296
+ built = fuser.fuse(other, "norm", default_vllm_config)
297
+ assert built.variance_epsilon == 1e-5
298
+
299
+
300
def test_fused_norm_is_gather_capable(default_vllm_config):
301
"""Every weighted fused norm is emitted gather-capable, so a norm on a head-sharded
302
projection (OLMoE-style) self-corrects at runtime with no QKV-specific
vllm/model_executor/models/transformers/fusers/mla.py
+4
-3
@@ -22,7 +22,6 @@ from vllm.model_executor.models.transformers.fx_utils import (
22
replace_expr,
23
returned_linear,
24
single_self_call,
25
- trace,
25
upstream_linear,
26
)
27
from vllm.model_executor.models.transformers.utils import (
@@ -57,8 +56,10 @@ def _norm_size(norm: nn.Module) -> int:
56
57
def _is_rms_norm(module: nn.Module) -> bool:
58
"""Whether `module` computes an RMSNorm, verified by `RMSNormFuser`'s matcher."""
60
- graph = trace(module)
61
- return graph is not None and RMSNormFuser.match(graph, module) is not None
59
+ # Go via `get_fuser` because it caches
60
+ from vllm.model_executor.models.transformers.fuser import get_fuser
61
+
62
+ return isinstance(get_fuser(module), RMSNormFuser)
63
64
65
def _top_level_index(funcdef: ast.FunctionDef, node: ast.AST) -> int:
vllm/model_executor/models/transformers/fusers/rms_norm.py
+63
-7
@@ -15,6 +15,7 @@ from vllm.distributed import (
15
)
16
from vllm.distributed.parallel_state import model_parallel_is_initialized
17
from vllm.distributed.utils import split_tensor_along_last_dim
18
+from vllm.logger import init_logger
19
from vllm.model_executor.models.transformers.fusers.base import BaseFuser
20
from vllm.model_executor.models.transformers.fx_utils import (
21
find_node,
@@ -32,6 +33,8 @@ from vllm.model_executor.models.transformers.layers import (
33
if TYPE_CHECKING:
34
from vllm.config import VllmConfig
35
36
+logger = init_logger(__name__)
37
+
38
39
def _is_squared(node: object, x: fx.Node) -> bool:
40
"""`x**2`, `x.square()` or `x * x`, through any dtype casts."""
@@ -122,6 +125,10 @@ class RMSNormFuser(BaseFuser):
125
"""Gemma-style `(1 + weight)` scaling (weight initialised at zero)."""
126
source_cls: str
127
"""Class name of the norm this was matched from (for logging)."""
128
+ eps_attr: str | None = None
129
+ """Attribute holding `eps`, read per instance in `fuse`."""
130
+ eps: float | None = None
131
+ """`eps` itself, when it is not held in an attribute (see `_eps_source`)."""
132
133
def info(self, name: str) -> str:
134
norm = "GemmaRMSNorm" if self.zero_centered else "RMSNorm"
@@ -140,7 +147,13 @@ class RMSNormFuser(BaseFuser):
147
if rms_norm is not None and rms_norm.args and peel(rms_norm.args[0]) is x:
148
if _has_trailing_compute(graph, rms_norm):
149
return None
143
- return cls(zero_centered=False, source_cls=type(module).__name__)
150
+ eps_attr, eps = cls._eps_source(graph, module)
151
+ return cls(
152
+ zero_centered=False,
153
+ source_cls=type(module).__name__,
154
+ eps_attr=eps_attr,
155
+ eps=eps,
156
+ )
157
# Handle explicit `x * rsqrt(mean(x**2, -1) + eps)` pattern.
158
# The rsqrt over the mean-square variance is the spine of the norm.
159
rsqrt = None
@@ -169,7 +182,51 @@ class RMSNormFuser(BaseFuser):
182
# The norm must be the last compute in forward, or it is not a pure norm.
183
if _has_trailing_compute(graph, tail):
184
return None
172
- return cls(zero_centered=zero_centered, source_cls=type(module).__name__)
185
+ eps_attr, eps = cls._eps_source(graph, module)
186
+ return cls(
187
+ zero_centered=zero_centered,
188
+ source_cls=type(module).__name__,
189
+ eps_attr=eps_attr,
190
+ eps=eps,
191
+ )
192
+
193
+ @classmethod
194
+ def _eps_source(
195
+ cls, graph: fx.Graph, module: nn.Module
196
+ ) -> tuple[str | None, float | None]:
197
+ """Where `fuse` should read `eps` from, resolved once per class."""
198
+ eps = cls._eps_from_graph(graph)
199
+ if eps is None:
200
+ return None, None
201
+ # Whatever supplied the constant must still equal it.
202
+ candidates = {
203
+ name: value
204
+ for name, value in vars(module).items()
205
+ if isinstance(value, float) and value == eps
206
+ }
207
+ # Use unique markers and retrace to verify exactly which attribute is eps.
208
+ markers = {float(-index - 1): name for index, name in enumerate(candidates)}
209
+ marked = None
210
+ if markers:
211
+ try:
212
+ for marker, name in markers.items():
213
+ setattr(module, name, marker)
214
+ if (remarked := trace(module)) is not None:
215
+ marked = cls._eps_from_graph(remarked)
216
+ finally:
217
+ for name, value in candidates.items():
218
+ setattr(module, name, value)
219
+ if (name := markers.get(marked)) is not None:
220
+ return name, None
221
+ logger.debug_once(
222
+ "%s does not hold its eps (%s) in an attribute. Every instance in this "
223
+ "model will use the value traced from this instance. If this is not "
224
+ "desired, consider storing and reading eps using attribute of %s.",
225
+ type(module).__name__,
226
+ eps,
227
+ type(module).__name__,
228
+ )
229
+ return None, eps
230
231
@staticmethod
232
def _eps_from_graph(graph: fx.Graph) -> float | None:
@@ -180,7 +237,7 @@ class RMSNormFuser(BaseFuser):
237
if fused is not None and fused.args and peel(fused.args[0]) is x:
238
args, kwargs = fused.args, fused.kwargs
239
eps = args[3] if len(args) > 3 else kwargs.get("eps")
183
- return eps if isinstance(eps, (int, float)) else None
240
+ return float(eps) if isinstance(eps, (int, float)) else None
241
for node in graph.nodes:
242
if is_op(node, "rsqrt") and (eps := _variance_eps(node, x)) is not None:
243
return eps
@@ -196,10 +253,9 @@ class RMSNormFuser(BaseFuser):
253
weight = getattr(module, "weight", None)
254
has_weight = weight is not None
255
hidden_size = weight.size(0) if has_weight else 0
199
- graph = trace(module)
200
- eps = self._eps_from_graph(graph) if graph is not None else None
201
- if eps is None:
202
- # If eps not in graph, match torch behaviour.
256
+ eps = getattr(module, self.eps_attr, None) if self.eps_attr else self.eps
257
+ if not isinstance(eps, (int, float)):
258
+ # If eps was not detected, match torch behaviour.
259
dtype = weight.dtype if has_weight else vllm_config.model_config.dtype
260
eps = torch.finfo(dtype).eps
261
if self.zero_centered: