[shoumikhin]

What is wrong today

RecomposeRmsNorm rebuilds a decomposed RMSNorm into one rms_norm node. It walked back to the norm's input over aten.to.dtype casts only, but edge graphs use _to_copy or _to_dim_order_copy. So for a Qwen- or Llama-style norm (upcast to float32, normalize, cast back), the rebuilt rms_norm read the float32 cast and lost the cast back.

A float16 norm did not lower, and bfloat16 models failed later:

Node aten_rms_norm_default ... was not decomposed or delegated.
masked_scatter: expected self and source to have same dtypes
RuntimeError: Missing out variants: {'aten::rms_norm'}

What this change does

The first commit rebuilds the whole norm only if its input and output share a dtype and no cast in between rounds below it:

• It skips dtype-only casts, back to the nearest tensor in the output dtype.
• Only such casts may sit between the normalization and the weight.
• It must scale the tensor it squares, past such casts (T5 squares a float32 copy), and square it with pow(2) or h * h.
• Outside quantization capture, the backend must support the dtype (not bfloat16).

So a float16 Qwen- or Llama-style norm is rebuilt on its float16 input.

A bfloat16 norm, or a float32 norm with a bfloat16 weight, still stays decomposed, and QNN's float16 mode then overflows when it squares values above 255.9. The third commit rebuilds its float32 normalization alone, as an rms_norm without a weight. On the QNN SDK's x86 simulator for SM8650 (float16), the worst relative error on rows with values from 300 to 2000 drops from 4.5 to 0.005.

Since rms_norm has no CPU kernel, every rebuild now needs what QNN's RmsNorm takes: the last of 1 to 4 dims, static sizes, a constant eps. A rank-5 or dynamic-shape float32 norm now lowers instead of failing.

The second commit keeps a dtype-changing _to_copy with memory_format, which RemoveRedundancy dropped without dim order.

Tests

test_recompose_rms_norm_keeps_the_norm_dtype covers the Qwen, Llama and T5 forms across all three dtypes, with and without dim order. Two more tests check that a different scaled tensor, a cube, a product of two tensors, a rank-5 input, a dynamic dim and an eps buffer stay decomposed.

The new tests pass on Linux x86_64 with torch 2.14 and 2.16 nightly. With the QNN SDK (SM8650), a float16 Qwen-style norm is delegated as one rms_norm, and a bfloat16 norm before masked_scatter lowers.

cc @cbilgin @psiddh

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ You can merge normally! (1 Unrelated Failure)

As of commit b8fd5ec with merge base a057bb4 ():

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

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

• pull / nxp-mcuxpresso-test / linux-job (gh) (trunk failure)
  undefined reference to SYS_INFO_PrintAll'`

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

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[shoumikhin]

@claude /pr-review

[claude]

Claude finished @shoumikhin's task in 25m 39s —— View job

---

PR Review: #24428

No blockers, but in rare mixed-dtype graphs the pass can still build an rms_norm that QNN rejects, or rebuild a later mul with the wrong weight; the float32-weight handling and the test tolerances also need another look.

Detailed review

Summary

The PR makes RecomposeRmsNorm walk past _to_copy/_to_dim_order_copy casts to the nearest input in the norm's output dtype. It leaves the norm decomposed when that dtype has no QNN type.

Nothing blocks the PR. The problems are:

• two gaps in mixed-dtype graphs, one of which gives silently wrong output;
• a design question about float32 weights;
• missing test coverage.

All findings come from reading the code. Python can't run in this environment, so none of them were executed.

Code Quality

1. A skipped norm lets a later mul be rebuilt with the wrong weight (backends/qualcomm/_passes/recompose_rms_norm.py:151-154).

   • When the new gate skips a norm, its weight mul y stays in the graph.
   • The "*" in the patterns spans up to three nodes (max_wildcard_life=3, _passes/utils.py:278), so it can cover y and the cast back.
   • A following z = y * s then matches with s as gamma, when s is a [D]-shaped placeholder, get_attr or dq.
   • Example: a bfloat16 Qwen-style norm followed by * s, with s in float32. z is float32, the walk stops at the float32 upcast and the gate passes. z becomes rms_norm(x.to(float32), [D], s), and the real weight is silently dropped.
   • Before this PR, y was always rebuilt first, so this graph failed loudly at lowering. The older skips (no gamma found, the shape checks) already had this hole; the new gate opens it for bfloat16 norms.
   • Suggested fix: recompose only when the nodes between last_mul_node and rms_norm_pattern[2] are dtype-only casts.
   • Fix this →

