[quic-xiyushi]

Summary

• Extends RecomposeRmsNorm to support Gemma4-style pow(-0.5) normalization in addition to the existing rsqrt-based Qwen/Llama patterns.
• Adds mathematical-equivalence validation, including:
  • square exponent 2;
  • last-axis reduction with keepdim=True;
  • matching normalization input;
  • add(alpha=1);
  • finite, non-negative epsilon;
  • only known FP32 upcast/output-cast wildcard paths;
  • a static one-dimensional weight matching the input’s final dimension.
• Introduces _get_scalar to resolve compile-time scalar values through lifted constants and wrapper operations such as detach and lift_fresh_copy.
• Canonicalizes static, shape-preserving RMSNorm weight views:
  • RemoveRedundancy handles preserved/native RMSNorm operations;
  • RecomposeRmsNorm handles decomposed RMSNorm patterns;
  • runtime and shape-changing weight views remain unfused.
• Uses the logical floating-point input dtype when deriving default epsilon, including quantized graphs where get_node() exposes integer storage behind a dequantize operation.
• Adds isolated annotation-pipeline coverage for RMSNorm recomposition.

Test plan

• python -m pytest backends/qualcomm/tests/test_passes.py -k "rms_norm" -q
• python -m pytest backends/qualcomm/tests/rework/passes/test.py -k "(test_layout_transform or test_recompose_rms_norm) and not gpu" -q
• python -m pytest backends/qualcomm/tests/rework/htp/op/v69/test.py -k test_rms_norm -q
• HTP V69 tuple-output layout-sensitive operators: python -m pytest backends/qualcomm/tests/rework/htp/op/v69/test.py -k "test_batchnorm_2d or test_maxpool_2d or test_group_norm or test_instance_norm_2d or test_adaptive_max_pool_2d_with_indices" -q
• On-device FP32-accumulation coverage:
  • TestQNNFloatingPointOperator.test_qnn_backend_fp32_rms_norm
  • Passed on SM8850 (morph-drental-vm-16, device 5e712b29), including export, QNN compilation, deployment, and execution.
  • Validates that fused RMSNorm remains numerically correct for FP16 input when the decomposed model performs RMS accumulation in FP32.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24166

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 2 New Failures, 2 Pending, 1 Unrelated Failure

As of commit 453116b with merge base 14b5e1d ():

NEW FAILURES - The following jobs have failed:

• pull / nxp-mcuxpresso-test / linux-job (gh)
  undefined reference to SYS_INFO_PrintAll'`
• pull / unittest-buck / macos / macos-job (gh)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 3

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / unittest-buck / linux / linux-job (gh) (trunk failure)
  RuntimeError: Command bash /__w/_temp/exec_script failed with exit code 3

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[meta-cla]

Hi @quic-xiyushi!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: Tanya-Motwani / name: Tanya Motwani (0a06661)
• ✅ login: quic-xiyushi / name: Xiyue Shi (0a06661)

[quic-xiyushi]

/easycla

[quic-xiyushi]

@pytorchbot label "release notes: qualcomm"

[psiddh]

@claude review this pr thoroughlly

[claude]

Claude finished @psiddh's task in 4m 27s —— View job

---

Summary: This PR adds Gemma-style pow(-0.5) RMSNorm recomposition and new equivalence guards, but the guards never look at the nodes the "*" wildcard skips or at the shape and origin of the chosen weight, so wrong fusions are still possible; it also changes which operand is picked as gamma, which can regress quantized graphs. Request Changes.

Detailed review

PR Review: #24166

Summary

The PR extends RecomposeRmsNorm to match Gemma4 pow(x, -0.5) normalization, adds equivalence guards, and moves view_copy off the RMSNorm weight in LayoutTransform. The guards leave two holes that can still produce wrong fusions, the change to how gamma is picked can regress quantized graphs, and the view_copy bypass is in the wrong pass.

Code Quality

• recompose_rms_norm.py:179-218, :249-253: the guards ignore the "*" wildcard.

  • find_pattern lets "*" skip up to 3 nodes of any kind between last_mul and outer_mul (utils.py:278, max_wildcard_life=3). _is_valid_rms_norm throws away pattern[1], and the rewrite drops whatever was skipped.
  • Examples that still fuse into a plain rms_norm(x, [D], w, eps):
    • (norm + b) * w
    • torch.relu(norm) * w
    • torch.tanh(x * torch.pow(x.pow(2).mean(-1, keepdim=True) + eps, -0.5)) * w
  • The hole existed before this PR. The PR now runs a third pattern through it and describes the guards as "strict mathematical equivalence guards", so it should close it: allow only to.dtype / _to_dim_order_copy / q/dq nodes in the wildcard slot, or reject the match.
  • All three reviewers flagged this.

• recompose_rms_norm.py:233-241: gamma is not checked.

  • _get_rms_norm_weight takes the first get_attr/placeholder/dq operand of rank ≥ 1 and uses its full shape as normalized_shape. Nothing checks that the operand is a parameter, or that its shape is [input.shape[-1]].
  • Wrong-shape gain: nn.Parameter(torch.ones(1)) gives normalized_shape=[1] for an input whose last dim is 16. A (1,1,D) or (S,D) weight gives a multi-dim normalized_shape, even though the mean reduced only the last axis. Either way the result is an invalid or different rms_norm. The builder guard at op_rms_norm.py:52-54 uses and, so it can let some of these through.
  • Graph input as weight: _is_placeholder also accepts user inputs. norm * y, where y is a (D,) graph input, fuses with a runtime gamma, and op_rms_norm.py then calls get_parameter on it, which returns None.
  • Suggested fix: require is_parameter (or a dq of a parameter) and gamma_shape == [square.args[0].meta["val"].shape[-1]].

