Make ScalelessRMSNorm a torch.nn.RMSNorm; fix SDPACustom view -> reshape (#19376)

[navsud]

Summary:

Two related changes that together unblock the QNN export path for VLM/STITO:

(1) ScalelessRMSNorm: re-implement as torch.nn.RMSNorm subclass

ScalelessRMSNorm was previously implemented as a hand-rolled RMS normalization
(decomposed into mean / rsqrt / mul). On the QNN export path, this decomposition
fails to lower for an LLM. Using torch.nn.RMSNorm() directly works.

Re-implement ScalelessRMSNorm as a torch.nn.RMSNorm subclass whose weight is
hardcoded to ones and frozen (requires_grad=False). This keeps the public
interface (ScalelessRMSNorm(dim, eps)) unchanged while letting backends see a
proper RMSNorm op so it composes/decomposes cleanly for QNN.

(2) SDPACustom / QuantizedSDPA: replace .view() with .reshape()

Switching to torch.nn.RMSNorm changes how strides propagate through the export
graph compared to the hand-rolled decomposition, exposing a latent bug in
source_transformation/sdpa.py. The output of torch.ops.llama.custom_sdpa retains
the non-contiguous (transposed) strides of its inputs, so
output.view(bsz, seqlen, self.dim) — which merges the last two dims
(n_heads, head_dim) — fails during torch.export with:

Cannot view a tensor with shape (1, s0, 32, 64) and strides
(2048*s0, 64, 64*s0, 1) as a tensor with shape (1, s0, 2048)

Switching to .reshape() inserts .contiguous() only when needed and matches the pattern already used elsewhere in this file (SDPASimple, SDPAFlex, SDPACoreML, and attention.py).

Reviewed By: billmguo, telgamal-1

Differential Revision: D104258950

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 3 New Failures, 3 Pending

As of commit 07aa7e5 with merge base 91aef57 ():

NEW FAILURES - The following jobs have failed:

• Docathon Labels Sync / check-labels (gh)
  Process completed with exit code 1.
• pull / android / run-emulator (gh)
  The process '/usr/local/lib/android/sdk/platform-tools/adb' failed with exit code 224
• pull / unittest-editable / macos / macos-job (gh)
  export/tests/test_target_recipes.py::TestTargetRecipes::test_mv3_model

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

[meta-codesync]

@navsud has exported this pull request. If you are a Meta employee, you can view the originating Diff in D104258950.