Add rank-4 RMSNormalization reshape rule for QNN HTP - #372
Merged
Conversation
The QNN HTP miscomputes RMSNormalization over the last axis of a rank-4 tensor (query/key norm, shaped (batch, seq, heads, head_dim)) -- the result is numerically wrong even though the graph finalizes. Rank-3 RMSNormalization (over the last axis of a rank-3 tensor) runs correctly. ReshapeRank4RMSNorm reshapes the rank-4 input to (batch, seq*heads, head_dim), applies the identical RMSNormalization over the last axis, and reshapes back. Both reshapes use constant shapes ([0, -1, head_dim] and [0, -1, heads, head_dim]) so no Shape op is introduced; the transform is numerically exact. Gated on a new supports_rank4_rmsnorm capability, set False only for the qnn EP, applied in the lowering stage.
Performance Comparison
|
|
The author of this PR, shreyshah-microsoft, is not an activated member of this organization on Codecov. |
Contributor
There was a problem hiding this comment.
Pull request overview
Adds an EP-gated lowering rewrite to make query/key RMSNormalization numerically correct on Qualcomm Hexagon HTP by reshaping rank-4 RMSNorm inputs to rank-3 and back during the optimization pipeline.
Changes:
- Introduces
Rank4RMSNormToRank3rewrite rule (reshape_rank4_rmsnorm_rules) and exports it viamobius.rewrite_rules. - Wires the rule into the lowering stage in
optimize_model()whenEpCapabilities.supports_rank4_rmsnormisFalse. - Adds a unit test validating the rewrite adds exactly 2 reshapes and is numerically exact on CPU.
Show a summary per file
| File | Description |
|---|---|
src/mobius/rewrite_rules/_reshape_rmsnorm.py |
New rewrite rule: rank-4 RMSNorm → reshape-to-rank-3 RMSNorm → reshape-back. |
src/mobius/rewrite_rules/_reshape_rmsnorm_test.py |
New test for structural + numerical equivalence of the rewrite. |
src/mobius/rewrite_rules/__init__.py |
Exports reshape_rank4_rmsnorm_rules. |
src/mobius/_optimizations.py |
Adds ReshapeRank4RMSNorm lowering stage gated by EP capability. |
src/mobius/_execution_providers.py |
Adds supports_rank4_rmsnorm capability and disables it for qnn. |
Copilot's findings
- Files reviewed: 5/5 changed files
- Comments generated: 1
Build the fixture model via ir.Value/ir.Graph/GraphBuilder and serialize with ir.serde.serialize_model for ORT, per the repo's no-explicit-protobuf guideline (review feedback). Behavior unchanged: rewrite adds exactly two Reshapes and stays numerically exact on CPU.
kunal-vaishnavi
approved these changes
Jun 25, 2026
kunal-vaishnavi
enabled auto-merge (squash)
June 25, 2026 19:05
justinchuby
reviewed
Jun 25, 2026
This is a Qualcomm Hexagon HTP hardware workaround, not a generic transform, so make the names say so: HtpRank4RMSNormToRank3 / htp_rank4_rmsnorm_rules, the _htp_rank4_rmsnorm module, and the HtpRank4RMSNorm lowering stage. The supports_rank4_rmsnorm EP capability stays generic (it is per-EP by design; only the QNN HTP sets it False).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
A rewrite rule that reshapes rank-4
RMSNormalization(query/key norm) to rank-3 for the QNN HTP, which miscomputes RMSNormalization over the last axis of a rank-4 tensor.Why
Gemma-4 (and Qwen3) apply RMSNormalization over the head dimension after reshaping q/k to
(batch, seq, heads, head_dim), so the norm runs on a rank-4 tensor. On the Hexagon HTP this computes incorrectly -- the graph finalizes but the q/k-norm output is numerically wrong (uncorrelated with CPU); rank-3 RMSNormalization runs correctly. Pairs with the QNN EP profile (#370): that targets the HTP, this makes the q/k-norm graph compute correctly there.How
ReshapeRank4RMSNormreshapes(B, S, H, Dh)->(B, S*H, Dh), applies the identical RMSNormalization over the last axis, and reshapes back. Both reshapes use constant shapes ([0, -1, head_dim],[0, -1, heads, head_dim]) so noShapeop is introduced and the graph stays static for the HTP. Numerically exact -- the normalized axis is unchanged. Gated on a newsupports_rank4_rmsnormcapability (defaultTrue;Falseonly forqnn), applied in the lowering stage alongside SeparateRoPE.Test
_reshape_rmsnorm_test.py: the rule reshapes a rank-4 RMSNorm to rank-3 (2 Reshapes added) and the rewritten graph matches the original on CPU (where rank-4 RMSNorm is correct); a rank-3 RMSNorm is left untouched.