• recompose_rms_norm.py:235-240: possible regression in which operand becomes gamma.

  • The old code took the last matching argument of last_mul. The new code returns the first.
  • When q/dq nodes are present, norm * self.weight becomes mul(dq(norm), dq(weight)). args[0] is then a rank-4 dq and is chosen as gamma.
  • I have not confirmed this happens in the current pipelines. It requires the annotation-time fusion to have been missed, for example because _get_scalar returned None during capture. Excluding the operand that leads back to outer_mul makes the choice independent of argument order.

• layout_transform.py:341-377: bypass_view_copy_for_rms_norm is in the wrong pass, and its comment is wrong.

  • Comment: it says "RecomposeRmsNorm may reshape gamma for broadcasting". RecomposeRmsNorm passes gamma_node through unchanged (recompose_rms_norm.py:287-296). The view_copy comes from model code, as in the new view_as tests.
  • Placement: removing a shape-preserving view_copy is a redundancy cleanup, not a layout concern. Adding exir_ops.edge.aten.view_copy.default: self._same_shape_condition to RemoveRedundancy.redundant_ops_general (remove_redundancy.py:19-28) would cover it.
    • It would apply to every op that reads a static weight, not only rms_norm.
    • RecomposeRmsNorm already depends on RemoveRedundancy (qnn_pass_manager.py:343), so it would run first.
    • LayoutTransform would no longer need its own DCE.
  • All three reviewers flagged this.

• recompose_rms_norm.py:220-231: only the first match per pattern is checked.

  • find_patterns can return several wildcard alignments for one pattern, but only matched_patterns[0] is validated. A later valid alignment is never tried, so the fusion is missed.

• recompose_rms_norm.py:217, :275: epsilon is read only from add.args[1], and the check runs twice.

  • torch.tensor(eps) + mean puts the scalar in args[0], so _get_scalar(mean) returns None and the fusion is rejected.
  • call() also recomputes eps and checks for None again, although _is_valid_rms_norm has already guaranteed it is valid. Returning eps from the validator would remove the duplicate.

• recompose_rms_norm.py:151-177: folding the fp32 upcast drops the fp32 accumulation.

  • Folding the fp16→fp32 _to_dim_order_copy produces rms_norm(fp16 input, fp32 weight) -> fp32; the new pattern.py test asserts exactly this.
  • Models upcast so that mean(x²) is accumulated in fp32, which avoids fp16 overflow for large activations.
  • Whether HTP accepts mixed fp16/fp32 tensors for RmsNorm, and accumulates in fp32, is not tested. Only graph metadata is checked.
  • Confidence is low: this needs an on-device check.

• recompose_rms_norm.py:100: getattr(node.meta.get("val"), "constant", None) is the dynamic getattr pattern that CLAUDE.md discourages. An explicit isinstance check on the FakeTensor would be clearer.

• recompose_rms_norm.py:116-117: _is_pow_neg_half is a one-line helper used once. Call _is_pow_with_exponent(..., -0.5, ...) directly.

• recompose_rms_norm.py:132-142: lifted constants are copied by placeholder name in __init__.

  • Edge passes are created before any of them run (qnn_pass_manager.py:412-427), so the lookup depends on placeholder names staying the same across earlier passes.
  • If a name changes, the fusion is skipped with only a debug log. Low confidence that this happens in practice.

• op_rms_norm.py:114-117: default epsilon ignores the input dtype.

  • With eps=None, aten uses torch.finfo(input.dtype).eps, but the builder always uses the fp32 value. This predates the PR.
  • The new kwargs path makes eps=None reachable, and test_rms_norm_builder_reads_keyword_weight_and_eps locks the fp32 value in.

Testing

• No negative tests for the risky cases:
  • an op in the wildcard slot ((norm + b) * w, relu(norm) * w);
  • a gamma whose shape doesn't match the last dim ((1,), (1,1,D));
  • a graph-input gamma;
  • x * y squaring with different operands.
• The Gemma subtests check only node counts. None checks the fused node's eps or which node became args[2], so the gamma-selection regression above would not be caught.
• The view_copy bypass tests run only when quantizer is None, so the quantized flow is not covered.
• test_rms_norm_builder_reads_keyword_weight_and_eps (test_passes.py:176) mocks get_node/get_tensor/define_tensor with fixed-length side_effect lists, so it will break on any internal refactor. Nothing in the pipeline produces kwargs-form rms_norm either, since RecomposeRmsNorm emits positional args.
• test_rms_norm_recomposition_unwraps_literal_scalar passes a real torch.tensor to lift_fresh_copy, which export never produces. The lifted-constants lookup and the FakeTensor .constant path are only exercised indirectly.

Recommendation

Request Changes. The guards should check the wildcard-skipped nodes and gamma's shape and origin; until they do, wrong fusions can still happen. Picking gamma should not depend on operand order. The view_copy bypass belongs in RemoveRedundancy, and its comment needs to be corrected.

[psiddh]

@quic-xiyushi can you look at claude review ? I looked at it and it seems reasonable suggestion to me

[quic-xiyushi]

@quic-xiyushi can you look at claude review ? I looked at it and it seems reasonable suggestion to me

@psiddh , I’ve addressed Claude’s review suggestions. Thank you!