[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>

tests/models/transformers/fusers/test_rms_norm.py

@@ -50,6 +50,55 @@ class WeightlessRMSNorm(RMSNorm):
         return self._rms(x.to(torch.float32)).to(x.dtype)
 
 
+class Gemma4RMSNorm(RMSNorm):
+    """`pow(v, -0.5)` where the others write `rsqrt(v)` (HF `Gemma4RMSNorm._norm`).
+
+    Gemma 4 spells the reciprocal square root as `torch.pow(..., -0.5)` to keep
+    Torch and JAX in agreement, and -- unlike Gemma 1-3 -- its weight is *not*
+    zero-centered, so this fuses to a plain `RMSNorm`.
+    """
+
+    def _rms(self, x):
+        mean_squared = x.pow(2).mean(-1, keepdim=True) + self.variance_epsilon
+        return x * torch.pow(mean_squared, -0.5)
+
+    def forward(self, x):
+        return (self._rms(x.float()) * self.weight.float()).type_as(x)
+
+
+class PowOperatorRMSNorm(Gemma4RMSNorm):
+    """The `v ** -0.5` operator spelling of the same reciprocal square root."""
+
+    def _rms(self, x):
+        return x * (x.pow(2).mean(-1, keepdim=True) + self.variance_epsilon) ** -0.5
+
+
+class KwargPowRMSNorm(Gemma4RMSNorm):
+    """The keyword spelling `torch.pow(v, exponent=-0.5)` of the same rsqrt."""
+
+    def _rms(self, x):
+        mean_squared = x.pow(2).mean(-1, keepdim=True) + self.variance_epsilon
+        return x * torch.pow(mean_squared, exponent=-0.5)
+
+
+class KwargBasePowRMSNorm(Gemma4RMSNorm):
+    """The fully functional spelling `torch.pow(input=v, exponent=-0.5)`."""
+
+    def _rms(self, x):
+        mean_squared = x.pow(2).mean(-1, keepdim=True) + self.variance_epsilon
+        return x * torch.pow(input=mean_squared, exponent=-0.5)
+
+
+class WeightlessGemma4RMSNorm(Gemma4RMSNorm):
+    """Gemma 4 with `with_scale=False`: the `pow` spelling and no scale parameter."""
+
+    def __init__(self, hidden: int = 16, eps: float = 1e-6):
+        super().__init__(hidden, eps, weight=False)
+
+    def forward(self, x):
+        return self._rms(x.float()).type_as(x)
+
+
 class LayerNorm(RMSNorm):
     """An RMSNorm not named `*RMSNorm`, keeping the input dtype (no upcast)."""
 
@@ -156,6 +205,11 @@ class UntraceableGatedRMSNorm(RMSNorm):
         (RMSNorm, 1e-5, False),
         (GemmaRMSNorm, 1e-6, True),
         (WeightlessRMSNorm, 1e-6, False),
+        (Gemma4RMSNorm, 1e-6, False),
+        (PowOperatorRMSNorm, 1e-6, False),
+        (KwargPowRMSNorm, 1e-6, False),
+        (KwargBasePowRMSNorm, 1e-6, False),
+        (WeightlessGemma4RMSNorm, 1e-6, False),
         (LayerNorm, 1e-6, False),
         (torch.nn.RMSNorm, 1e-5, False),  # fused `F.rms_norm` op
     ],
@@ -188,6 +242,10 @@ def test_gated_rms_norm_is_not_fused(cls):
         (RMSNorm, "RMSNorm", False),
         (GemmaRMSNorm, "GemmaRMSNorm", True),
         (WeightlessRMSNorm, "RMSNorm", False),
+        (Gemma4RMSNorm, "RMSNorm", False),
+        (KwargPowRMSNorm, "RMSNorm", False),
+        (KwargBasePowRMSNorm, "RMSNorm", False),
+        (WeightlessGemma4RMSNorm, "RMSNorm", False),
     ],
 )
 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
     assert built.weight.shape[0] == (weight.size(0) if weight is not None else 0)
 
 