2. The new dtype gate does not check the weight (recompose_rms_norm.py:151-153).

   • The partitioner's has_unsupported_dtype (partition/utils.py:149-159, called at qnn_partitioner.py:177) rejects a node if any unquantized input has no QNN type. The pass checks only the output dtype.
   • Example: UpcastNorm(torch.bfloat16) with a float32 x has a float32 output, so the norm is rebuilt with a bfloat16 gamma. The partitioner keeps it on CPU, and lowering fails with the same Missing out variants: {'aten::rms_norm'} error this PR fixes.
   • A float16 x with a bfloat16 weight fails the same way, since it promotes to float32.
   • Suggested fix: also require gamma_node.meta["val"].dtype to have a QNN type.
   • Fix this →

3. With a float32 weight, the target dtype is taken after type promotion (recompose_rms_norm.py:146-150).

   • The comment says rms_norm computes in its input's dtype, which "has to be the pattern's output dtype". That doesn't match eager: aten.rms_norm upcasts float16/bfloat16 internally and applies the weight after casting back to the input dtype. Its output dtype is promote_types(input, weight).
   • Because the target is last_mul_node's dtype, a Qwen-style float16 x with a float32 weight is rebuilt on x.to(float32). That drops the model's h.to(float16) rounding before the weight multiply, which is the kind of rounding the comment at lines 103-105 is meant to preserve. The float32-weight test case only passes because of atol=rtol=1e-2 at test_passes.py:561.
   • A cast-free float16 norm with a float32 weight is now left decomposed: x is float16 and the target is float32. The old pass rebuilt it, and the PR description doesn't mention this change.
   • If QNN's RmsNorm needs a single dtype for input, weight and output, the comment should say so, and the weight then has to match too (see item 2).
   • Otherwise, target the dtype of the weight mul's other operand (the cast back). That rebuilds both cases on the float16 x and matches what eager computes for the original module.

4. The walk still skips narrowing casts (recompose_rms_norm.py:106-109).

   • The loop walks past every cast whose output is not the target dtype, including narrowing ones. So the guarantee in the comment above it only covers round trips that end in the target dtype.
   • Example: x32 → .to(float16) → norm → .to(float32) → w32 * is rebuilt as a float32 norm on x32, which drops the float16 rounding and compute.
   • Suggested fix: walk only past casts that don't narrow, or narrow the comment's claim.
   • Fix this →

5. _is_cast compares dim_order with the contiguous order instead of the input's (recompose_rms_norm.py:99).

   • MemoryFormatOpsPass gives a preserve-format cast the input's dim order (exir/passes/memory_format_ops_pass.py:114-117).
   • So a dtype-only cast of a channels-last tensor isn't treated as a cast, and its norm is not recomposed.
   • A cast that also turns a channels-last input contiguous is skipped, which contradicts "Only casts that change nothing but the dtype are skipped".
   • Suggested fix: compare with list(node.args[0].meta["val"].dim_order()).
   • Fix this →

Testing

• The test is too lenient to catch numeric bugs (backends/qualcomm/tests/test_passes.py:509, :561).
  • UpcastNorm's weight is torch.ones, and the comparison uses atol=rtol=1e-2. A wrong or dropped gamma (item 1) or a dropped rounding (item 3) still passes.
  • ScaledNorm above uses torch.randn and default tolerances.
  • Only the float32-weight case should need a looser bound. I haven't run the other cases to confirm they pass at default tolerances.
• Cases with no test:
  • the Llama-style pattern (x * x, x.float(), .type_as(x); examples/models/llama/norm.py:34-41), which the description says was checked on device;
  • the round trip described in the _get_input_node comment (float32 → float16 → float32 before the norm);
  • _is_cast rejecting a cast with extra kwargs or a non-contiguous dim_order. Dropping those two conditions (lines 98-99 of the pass) still passes every case.
• The new cases are in the wrong test. They are appended to test_recompose_rms_norm_only_takes_last_dim_with_weight, which is about shape gating. A separate test would make a failure point at the dtype logic.
• Fix this →

Recommendation

Needs Discussion

Nothing blocks. Item 3 needs an answer on whether QNN's RmsNorm requires one dtype throughout, because that decides how items 2 and 3 should be fixed. Items 1 and 2 are small enough to close in this PR.

