Cortex-M: fix quantized lowering for decomposed attention

[amacharla15]

Summary

Fix Cortex-M lowering for explicitly decomposed Transformer attention so that quantized bmm, softmax, and scalar-multiply operations can lower to Cortex-M operators.

This does not add general Transformer support or automatically decompose Transformer operators. The regression test explicitly applies the existing SDPA and LayerNorm decomposition passes before PT2E quantization.

Problem

After decomposing Transformer attention before quantization, _softmax and bmm could be quantized, but the attention scaling mul did not lower correctly.

There were three related issues in the Cortex-M lowering pipeline:

1. Scalar multiplication is represented as an aten.full tensor after scalar-to-tensor replacement. For attention scaling, this rank-1 scalar tensor needs to be materialized at the MUL output shape before quantization so that Cortex-M quantized MUL can be selected.

2. After FoldAndAnnotateQParamsPass, output requantization information is stored in node metadata while the graph still contains aten.linear. Generic ATen fake/meta retracing therefore reconstructs the linear output as int32, even though the Cortex-M quantized linear output is requantized to int8. This later causes an int8/int32 mismatch when retracing BMM.

3. Quantized scalar constants needed by MUL must be materialized after qparams are attached, without globally folding unrelated constant operations.

Changes

• Add Cortex-M-specific retracing behavior so folded quantized aten.linear nodes use the dtype described by their output qparams.
• Use this retracing behavior in scalar replacement and quant-node replacement passes.
• Materialize only scalar-generated aten.full operands when their shape must match a MUL output.
• Leave general tensor broadcasting unchanged so existing fallback behavior is preserved.
• Materialize only quantized aten.full constants that eventually feed a quantized MUL.
• Add a regression test using a small torch.nn.Transformer with explicit SDPA and LayerNorm decomposition.

The regression verifies that decomposed attention lowers to:

• cortex_m.quantized_batch_matmul
• cortex_m.softmax
• cortex_m.quantized_mul

and that decomposed attention no longer retains the corresponding ATen BMM or softmax operations.

Validation

Local validation on the updated implementation:

• Transformer regression passes with numerical comparison at qtol=2.
• Cortex-M MUL dialect tests: 9 passed, 4 xfailed.
• Silero VAD regression passes.
• Shared-qspec regression preserves the expected xfail.
• Full Cortex-M dialect suite: 240 passed, 1 skipped, 24 xfailed.
• lintrunner -a passes, including mypy.

Cortex-M55 FVP validation

As additional validation, I serialized the decomposed Transformer and tested it end-to-end on the Cortex-M55 Corstone-300 FVP.

The stock Cortex-M runner initially could not load the full decomposed Transformer because its selective operator list does not include several portable operators present in the resulting program:

• aten::permute_copy.out
• aten::squeeze_copy.dims_out
• aten::expand_copy.out
• aten::mean.out
• aten::var.correction_out

I also checked the remaining expand_copy operations explicitly. The decomposed Edge graph contains 12 aten.expand_copy nodes before the Cortex-M backend passes run, and the lowered graph contains the same 12 afterward. The narrowed MaterializeMulBroadcastPass in this PR therefore does not introduce additional expand_copy operations; these shape operations are already present before Cortex-M lowering.

For local runtime validation only, I rebuilt the Cortex-M runner with those five existing portable kernels enabled. No PR source files were changed, and the working tree remained clean.

With that runner, the full path completed successfully:

• PT2E quantization
• Cortex-M lowering
• ExecuTorch serialization
• Cortex-M55 Corstone-300 FVP execution
• numerical output comparison with qtol=2

The end-to-end FVP run passed.

Scope / limitations

This PR intentionally does not:

• add automatic Transformer decomposition to the Cortex-M backend,
• claim general or official Transformer support,
• force every decomposed Transformer MUL to use a Cortex-M quantized kernel,
• add general quantized tensor broadcasting support,
• or change the stock Cortex-M runner's selected portable operator set.

