[transformer] RMSNorm matching for alternative rsqrt

[bohnstingl]

Purpose

Two gaps in the Transformers backend's RMSNorm fuser, both surfaced by Gemma 4.

1. The pow(v, -0.5) spelling of the reciprocal square root was not matched.
The fuser identifies a norm structurally by finding rsqrt over the mean-square
variance. Gemma4RMSNorm._norm writes hidden_states * torch.pow(mean_squared, -0.5)
instead — deliberately, per its own comment, "to address compiler differences between
Torch and JAX". Every Gemma 4 norm therefore failed to match and was left as the
unfused HF module. _is_inverse_sqrt now accepts rsqrt(v), pow(v, -0.5) and the
v ** -0.5 operator spelling, in both the detection scan and the eps extraction.
Gemma 4's weight is not zero-centered (unlike Gemma 1-3), so these fuse to a plain
RMSNorm.

The OOT part has been moved to a separate PR, currently either #57715 or #53153

Test Plan

tests/models/transformers/fusers/test_rms_norm.py:

• Three Gemma-4-shaped norms added to the detection and build parametrizations:
  Gemma4RMSNorm (the torch.pow spelling), PowOperatorRMSNorm (** -0.5), and
  WeightlessGemma4RMSNorm (with_scale=False).
• test_pow_spelled_norm_fuses_to_equivalent_math compares the fused module's output
  against the unfused forward — matching pow(v, -0.5) is only sound if it really is
  the reciprocal square root, so this asserts the math rather than just the swap.

pytest tests/models/transformers/fusers/test_rms_norm.py

Beyond the unit tests, the real classes and a real engine:

• Every *RMSNorm submodule of the Gemma-4 checkpoints (31B, 12B, E4B, E2B-it), built on
  the meta device from their own configs, run through match().
• transformers' own Gemma4RMSNorm and Gemma4UnifiedRMSNorm, with and without scale,
  fused and compared against the unfused forward in fp32 and fp16.
• A Gemma-4 engine on CPU via --model-impl transformers, greedy decode, before and after.

Test Result

	result
upstream baseline	23 passed
with this change	37 passed
the 14 new cases against the unmodified fuser	14 failed, 23 passed

The third row is the one that matters: every new case fails without the change, so none
of them passes incidentally, and no pre-existing test changes outcome.

On the real checkpoints, 611 / 339 / 524 / 464 norms respectively go from 0 matched to all
matched, every one detected as zero_centered=False, eps_attr='eps'. Fused and unfused
output for the four real-class configurations is bit-identical (max|diff| = 0.0) in both
dtypes. In the engine, all ten distinct norm sites log
Fused: <name> (Gemma4RMSNorm) -> RMSNorm (CustomOp) where the baseline logs none, and
both runs emit the same token ids.

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[hmellor]

Thanks for this, the pow half is a clean catch and the equivalence test is the right way to prove it.

Two things overall:

1. There are two independent changes here: the pow(v, -0.5) matching (self-contained, clearly correct, well tested) and the OOT override derivation (new machinery, no in-tree consumer). I'd land the first on its own and split the second, mainly because the second has a gap that needs resolving first (see the comment on the op_registry_oot import).
2. Given the branch name, the OOT half is presumably aimed at the hw-agnostic series, which makes the registry it reads the central question rather than a detail.

Also, which platform plugin actually registers "RMSNorm" out of tree today? The tests fabricate one, so nothing here exercises the real path end to end.

[hmellor]

This hard-codes vllm.model_executor.custom_op's registry, but the RMSNorm / GemmaRMSNorm this file imports come from transformers/layers.py::_resolve, which returns vllm.model_executor.hw_agnostic.layers.layernorm.RMSNorm when VLLM_USE_HW_AGNOSTIC is set. That class inherits hw_agnostic/custom_op.py's CustomOp, which has its own module-level op_registry_oot (vllm/model_executor/hw_agnostic/custom_op.py:22), a different dict from the one imported here.

