Cortex-M: Fix pass ordering for quantized attention (#23923)

Convert scalar arguments before folding quantization nodes, and replace
quantization nodes after lowering ATen ops to Cortex-M kernels. Avoid
retracing partially lowered attention graphs whose matrix inputs have
incompatible int8/int32 types.

Unblock default Cortex-M lowering of attention patterns used by DeiT
Tiny and similar models. Add M55/M85 coverage for numerical accuracy and
PTE serialization.

AI-assisted: OpenAI Codex

Signed-off-by: Zingo Andersen <Zingo.Andersen@arm.com>

Author: zingo

backends/cortex_m/passes/cortex_m_pass_manager.py

@@ -111,13 +111,15 @@ class CortexMPassManager(ExportedProgramPassManager):
     legacy_pass_list: list[PassClass] = [
         # Run before folding so qparams attach to max_pool2d values, not tuple + getitem.
         RemoveGetItemPass,
-        FoldAndAnnotateQParamsPass,
+        # Retrace scalar ops while the graph still has floating-point Q/DQ inputs.
         ReplaceScalarWithTensorArgPass,
-        ReplaceQuantNodesPass,
+        FoldAndAnnotateQParamsPass,
         ActivationFusionPass,
         QuantizedClampActivationPass,
         DecomposeHardswishPass,
         AtenToCortexMPass,
+        # Retrace only after quantized ATen ops have become Cortex-M kernels.
+        ReplaceQuantNodesPass,
         FuseConvPaddingPass,
         InitializeScratchBuffersPass,
         LiftConstantTensorsPass,

backends/cortex_m/test/build_test_runner.sh

@@ -51,8 +51,11 @@ ops_list=(
     aten::cat.out
     aten::full.out
     aten::ge.Tensor_out
+    aten::permute_copy.out
+    aten::slice_copy.Tensor_out
     aten::unsqueeze_copy.out
     aten::squeeze_copy.dim_out
+    aten::squeeze_copy.dims_out
     aten::select_copy.int_out
     aten::amax.out
     cortex_m::quantize_per_tensor.out

backends/cortex_m/test/ops/test_attention.py

@@ -0,0 +1,58 @@
+# Copyright 2026 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+from executorch.backends.arm.test.common import parametrize
+from executorch.backends.cortex_m.test.tester import CortexMTester, McuTestCase
+
+
+class CortexMAttention(torch.nn.Module):
+    ops_before_transforms = {
+        "executorch_exir_dialects_edge__ops_aten_linear_default": 1,
+        "executorch_exir_dialects_edge__ops_aten_bmm_default": 2,
+        "executorch_exir_dialects_edge__ops_aten_mul_Tensor": 1,
+        "executorch_exir_dialects_edge__ops_aten__softmax_default": 1,
+    }
+
+    ops_after_transforms = {
+        "executorch_exir_dialects_edge__ops_cortex_m_quantized_linear_default": 1,
+        "executorch_exir_dialects_edge__ops_cortex_m_quantized_batch_matmul_default": 2,
+        "executorch_exir_dialects_edge__ops_cortex_m_softmax_default": 1,
+    }
+
+    def __init__(self):
+        super().__init__()
+        self.qkv = torch.nn.Linear(12, 36)
+
+    def forward(self, x):
+        qkv = self.qkv(x).reshape(1, 5, 3, 3, 4).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv.unbind(0)
+        scores = (q * 0.5) @ k.transpose(-2, -1)
+        return (scores.softmax(dim=-1) @ v).transpose(1, 2).reshape(1, 5, 12)
+
+
+test_cases = {
+    "qkv_attention": McuTestCase(CortexMAttention(), (torch.randn(1, 5, 12),)),
+}
+
+
+@parametrize("test_case", test_cases)
+def test_dialect_attention(test_case, cortex_m_target):
+    tester = CortexMTester(
+        test_case.model, test_case.example_inputs, target_config=cortex_m_target
+    )
+    tester.test_dialect(
+        test_case.model.ops_before_transforms,
+        test_case.model.ops_after_transforms,
+        qtol=2,
+    )
+
+
+@parametrize("test_case", test_cases)
+def test_implementation_attention(test_case, cortex_m_target):
+    tester = CortexMTester(
+        test_case.model, test_case.example_inputs, target_config=cortex_m_target
+    )
+    tester.test_implementation(qtol=2)

backends/cortex_m/test/ops/test_sdpa.py

@@ -129,9 +129,6 @@ def _qkv(query_shape, key_shape, value_shape, dtype=torch.float32):
     "explicit_scale": (_SCALE_NOT_FOLDED, RuntimeError),
     "multihead_attention": (_LINEAR_RETRACED_AS_INT32, Exception),
 }
-xfail_cases_impl: dict[str, xfail_type] = {
-    "multihead_attention": (_LINEAR_RETRACED_AS_INT32, Exception),
-}
 
 
 def _transform_for_annotation(model, example_inputs, dynamic_shapes=None):
@@ -161,7 +158,7 @@ def test_dialect_sdpa(test_case, cortex_m_target):
     )
 
 
-@parametrize("test_case", test_cases, xfails=xfail_cases_impl)
+@parametrize("test_case", test_cases)
 def test_implementation_sdpa(test_case, cortex_m_target):
     tester = CortexMTester(
         test_case.model, test_case.example_inputs, target_config=cortex_m_target

backends/cortex_m/test/targets.bzl

@@ -1,5 +1,6 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
+# Copyright 2026 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
@@ -104,6 +105,23 @@ def define_common_targets(is_fbcode = False):
             ],
         )
 
+        python_pytest(
+            name = "test_attention",
+            srcs = ["ops/test_attention.py"],
+            resources = ["conftest.py"],
+            compile = "with-source",
+            typing = False,
+            # Implementation tests require the Cortex-M FVP runner, which Buck does not provide.
+            env = {"PYTEST_ADDOPTS": "-k 'not test_implementation'"},
+            deps = [
+                "//caffe2:torch",
+                "//executorch/backends/arm/test:common",
+                "//executorch/backends/cortex_m:target_config",
+                ":tester",
+                "fbsource//third-party/pypi/pytest:pytest",
+            ],
+        )
+
         python_pytest(
             name = "test_pass_manager",
             srcs = ["test_pass_manager.py"],