Cortex-M: fix quantized lowering for decomposed attention

backends/cortex_m/passes/BUCK

@@ -10,6 +10,18 @@ load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
 
 oncall("executorch")
 
+fbcode_target(_kind = runtime.python_library,
+    name="cortex_m_quantized_pass",
+    srcs=[
+        "cortex_m_quantized_pass.py",
+    ],
+    deps=[
+        "//caffe2:torch",
+        "//executorch/exir:pass_base",
+        "//executorch/exir/dialects:lib",
+    ],
+)
+
 fbcode_target(_kind = runtime.python_library,
     name="replace_quant_nodes_pass",
     srcs=[
@@ -21,6 +33,7 @@ fbcode_target(_kind = runtime.python_library,
         "//executorch/exir:pass_base",
         "//executorch/exir/dialects:lib",
         "//executorch/backends/cortex_m/ops:ops",
+        "//executorch/backends/cortex_m/passes:cortex_m_quantized_pass",
         "//executorch/backends/cortex_m/passes:passes_utils",
     ],
 )
@@ -37,7 +50,10 @@ fbcode_target(_kind = runtime.python_library,
         "decompose_mean_pass.py",
         "explicit_layout_pass.py",
         "matmul_to_bmm_pass.py",
+        "materialize_mul_broadcast_pass.py",
+        "materialize_quantized_mul_constants_pass.py",
         "quantized_clamp_activation_pass.py",
+        "replace_scalar_with_tensor_pass.py",
         "scratch_buffer_sizes.py",
     ],
     deps=[
@@ -46,6 +62,7 @@ fbcode_target(_kind = runtime.python_library,
         "//executorch/backends/cortex_m:cmsis_nn",
         "//executorch/backends/cortex_m:target_config",
         "//executorch/backends/cortex_m/ops:ops",
+        "//executorch/backends/cortex_m/passes:cortex_m_quantized_pass",
         "//executorch/backends/cortex_m/passes:passes_utils",
         "//executorch/backends/cortex_m/passes:replace_quant_nodes_pass",
         "//executorch/backends/transforms:aten_to_dialect_pass",

backends/cortex_m/passes/cortex_m_pass_manager.py

@@ -41,9 +41,14 @@
     CortexMReplaceOpsWithChannelsLastVariants,
     ValidateCortexMExplicitLayoutPass,
 )
+from .materialize_mul_broadcast_pass import MaterializeMulBroadcastPass
+from .materialize_quantized_mul_constants_pass import (
+    MaterializeQuantizedMulConstantsPass,
+)
 from .matmul_to_bmm_pass import MatmulToBmmPass
 from .quantized_clamp_activation_pass import QuantizedClampActivationPass
 from .replace_quant_nodes_pass import ReplaceQuantNodesPass
+from .replace_scalar_with_tensor_pass import CortexMReplaceScalarWithTensorArgPass
 
 PassClass = Type[ExportPass]
 
@@ -53,7 +58,9 @@ class CortexMPassManager(PassManager):
         # Run before folding so qparams attach to max_pool2d values, not tuple + getitem.
         RemoveGetItemPass,
         FoldAndAnnotateQParamsPass,
-        ReplaceScalarWithTensorArgPass,
+        # Materialize only quantized MUL constants after qparams are attached.
+        MaterializeQuantizedMulConstantsPass,
+        CortexMReplaceScalarWithTensorArgPass,
         ReplaceQuantNodesPass,
         ActivationFusionPass,
         QuantizedClampActivationPass,
@@ -84,6 +91,7 @@ class CortexMPassManager(PassManager):
     pass_list_transform_for_annotation: list[PassClass] = [
         ScalarsToAttributePass,
         ReplaceScalarWithTensorArgPass,
+        MaterializeMulBroadcastPass,
         ClampHardswishPass,
         DecomposeMeanPass,
         MatmulToBmmPass,

backends/cortex_m/passes/cortex_m_quantized_pass.py

@@ -0,0 +1,83 @@
+# 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.
+
+from typing import Any
+
+import torch
+
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.pass_base import ExportPass, NodeMetadata, ProxyValue
+from torch.utils import _pytree as pytree
+
+
+class CortexMQuantizedPass(ExportPass):
+    """Base pass for retracing folded Cortex-M quantized graphs.
+
+    Cortex-M's quantized linear kernel includes requantization and therefore
+    produces the dtype described by its output qparams, normally int8.
+
+    Once Q/DQ nodes have been folded, retracing ``aten.linear`` through the
+    regular PyTorch fake kernel would incorrectly recreate an int32 output.
+    Preserve the Cortex-M output dtype instead.
+    """
+
+    @staticmethod
+    def _is_quantized_meta(meta: NodeMetadata) -> bool:
+        input_qparams = meta.data.get("input_qparams", {})
+        output_qparams = meta.data.get("output_qparams", {})
+        return bool(input_qparams) and bool(output_qparams)
+
+    def _call_quantized_op_without_fake_kernel(
+        self,
+        op,
+        args: tuple[ProxyValue, ...],
+        kwargs: dict[str, Any],
+        meta: NodeMetadata,
+    ) -> ProxyValue:
+        old_val = meta.data["val"]
+        output_qparams = meta.data.get("output_qparams", {})
+        dtype = (
+            next(iter(output_qparams.values())).dtype
+            if output_qparams
+            else old_val.dtype
+        )
+
+        result_data = torch.empty_like(old_val, dtype=dtype)
+
+        args_proxy, kwargs_proxy = pytree.tree_map_only(
+            ProxyValue,
+            lambda value: value.proxy,
+            (args, kwargs),
+        )
+
+        result_proxy = self.tracer.create_proxy(
+            "call_function",
+            op,
+            args_proxy,
+            kwargs_proxy,
+        )
+        result_proxy.node.meta.update(meta.data)
+        self.tracer.set_metadata(result_proxy.node, result_data)
+
+        return ProxyValue(result_data, result_proxy)
+
+    def call_operator(self, op, args, kwargs, meta):
+        if (
+            op
+            in {
+                exir_ops.edge.aten.linear.default,
+                exir_ops.edge.aten.bmm.default,
+            }
+            and isinstance(meta, NodeMetadata)
+            and self._is_quantized_meta(meta)
+        ):
+            return self._call_quantized_op_without_fake_kernel(
+                op,
+                args,
+                kwargs,
+                meta,
+            )
+
+        return super().call_operator(op, args, kwargs, meta)

backends/cortex_m/passes/materialize_mul_broadcast_pass.py

@@ -0,0 +1,82 @@
+# 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._passes.arm_pass_utils import get_first_fake_tensor
+from executorch.backends.cortex_m.passes.passes_utils import (
+    is_channel_broadcast,
+    is_channels_last,
+)
+from executorch.exir.pass_base import ExportPass, PassResult
+from torch.fx import GraphModule, Node
+
+
+class MaterializeMulBroadcastPass(ExportPass):
+    """Materialize scalar MUL broadcasts unsupported by the Cortex-M kernel.
+
+    ``ReplaceScalarWithTensorArgPass`` represents scalar multiplication as a
+    rank-1 ``aten.full`` tensor. When that broadcast is unsupported by the
+    Cortex-M quantized MUL kernel, materialize the scalar constant directly at
+    the MUL output shape. Leave general tensor broadcasting unchanged so those
+    graphs keep their existing fallback behavior.
+    """
+
+    @staticmethod
+    def _is_scalar_full(node: Node) -> bool:
+        if node.op != "call_function" or node.target != torch.ops.aten.full.default:
+            return False
+        return get_first_fake_tensor(node).numel() == 1
+
+    def call(self, graph_module: GraphModule) -> PassResult:
+        modified = False
+
+        for node in list(graph_module.graph.nodes):
+            if node.op != "call_function" or node.target != torch.ops.aten.mul.Tensor:
+                continue
+
+            lhs, rhs = node.args[:2]
+            if not isinstance(lhs, Node) or not isinstance(rhs, Node):
+                continue
+
+            lhs_tensor = get_first_fake_tensor(lhs)
+            rhs_tensor = get_first_fake_tensor(rhs)
+            output_tensor = get_first_fake_tensor(node)
+
+            if lhs_tensor.shape == rhs_tensor.shape:
+                continue
+
+            if is_channel_broadcast(lhs_tensor, rhs_tensor) and is_channels_last(
+                lhs_tensor
+            ):
+                continue
+
+            output_shape = list(output_tensor.shape)
+            new_args = list(node.args)
+
+            for index, arg in enumerate((lhs, rhs)):
+                tensor = get_first_fake_tensor(arg)
+                if tensor.shape == output_tensor.shape or not self._is_scalar_full(arg):
+                    continue
+
+                full_args = list(arg.args)
+                full_args[0] = output_shape
+                with graph_module.graph.inserting_before(node):
+                    materialized = graph_module.graph.call_function(
+                        torch.ops.aten.full.default,
+                        args=tuple(full_args),
+                        kwargs=dict(arg.kwargs),
+                    )
+
+                new_args[index] = materialized
+                modified = True
+
+            node.args = tuple(new_args)
+
+        if modified:
+            graph_module.graph.eliminate_dead_code()
+            graph_module.recompile()
+            graph_module = super().call(graph_module).graph_module
+
+        return PassResult(graph_module, modified)

backends/cortex_m/passes/materialize_quantized_mul_constants_pass.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.
+
+from executorch.backends.arm._passes import ComputeConstantOpsAOTPass
+from executorch.exir.dialects._ops import ops as exir_ops
+from torch.fx import Node
+
+
+class MaterializeQuantizedMulConstantsPass(ComputeConstantOpsAOTPass):
+    """Precompute only quantized scalar constants that feed Cortex-M MUL.
+
+    The generic constant-materialization pass also folds unrelated operators
+    such as ``arange``. Cortex-M only needs the ``aten.full`` constants created
+    for scalar MULs here, after qparams have been attached, so keep the scope
+    limited to those nodes.
+    """
+
+    targeted_ops = [exir_ops.edge.aten.full.default]
+
+    _passthrough_ops = {
+        exir_ops.edge.aten.clone.default,
+        exir_ops.edge.dim_order_ops._clone_dim_order.default,
+    }
+
+    @classmethod
+    def _feeds_single_quantized_mul(cls, node: Node) -> bool:
+        current = node
+        visited: set[Node] = set()
+
+        while True:
+            users = list(current.users)
+            if len(users) != 1:
+                return False
+
+            user = users[0]
+            if user in visited:
+                return False
+            visited.add(user)
+
+            if user.target == exir_ops.edge.aten.mul.Tensor:
+                return bool(user.meta.get("input_qparams")) and bool(
+                    user.meta.get("output_qparams")
+                )
+
+            if user.target not in cls._passthrough_ops:
+                return False
+
+            current = user
+
+    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):
+            return False
+        return super().compute_node_aot(node)

backends/cortex_m/passes/replace_quant_nodes_pass.py

@@ -11,10 +11,12 @@
 import torch
 
 from executorch.exir.dialects._ops import ops as exir_ops
-from executorch.exir.pass_base import ExportPass, NodeMetadata, ProxyValue
+from executorch.exir.pass_base import NodeMetadata, ProxyValue
 
+from .cortex_m_quantized_pass import CortexMQuantizedPass
 
-class ReplaceQuantNodesPass(ExportPass):
+
+class ReplaceQuantNodesPass(CortexMQuantizedPass):
     """
     Replace quantize and dequantize nodes with the corresponding
     cortex_m.quantize_per_tensor and cortex_m.dequantize_per_tensor nodes.

backends/cortex_m/passes/replace_scalar_with_tensor_pass.py

@@ -0,0 +1,17 @@
+# 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.
+
+from executorch.backends.transforms.replace_scalar_with_tensor import (
+    ReplaceScalarWithTensorArgPass,
+)
+
+from .cortex_m_quantized_pass import CortexMQuantizedPass
+
+
+class CortexMReplaceScalarWithTensorArgPass(
+    ReplaceScalarWithTensorArgPass,
+    CortexMQuantizedPass,
+):
+    """Scalar-to-tensor replacement with Cortex-M quantized retracing."""