---

• Gather context (diff, changed files, CLAUDE.md)
• Run three independent reviewers
• Consolidate findings
• Fact-check findings
• Post review

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[shoumikhin]

Thanks for the review. The latest push changes this:

1. Wrong weight after a skipped norm: fixed. Only casts that change nothing but the dtype may now sit between the normalization and the weight multiply, so a later * s is no longer rebuilt as the norm. The new test has a bfloat16 norm followed by a float32 scale.
2. Weight dtype: not changed, because this case is never rebuilt. When the weight's dtype differs from the multiply's output, the weight reaches the multiply through a cast in the edge graph. The pass then finds no parameter or constant as gamma, so the norm stays decomposed. I checked float32 input with a bfloat16 weight, float16 input with a bfloat16 weight, and float32 input with a float16 weight: no rms_norm is built in any of them.
3. One dtype or mixed: I found nothing that says QNN's RmsNorm takes different dtypes for input, weight and output, so I went with one dtype throughout. The norm is rebuilt only when its input and its output share a dtype, and no cast in between rounds below that dtype. So a float16 Qwen-style norm with a float32 weight now stays decomposed. Before, it was rebuilt in float32 and lost the float16 rounding. The description now says this.
4. Narrowing casts: fixed. The walk stops at a cast to a dtype narrower than the output.
5. dim_order: fixed. A cast is compared with its input's dim order. The test has a channels-last input, once with a cast that keeps its dim order and once with a cast to contiguous, which stops the rebuild.

Tests: the dtype cases are now in their own test, test_recompose_rms_norm_keeps_the_norm_dtype, with random weights and default tolerances. It adds a Llama-style norm (type_as), float16 rounding before a float32 norm (rebuilt on the rounded tensor) and inside one (left decomposed), and the channels-last cases.

[shoumikhin]

@claude /pr-review

[claude]

Claude finished @shoumikhin's task in 36m 11s —— View job

---

PR Review: #24428

During quantization capture, the new check between the normalization and the weight stops a norm that casts back with type_as from being recomposed, which is a regression that needs a fix before merge; the other two items are older bugs next to the change.

Detailed review

Summary

The PR makes RecomposeRmsNorm rebuild a norm only on an input in the norm's output dtype. It finds that input by walking past casts that change only the dtype, and requires that only such casts sit between the normalization and the weight.

• Item 1 is a regression in the quantization flow, caused by the new check on the weight side.
• Items 2 and 3 are older bugs next to the change.

I covered both changed files. Python can't run in this environment, so nothing was executed; every finding comes from reading the code.

Code Quality

1. During quantization capture, a norm that casts back with type_as is no longer recomposed (backends/qualcomm/_passes/recompose_rms_norm.py:91, :133-140).

   • Why it misses:
     • In capture mode, skip_targets is only torch.ops.aten.to.dtype.
     • A captured graph keeps h.type_as(x) as aten.type_as.default, even when the dtype doesn't change.
       • Arm's annotation pipeline runs a pass that removes such no-op calls on parameters (backends/arm/_passes/remove_redundant_type_as_pass.py:24, arm_pass_manager.py:808).
       • _cast_source in backends/transforms/fuse_rms_norm.py:210 matches the same op.
     • QNN's capture-mode RemoveRedundancy removes only _assert_tensor_metadata, so the type_as node stays.
     • The new walk back from the weight multiply stops at that node, normalized is pattern[2] fails, and _get_norm_input returns None.
     • The old pass let the "*" take the type_as and built one rms_norm on x.
   • Trigger: the Meta reference Llama RMSNorm (self._norm(x.float()).type_as(x) * self.weight) or the PR's TypeAsNorm, in float32 or bfloat16, run through RecomposeRmsNorm(quantization_capture=True).
     • The norm stays decomposed through annotation, so its ops are quantized one by one instead of as one rms_norm.
     • For bfloat16, this contradicts the description, which says a bfloat16 norm stays decomposed unless it is about to be quantized.
   • Test gap: TypeAsNorm runs only on the edge path (backends/qualcomm/tests/test_passes.py:561). The capture rows (:567-568) use only UpcastNorm, which casts back with to.dtype.
   • Fix: add torch.ops.aten.type_as.default to the capture skip_targets. Add capture rows such as (TypeAsNorm(f32), f32, c, True, True, [(f32, True)]) and the bfloat16 version.
   • Related:
     • _is_cast reads layout, device and dim order only from node.kwargs (:102-111).
     • aten.to.dtype takes memory_format positionally. So in capture, a to.dtype that makes a channels-last input contiguous still counts as changing only the dtype. The values are unaffected.
     • Comparing node.args[0].meta["val"] with node.meta["val"] instead would apply one rule to all four cast ops. _cast_source already compares device and layout this way (fuse_rms_norm.py:194).
   • Basis: I confirmed the mechanism by reading the code. That captured graphs keep aten.type_as.default is inferred from those two passes, not from a run.
   • Fix this →