+@pytest.mark.parametrize(
+    "cls", [Gemma4RMSNorm, PowOperatorRMSNorm, KwargPowRMSNorm, KwargBasePowRMSNorm]
+)
+def test_pow_spelled_norm_fuses_to_equivalent_math(cls, default_vllm_config):
+    """Matching `pow(v, -0.5)` is only sound if it really is the reciprocal square
+    root: the fused norm must reproduce the unfused forward, not just replace it.
+    `pow` and `rsqrt` take different paths to the same value, so compare at float32
+    tolerance rather than bit-for-bit."""
+    torch.manual_seed(0)
+    module = cls(16)
+    with torch.no_grad():
+        module.weight.copy_(torch.randn(16))
+    built = get_fuser(module, RMSNormFuser).fuse(module, "norm", default_vllm_config)
+    with torch.no_grad():
+        built.weight.copy_(module.weight)
+    x = torch.randn(4, 16)
+    torch.testing.assert_close(built(x), module(x), rtol=1e-5, atol=1e-6)
+
+
 def test_weightless_norm_has_no_hidden_size(default_vllm_config):
     """A weightless norm states no hidden size, and the model's (LM) hidden size
     would be wrong for one on a sub-dimension: Llama 4's `qk_norm` normalizes

vllm/model_executor/models/transformers/fusers/rms_norm.py

@@ -36,33 +36,53 @@ if TYPE_CHECKING:
 logger = init_logger(__name__)
 
 
+def _operand(node: fx.Node, index: int, name: str) -> object | None:
+    """Operand `index` of `node`, whether it was passed positionally or as `name`."""
+    if len(node.args) > index:
+        return node.args[index]
+    return node.kwargs.get(name)
+
+
 def _is_squared(node: object, x: fx.Node) -> bool:
     """`x**2`, `x.square()` or `x * x`, through any dtype casts."""
     node = peel(node)
     if is_op(node, "pow"):
-        base, exp = node.args
-        return peel(base) is x and exp == 2
+        return (
+            peel(_operand(node, 0, "input")) is x and _operand(node, 1, "exponent") == 2
+        )
     if is_op(node, "square"):
-        return peel(node.args[0]) is x
+        return peel(_operand(node, 0, "input")) is x
     if is_op(node, "mul"):
-        a, b = node.args
-        return peel(a) is x and peel(b) is x
+        return (
+            peel(_operand(node, 0, "input")) is x
+            and peel(_operand(node, 1, "other")) is x
+        )
+    return False
+
+
+def _is_inverse_sqrt(node: object) -> bool:
+    """Detect `rsqrt(v)`, or the `pow(v, -0.5)` / `v ** -0.5` spelling of it."""
+    if is_op(node, "rsqrt"):
+        return True
+    if is_op(node, "pow"):
+        return _operand(node, 1, "exponent") == -0.5
     return False
 
 
 def _variance_eps(rsqrt: fx.Node, x: fx.Node) -> float | None:
-    """Eps from `rsqrt(mean(x**2, -1) + eps)`, or `None` if not that shape."""
-    add = peel(rsqrt.args[0])
+    """`eps` from `rsqrt(mean(x**2, -1) + eps)`, or `None` if not that shape."""
+    add = peel(_operand(rsqrt, 0, "input"))
     if not is_op(add, "add"):
         return None
-    consts = [a for a in add.args if isinstance(a, (int, float))]
-    nodes = [a for a in add.args if isinstance(a, fx.Node)]
+    operands = [_operand(add, 0, "input"), _operand(add, 1, "other")]
+    consts = [a for a in operands if isinstance(a, (int, float))]
+    nodes = [a for a in operands if isinstance(a, fx.Node)]
     if len(consts) != 1 or len(nodes) != 1:
         return None
     mean = peel(nodes[0])
     if not is_op(mean, "mean"):
         return None
-    if not _is_squared(mean.args[0], x):
+    if not _is_squared(_operand(mean, 0, "input"), x):
         return None
     return float(consts[0])
 
@@ -72,7 +92,10 @@ def _is_one_plus(node: object) -> bool:
     node = peel(node)
     if not is_op(node, "add"):
         return False
-    return any(isinstance(a, (int, float)) and a == 1 for a in node.args)
+    return any(
+        isinstance(a, (int, float)) and a == 1
+        for a in (_operand(node, 0, "input"), _operand(node, 1, "other"))
+    )
 
 
 def _has_trailing_compute(graph: fx.Graph, node: fx.Node) -> bool:
@@ -144,7 +167,7 @@ class RMSNormFuser(BaseFuser):
             return None
         # Handle native torch `rms_norm` op.
         rms_norm = find_node(graph, lambda n: is_op(n, "rms_norm"))
-        if rms_norm is not None and rms_norm.args and peel(rms_norm.args[0]) is x:
+        if rms_norm is not None and peel(_operand(rms_norm, 0, "input")) is x:
             if _has_trailing_compute(graph, rms_norm):
                 return None
             eps_attr, eps = cls._eps_source(graph, module)
@@ -158,14 +181,16 @@ class RMSNormFuser(BaseFuser):
         # The rsqrt over the mean-square variance is the spine of the norm.
         rsqrt = None
         for node in graph.nodes:
-            if is_op(node, "rsqrt") and _variance_eps(node, x) is not None:
+            if _is_inverse_sqrt(node) and _variance_eps(node, x) is not None:
                 rsqrt = node
                 break
         if rsqrt is None:
             return None
         # The `x * rsqrt(...)` normalize multiply.
         normalize = find_node(
-            graph, lambda n: is_op(n, "mul") and rsqrt in map(peel, n.args)
+            graph,
+            lambda n: is_op(n, "mul")
+            and rsqrt in map(peel, (_operand(n, 0, "input"), _operand(n, 1, "other"))),
         )
         if normalize is None:
             return None
@@ -174,7 +199,11 @@ class RMSNormFuser(BaseFuser):
         for node in graph.nodes:
             if not is_op(node, "mul") or node is normalize:
                 continue
-            operands = [peel(a) for a in node.args if isinstance(a, fx.Node)]
+            operands = [
+                peel(a)
+                for a in (_operand(node, 0, "input"), _operand(node, 1, "other"))
+                if isinstance(a, fx.Node)
+            ]
             if len(operands) == 2 and normalize in operands:
                 weight = next(o for o in operands if o is not normalize)
                 tail, zero_centered = node, _is_one_plus(weight)
@@ -234,12 +263,12 @@ class RMSNormFuser(BaseFuser):
         if (x := find_node(graph, lambda n: n.op == "placeholder")) is None:
             return None
         fused = find_node(graph, lambda n: is_op(n, "rms_norm"))
-        if fused is not None and fused.args and peel(fused.args[0]) is x:
+        if fused is not None and peel(_operand(fused, 0, "input")) is x:
             args, kwargs = fused.args, fused.kwargs
             eps = args[3] if len(args) > 3 else kwargs.get("eps")
             return float(eps) if isinstance(eps, (int, float)) else None
         for node in graph.nodes:
-            if is_op(node, "rsqrt") and (eps := _variance_eps(node, x)) is not None:
+            if _is_inverse_sqrt(node) and (eps := _variance_eps(node, x)) is not None:
                 return eps
         return None