[transformer] RMSNorm matching for alternative rsqrt (#54461)
Signed-off-by: Thomas Ortner <boh@zurich.ibm.com> Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Thomas Ortner committed
Sep 24, 2026 at 18:33 UTC
8b84e15066b07cf0064a0ff05f153cc85cb42817
2 files changed
+123
-17
tests/models/transformers/fusers/test_rms_norm.py
+77
@@ -50,6 +50,55 @@ class WeightlessRMSNorm(RMSNorm):
50
return self._rms(x.to(torch.float32)).to(x.dtype)
51
52
53
+class Gemma4RMSNorm(RMSNorm):
54
+ """`pow(v, -0.5)` where the others write `rsqrt(v)` (HF `Gemma4RMSNorm._norm`).
55
+
56
+ Gemma 4 spells the reciprocal square root as `torch.pow(..., -0.5)` to keep
57
+ Torch and JAX in agreement, and -- unlike Gemma 1-3 -- its weight is *not*
58
+ zero-centered, so this fuses to a plain `RMSNorm`.
59
+ """
60
+
61
+ def _rms(self, x):
62
+ mean_squared = x.pow(2).mean(-1, keepdim=True) + self.variance_epsilon
63
+ return x * torch.pow(mean_squared, -0.5)
64
+
65
+ def forward(self, x):
66
+ return (self._rms(x.float()) * self.weight.float()).type_as(x)
67
+
68
+
69
+class PowOperatorRMSNorm(Gemma4RMSNorm):
70
+ """The `v ** -0.5` operator spelling of the same reciprocal square root."""
71
+
72
+ def _rms(self, x):
73
+ return x * (x.pow(2).mean(-1, keepdim=True) + self.variance_epsilon) ** -0.5
74
+
75
+
76
+class KwargPowRMSNorm(Gemma4RMSNorm):
77
+ """The keyword spelling `torch.pow(v, exponent=-0.5)` of the same rsqrt."""
78
+
79
+ def _rms(self, x):
80
+ mean_squared = x.pow(2).mean(-1, keepdim=True) + self.variance_epsilon
81
+ return x * torch.pow(mean_squared, exponent=-0.5)
82
+
83
+
84
+class KwargBasePowRMSNorm(Gemma4RMSNorm):
85
+ """The fully functional spelling `torch.pow(input=v, exponent=-0.5)`."""
86
+
87
+ def _rms(self, x):
88
+ mean_squared = x.pow(2).mean(-1, keepdim=True) + self.variance_epsilon
89
+ return x * torch.pow(input=mean_squared, exponent=-0.5)
90
+
91
+
92
+class WeightlessGemma4RMSNorm(Gemma4RMSNorm):
93
+ """Gemma 4 with `with_scale=False`: the `pow` spelling and no scale parameter."""
94
+
95
+ def __init__(self, hidden: int = 16, eps: float = 1e-6):
96
+ super().__init__(hidden, eps, weight=False)
97
+
98
+ def forward(self, x):
99
+ return self._rms(x.float()).type_as(x)
100
+
101
+
102
class LayerNorm(RMSNorm):
103
"""An RMSNorm not named `*RMSNorm`, keeping the input dtype (no upcast)."""
104
@@ -156,6 +205,11 @@ class UntraceableGatedRMSNorm(RMSNorm):
205
(RMSNorm, 1e-5, False),
206
(GemmaRMSNorm, 1e-6, True),
207
(WeightlessRMSNorm, 1e-6, False),
208
+ (Gemma4RMSNorm, 1e-6, False),
209
+ (PowOperatorRMSNorm, 1e-6, False),
210
+ (KwargPowRMSNorm, 1e-6, False),
211
+ (KwargBasePowRMSNorm, 1e-6, False),
212
+ (WeightlessGemma4RMSNorm, 1e-6, False),
213
(LayerNorm, 1e-6, False),
214
(torch.nn.RMSNorm, 1e-5, False), # fused `F.rms_norm` op
215
],
@@ -188,6 +242,10 @@ def test_gated_rms_norm_is_not_fused(cls):
242
(RMSNorm, "RMSNorm", False),
243
(GemmaRMSNorm, "GemmaRMSNorm", True),
244
(WeightlessRMSNorm, "RMSNorm", False),
245
+ (Gemma4RMSNorm, "RMSNorm", False),
246
+ (KwargPowRMSNorm, "RMSNorm", False),
247
+ (KwargBasePowRMSNorm, "RMSNorm", False),
248
+ (WeightlessGemma4RMSNorm, "RMSNorm", False),
249
],
250
)
251
def test_rms_norm_builds_vllm_class(cls, expected, zero_centered, default_vllm_config):
@@ -213,6 +271,25 @@ def test_rms_norm_builds_vllm_class(cls, expected, zero_centered, default_vllm_c
271
assert built.weight.shape[0] == (weight.size(0) if weight is not None else 0)
272
273
274
+@pytest.mark.parametrize(
275
+ "cls", [Gemma4RMSNorm, PowOperatorRMSNorm, KwargPowRMSNorm, KwargBasePowRMSNorm]
276
+)
277
+def test_pow_spelled_norm_fuses_to_equivalent_math(cls, default_vllm_config):
278
+ """Matching `pow(v, -0.5)` is only sound if it really is the reciprocal square
279
+ root: the fused norm must reproduce the unfused forward, not just replace it.
280
+ `pow` and `rsqrt` take different paths to the same value, so compare at float32
281
+ tolerance rather than bit-for-bit."""
282
+ torch.manual_seed(0)
283
+ module = cls(16)
284
+ with torch.no_grad():
285
+ module.weight.copy_(torch.randn(16))
286
+ built = get_fuser(module, RMSNormFuser).fuse(module, "norm", default_vllm_config)
287
+ with torch.no_grad():
288
+ built.weight.copy_(module.weight)
289
+ x = torch.randn(4, 16)
290
+ torch.testing.assert_close(built(x), module(x), rtol=1e-5, atol=1e-6)
291
+
292
+
293
def test_weightless_norm_has_no_hidden_size(default_vllm_config):
294
"""A weightless norm states no hidden size, and the model's (LM) hidden size
295
would be wrong for one on a sub-dimension: Llama 4's `qk_norm` normalizes
vllm/model_executor/models/transformers/fusers/rms_norm.py
+46
-17
@@ -36,33 +36,53 @@ if TYPE_CHECKING:
36
logger = init_logger(__name__)
37
38
39
+def _operand(node: fx.Node, index: int, name: str) -> object | None:
40
+ """Operand `index` of `node`, whether it was passed positionally or as `name`."""
41
+ if len(node.args) > index:
42
+ return node.args[index]
43
+ return node.kwargs.get(name)
44
+
45
+
46
def _is_squared(node: object, x: fx.Node) -> bool:
47
"""`x**2`, `x.square()` or `x * x`, through any dtype casts."""
48
node = peel(node)
49
if is_op(node, "pow"):
43
- base, exp = node.args
44
- return peel(base) is x and exp == 2
50
+ return (
51
+ peel(_operand(node, 0, "input")) is x and _operand(node, 1, "exponent") == 2
52
+ )
53
if is_op(node, "square"):
46
- return peel(node.args[0]) is x
54
+ return peel(_operand(node, 0, "input")) is x
55
if is_op(node, "mul"):
48
- a, b = node.args
49
- return peel(a) is x and peel(b) is x
56
+ return (
57
+ peel(_operand(node, 0, "input")) is x
58
+ and peel(_operand(node, 1, "other")) is x
59
+ )
60
+ return False
61
+
62
+
63
+def _is_inverse_sqrt(node: object) -> bool:
64
+ """Detect `rsqrt(v)`, or the `pow(v, -0.5)` / `v ** -0.5` spelling of it."""
65
+ if is_op(node, "rsqrt"):
66
+ return True
67
+ if is_op(node, "pow"):
68
+ return _operand(node, 1, "exponent") == -0.5
69
return False
70
71
72
def _variance_eps(rsqrt: fx.Node, x: fx.Node) -> float | None:
54
- """Eps from `rsqrt(mean(x**2, -1) + eps)`, or `None` if not that shape."""
55
- add = peel(rsqrt.args[0])
73
+ """`eps` from `rsqrt(mean(x**2, -1) + eps)`, or `None` if not that shape."""
74
+ add = peel(_operand(rsqrt, 0, "input"))
75
if not is_op(add, "add"):
76
return None
58
- consts = [a for a in add.args if isinstance(a, (int, float))]
59
- nodes = [a for a in add.args if isinstance(a, fx.Node)]
77
+ operands = [_operand(add, 0, "input"), _operand(add, 1, "other")]
78
+ consts = [a for a in operands if isinstance(a, (int, float))]
79
+ nodes = [a for a in operands if isinstance(a, fx.Node)]
80
if len(consts) != 1 or len(nodes) != 1:
81
return None
82
mean = peel(nodes[0])
83
if not is_op(mean, "mean"):
84
return None
65
- if not _is_squared(mean.args[0], x):
85
+ if not _is_squared(_operand(mean, 0, "input"), x):
86
return None
87
return float(consts[0])
88
@@ -72,7 +92,10 @@ def _is_one_plus(node: object) -> bool:
92
node = peel(node)
93
if not is_op(node, "add"):
94
return False
75
- return any(isinstance(a, (int, float)) and a == 1 for a in node.args)
95
+ return any(
96
+ isinstance(a, (int, float)) and a == 1
97
+ for a in (_operand(node, 0, "input"), _operand(node, 1, "other"))
98
+ )
99
100
101
def _has_trailing_compute(graph: fx.Graph, node: fx.Node) -> bool:
@@ -144,7 +167,7 @@ class RMSNormFuser(BaseFuser):
167
return None
168
# Handle native torch `rms_norm` op.
169
rms_norm = find_node(graph, lambda n: is_op(n, "rms_norm"))
147
- if rms_norm is not None and rms_norm.args and peel(rms_norm.args[0]) is x:
170
+ if rms_norm is not None and peel(_operand(rms_norm, 0, "input")) is x:
171
if _has_trailing_compute(graph, rms_norm):
172
return None
173
eps_attr, eps = cls._eps_source(graph, module)
@@ -158,14 +181,16 @@ class RMSNormFuser(BaseFuser):
181
# The rsqrt over the mean-square variance is the spine of the norm.
182
rsqrt = None
183
for node in graph.nodes:
161
- if is_op(node, "rsqrt") and _variance_eps(node, x) is not None:
184
+ if _is_inverse_sqrt(node) and _variance_eps(node, x) is not None:
185
rsqrt = node
186
break
187
if rsqrt is None:
188
return None
189
# The `x * rsqrt(...)` normalize multiply.
190
normalize = find_node(
168
- graph, lambda n: is_op(n, "mul") and rsqrt in map(peel, n.args)
191
+ graph,
192
+ lambda n: is_op(n, "mul")
193
+ and rsqrt in map(peel, (_operand(n, 0, "input"), _operand(n, 1, "other"))),
194
)
195
if normalize is None:
196
return None
@@ -174,7 +199,11 @@ class RMSNormFuser(BaseFuser):
199
for node in graph.nodes:
200
if not is_op(node, "mul") or node is normalize:
201
continue
177
- operands = [peel(a) for a in node.args if isinstance(a, fx.Node)]
202
+ operands = [
203
+ peel(a)
204
+ for a in (_operand(node, 0, "input"), _operand(node, 1, "other"))
205
+ if isinstance(a, fx.Node)
206
+ ]
207
if len(operands) == 2 and normalize in operands:
208
weight = next(o for o in operands if o is not normalize)
209
tail, zero_centered = node, _is_one_plus(weight)
@@ -234,12 +263,12 @@ class RMSNormFuser(BaseFuser):
263
if (x := find_node(graph, lambda n: n.op == "placeholder")) is None:
264
return None
265
fused = find_node(graph, lambda n: is_op(n, "rms_norm"))
237
- if fused is not None and fused.args and peel(fused.args[0]) is x:
266
+ if fused is not None and peel(_operand(fused, 0, "input")) is x:
267
args, kwargs = fused.args, fused.kwargs
268
eps = args[3] if len(args) > 3 else kwargs.get("eps")
269
return float(eps) if isinstance(eps, (int, float)) else None
270
for node in graph.nodes:
242
- if is_op(node, "rsqrt") and (eps := _variance_eps(node, x)) is not None:
271
+ if _is_inverse_sqrt(node) and (eps := _variance_eps(node, x)) is not None:
272
return eps
273
return None
274