2. Nothing checks that the tensor being scaled is the norm's input (recompose_rms_norm.py:123, :140). This bug predates the PR.

   • Problem: the input is taken from pattern[6].args[0], the base of the square, and the new check verifies only the weight side. Nothing checks that:
     • the other operand of pattern[2] (the h in h * rsqrt(...)) is that same tensor;
     • both operands of the Llama-form h * h are the same tensor;
     • the pow exponent is 2.
   • Trigger: y * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-6) * w is rebuilt as rms_norm(x, [D], w), and y is silently dropped.
   • Why fix it here: _match_rms_norm_core in backends/transforms/fuse_rms_norm.py:104-108 already checks the scaled tensor and the exponent. This PR rewrote how the input is chosen, and its test asserts which tensor the norm reads.
   • Fix this →

3. With _skip_dim_order=True, the dim-order rule never sees the cast it targets (recompose_rms_norm.py:102-111, test_passes.py:559). The cause predates the PR and is outside the diff.

   • Cause: RemoveRedundancy runs before this pass in the to_edge pipeline (qnn_pass_manager.py:344). Its _to_copy_op_condition (remove_redundancy.py:51-52) deletes every _to_copy that has a memory_format kwarg, including one that changes the dtype. It doesn't re-trace, so the metadata downstream stays stale.
   • Trigger: ContiguousUpcastNorm(f16) lowered with _skip_dim_order=True.
     • _to_copy(x, dtype=torch.float32, memory_format=torch.contiguous_format) is deleted.
     • The pow then reads the float16 x, while its metadata still says float32.
     • The norm is rebuilt on x, although under the PR's rule this cast keeps the norm decomposed.
   • Wider effect: any x.to(dtype, memory_format=...) in that flow loses its dtype change the same way.
   • Why the test misses it: it calls the pass on its own and checks the rule only with dim order on.
   • Fix: make _to_copy_op_condition also require an unchanged dtype, as _dim_order_op_condition does.
   • Fix this →

Recommendation

Request Changes

Item 1 makes the quantization flow worse for norms written with type_as. The fix is one skip_targets entry and two test rows. Items 2 and 3 predate the PR and can be follow-ups.

---

• Gather context (diff, changed files, CLAUDE.md)
• Run three independent reviewers
• Consolidate findings
• Fact-check findings
• Post review

[shoumikhin]

Thanks. The latest push changes this:

1. type_as in quantization capture: fixed. aten.type_as now counts as a cast in captured graphs. _is_cast also compares the cast's output with its input (layout, device, dim order) instead of reading kwargs, so one rule covers every cast op, however it passes its arguments. The test runs TypeAsNorm in capture for float32 and bfloat16.
2. The scaled tensor and the exponent: fixed here. The norm is rebuilt only when the tensor it scales is the tensor it squares, and the square is pow(2) or h * h. The comparison looks past dtype-only casts, because T5's norm squares a float32 copy of the tensor it scales; the test covers T5's norm too. A new test checks that a different scaled tensor, a cube and a product of two tensors stay decomposed.
3. RemoveRedundancy dropping a dtype change: fixed in a second commit. A _to_copy with memory_format is now removed only when it keeps the dtype, as _to_dim_order_copy already was. A new test exports such casts without dim order: the layout-only one is removed, and the float16 one stays.

[shoumikhin]

@claude /pr-review

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[claude]

Claude finished @shoumikhin's task in 42m 16s —— View job

---

PR Review: #24428

Nothing blocks this PR. Quantization capture can still rebuild a norm whose weight has another dtype, no test runs the lifted pow exponent that real edge-flow norms go through, and recompose_rms_norm.py needs two small cleanups.

Detailed review

Summary

The PR changes two passes:

• RecomposeRmsNorm now rebuilds a norm only when its input and output share a dtype and no cast between them rounds to a narrower one.
• RemoveRedundancy now keeps a _to_copy that changes the dtype.

I found no blockers: three minor code problems and one test gap. All three changed files were covered. Python can't run in this environment, so every finding comes from reading the code.

Code Quality

1. Quantization capture still rebuilds a norm whose weight has another dtype (backends/qualcomm/_passes/recompose_rms_norm.py:158-161; gamma is picked at :198-200).

   • Cause: The earlier reply said the weight reaches the multiply through a cast. That is only true for edge graphs.
     • The cast comes from RemoveMixedTypeOperators (exir/passes/remove_mixed_type_operators.py:25, :87), which only runs inside to_edge (exir/program/_program.py:649).
     • The annotation pipeline (qnn_pass_manager.py:173-211) has no such pass, so in a captured graph the weight get_attr feeds the multiply directly.
   • Trigger: UpcastNorm(torch.bfloat16), or a float16 weight, with a float32 x, run through RecomposeRmsNorm(quantization_capture=True).
     • The weight multiply outputs float32, the walk returns x, and _get_norm_input never looks at the weight.
     • The pass then builds aten.rms_norm(x_f32, [32], weight_bf16). That breaks the description's assumption that QNN's RmsNorm takes one dtype throughout.
     • No capture row in test_recompose_rms_norm_keeps_the_norm_dtype uses a weight in another dtype.
   • Scope: Only a captured model that mixes dtypes inside the norm is affected. I didn't trace what QNN does with such a node when it stays in floating point.
   • Fix: In _get_norm_input, also return None when gamma_node.meta["val"].dtype != dtype, so both flows apply the same rule.
   • Fix this →

2. The _get_norm_input docstring misstates what rms_norm computes (recompose_rms_norm.py:143).

   • Problem: The docstring says "rms_norm computes in its input's dtype". ATen's rms_norm actually upcasts float16 and bfloat16 inputs to float32 for the normalization, then casts back before applying the weight.
   • Why it matters:
     • That upcast is why the float16 Qwen, Llama and T5 rows match eager at default tolerances.
     • Read literally, the docstring says rebuilding those norms on x_f16 drops the float32 math the model asks for. That is the opposite of the reason the pass can rebuild them.
   • Fix: State the constraint the code actually enforces: rms_norm returns its input's dtype, so the input has to be in the pattern's output dtype. Separately, state the assumption that QNN's RmsNorm takes one dtype throughout.
   • Basis: ATen's rms_norm implementation, which isn't in this checkout.
   • Fix this →

3. getattr(exponent.meta["val"], "constant", None) is the dynamic getattr pattern CLAUDE.md rules out (recompose_rms_norm.py:127).

   • The default is never used: FakeTensor always defines .constant. lift_constant_tensor_pass sets it on every lifted constant (exir/program/_program.py:289), and transform_for_export_pipeline calls that pass (qnn_pass_manager.py:476).
   • Same file: Line 225 already reads eps.meta["val"].constant.item() directly.
   • Fix: Use exponent.meta["val"].constant, and keep the None check for an exponent that isn't a constant.
   • Fix this →

Testing

• No test runs the exponent branch that real edge-flow norms go through (recompose_rms_norm.py:125-128; test_passes.py:599-602, :620).
  • Why it's the main path:
    • to_edge_transform_and_lower_to_qnn runs transform_for_export_pipeline before to_edge (backends/qualcomm/utils/utils.py:524).
    • That pipeline includes LiftConstantScalarOperands (qnn_pass_manager.py:245). It rewrites pow.Tensor_Scalar(h, 2) as pow.Tensor_Tensor(h, <lifted constant>) (lift_constant_scalar_operands.py:57).
    • So in the real flow, every pow-form norm (Qwen, T5) reaches this pass through lines 125-128.
  • What the tests run instead: _check_recomposed_rms_norm and test_recompose_rms_norm_needs_the_norm_of_one_tensor call to_edge on the raw export. Only the scalar-exponent branch runs there.
  • What's left unchecked: Neither the pow(2) acceptance nor the pow(3) rejection is tested in the form the pipeline produces. A regression there would only show up on device.
  • Fix: Run transform_for_export_pipeline(ep) before to_edge, at least for the float16 UpcastNorm row and the pow3 case. The tests at test_passes.py:1006-1011 already do this.
  • Fix this →

Recommendation

Approve