The changes are limited to making the existing explicit decomposition workflow lower the targeted quantized attention operations correctly while preserving existing behavior for unrelated Cortex-M graphs.

• The experimental use_explicit_layout=True pipeline is not covered by this fix. I reproduced the same folded aten.linear int8/int32 retracing failure there. Mirroring the legacy-path fixes gets past the initial failure, but the same issue then appears while retracing through ReplaceSqueezeAndUnsqueezeWithViewPass and CortexMCanonicalizeViewCopyPermutePass. Supporting decomposed Transformer attention in explicit-layout mode therefore requires a broader Cortex-M-aware retracing change.

Additional validation after review fixes

• Cortex-M test suite, excluding FVP/implementation tests and the MobileNet v2, MobileNet v3, and Wav2Letter model files whose optional torchvision/torchaudio dependencies are not installed locally: 466 passed, 5 skipped, 36 xfailed.
• Focused decomposed-SDPA regression passes numerical comparison at qtol=2.
• The focused SDPA model also serializes and executes end-to-end on the stock Cortex-M55 Corstone-300 FVP runner at qtol=2.
• lintrunner -a passes, including MYPY.

Addresses #21943

cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• linux.dgx.b200 runner queuing exceeds 4 hours and dgxb200-05 "CUDA not available"

❌ 14 Awaiting Approval, 1 Unclassified Failure

As of commit c09864c with merge base 9918b00 ():

AWAITING APPROVAL - The following workflows need approval before CI can run:

• Build documentation (gh)
• Build Presets (gh)
• Lint (gh)
• periodic (gh)
• pull (gh)
• Test ARM Backend (gh)
• Test CoreML Backend (gh)
• Test Cortex-M Backend (gh)
• Test NXP Backend (gh)
• Test QNN Backend (gh)
• Test QNN Windows MSVC build (gh)
• Test Vulkan Backend (gh)
• Test WebGPU Backend (gh)
• Test XNNPACK Backend (gh)

UNCLASSIFIED FAILURE - DrCI could not classify the following job because the workflow did not run on the merge base. The failure may be pre-existing on trunk or introduced by this PR:

• Cadence Build & Test / Resolve CI docker image / resolve (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  ##[error]Refusing to check out fork pull request code from a 'pull_request_target' workflow. This workflow runs with the base repository's GITHUB_TOKEN, secrets, default-branch cache scope, and runner access. Fetching and executing a fork's code in that trusted context commonly leads to "pwn request" vulnerabilities. To opt in, review the risks at https://gh.io/securely-using-pull_request_target and set 'allow-unsafe-pr-checkout: true' on the actions/checkout step.

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

[meta-cla]

Hi @amacharla15!

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: amacharla15 / name: amacharla15 (60ae655)

[amacharla15]

@pytorchbot label "release notes: none"

[meta-cla]

Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks!

[pytorch-bot]

The following ciflow label(s) have been added but CI has not been triggered yet because the workflows are awaiting approval:

• ciflow/trunk

Once a maintainer approves the workflows (scroll to the bottom of the PR page), the corresponding CI jobs will be triggered automatically. Please ping one of the reviewers if you do not have access to approve and run workflows.

[amacharla15]

Thanks for triggering the Cortex-M CI.

The first CI run exposed a few regressions from the initial implementation that were broader than the intended fix:

• the MUL broadcast pass inserted expand_copy for existing dynamic broadcast cases, which caused Cortex-M runtime implementation tests to fail because that operator was not registered by the runner;
• running the generic constant-materialization pass globally also changed unrelated graphs, including folding the arange in the Silero VAD test and changing an existing shared-qspec expected-xfail case.

I narrowed both parts in the latest commit:

• MaterializeMulBroadcastPass now only materializes the scalar-generated aten.full operand needed for the decomposed-attention MUL and leaves general tensor broadcasting unchanged;
• constant materialization is now limited to quantized aten.full constants that feed quantized MUL instead of running the generic constant pass over the whole graph.

I also narrowed the Transformer regression so it checks the relevant decomposed attention lowering without requiring unrelated LayerNorm MULs to lower.

After these changes, locally:

• Transformer regression: 1 passed
• Cortex-M MUL dialect tests: 9 passed, 4 xfailed
• Silero VAD regression: 1 passed
• shared-qspec case: restored to its expected xfail
• full Cortex-M dialect suite: 240 passed, 1 skipped, 24 xfailed
• lintrunner -a: all checks passed

I updated the PR description to reflect the narrowed implementation and scope.

While the new CI run is waiting/running, I am also checking serialization/FVP loading of the lowered Transformer as an additional validation step.

[amacharla15]

Additional runtime validation on the latest commit (e1171a6):

I serialized the decomposed Transformer and tested it on the Cortex-M55 Corstone-300 FVP.

The stock Cortex-M runner initially could not load the full program because its selective operator list omits five portable operators present in the decomposed Transformer: aten::permute_copy.out, aten::squeeze_copy.dims_out, aten::expand_copy.out, aten::mean.out, and aten::var.correction_out.

The narrowed MaterializeMulBroadcastPass in this PR does not introduce expand_copy; it materializes the scalar aten.full directly at the MUL output shape. The remaining expand_copy nodes come from the broader decomposed Transformer graph.

For local validation only, I rebuilt the runner with those five existing portable kernels enabled. No PR source files were changed, and the working tree remained clean.

With that runner, the full path passed:

• PT2E quantization
• Cortex-M lowering
• ExecuTorch serialization
• Cortex-M55 Corstone-300 FVP execution
• numerical comparison with qtol=2

So the end-to-end FVP validation passes with the required portable operators registered.

The CI checks for e1171a6 are currently running.

[amacharla15]

I also ran a few targeted ablations on the latest implementation:

• Without MaterializeMulBroadcastPass, the attention MUL no longer lowers to cortex_m.quantized_mul.
• Without MaterializeQuantizedMulConstantsPass, lowering still happens but numerical comparison fails (mean abs error ~0.107).
• Reverting either Cortex-M-specific retracing site reproduces the original int8/int32 BMM mismatch.

So the two materialization passes are independently needed at different stages, and the retracing fix is required in both passes.

These were local-only checks; the working tree was restored afterward.

[Copilot]

🟡 Changes recommended

Critical issues remain for explicit-layout integration and dynamic-shape scalar broadcasting, with additional constant-materialization and test-coverage concerns.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Fixes Cortex-M lowering for explicitly decomposed, quantized Transformer attention.

Changes:

• Adds Cortex-M-aware quantized retracing.
• Materializes scalar broadcasts and quantized MUL constants.
• Adds Transformer lowering regression coverage.

File summaries

File	Reviewed changes
backends/cortex_m/test/models/test_transformer.py	Adds decomposed Transformer regression coverage.
backends/cortex_m/passes/replace_scalar_with_tensor_pass.py	Uses Cortex-M-aware scalar replacement.
backends/cortex_m/passes/replace_quant_nodes_pass.py	Uses Cortex-M-aware retracing.
backends/cortex_m/passes/materialize_quantized_mul_constants_pass.py	Materializes quantized MUL constants.
backends/cortex_m/passes/materialize_mul_broadcast_pass.py	Materializes scalar MUL broadcasts.
backends/cortex_m/passes/cortex_m_quantized_pass.py	Preserves quantized linear output dtypes.
backends/cortex_m/passes/cortex_m_pass_manager.py	Integrates the lowering passes.

Review details

Suppressed comments (2)

backends/cortex_m/passes/materialize_quantized_mul_constants_pass.py:41

• This predicate succeeds when any user path reaches a quantized MUL, but ComputeConstantOpsAOTPass.compute_node_aot then replaces all users of the aten.full node with one quantized placeholder. If a shared full constant also feeds an unrelated float operation, that consumer receives the int8 buffer and its dtype/value changes. Clone/materialize only the eligible MUL use, or require all users to have compatible quantization.

            if user.target == exir_ops.edge.aten.mul.Tensor:
                return bool(user.meta.get("input_qparams")) and bool(
                    user.meta.get("output_qparams")
                )

backends/cortex_m/test/models/test_transformer.py:97

• This only proves that some mul lowered to cortex_m.quantized_mul; it does not prove that the decomposed attention scaling multiply was the one lowered. The Transformer/decomposition graph contains multiple MULs, and check() only checks for presence, so the test can pass while the problematic aten.mul.Tensor remains. Please add a targeted assertion for the scale node (or a smaller graph where that MUL is unambiguous), while keeping the intentional allowance for unrelated Transformer MULs.

            "executorch_exir_dialects_edge__ops_cortex_m_quantized_mul_default",

• Files reviewed: 7/7 changed files
• Comments generated: 2
• Review effort level: Lite

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

🟡 Changes recommended

Unresolved pipeline coverage, packaging, and quantized lowering issues block approval.

Get a fresh assessment by requesting another Copilot review.

Review details

Suppressed comments (5)

backends/cortex_m/passes/cortex_m_pass_manager.py:63

• These new passes are wired only into legacy_pass_list. The explicit_layout_pass_list still uses the generic scalar-replacement pass and omits MaterializeQuantizedMulConstantsPass, so use_explicit_layout=True does not get the post-qparam constant materialization or dtype-preserving retrace needed by the same decomposed attention graph. Add the equivalent passes to that list or explicitly constrain this fix to legacy lowering.

        MaterializeQuantizedMulConstantsPass,
        CortexMReplaceScalarWithTensorArgPass,

backends/cortex_m/passes/cortex_m_quantized_pass.py:79

• This base pass retraces the entire graph, but the fallback here only preserves quantized linear. A quantized aten.bmm also has requantized output qparams after FoldAndAnnotateQParamsPass; falling through to the generic fake aten.bmm recreates its accumulator dtype, so the scalar-replacement retrace can feed later attention ops with metadata that no longer matches the Cortex-M BMM output. Include BMM in this no-fake-kernel path (the existing Arm pass already treats it this way).

        return super().call_operator(op, args, kwargs, meta)

backends/cortex_m/passes/materialize_mul_broadcast_pass.py:30

• numel() == 1 is not enough to identify a scalar-generated constant: a model can also contain an ordinary aten.full((1,), ...) tensor used by a MUL. This pass will expand that unrelated constant to the full output shape before quantization, potentially creating large persistent buffers and changing the fallback behavior the PR says it is preserving. Mark full nodes created by scalar replacement and require that marker here.

        return get_first_fake_tensor(node).numel() == 1

backends/cortex_m/passes/materialize_quantized_mul_constants_pass.py:58

• This path explicitly accepts full -> clone -> mul, but the AOT replacement only turns the full into a precomputed buffer; it never changes the clone's shape. When the full is still the rank-1 scalar created by scalar replacement, the clone remains rank-1 and the later Cortex-M quantized_mul still receives an unsupported broadcast. Make the shape-materialization pass follow the same passthrough chain (or reject it here) so every accepted path has matching MUL input shapes.

        if not self._feeds_single_quantized_mul(node):
            return False
        return super().compute_node_aot(node)

backends/cortex_m/passes/materialize_quantized_mul_constants_pass.py:45

• This predicate accepts a partially annotated MUL because it only checks that input_qparams is non-empty, but _get_mul_replacement later unconditionally reads input_qparams[0] and [1]. If only one operand is quantized, this pass can materialize the constant and leave a graph that fails during lowering; require both operand entries (and make the positive unit test model both).

                return bool(user.meta.get("input_qparams")) and bool(
                    user.meta.get("output_qparams")
                )

• Files reviewed: 8/8 changed files
• Comments generated: 1
• Review effort level: Lite

[Copilot]

🔵 Needs a closer look

Quantized BMM handling and constant-size restrictions need correction, and pass-level behavior remains insufficiently tested.

Review details

Suppressed comments (3)

backends/cortex_m/passes/cortex_m_quantized_pass.py:70

• This override only avoids fake-kernel execution for aten.linear, but both new subclasses retrace the entire folded graph and still send quantized aten.bmm nodes through the generic ExportPass.call_operator. The existing ArmPass explicitly handles quantized BMM because the ATen fake kernel cannot propagate int8 inputs (backends/arm/_passes/arm_pass.py:89-98, backends/cortex_m/passes/matmul_to_bmm_pass.py:82-84), so the same int32/error path can remain after the linear fix. Preserve that no-fake-kernel BMM handling as well.

    def call_operator(self, op, args, kwargs, meta):
        if (
            op == exir_ops.edge.aten.linear.default
            and isinstance(meta, NodeMetadata)
            and self._is_quantized_meta(meta)

backends/cortex_m/passes/materialize_quantized_mul_constants_pass.py:57

• This pass is documented as handling scalar aten.full constants, but this predicate never checks the tensor size. Any quantized aten.full of arbitrary shape with a single passthrough path into a quantized MUL will therefore be converted into a persistent buffer by ComputeConstantOpsAOTPass, which can unnecessarily inflate the Cortex-M PTE for large masks or lookup tensors. Restrict this pass to one-element constants (or make the intended size policy explicit).

        output_qparams = node.meta.get("output_qparams", {})
        if set(output_qparams) != {0}:
            return False
        if not self._feeds_single_quantized_mul(node):
            return False

backends/cortex_m/passes/materialize_quantized_mul_constants_pass.py:56

• The new tests cover only _feeds_single_quantized_mul; they never run compute_node_aot()/call() on an exported graph. The behavior introduced here—quantizing and replacing an eligible aten.full with a buffer while rejecting unrelated constants—therefore remains unverified, and these predicate tests would pass if the actual pass were a no-op. Add a pass-level graph test that inspects the transformed graph and state dict.

    def compute_node_aot(self, node: Node) -> bool:
        output_qparams = node.meta.get("output_qparams", {})
        if set(output_qparams) != {0}:
            return False
        if not self._feeds_single_quantized_mul(node):

• Files reviewed: 9/9 changed files
• Comments generated: 0 new
• Review effort level: Lite

[Copilot]

🔵 Needs a closer look

Two unresolved comments request end-to-end pass-level coverage for quantized constant materialization.

Review details

Suppressed comments (2)

backends/cortex_m/passes/materialize_quantized_mul_constants_pass.py:56

• Please add a pass-level test for this new compute_node_aot path. The added unit tests only exercise _feeds_single_quantized_mul, so a regression that stops replacing the aten.full with a quantized persistent constant (or mishandles its qparams) would still pass; the transformer test validates lowering and numerics but not this materialization. Assert the pass result and the resulting graph/buffer state after invoking the pass.

    def compute_node_aot(self, node: Node) -> bool:
        output_qparams = node.meta.get("output_qparams", {})
        if set(output_qparams) != {0}:
            return False
        if not self._feeds_single_quantized_mul(node):

backends/cortex_m/test/test_materialize_quantized_mul_constants.py:36

• This only tests the private path predicate; it never runs MaterializeQuantizedMulConstantsPass.call() or compute_node_aot(), so regressions in the actual aten.full replacement/quantization would still pass. Add an end-to-end pass test that verifies the full node is replaced with the expected quantized constant while preserving the clone path.

    assert MaterializeQuantizedMulConstantsPass._feeds_single_quantized_mul(full)

• Files reviewed: 9/9 changed files
• Comments generated: 0 new
• Review effort level: Lite

[amacharla15]

@zingo the workflows on c09864c are still awaiting approval, so CI hasn’t run on the current head yet. Could you approve them when you get a chance?

[AdrianLundell]

Hello @amacharla15 and welcome with your first contribution to Executorch!

There are a number of things going on in this patch, as a general tip I strongly recommend splitting it up into several PR:s with one pass each for example, with briefer commit messages. This will significantly increase the chance of getting the changes merged since it is much faster to review and an issue in one pass does not block all other passes for example.

In regards to the patch:
Linear Q/DQ change: I happen to know that there will be a prototype for changing how the folded q/dq parameters are handled in the graph coming up just this week, so it makes sense to wait for that before merging a new solution, bad timing unfortunately.

Matrialize mul change: Introducing full shape constants for a single scalar mul is not so efficient, would prefer to see this by modifying the qparams, or a new accelerated kernel if there are cases where that does not work.

Matrialize constant change: Likely not need with other mul approach.

Let me know if you have any questions! Best regards,
Adrian

[amacharla15]

Thanks @AdrianLundell ;

Here is my observation and work:

I split the linear and scalar-MUL issues apart and followed up on your suggestion to handle the scalar scale through qparams rather than materializing a full-shape constant.

While investigating that path, I found the existing draft PR #21027, which removes a DQ -> MUL -> Q sequence when the surrounding qparams already make the MUL an integer-preserving identity. I tested that approach against the decomposed nn.Transformer reproducer from #21943.

The SDPA decomposition splits the attention scale across Q and K, with each branch using head_dim^(-1/4). On the full Transformer there are 6 such attention-scale branches. With independently observed qparams, only 2/6 satisfy the relation required by the fold:

S_out ~= S_in * c and Z_out == Z_in

The other 4 are correctly rejected by the fold, so those remaining MULs originate from the qparam choices made earlier during quantization rather than from the folding logic itself.

I then prototyped a small Cortex-M quantizer stage that derives the scaled activation qparams from the input:

S_out = S_in * c and Z_out = Z_in

for the positive scalar-scaled BMM-input pattern emitted by the existing SDPA decomposition.

One detail from the real graph: the scaled value reaches BMM through expand / view, so the derived qspec goes on the qspec-owning edge from SharedQspecQuantizer and the BMM input shares it. The stage only replaces an independently specified QuantizationSpec — existing shared or derived relationships are untouched.

With that change, all 6/6 attention-scale branches in the Transformer reproducer satisfy the identity relation needed by the later fold. I also tested the actual decomposed SDPA path with head dimensions 16 and 64, including the non-exact 64^(-1/4) scale, and verified that the derived qparams reach both score-BMM operands.

I have kept the linear int8/int32 metadata issue separate and stopped pursuing the Cortex-M-specific retracing workaround there, since you mentioned the upcoming folded-Q/DQ prototype.

Does this derived-qparam approach look like the direction you had in mind for the scalar-MUL issue? If so, I can clean it up and send it as a separate focused PR.

[AdrianLundell]

Sounds like good progress, I like this direction so definitely worth opening a PR!

As a general heads up I prefer to keep the quantizer simple and predictable, do any decomps needed to get correct qparams in transform_for_annotation, and then do fusings in the cortex-m lowering pass pipeline, if this is possible rather than the derived-qparam approach. Having a clear split of responsibilities across the different stages of the lowering makes it easier to understand.

[amacharla15]

Thanks @AdrianLundell , I implemented the transform_for_annotation approach you suggested as #22892.

It moves the split default-SDPA score scale after the score BMM, keeps CortexMQuantizer unchanged, and leaves downstream fusion separate in the lowering pipeline. With the stock quantizer, all three attention score scales in the Transformer reproducer now get the proportional qparams needed for a later fold; I also checked this across head dimensions 16/32/64/128.

I’ve kept the linear folded-Q/DQ issue separate as discussed.