So with the env var set, _norm_impl reads a registry the plugin never wrote to, returns base, and the override is silently skipped: exactly the failure this PR is fixing, in the configuration where OOT layers are most likely to be in play.

_resolve also falls back per symbol and logs a warning, so RMSNorm can be hw-agnostic while GemmaRMSNorm falls back to the vLLM one in the same process. A single hard-coded registry cannot be correct for both.

Suggest resolving the registry from the base class itself, by walking base.__mro__ to its CustomOp ancestor and using that module's registry (both modules already export maybe_get_oot_by_class for this shape of lookup), plus a test case with VLLM_USE_HW_AGNOSTIC set so the two paths stay honest.

[bohnstingl]

This is a good catch in principle. However, there is a related, but separate effort to consolidate the two op_registry_oot into a single one, see #53153. In particular, the registry will remain in vllm/model_executor/custom_op.py as today and the vllm/model_executor/hw_agnostic/custom_op.py just re-exports the same registry. This is also needed so that existing OOT plugins don't need to adjust in any way, but are still able to use the hw_agnostic path via VLLM_USE_HW_AGNOSTIC

[hmellor]

Worth saying in this docstring that the cache is what keeps the class count bounded, not just what keeps it fresh.

Fusers caches the fuser per module class, but recursive_replace calls fuse() per instance, so this runs once per norm site: 611 times on the Gemma-4 31B in your own results. Without the memo that is 611 structurally identical classes, which is the pathological input for dynamo guards and per-class compile caches, i.e. the opposite of what #47187 was for. I suspect that invariant (one class per norm kind, not one per instance) is the main thing the fixed TPAwareRMSNorm / TPAwareGemmaRMSNorm classes were buying, and the memo is what preserves it once they become dynamic.

The current wording ("a later registration change is picked up rather than stale") reads as an optimisation note, and someone will eventually simplify it away.

[bohnstingl]

I updated the docstring of the function to explain the caching behavior.

[hmellor]

del base # only part of the cache key resolves cleanly here: RMSNorm and GemmaRMSNorm are siblings under CustomOp (vllm/model_executor/layers/layernorm.py:37 and :132), neither subclasses the other, so _norm_impl's issubclass(impl, base) guard means no single class can be a usable override for both. impl alone determines the result, so the parameter and the del can both go.

[bohnstingl]

I removed.

[hmellor]

This silently returns the RMSNorm variant for any base that is neither of the two. Only two call sites today, but {RMSNorm: TPAwareRMSNorm, GemmaRMSNorm: TPAwareGemmaRMSNorm}[base] fails loudly instead and reads better.

[bohnstingl]

I adopted the suggestion.

[hmellor]

info() and fuse() now resolve the override independently, so they can disagree if the registry changes between the two calls, and info() reports impl.__name__ while the class actually built is TPAware{impl.__name__}. Both cosmetic, but having fuse() stash the resolved class on the fuser and info() read that would make the log provably describe what was installed.

[bohnstingl]

Fix applied: fuse() stashes and returns the tp_aware + classname and info() returns that.

[hmellor]

This asserts identity across two _tp_aware calls, but the path that actually runs hundreds of times per model is fuse(). Worth asserting type(fuser.fuse(m1, ...)) is type(fuser.fuse(m2, ...)) as well, so the one-class-per-norm-kind invariant is guarded where it can actually regress.

[bohnstingl]

I adjusted the tests.

[hmellor]

Re-reviewed at 86acd690b. The pow(v, -0.5) / v ** -0.5 detection looks right to me: I checked it against the real transformers Gemma4RMSNorm (5.18.0.dev0) and it matches with zero_centered=False, both with and without with_scale. The _operand refactor is a genuine improvement, the old _variance_eps would have raised IndexError on a kwarg-spelled rsqrt(input=v). tests/models/transformers/fusers/ passes and lint is clean.