Nothing blocks. Item 1 and the test gap are small and worth folding into this PR. Items 2 and 3 are one-line cleanups.

---

• Gather context (diff, changed files, CLAUDE.md)
• Run three independent reviewers
• Consolidate findings
• Fact-check findings
• Post review

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[shoumikhin]

About a weight in another dtype during quantization capture: I kept the rebuild on purpose. The quantizer quantizes the weight with the same spec as the input, so QNN gets the input, the weight and the output of RmsNorm as one quantized type (8-bit with 8a8w, 16-bit with 16a16w).

I ran a float32 input through PTQ and lowered it for SM8650, with a bfloat16 weight (8a8w and 16a16w) and a float16 weight (8a8w). With the rebuild, the whole norm becomes one delegated rms_norm, and nothing stays on the CPU. With the suggested check, the norm stays decomposed, and QNN rejects the weight multiply because its two inputs get different types, so it runs on the CPU. The model after convert_pt2e is also usually a little less accurate (for example, 3.1e-2 instead of 2.8e-2 relative error with 8a8w and a bfloat16 weight). ATen's rms_norm with a mixed weight gives the same output as the original module.

I added a capture case with a bfloat16 weight and a float32 input that expects the rebuild, so this choice is now tested.

[shoumikhin]

@claude /pr-review

[claude]

Claude finished @shoumikhin's task in 49m 38s —— View job

---

PR Review: #24428

Nothing blocks this PR: three minor gaps leave norms decomposed that could be rebuilt, mostly during quantization capture, and one edge-flow case is untested.

Detailed review

Summary

The PR makes RecomposeRmsNorm rebuild a norm only when its input and output share a dtype and only dtype-only casts sit inside it. It also makes RemoveRedundancy keep a _to_copy that changes the dtype.

There are no blockers:

• Three performance findings: each leaves a norm decomposed where it could be one rms_norm. None gives wrong output.
• One testing finding: it leaves a design choice unpinned.

All three changed files were covered. Python needs approval in this environment, so nothing was executed. Every finding comes from reading the code.

Testing

• No edge row has a weight narrower than its input (backends/qualcomm/tests/test_passes.py:639-657).
  • What's pinned: row 653 pins that capture rebuilds UpcastNorm(bf16) on a float32 input.
  • What isn't: the edge flow leaves the same model decomposed for an indirect reason. The pass never checks the weight's dtype (recompose_rms_norm.py:198-200). It finds no gamma only because to_edge runs RemoveMixedTypeOperators (exir/program/_program.py:649), which casts the weight before the multiply.
  • Risk: if gamma detection ever looks past that cast, the edge flow builds rms_norm(x_f32, [32], w_bf16). bfloat16 has no QNN type, so this brings back the Missing out variants: {'aten::rms_norm'} failure the PR fixes. No test would catch it.
  • Fix: add the edge row (UpcastNorm(bf16), f32, c, False, True, []) next to row 653.
  • Fix this →

Performance

1. Capture mode now leaves a norm decomposed when its weight is wider than its input (backends/qualcomm/_passes/recompose_rms_norm.py:148, :164-169).

   • Trigger: a Qwen-style norm with bfloat16 or float16 activations and a float32 weight, run through RecomposeRmsNorm(quantization_capture=True). For example, UpcastNorm(f32) with a bfloat16 x.
   • Why:
     • The target dtype is the weight multiply's output, which type promotion makes float32.
     • Against float32, the cast back to bfloat16 counts as narrowing. So _strip_casts(normalized, dtype) stops before pattern[2], and _get_norm_input returns None.
   • Before this PR: the pass rebuilt it as rms_norm(x_bf16, [D], w_f32). ATen's rms_norm casts back to the input dtype before it applies the weight, so that node computes what the module computes.
   • Reply 3: its reasoning covers this direction too. The quantizer gives the weight the input's quantized type, and QNN gets one quantized RmsNorm.
     • Reply 3 found that leaving the mirror case decomposed sent the weight multiply to CPU. I didn't check whether this case does the same.
     • No capture row has a wider weight.
   • Fix, if this should be rebuilt: take dtype from the weight multiply's other operand (normalized) instead of its output.
     • Edge graphs stay as they are. There, that operand is the cast RemoveMixedTypeOperators adds to the promoted dtype (exir/passes/remove_mixed_type_operators.py:32, :75-80, :100), so row 646 stays decomposed.
     • Add the capture row (UpcastNorm(f32), bf16, c, True, True, [(bf16, True)]).
   • If it should stay decomposed: a capture row expecting [] would pin that choice.
   • Basis:
     • The old and new behavior are traced in the code.
     • The order of the cast back and the weight inside ATen's rms_norm comes from its implementation, which isn't in this checkout.
   • Fix this →

