diff --git a/backends/cortex_m/passes/BUCK b/backends/cortex_m/passes/BUCK index 2ae348a57e0..297bcf01cad 100644 --- a/backends/cortex_m/passes/BUCK +++ b/backends/cortex_m/passes/BUCK @@ -35,6 +35,7 @@ fbcode_target(_kind = runtime.python_library, "cortex_m_pass_manager.py", "decompose_hardswish_pass.py", "decompose_mean_pass.py", + "decompose_sdpa_pass.py", "explicit_layout_pass.py", "fuse_conv_padding_pass.py", "initialize_scratch_buffers_pass.py", diff --git a/backends/cortex_m/passes/cortex_m_pass_manager.py b/backends/cortex_m/passes/cortex_m_pass_manager.py index 8d44cd9b298..93ff439933f 100644 --- a/backends/cortex_m/passes/cortex_m_pass_manager.py +++ b/backends/cortex_m/passes/cortex_m_pass_manager.py @@ -40,6 +40,7 @@ from .clamp_hardswish_pass import ClampHardswishPass from .decompose_hardswish_pass import DecomposeHardswishPass from .decompose_mean_pass import DecomposeMeanPass +from .decompose_sdpa_pass import DecomposeSDPAPass from .explicit_layout_pass import ( CortexMCanonicalizeViewCopyPermutePass, CortexMReplaceOpsWithChannelsLastVariants, @@ -121,6 +122,7 @@ class CortexMPassManager(ExportedProgramPassManager): pass_list = legacy_pass_list pass_list_transform_for_annotation: list[PassClass] = [ + DecomposeSDPAPass, ScalarsToAttributePass, ReplaceScalarWithTensorArgPass, ClampHardswishPass, diff --git a/backends/cortex_m/passes/decompose_sdpa_pass.py b/backends/cortex_m/passes/decompose_sdpa_pass.py new file mode 100644 index 00000000000..97d3848d50a --- /dev/null +++ b/backends/cortex_m/passes/decompose_sdpa_pass.py @@ -0,0 +1,92 @@ +# 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 math +from typing import Dict + +import torch +from executorch.exir.pass_base import ExportPass, NodeMetadata, ProxyValue +from torch._ops import OpOverload +from torch.fx.node import Argument + + +class DecomposeSDPAPass(ExportPass): + """ + Decomposes float32 scaled_dot_product_attention without mask, causal + masking, dropout or GQA into matmul -> mul -> softmax -> matmul. Only + static, non-empty shapes where query, key and value have the same leading + dimensions are decomposed, since MatmulToBmmPass turns exactly those + matmuls into bmm. + + The default SDPA decomposition splits the scale over query and key before + the score matmul. Applying it once to the scores instead means the + quantizer observes the scores before and after the scale, so their qparams + differ by the scale and the mul can be folded during lowering. + + Without a mask a score row can only be all -inf for non-finite inputs, so + regular softmax is used rather than _safe_softmax. + """ + + _known_kwargs = {"attn_mask", "dropout_p", "is_causal", "scale", "enable_gqa"} + + def call_operator( + self, + op: OpOverload, + args: tuple[Argument, ...], + kwargs: Dict[str, Argument], + meta: NodeMetadata, + ) -> ProxyValue: + if op != torch.ops.aten.scaled_dot_product_attention.default: + return super().call_operator(op, args, kwargs, meta) + + query, key, value = args[:3] + attn_mask = args[3] if len(args) > 3 else kwargs.get("attn_mask") + dropout_p = args[4] if len(args) > 4 else kwargs.get("dropout_p", 0.0) + is_causal = args[5] if len(args) > 5 else kwargs.get("is_causal", False) + scale = kwargs.get("scale") + + if ( + not isinstance(query, ProxyValue) + or not isinstance(key, ProxyValue) + or not isinstance(value, ProxyValue) + or not kwargs.keys() <= self._known_kwargs + or attn_mask is not None + or dropout_p != 0.0 + or is_causal + or kwargs.get("enable_gqa", False) + ): + return super().call_operator(op, args, kwargs, meta) + + shapes = [t.to_tensor().shape for t in (query, key, value)] + if ( + # Scaling after the score matmul could overflow in half precision, + # which is why the default decomposition splits the scale. + query.to_tensor().dtype != torch.float32 + # MatmulToBmmPass reshapes with -1, which is ambiguous for empty + # tensors. + or not all(isinstance(dim, int) and dim > 0 for s in shapes for dim in s) + or len(shapes[0]) < 3 + or not shapes[0][:-2] == shapes[1][:-2] == shapes[2][:-2] + ): + return super().call_operator(op, args, kwargs, meta) + + if scale is None: + scale = 1.0 / math.sqrt(shapes[0][-1]) + + key_t = super().call_operator( + torch.ops.aten.transpose.int, (key, -2, -1), {}, meta + ) + scores = super().call_operator( + torch.ops.aten.matmul.default, (query, key_t), {}, meta + ) + # A Python scalar, unlike mul.Scalar, becomes a constant buffer in + # ScalarsToAttributePass rather than a full op evaluated at runtime. + scores = super().call_operator( + torch.ops.aten.mul.Tensor, (scores, scale), {}, meta + ) + attn = super().call_operator(torch.ops.aten.softmax.int, (scores, -1), {}, meta) + return super().call_operator( + torch.ops.aten.matmul.default, (attn, value), {}, meta + ) diff --git a/backends/cortex_m/test/ops/test_sdpa.py b/backends/cortex_m/test/ops/test_sdpa.py new file mode 100644 index 00000000000..b17cc54fd73 --- /dev/null +++ b/backends/cortex_m/test/ops/test_sdpa.py @@ -0,0 +1,239 @@ +# 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 math +from collections import Counter + +import torch +from executorch.backends.arm.test.common import parametrize, xfail_type +from executorch.backends.cortex_m.quantizer.quantizer import CortexMQuantizer +from executorch.backends.cortex_m.test.tester import CortexMTester, McuTestCase + + +class CortexMSDPA(torch.nn.Module): + ops_before_transforms = { + "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_batch_matmul_default": 2, + "executorch_exir_dialects_edge__ops_cortex_m_softmax_default": 1, + "executorch_exir_dialects_edge__ops_aten_mul_Tensor": 0, + } + + def __init__(self, attn_mask: torch.Tensor | None = None, **sdpa_kwargs): + super().__init__() + self.register_buffer("attn_mask", attn_mask) + self.sdpa_kwargs = sdpa_kwargs + + def forward(self, query, key, value): + return torch.nn.functional.scaled_dot_product_attention( + query, key, value, attn_mask=self.attn_mask, **self.sdpa_kwargs + ) + + +class CortexMMultiheadAttention(torch.nn.Module): + ops_before_transforms = { + "executorch_exir_dialects_edge__ops_aten_linear_default": 2, + "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": 2, + "executorch_exir_dialects_edge__ops_cortex_m_quantized_batch_matmul_default": 2, + "executorch_exir_dialects_edge__ops_cortex_m_softmax_default": 1, + "executorch_exir_dialects_edge__ops_aten_mul_Tensor": 0, + } + + def __init__(self, embed_dim: int, num_heads: int): + super().__init__() + self.attention = torch.nn.MultiheadAttention( + embed_dim, num_heads, batch_first=True + ) + + def forward(self, x): + return self.attention(x, x, x, need_weights=False)[0] + + +def _qkv(query_shape, key_shape, value_shape, dtype=torch.float32): + return ( + torch.randn(query_shape, dtype=dtype), + torch.randn(key_shape, dtype=dtype), + torch.randn(value_shape, dtype=dtype), + ) + + +test_cases = { + "self_attention": McuTestCase( + CortexMSDPA(), _qkv((1, 2, 8, 16), (1, 2, 8, 16), (1, 2, 8, 16)) + ), + "cross_attention": McuTestCase( + CortexMSDPA(), _qkv((1, 2, 4, 32), (1, 2, 12, 32), (1, 2, 12, 8)) + ), + "rank3": McuTestCase(CortexMSDPA(), _qkv((2, 8, 64), (2, 8, 64), (2, 8, 64))), + "explicit_scale": McuTestCase( + CortexMSDPA(scale=0.3), _qkv((1, 2, 8, 16), (1, 2, 8, 16), (1, 2, 8, 16)) + ), + "multihead_attention": McuTestCase( + CortexMMultiheadAttention(embed_dim=32, num_heads=4), + (torch.randn(1, 8, 32),), + ), +} + +not_decomposed_test_cases = { + "attn_mask": McuTestCase( + CortexMSDPA(attn_mask=torch.randn(8, 8)), + _qkv((1, 2, 8, 16), (1, 2, 8, 16), (1, 2, 8, 16)), + ), + "causal": McuTestCase( + CortexMSDPA(is_causal=True), _qkv((1, 2, 8, 16), (1, 2, 8, 16), (1, 2, 8, 16)) + ), + "dropout": McuTestCase( + CortexMSDPA(dropout_p=0.5), _qkv((1, 2, 8, 16), (1, 2, 8, 16), (1, 2, 8, 16)) + ), + "gqa": McuTestCase( + CortexMSDPA(enable_gqa=True), _qkv((1, 4, 8, 16), (1, 2, 8, 16), (1, 2, 8, 16)) + ), + "float16": McuTestCase( + CortexMSDPA(), + _qkv((1, 2, 8, 16), (1, 2, 8, 16), (1, 2, 8, 16), dtype=torch.float16), + ), + "empty_key": McuTestCase( + CortexMSDPA(), _qkv((1, 2, 8, 16), (1, 2, 0, 16), (1, 2, 0, 16)) + ), + "broadcast_batch": McuTestCase( + CortexMSDPA(), _qkv((2, 2, 8, 16), (1, 2, 8, 16), (1, 2, 8, 16)) + ), + "rank2": McuTestCase(CortexMSDPA(), _qkv((8, 16), (8, 16), (8, 16))), +} + +_SCALE_NOT_FOLDED = ( + "The attention scale mul is left as an fp32 aten.mul between dequantize " + "and quantize, no pass folds it into the quantize yet" +) +_LINEAR_RETRACED_AS_INT32 = ( + "Retracing a folded quantized linear gives an int32 output, which fails " + "the dtype check of the bmm consuming it" +) + +xfail_cases_dialect: dict[str, xfail_type] = { + "self_attention": (_SCALE_NOT_FOLDED, RuntimeError), + "cross_attention": (_SCALE_NOT_FOLDED, RuntimeError), + "rank3": (_SCALE_NOT_FOLDED, RuntimeError), + "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): + # Without check_guards=False, whether export adds a _guards_fn call_module, + # which ExportPass rejects, depends on the path of the calling file. + graph_module = torch.export.export( + model, example_inputs, dynamic_shapes=dynamic_shapes, strict=True + ).module(check_guards=False) + return CortexMQuantizer().transform_for_annotation(graph_module) + + +def _op_counts(graph_module) -> Counter: + return Counter( + node.target for node in graph_module.graph.nodes if node.op == "call_function" + ) + + +@parametrize("test_case", test_cases, xfails=xfail_cases_dialect) +def test_dialect_sdpa(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, xfails=xfail_cases_impl) +def test_implementation_sdpa(test_case, cortex_m_target): + tester = CortexMTester( + test_case.model, test_case.example_inputs, target_config=cortex_m_target + ) + tester.test_implementation(qtol=2) + + +@parametrize("test_case", test_cases) +def test_sdpa_decomposition(test_case): + inputs = test_case.get_example_inputs() + graph_module = _transform_for_annotation(test_case.model, inputs) + + op_counts = _op_counts(graph_module) + assert op_counts[torch.ops.aten.scaled_dot_product_attention.default] == 0 + assert op_counts[torch.ops.aten.matmul.default] == 0 + assert op_counts[torch.ops.aten.bmm.default] == 2 + assert op_counts[torch.ops.aten.mul.Tensor] == 1 + assert op_counts[torch.ops.aten.softmax.int] == 1 + torch.testing.assert_close(graph_module(*inputs), test_case.model(*inputs)) + + +@parametrize("test_case", not_decomposed_test_cases) +def test_sdpa_not_decomposed(test_case): + graph_module = _transform_for_annotation( + test_case.model, test_case.get_example_inputs() + ) + + assert ( + _op_counts(graph_module)[torch.ops.aten.scaled_dot_product_attention.default] + == 1 + ) + + +def test_sdpa_with_dynamic_shape_not_decomposed(): + seq_len = torch.export.Dim("seq_len", min=2, max=64) + graph_module = _transform_for_annotation( + CortexMSDPA(), + _qkv((1, 2, 8, 16), (1, 2, 8, 16), (1, 2, 8, 16)), + dynamic_shapes=({2: seq_len}, {2: seq_len}, {2: seq_len}), + ) + + assert ( + _op_counts(graph_module)[torch.ops.aten.scaled_dot_product_attention.default] + == 1 + ) + + +@parametrize("test_case", test_cases) +def test_sdpa_scale_mul_qparams_allow_folding(test_case): + """The scale mul can only be folded into the quantize after it without + changing any int8 value if that quantize uses the scale of the dequantize + before it times the attention scale, and the same zero point.""" + tester = CortexMTester(test_case.model, test_case.example_inputs) + tester.quantize() + graph_module = tester.get_artifact() + + muls = [ + node + for node in graph_module.graph.nodes + if node.target == torch.ops.aten.mul.Tensor + ] + assert len(muls) == 1 + dequantize, scale_constant = muls[0].args + (quantize,) = muls[0].users + assert ( + dequantize.target + == torch.ops.quantized_decomposed.dequantize_per_tensor.default + ) + assert quantize.target == torch.ops.quantized_decomposed.quantize_per_tensor.default + + scale = getattr(graph_module, scale_constant.target).item() + in_scale, in_zero_point = dequantize.args[1:3] + out_scale, out_zero_point = quantize.args[1:3] + assert out_zero_point == in_zero_point + assert math.isclose(out_scale, in_scale * scale, rel_tol=1e-6)