Cortex-M: decompose SDPA with score scaling after matmul

[amacharla15]

Summary

Decompose eligible aten.scaled_dot_product_attention directly in the
Cortex-M transform_for_annotation pipeline.

The Cortex-M decomposition is:

scores = matmul(query, key.transpose(-2, -1))
scores = scores * scale
attn = softmax(scores, dim=-1)
output = matmul(attn, value)

The existing MatmulToBmmPass then converts the matmuls to BMMs.

This replaces the previous approach that rewrote the output of
DecomposeSDPAWithRegularSoftmaxPass. That Arm decomposition is not part of
the Cortex-M pipeline, so the previous pass did not run in the real flow.

Keeping the attention scale as a single MUL on the score tensor allows the
stock CortexMQuantizer to produce proportional qparams around that MUL,
making it suitable for later folding without adding Cortex-M-specific qspec
logic to the quantizer.

The decomposition is intentionally limited to fp32, unmasked, non-causal,
zero-dropout, non-GQA SDPA with static, non-empty shapes of rank 3 or higher
and matching leading dimensions. Unsupported forms are left unchanged.

Addresses the scale/decomposition part of #21943.

Tests

Added a full SDPA operator test starting from
torch.nn.functional.scaled_dot_product_attention, with no manual SDPA
decomposition.

The tests verify that:

• raw SDPA is removed by the Cortex-M transform-for-annotation pipeline;
• the two matmuls become BMMs;
• exactly one scale MUL and one softmax remain, and the decomposed graph
  matches eager SDPA;
• the MUL sits between a dequantize and quantize with
  out_scale ~= in_scale * scale and equal zero points, so it can be folded
  later;
• unsupported SDPA forms, including dynamic shapes, are left unchanged.

The SDPA dialect cases are xfailed because no lowering pass folds the scale
MUL yet, using the existing xfails mechanism of parametrize with its
default strict behavior.

nn.MultiheadAttention(need_weights=False) regresses with this change. On
main it lowers with the attention left in fp32. With this change its
attention BMMs are quantized, and a BMM then gets an int8 input together with
the int32 output of a retraced folded linear, the separate issue from #21943.
That case is also xfailed.

Local verification:

• SDPA tests without implementation/FVP coverage: 19 passed, 5 xfailed,
  5 deselected;
• existing BMM/matmul/softmax tests passed;
• lintrunner passed;
• full Cortex-M suite comparison introduced no new failures versus the parent
  commit.

FVP implementation tests were not run locally and are covered by the
Cortex-M CI job.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

⏳ No Failures, 2 Pending

As of commit 6864d71 with merge base ed72896 ():
💚 Looks good so far! There are no failures yet. 💚

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 reached their quota limit.

[amacharla15]

@pytorchbot label "release notes: none"

[nil-is-all]

Hi @AdrianLundell, CI looks good on this PR. Could you give a final review?

[AdrianLundell]

Thanks for the update, this is getting closer now! The new pass is targeting the decomposition from the Arm DecomposeSDPAWithRegularSoftmaxPass and then rewrites it, I see two things wrong with this:

1. As currently written the pass will never do anything since DecomposeSDPAWithRegularSoftmaxPass is not in the Cortex-M pass pipeline.
2. Rewriting the decomposition of one pass with another one is a bit convoluted, it would make more sense to me to reimplement the SDPA decomposition in a way which suits the Cortex-M backend better directly instead.

Please also add a full SDPA operator test with these changes to make sure the full flow makes sense, if there are more changes required to make them pass fully that is fine, just make sure that the intended change is happening (muls are annotated correctly for folding) and then xfail the tests with a comment about remaining work needed.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[amacharla15]

Thanks for the review @AdrianLundell. I updated the PR to match the points you raised.

• You pointed out that the previous pass depended on DecomposeSDPAWithRegularSoftmaxPass, which is not part of the Cortex-M pipeline, so it never ran in the real flow. I removed that dependency and replaced the old rewrite pass with DecomposeSDPAPass, which runs directly in transform_for_annotation.

• You also suggested decomposing SDPA directly in a Cortex-M-friendly form instead of rewriting another decomposition. The new pass lowers eligible SDPA directly to:

  matmul -> mul(scale) -> softmax -> matmul
  

  with the scale applied once to the scores. The existing MatmulToBmmPass then converts the supported matmuls to BMMs. Unsupported SDPA forms are left unchanged.

• SDPA operator test: ops/test_sdpa.py now starts from F.scaled_dot_product_attention and runs through the Cortex-M flow without manually invoking any SDPA decomposition pass.

• You also asked to make sure the intended MUL behavior is present for folding, and to xfail the remaining work if needed. The test verifies that the score MUL sits between DQ and Q with out_scale ~= in_scale * scale and equal zero points. The dialect cases that still need the MUL-folding pass are xfailed only on the specific expected failure message, so unrelated failures still fail normally.

One remaining issue: nn.MultiheadAttention(need_weights=False) lowers on main today with attention in float. With this change its attention BMMs become quantized, and lowering reaches the separate folded-linear retrace issue from #21943 (int8 vs int32).

[AdrianLundell]

Please re-use the existing xfail mechanism in the parametrize decorator instead.

[AdrianLundell]

Looks good overall to me now, just fix the tests to use the regular xfailing mechanism instead of a special function and it is an approve from my side. @rascani Do you have any input also?

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[amacharla15]

Thanks @AdrianLundell , switched to the xfails mechanism in parametrize.

[rascani]

This LGTM, I think I'd just ask that we follow-up with a fix for the nn.MultiheadAttention(need_weights=False) regression soon. @amacharla15 - let us know if you're planning on tackling that.

Thanks for this! I really appreciate your contribution.

[amacharla15]

Thanks @rascani, really appreciate the review! Yes, I’d like to take the nn.MultiheadAttention(need_weights=False) follow-up.

@AdrianLundell, on #22796 you mentioned an upcoming prototype for changing how folded Q/DQ parameters are handled in the graph. Is that work available now, and should I base this follow-up on it rather than adding a separate Cortex-M-specific workaround?