2. Capture mode doesn't count _to_copy as a cast, though the annotation pipeline emits it (recompose_rms_norm.py:91-94).

   • Source:
     • DecomposeWrapWithAutocast runs just before RecomposeRmsNorm in the annotation pipeline (qnn_pass_manager.py:176-180).
     • It retraces enabled autocast regions with make_fx (decompose_wrap_with_autocast.py:64-67). That traces below the composite ops, so to.dtype and type_as come out as aten._to_copy.default.
   • Trigger: a bfloat16 or float16 Qwen-style norm inside an enabled torch.autocast region, quantized with QnnQuantizer.
     • The walk stops at the float32 _to_copy, so the norm stays decomposed and is quantized op by op.
     • The same norm outside an autocast region is rebuilt (rows 651, 654).
   • Fix: add torch.ops.aten._to_copy.default to the capture skip_targets. _is_cast compares tensor metadata rather than kwargs, so it handles _to_copy as is.
   • Basis: read the pass order and the retrace. That the retraced casts are _to_copy follows from how make_fx traces without pre_dispatch; I didn't run it.
   • Fix this →

3. The mul-form square compares its operands by identity, but the scaled tensor is compared past casts (recompose_rms_norm.py:131 vs :137).

   • Trigger: (x.float() * x.float()).mean(-1, keepdim=True) with a half-precision x.
     • The two casts are separate nodes, so square.args[1] is base fails and the norm stays decomposed.
     • In capture mode the pass rebuilt this form before.
   • Fix: compare self._strip_casts(square.args[1], dtype) is self._strip_casts(base, dtype), as line 137 does for the scaled tensor.
   • Confidence: low that this spelling shows up in real models. It also assumes export keeps the two casts as two nodes, which I didn't run.
   • Fix this →

Recommendation

Approve

Nothing blocks. Performance item 1 needs an answer: should a weight wider than the input get the same capture-mode rebuild that reply 3 keeps for a narrower one? The other items are small.

---

• Gather context (diff, changed files, CLAUDE.md)
• Run three independent reviewers
• Consolidate findings
• Fact-check findings
• Post review

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[shoumikhin]

Thanks for the review. The new push adds a third commit and answers the 4 points:

• Edge row with a narrower weight: added. A float32 input with a bfloat16 weight now rebuilds only the float32 normalization, as rms_norm without a weight, and keeps the weight multiply. The test now also checks whether each rebuilt norm keeps its weight.
• Capture with a weight wider than the input: I kept it decomposed on purpose. ATen's rms_norm returns the input's dtype, also with a wider weight. For a bfloat16 input and a float32 weight, the rebuilt node returns bfloat16 while the module returns float32. I tried the suggested change and this capture case failed exactly that way. So main's rebuild here changed the output dtype. A capture row now pins it as decomposed.
• _to_copy casts from an autocast region during capture: confirmed. A norm inside torch.autocast reaches this pass with _to_copy casts. Main rebuilds it on the float32 cast, so the graph returns float32 where the module returns bfloat16. Here it stays decomposed and the output is right. I would rather rebuild it in a separate change with its own test.
• x.float() * x.float() as the square: fixed. The square is now compared past dtype-only casts, so this spelling is rebuilt like x.pow(2). Left decomposed, it would overflow in float16 mode the same way. The type_as case in the test now squares two casts.

[shoumikhin]

@claude /pr-review

[claude]

PR Review: #24428

• Gather context
• Run three independent reviewers
• Consolidate findings
• Fact-check findings
• Post review

Fact-check finished. I confirmed that the base branch rebuilt weighted norms whose dims are dynamic, and that the new static-size check at recompose_rms_norm.py:247 now leaves those norms decomposed. I'm writing up the review now.

View job run

[shoumikhin]

@claude /pr-review

[claude]

PR Review: #24428

• Gather context (diff, changed files, CLAUDE.md)
• Run three independent reviewers
• Consolidate findings
• Fact-check findings
• Post review

Three reviewers are now going through the three commits, each starting from a different part of the change. Python needs approval in this environment, so all findings will come from reading the code.

View job run