One substantive thing left (the OOT registry on the hw-agnostic path), plus three smaller ones inline.

[hmellor]

Small accuracy nit on this docstring: it says every operand is read through _operand rather than off node.args directly, but a few reads still index node.args directly, _is_one_plus (L105), the rsqrt in map(peel, n.args) and [peel(a) for a in node.args ...] searches in match, and _eps_from_graph's fused.args[0].

The consequence is small and fails safe: a scale multiply spelled torch.mul(input=normalized, other=1 + self.weight) simply doesn't match, so the norm is left unfused rather than fused without the 1 +. Worth either narrowing the claim or converting the remaining reads.

[bohnstingl]

I removed the inaccuracy docstring.

[hmellor]

This reads op_registry_oot from vllm.model_executor.custom_op, but base here comes from transformers/layers.py, which under VLLM_USE_HW_AGNOSTIC=1 resolves RMSNorm to hw_agnostic.layers.layernorm.RMSNorm. That class derives from a different CustomOp with its own registry (hw_agnostic/custom_op.py:22); the two dicts are not the same object.

So on the hw-agnostic path an override registered where hw_agnostic.CustomOp.__new__ actually looks is invisible here, and one registered in the in-tree registry subclasses the in-tree RMSNorm and therefore fails the issubclass guard below. Either way the plugin's kernels are silently dropped, which is the failure this PR sets out to fix. Given the branch is transformer_hw-agnostic_p2, that path seems like the intended target.

It also splits within a single model: hw-agnostic defines RMSNorm but not GemmaRMSNorm, so the two norms resolve from different registries.

The fix is to derive the registry from base rather than hard-importing it: walk base.__mro__ to the CustomOp / PluggableLayer it actually derives from and read op_registry_oot off that class's module. Alternatively this could wait on #53153, which collapses the two registries at source.

Separately, the else base fallback drops an unusable registration silently. A logger.warning_once there would make a mis-registered plugin visible rather than reintroducing the silent-skip mode this PR is trying to remove.

[bohnstingl]

I isolated the current PR from the OOT plugin topic, which will be handled in a separate PR.

[hmellor]

This reason isn't quite what the cited code does. CustomOp.__new__ doesn't refuse a non-subclass: custom_op.py:118-128 calls object.__new__(override) unconditionally, and Python then skips __init__ because the result isn't an instance of cls, leaving an uninitialised module. Keeping the in-tree class is the right behaviour, it's just the justification that needs rewording.

[bohnstingl]

See comment above.

[hmellor]

fused_cls is None until fuse() runs, so RMSNormFuser(zero_centered=False, source_cls="X").info("norm") raises AttributeError: 'NoneType' object has no attribute '__module__'.

Nothing breaks today because the only caller (transformers/base.py:560) calls fuse immediately before, but BaseFuser.info is public with no documented ordering constraint, and the other implementations (StackedFuser.info, MLAFuser.info) are pure. A debug dump of get_fusers, or anything logging a matched-but-not-yet-applied candidate, would crash model load instead of printing a string.

Worth noting too that get_fusers is cached process-globally on (type(module), children), so fused_cls is shared mutable state on a fuser reused across every model in the process. Guarding with a None fallback, or computing the class in info() via _tp_aware(...) instead of stashing it, removes the ordering dependency.

[bohnstingl]

The PR should no longer touch info()

[bohnstingl]

@hmellor thank you for your comments. I went through them and addressed them.
I removed the touching points to the hardware-agnostic work stream, as this will be handled in a separate PR

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @bohnstingl.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

Hi @bohnstingl, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

[hmellor]

/ci run

[github-actions]

✅ Triggered Buildkite CI #90940 for commit b8d19154baf4.

[hmellor]

/ci retry

[github-actions]

✅ Queued 6 failed job(s) for retry in Buildkite CI #90940.