From 485ab0bd42a47c8f08cc018c927033c1331c8771 Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Thu, 16 Jul 2026 15:42:58 -0700 Subject: [PATCH 1/5] Cortex-M backend: fold constant attention-score scale into the softmax-input quantize A constant elementwise scale between a quantized producer and a per-tensor quantize -- e.g. the attention scores/sqrt(d) before softmax -- otherwise stays an fp32 aten.div/mul with dq/q boundaries. Since quantize(x/c, S) == quantize(x, S*c) (and *c divides S), the constant folds into the adjacent quantize with no numerical change. Add a cortex_m-scoped FoldScaleIntoQuantizePass (run before FoldAndAnnotateQParamsPass, while the softmax-input quantize is still explicit) that folds a constant scalar div/mul whose sole consumer is a quantize_per_tensor into that quantize's scale and drops the div/mul. Only constant (param, numel==1) scales fold, so data-dependent normalization (e.g. LayerNorm's divide by sqrt(var)) is correctly left untouched. Verified test-first (RED->GREEN) with div and mul scaled-attention cases (numerics at qtol=1 pin the fold direction); full cortex_m op dialect suite still passes (179). On a SAM mask decoder this removes the 7 attention /sqrt(d) divs so the attention chain is int8 through softmax. Authored with Claude Code. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../cortex_m/passes/cortex_m_pass_manager.py | 4 + .../passes/fold_scale_into_quantize_pass.py | 81 ++++++++++++++++++ .../cortex_m/test/ops/test_attention_scale.py | 85 +++++++++++++++++++ 3 files changed, 170 insertions(+) create mode 100644 backends/cortex_m/passes/fold_scale_into_quantize_pass.py create mode 100644 backends/cortex_m/test/ops/test_attention_scale.py diff --git a/backends/cortex_m/passes/cortex_m_pass_manager.py b/backends/cortex_m/passes/cortex_m_pass_manager.py index 892baf136ed..3e1ac7077f0 100644 --- a/backends/cortex_m/passes/cortex_m_pass_manager.py +++ b/backends/cortex_m/passes/cortex_m_pass_manager.py @@ -27,6 +27,7 @@ from .clamp_hardswish_pass import ClampHardswishPass from .decompose_hardswish_pass import DecomposeHardswishPass from .decompose_mean_pass import DecomposeMeanPass +from .fold_scale_into_quantize_pass import FoldScaleIntoQuantizePass from .matmul_to_bmm_pass import MatmulToBmmPass from .quantized_clamp_activation_pass import QuantizedClampActivationPass from .replace_quant_nodes_pass import ReplaceQuantNodesPass @@ -36,6 +37,9 @@ class CortexMPassManager(PassManager): pass_list: list[PassClass] = [ + # Fold constant scales (e.g. attention /sqrt(d)) into the adjacent + # quantize before its scale is folded into op meta. + FoldScaleIntoQuantizePass, # Run before folding so qparams attach to max_pool2d values, not tuple + getitem. RemoveGetItemPass, FoldAndAnnotateQParamsPass, diff --git a/backends/cortex_m/passes/fold_scale_into_quantize_pass.py b/backends/cortex_m/passes/fold_scale_into_quantize_pass.py new file mode 100644 index 00000000000..25f3f9b5549 --- /dev/null +++ b/backends/cortex_m/passes/fold_scale_into_quantize_pass.py @@ -0,0 +1,81 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# 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 Optional, Set, Type + +from executorch.backends.arm._passes.arm_pass_utils import ( + get_param_tensor, + is_param_node, +) +from executorch.exir import ExportedProgram +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass, PassResult +from torch.fx import GraphModule, Node + +_QUANTIZE = exir_ops.edge.quantized_decomposed.quantize_per_tensor.default +_DIV = exir_ops.edge.aten.div.Tensor +_MUL = exir_ops.edge.aten.mul.Tensor + + +class FoldScaleIntoQuantizePass(ExportPass): + """Fold a constant elementwise scale (``x / c`` or ``x * c``) into the scale + of the per-tensor quantize that consumes it, then drop the div/mul. + + Because ``quantize(x / c, scale=S) == quantize(x, scale=S*c)`` and + ``quantize(x * c, scale=S) == quantize(x, scale=S/c)`` produce identical int8 + values, the constant scale can be absorbed into the adjacent quantize with no + numerical change. This erases the attention-score ``/sqrt(d)`` scale -- an + fp32 div that otherwise stays between the QK^T bmm and softmax -- so the + attention chain is int8 through softmax. + + Runs before ``FoldAndAnnotateQParamsPass`` while the softmax-input quantize is + still an explicit ``quantized_decomposed.quantize_per_tensor`` node. + """ + + _passes_required_after: Set[Type[ExportPass]] = set() + + def __init__(self, exported_program: Optional[ExportedProgram] = None) -> None: + super().__init__() + self.exported_program = exported_program + + def call(self, graph_module: GraphModule) -> PassResult: + ep = self.exported_program + if ep is None: + return PassResult(graph_module, False) + + modified = False + for node in list(graph_module.graph.nodes): + if node.op != "call_function" or node.target not in (_DIV, _MUL): + continue + + scaled, const = node.args[0], node.args[1] + if not (isinstance(scaled, Node) and isinstance(const, Node)): + continue + if not is_param_node(ep, const): + continue + const_t = get_param_tensor(ep, const) + if const_t is None or const_t.numel() != 1: + continue + c = float(const_t.reshape(-1)[0]) + if c == 0.0: + continue + + users = list(node.users) + if len(users) != 1 or users[0].target != _QUANTIZE: + continue + + quantize = users[0] + scale = quantize.args[1] + new_scale = scale * c if node.target is _DIV else scale / c + quantize.update_arg(1, new_scale) + quantize.replace_input_with(node, scaled) + graph_module.graph.erase_node(node) + modified = True + + if modified: + graph_module.graph.eliminate_dead_code() + graph_module.recompile() + return PassResult(graph_module, modified) diff --git a/backends/cortex_m/test/ops/test_attention_scale.py b/backends/cortex_m/test/ops/test_attention_scale.py new file mode 100644 index 00000000000..1636b128331 --- /dev/null +++ b/backends/cortex_m/test/ops/test_attention_scale.py @@ -0,0 +1,85 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# 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 + +import torch +from executorch.backends.arm.test.common import parametrize +from executorch.backends.cortex_m.test.tester import CortexMTester, McuTestCase + + +class CortexMScaledAttentionDiv(torch.nn.Module): + """``softmax(bmm(q, k^T) / sqrt(d))`` -- the attention-score scale is an fp32 + ``aten.div.Tensor`` by a constant that otherwise stays between the QK^T bmm + and softmax. It must fold into the softmax-input quantize scale + (``quantize(x / c, S) == quantize(x, S*c)``) so no fp32 div remains and the + chain lowers to ``cortex_m.quantized_batch_matmul`` + ``cortex_m.softmax``. + """ + + ops_before_transforms = { + "executorch_exir_dialects_edge__ops_aten_bmm_default": 1, + "executorch_exir_dialects_edge__ops_aten_div_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": 1, + "executorch_exir_dialects_edge__ops_cortex_m_softmax_default": 1, + "executorch_exir_dialects_edge__ops_aten_div_Tensor": 0, + "executorch_exir_dialects_edge__ops_aten_mul_Tensor": 0, + } + + def forward(self, q, k): + scores = torch.bmm(q, k.transpose(-2, -1)) / math.sqrt(q.shape[-1]) + return torch.softmax(scores, dim=-1) + + +class CortexMScaledAttentionMul(torch.nn.Module): + """Same, but the scale is applied as ``* (1/sqrt(d))`` -- an + ``aten.mul.Tensor`` by a constant, folded via ``quantize(x*c, S) == + quantize(x, S/c)``. + """ + + ops_before_transforms = { + "executorch_exir_dialects_edge__ops_aten_bmm_default": 1, + "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": 1, + "executorch_exir_dialects_edge__ops_cortex_m_softmax_default": 1, + "executorch_exir_dialects_edge__ops_aten_div_Tensor": 0, + "executorch_exir_dialects_edge__ops_aten_mul_Tensor": 0, + } + + def forward(self, q, k): + scores = torch.bmm(q, k.transpose(-2, -1)) * (1.0 / math.sqrt(q.shape[-1])) + return torch.softmax(scores, dim=-1) + + +test_cases = { + "scaled_attn_div": McuTestCase( + CortexMScaledAttentionDiv(), + (torch.rand(1, 8, 16), torch.rand(1, 8, 16)), + ), + "scaled_attn_mul": McuTestCase( + CortexMScaledAttentionMul(), + (torch.rand(1, 8, 16), torch.rand(1, 8, 16)), + ), +} + + +@parametrize("test_case", test_cases) +def test_dialect_attention_scale(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=1, + ) From c8a910a857b8ed8dbe10e1cc2daeb14f58a73995 Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Mon, 20 Jul 2026 11:36:23 -0700 Subject: [PATCH 2/5] Cortex-M: add fold_scale_into_quantize_pass.py to the passes BUCK srcs This branch adds backends/cortex_m/passes/fold_scale_into_quantize_pass.py and imports it from cortex_m_pass_manager. Add the file to the cortex_passes python_library srcs so the internal buck build can resolve the import; OSS pytest imports by path and does not catch a missing buck src. The pass's imports (arm _passes, exir, exir.dialects._ops, exir.pass_base, torch.fx) are already covered by the target's existing deps and sibling srcs, so no new dep is needed. Co-Authored-By: Claude Opus 4.8 (1M context) --- backends/cortex_m/passes/BUCK | 1 + 1 file changed, 1 insertion(+) diff --git a/backends/cortex_m/passes/BUCK b/backends/cortex_m/passes/BUCK index d301e14823c..2a66ca3e80c 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", + "fold_scale_into_quantize_pass.py", "matmul_to_bmm_pass.py", "quantized_clamp_activation_pass.py", "scratch_buffer_sizes.py", From c0635f5032835bdf1ed8b0e7e67afc8ea5c57f12 Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Mon, 20 Jul 2026 15:48:10 -0700 Subject: [PATCH 3/5] Cortex-M: fold the attention scale after qparam annotation Address AdrianLundell's review of #21027: run FoldScaleIntoQuantizePass after FoldAndAnnotateQParamsPass and fold the constant scale into a rescaling op that owns quantization parameters, rather than into an explicit quantize node before annotation. Post-annotation the constant scale (e.g. attention /sqrt(d)) is already absorbed into the consuming op's input scale: the quantizer's observer saw the post-scale range, so the softmax-input scale equals the bmm-output scale divided by the constant (multiplied, for mul) with matching zero-points. The div/mul therefore survives inside a redundant dequantize -> op -> quantize sandwich that is the identity on the int8 values. The pass recognises that relation -- checking the scale, zero-point, and clamp range/dtype all line up -- and deletes the sandwich, wiring the producer's int8 straight into the consumer. It rewrites no quantization parameter, so a SharedQspec consumer that shares the scale (a view/reshape/pool) is never disturbed, the case Adrian raised. If the relation does not hold the fold is skipped. Tests cover the div and mul folds, bit-exactness over the int8 range, a constant scale feeding a SharedQspec pool (shared scale unchanged), a constant-first operand (must not crash), and the skip guards. Authored with Claude Code. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../cortex_m/passes/cortex_m_pass_manager.py | 6 +- .../passes/fold_scale_into_quantize_pass.py | 121 +++++++--- .../test/ops/test_fold_scale_into_quantize.py | 215 ++++++++++++++++++ 3 files changed, 303 insertions(+), 39 deletions(-) create mode 100644 backends/cortex_m/test/ops/test_fold_scale_into_quantize.py diff --git a/backends/cortex_m/passes/cortex_m_pass_manager.py b/backends/cortex_m/passes/cortex_m_pass_manager.py index 3e1ac7077f0..b66de52ff7d 100644 --- a/backends/cortex_m/passes/cortex_m_pass_manager.py +++ b/backends/cortex_m/passes/cortex_m_pass_manager.py @@ -37,12 +37,12 @@ class CortexMPassManager(PassManager): pass_list: list[PassClass] = [ - # Fold constant scales (e.g. attention /sqrt(d)) into the adjacent - # quantize before its scale is folded into op meta. - FoldScaleIntoQuantizePass, # Run before folding so qparams attach to max_pool2d values, not tuple + getitem. RemoveGetItemPass, FoldAndAnnotateQParamsPass, + # After qparam-folding: drop a constant scale (e.g. attention /sqrt(d)) + # the quantizer already absorbed into the adjacent quantize scale. + FoldScaleIntoQuantizePass, ReplaceScalarWithTensorArgPass, ReplaceQuantNodesPass, ActivationFusionPass, diff --git a/backends/cortex_m/passes/fold_scale_into_quantize_pass.py b/backends/cortex_m/passes/fold_scale_into_quantize_pass.py index 25f3f9b5549..055b2ac9e26 100644 --- a/backends/cortex_m/passes/fold_scale_into_quantize_pass.py +++ b/backends/cortex_m/passes/fold_scale_into_quantize_pass.py @@ -4,7 +4,8 @@ # 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 Optional, Set, Type +import math +from typing import cast, Optional, Set, Tuple, Type from executorch.backends.arm._passes.arm_pass_utils import ( get_param_tensor, @@ -16,23 +17,35 @@ from torch.fx import GraphModule, Node _QUANTIZE = exir_ops.edge.quantized_decomposed.quantize_per_tensor.default +_DEQUANTIZE = exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default _DIV = exir_ops.edge.aten.div.Tensor _MUL = exir_ops.edge.aten.mul.Tensor class FoldScaleIntoQuantizePass(ExportPass): - """Fold a constant elementwise scale (``x / c`` or ``x * c``) into the scale - of the per-tensor quantize that consumes it, then drop the div/mul. - - Because ``quantize(x / c, scale=S) == quantize(x, scale=S*c)`` and - ``quantize(x * c, scale=S) == quantize(x, scale=S/c)`` produce identical int8 - values, the constant scale can be absorbed into the adjacent quantize with no - numerical change. This erases the attention-score ``/sqrt(d)`` scale -- an - fp32 div that otherwise stays between the QK^T bmm and softmax -- so the - attention chain is int8 through softmax. - - Runs before ``FoldAndAnnotateQParamsPass`` while the softmax-input quantize is - still an explicit ``quantized_decomposed.quantize_per_tensor`` node. + """Drop a constant elementwise scale (``x / c`` or ``x * c``) that the + quantizer already absorbed into the surrounding quantization scale. + + Runs after ``FoldAndAnnotateQParamsPass``. An attention-score ``/sqrt(d)`` + survives that pass as an unannotated fp32 op left inside a + ``dequantize_per_tensor(S_in) -> div/mul(c) -> quantize_per_tensor(S_out)`` + sandwich: the observer that set ``S_out`` saw the post-scale range, so + ``S_out == S_in / c`` (div) or ``S_in * c`` (mul) and the two zero-points + match. Under that relation the sandwich is the identity on the int8 values + (``quantize(dequantize(q, S_in) / c, S_out) == q``), so it is removed and the + producer's int8 wired straight into the consumer. + + The pass rewrites no quantization parameters -- it only recognises that the + constant is already in ``S_out`` and deletes the redundant fp32 round-trip. + Because no scale changes, a downstream SharedQspec consumer (view/reshape/ + pool) that shares ``S_out`` is never disturbed. If the ``S_out == S_in/c`` + (or ``*c``) and matching-zero-point relation does not hold -- per-channel + qparams, a non-scalar constant, or a scale the quantizer did not absorb -- + the fold is skipped and the fp32 op is left in place. + + A constant add/sub cannot fold this way: an additive shift makes the affine + observer change the scale (not just the zero-point), so no such identity + holds -- a separate, lossy transform, and no current model needs it. """ _passes_required_after: Set[Type[ExportPass]] = set() @@ -48,34 +61,70 @@ 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 not in (_DIV, _MUL): - continue - - scaled, const = node.args[0], node.args[1] - if not (isinstance(scaled, Node) and isinstance(const, Node)): - continue - if not is_param_node(ep, const): - continue - const_t = get_param_tensor(ep, const) - if const_t is None or const_t.numel() != 1: - continue - c = float(const_t.reshape(-1)[0]) - if c == 0.0: + match = self._absorbed_scale_sandwich(node, ep) + if match is None: continue - - users = list(node.users) - if len(users) != 1 or users[0].target != _QUANTIZE: - continue - - quantize = users[0] - scale = quantize.args[1] - new_scale = scale * c if node.target is _DIV else scale / c - quantize.update_arg(1, new_scale) - quantize.replace_input_with(node, scaled) + dequantize, quantize, producer = match + quantize.replace_all_uses_with(producer) + graph_module.graph.erase_node(quantize) graph_module.graph.erase_node(node) + graph_module.graph.erase_node(dequantize) modified = True if modified: graph_module.graph.eliminate_dead_code() graph_module.recompile() return PassResult(graph_module, modified) + + def _absorbed_scale_sandwich( + self, node: Node, ep: ExportedProgram + ) -> Optional[Tuple[Node, Node, Node]]: + """Return ``(dequantize, quantize, producer)`` when ``node`` is a + constant div/mul the quantizer already folded into the surrounding + quantize scale, so the sandwich can be deleted bit-exactly; else None.""" + if node.op != "call_function" or node.target not in (_DIV, _MUL): + return None + scaled, const = node.args[0], node.args[1] + if not (isinstance(scaled, Node) and isinstance(const, Node)): + return None + # A dequantize -> op -> quantize sandwich, each side used only by the op. + if scaled.target is not _DEQUANTIZE or len(scaled.users) != 1: + return None + # A dequantize's input is always a Node. + producer = cast(Node, scaled.args[0]) + users = list(node.users) + if len(users) != 1 or users[0].target is not _QUANTIZE: + return None + quantize = users[0] + + c = self._scalar_constant(const, ep) + if c is None: + return None + + # (de)quantize args: (input, scale, zero_point, qmin, qmax, dtype). + s_in, zp_in = scaled.args[1], scaled.args[2] + s_out, zp_out = quantize.args[1], quantize.args[2] + if not (isinstance(s_in, float) and isinstance(s_out, float)): + return None + expected = s_in / c if node.target is _DIV else s_in * c + # Bit-exact only if the quantizer folded c into s_out, the zero-points + # match, and the quantize's clamp range/dtype (args 3..5) match the + # dequantize's -- so the removed requantize is a no-op on the producer int8. + if ( + zp_in != zp_out + or scaled.args[3:6] != quantize.args[3:6] + or not math.isclose(s_out, expected, rel_tol=1e-6) + ): + return None + return scaled, quantize, producer + + def _scalar_constant(self, const: Node, ep: ExportedProgram) -> Optional[float]: + if not is_param_node(ep, const): + return None + const_t = get_param_tensor(ep, const) + if const_t is None or const_t.numel() != 1: + return None + c = float(const_t.reshape(-1)[0]) + if c == 0.0 or not math.isfinite(c): + return None + return c diff --git a/backends/cortex_m/test/ops/test_fold_scale_into_quantize.py b/backends/cortex_m/test/ops/test_fold_scale_into_quantize.py new file mode 100644 index 00000000000..d40924b8ea2 --- /dev/null +++ b/backends/cortex_m/test/ops/test_fold_scale_into_quantize.py @@ -0,0 +1,215 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Pass-level tests for ``FoldScaleIntoQuantizePass``. + +These run the pass on the post-``FoldAndAnnotateQParamsPass`` graph (its real +input state) and check the graph rewrite and its guards directly, without the +full lowering that ``test_attention_scale.py`` exercises end to end. +""" + +import copy +import math + +import torch +from executorch.backends.arm._passes import FoldAndAnnotateQParamsPass +from executorch.backends.arm._passes.fold_qdq_with_annotated_qparams_pass import ( + get_input_qparams, + get_output_qparams, +) +from executorch.backends.cortex_m.passes.cortex_m_pass_manager import CortexMPassManager +from executorch.backends.cortex_m.passes.fold_scale_into_quantize_pass import ( + FoldScaleIntoQuantizePass, +) +from executorch.backends.cortex_m.quantizer.quantizer import CortexMQuantizer +from executorch.backends.transforms.remove_getitem_op import RemoveGetItemPass +from executorch.exir import EdgeCompileConfig, to_edge +from executorch.exir.dialects._ops import ops as exir_ops +from torch.export import export +from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e + +_QD = torch.ops.quantized_decomposed +_QUANTIZE = exir_ops.edge.quantized_decomposed.quantize_per_tensor.default +_EDGE_CFG = EdgeCompileConfig( + _check_ir_validity=False, + _core_aten_ops_exception_list=[torch.ops.aten.max_pool2d.default], +) + + +class _AttnDiv(torch.nn.Module): + def forward(self, q, k): + return torch.softmax( + torch.bmm(q, k.transpose(-2, -1)) / math.sqrt(q.shape[-1]), dim=-1 + ) + + +class _AttnMul(torch.nn.Module): + def forward(self, q, k): + return torch.softmax( + torch.bmm(q, k.transpose(-2, -1)) * (1.0 / math.sqrt(q.shape[-1])), dim=-1 + ) + + +class _PoolScale(torch.nn.Module): + # A constant scale feeding a SharedQspec pool -- the passthrough case where + # folding must not disturb the shared scale. + def forward(self, x): + return torch.nn.functional.max_pool2d(x / 4.0, 2) + + +class _ScaleFirstMul(torch.nn.Module): + # Constant as the FIRST operand: mul(const, x). The constant lifts to a + # placeholder (empty args), so the pass must not read its args[0] before + # confirming the operand is a dequantize. + def __init__(self): + super().__init__() + self.register_buffer("scale", torch.tensor(0.125)) + + def forward(self, q, k): + return torch.softmax(self.scale * torch.bmm(q, k.transpose(-2, -1)), dim=-1) + + +_ATTN_INPUTS = (torch.rand(1, 8, 16), torch.rand(1, 8, 16)) +_POOL_INPUTS = (torch.rand(1, 1, 8, 8),) + + +def _to_edge(model, inputs): + gm = export(model.eval(), inputs, strict=True).module() + gm = prepare_pt2e(gm, CortexMQuantizer()) + gm(*inputs) + gm = convert_pt2e(gm) + return to_edge(export(gm, inputs, strict=True), compile_config=_EDGE_CFG) + + +def _run(edge, extra_passes): + passes = [RemoveGetItemPass, FoldAndAnnotateQParamsPass, *extra_passes] + return CortexMPassManager( + copy.deepcopy(edge).exported_program(), passes=passes + ).transform() + + +def _count(gm, needle): + return sum( + 1 for n in gm.graph.nodes if n.op == "call_function" and needle in str(n.target) + ) + + +def _find(gm, needle): + return next(n for n in gm.graph.nodes if needle in str(n.target)) + + +def test_div_folds_and_removes_the_sandwich(): + edge = _to_edge(_AttnDiv(), _ATTN_INPUTS) + annotated = _run(edge, []) + folded = _run(edge, [FoldScaleIntoQuantizePass]) + + assert _count(annotated.graph_module, "aten.div.Tensor") == 1 + assert _count(folded.graph_module, "aten.div.Tensor") == 0 + # The dequantize/quantize that wrapped the div go with it (one of each). + assert _count(folded.graph_module, "dequantize_per_tensor") == ( + _count(annotated.graph_module, "dequantize_per_tensor") - 1 + ) + # bmm now feeds softmax directly. + assert "bmm" in str(_find(folded.graph_module, "_softmax").args[0].target) + + +def test_mul_folds(): + edge = _to_edge(_AttnMul(), _ATTN_INPUTS) + assert _count(_run(edge, []).graph_module, "aten.mul.Tensor") == 1 + assert ( + _count(_run(edge, [FoldScaleIntoQuantizePass]).graph_module, "aten.mul.Tensor") + == 0 + ) + + +def test_fold_is_bit_exact(): + # Removing the sandwich is bit-exact iff quantize(dequantize(q) / c) == q for + # every int8 at the calibrated qparams; verify against the real kernels. + edge = _to_edge(_AttnDiv(), _ATTN_INPUTS) + div = _find(_run(edge, []).graph_module, "aten.div.Tensor") + dq, quant = div.args[0], next(iter(div.users)) + s_in, zp_in = dq.args[1], dq.args[2] + s_out, zp_out = quant.args[1], quant.args[2] + + c = math.sqrt(_ATTN_INPUTS[0].shape[-1]) + v = torch.arange(-128, 128, dtype=torch.int8) + deq = _QD.dequantize_per_tensor.default(v, s_in, zp_in, -128, 127, torch.int8) + requant = _QD.quantize_per_tensor.default( + deq / c, s_out, zp_out, -128, 127, torch.int8 + ) + assert torch.equal(requant, v) + + +def test_sharedqspec_consumer_scale_is_untouched(): + # Adrian's case: a constant scale feeding a SharedQspec pool. The fold + # rewrites no scale, so the pool's shared in/out qparams are identical before + # and after -- the shared cluster can never be corrupted -- and the divide + # still folds. + edge = _to_edge(_PoolScale(), _POOL_INPUTS) + annotated = _run(edge, []) + folded = _run(edge, [FoldScaleIntoQuantizePass]) + + assert _count(folded.graph_module, "aten.div.Tensor") == 0 + pool_before = _find(annotated.graph_module, "max_pool2d") + pool_after = _find(folded.graph_module, "max_pool2d") + # Pin the premise: max_pool2d is a SharedQspec op (input scale == output + # scale). Without this the test would silently stop guarding Adrian's case. + assert ( + get_input_qparams(pool_before)[0].scale + == get_output_qparams(pool_before)[0].scale + ) + assert ( + get_input_qparams(pool_before)[0].scale + == get_input_qparams(pool_after)[0].scale + ) + assert ( + get_output_qparams(pool_before)[0].scale + == get_output_qparams(pool_after)[0].scale + ) + + +def _perturb_softmax_quantize(program, arg_index, fn): + for node in program.graph_module.graph.nodes: + if ( + node.op == "call_function" + and node.target is _QUANTIZE + and any("softmax" in str(u.target) for u in node.users) + ): + args = list(node.args) + args[arg_index] = fn(args[arg_index]) + node.args = tuple(args) + + +def test_skips_when_scale_not_absorbed(): + program = _run(_to_edge(_AttnDiv(), _ATTN_INPUTS), []) + _perturb_softmax_quantize(program, 1, lambda s: s * 1.5) + result = FoldScaleIntoQuantizePass(exported_program=program).call( + program.graph_module + ) + assert not result.modified + assert _count(program.graph_module, "aten.div.Tensor") == 1 + + +def test_skips_on_zero_point_mismatch(): + program = _run(_to_edge(_AttnDiv(), _ATTN_INPUTS), []) + _perturb_softmax_quantize(program, 2, lambda z: z + 1) + result = FoldScaleIntoQuantizePass(exported_program=program).call( + program.graph_module + ) + assert not result.modified + assert _count(program.graph_module, "aten.div.Tensor") == 1 + + +def test_constant_first_operand_does_not_crash(): + # `const * tensor` puts the constant (a placeholder with empty args) at + # args[0]; the pass must not read its args[0] before confirming it is a + # dequantize. It should skip, not raise (regression). + program = _run(_to_edge(_ScaleFirstMul(), _ATTN_INPUTS), []) + result = FoldScaleIntoQuantizePass(exported_program=program).call( + program.graph_module + ) + assert not result.modified + assert _count(program.graph_module, "aten.mul.Tensor") == 1 From d90e409e47cf7396b9fe0fbcbc8bab9b313fb6da Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Tue, 21 Jul 2026 10:36:39 -0700 Subject: [PATCH 4/5] Cortex-M: add a matmul multi-head attention test for the scale-fold The div/mul cases (and test_attention_scale.py) use a bare rank-3 torch.bmm, whose output the quantizer quantizes directly, so they do not exercise the matmul->bmm dependency the fold actually relies on. Add a rank-4 multi-head q @ k^T attention like SAM's mask decoder: MatmulToBmmPass rewrites the matmul to a quantizable bmm at annotation, so the /sqrt(head_dim) scale lands in the requantize sandwich the fold removes. The test fails if matmul->bmm regresses (the scale would stay fp32 and never fold). Authored with Claude Code. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../test/ops/test_fold_scale_into_quantize.py | 42 +++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/backends/cortex_m/test/ops/test_fold_scale_into_quantize.py b/backends/cortex_m/test/ops/test_fold_scale_into_quantize.py index d40924b8ea2..2a36ef035e3 100644 --- a/backends/cortex_m/test/ops/test_fold_scale_into_quantize.py +++ b/backends/cortex_m/test/ops/test_fold_scale_into_quantize.py @@ -72,8 +72,35 @@ def forward(self, q, k): return torch.softmax(self.scale * torch.bmm(q, k.transpose(-2, -1)), dim=-1) +class _MultiHeadMatmulAttn(torch.nn.Module): + # Rank-4 multi-head attention with q @ k^T scaled by 1/sqrt(head_dim), like + # SAM's mask-decoder attention. MatmulToBmmPass rewrites the matmul to a + # quantizable bmm at annotation, so the scale lands in the requantize sandwich + # the fold removes -- the realistic path a bare rank-3 torch.bmm skips. + def __init__(self, dim=32, heads=4): + super().__init__() + self.heads = heads + self.head_dim = dim // heads + self.q_proj = torch.nn.Linear(dim, dim) + self.k_proj = torch.nn.Linear(dim, dim) + self.v_proj = torch.nn.Linear(dim, dim) + + def _heads(self, x): + b, n, c = x.shape + return x.reshape(b, n, self.heads, self.head_dim).transpose(1, 2) + + def forward(self, x): + q = self._heads(self.q_proj(x)) + k = self._heads(self.k_proj(x)) + v = self._heads(self.v_proj(x)) + attn = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim) + attn = torch.softmax(attn, dim=-1) + return (attn @ v).transpose(1, 2).reshape(x.shape) + + _ATTN_INPUTS = (torch.rand(1, 8, 16), torch.rand(1, 8, 16)) _POOL_INPUTS = (torch.rand(1, 1, 8, 8),) +_MHA_INPUT = (torch.rand(1, 8, 32),) def _to_edge(model, inputs): @@ -125,6 +152,21 @@ def test_mul_folds(): ) +def test_matmul_multihead_attention_scale_folds(): + # SAM-faithful path: the rank-4 q @ k^T matmul becomes a quantizable bmm via + # MatmulToBmmPass at annotation, so the /sqrt(head_dim) scale lands in a + # dequantize -> div -> quantize sandwich the fold removes. A bare rank-3 + # torch.bmm (the tests above) has its output quantized directly, so it does + # not exercise this matmul->bmm dependency -- if that pass regresses, the + # scale would stay fp32 and this test fails. + edge = _to_edge(_MultiHeadMatmulAttn(), _MHA_INPUT) + assert _count(_run(edge, []).graph_module, "aten.div.Tensor") == 1 + assert ( + _count(_run(edge, [FoldScaleIntoQuantizePass]).graph_module, "aten.div.Tensor") + == 0 + ) + + def test_fold_is_bit_exact(): # Removing the sandwich is bit-exact iff quantize(dequantize(q) / c) == q for # every int8 at the calibrated qparams; verify against the real kernels. From 7bde4b7cdf9f89e615425c41bc9534ca411cdcd6 Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Tue, 21 Jul 2026 16:36:50 -0700 Subject: [PATCH 5/5] Cortex-M: apply scale-fold review-panel fixes Add coverage for the two skip guards that had none -- a dtype/clamp-range mismatch (args 3..5, the int16-activation path) and a non-param runtime divisor; both are load-bearing (without them the pass would wrongly fold, or crash reading a param tensor for a non-param node). Remove the dead ArmPass _passes_required_after boilerplate (CortexMPassManager never reads it) and its now-unused typing imports. Correct the add/sub docstring: an additive shift moves the affine qparams by a generally non-integer amount, so no integer-preserving requantize identity exists. Reword two test comments that named a reviewer. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../passes/fold_scale_into_quantize_pass.py | 11 +++-- .../test/ops/test_fold_scale_into_quantize.py | 42 ++++++++++++++++++- 2 files changed, 45 insertions(+), 8 deletions(-) diff --git a/backends/cortex_m/passes/fold_scale_into_quantize_pass.py b/backends/cortex_m/passes/fold_scale_into_quantize_pass.py index 055b2ac9e26..3f3b742e9b7 100644 --- a/backends/cortex_m/passes/fold_scale_into_quantize_pass.py +++ b/backends/cortex_m/passes/fold_scale_into_quantize_pass.py @@ -5,7 +5,7 @@ # LICENSE file in the root directory of this source tree. import math -from typing import cast, Optional, Set, Tuple, Type +from typing import cast, Optional, Tuple from executorch.backends.arm._passes.arm_pass_utils import ( get_param_tensor, @@ -43,13 +43,12 @@ class FoldScaleIntoQuantizePass(ExportPass): qparams, a non-scalar constant, or a scale the quantizer did not absorb -- the fold is skipped and the fp32 op is left in place. - A constant add/sub cannot fold this way: an additive shift makes the affine - observer change the scale (not just the zero-point), so no such identity - holds -- a separate, lossy transform, and no current model needs it. + A constant add/sub cannot fold this way: an additive shift moves the affine + quantization parameters by a generally non-integer amount, so no + integer-preserving requantize identity exists -- a separate, lossy transform, + and no current model needs it. """ - _passes_required_after: Set[Type[ExportPass]] = set() - def __init__(self, exported_program: Optional[ExportedProgram] = None) -> None: super().__init__() self.exported_program = exported_program diff --git a/backends/cortex_m/test/ops/test_fold_scale_into_quantize.py b/backends/cortex_m/test/ops/test_fold_scale_into_quantize.py index 2a36ef035e3..fb16bf73a23 100644 --- a/backends/cortex_m/test/ops/test_fold_scale_into_quantize.py +++ b/backends/cortex_m/test/ops/test_fold_scale_into_quantize.py @@ -98,9 +98,21 @@ def forward(self, x): return (attn @ v).transpose(1, 2).reshape(x.shape) +class _RuntimeDiv(torch.nn.Module): + # Divisor is a runtime graph input, not a param/buffer, so is_param_node is + # False -- the scale is non-constant and must not fold. + def forward(self, q, k, d): + return torch.softmax(torch.bmm(q, k.transpose(-2, -1)) / d, dim=-1) + + _ATTN_INPUTS = (torch.rand(1, 8, 16), torch.rand(1, 8, 16)) _POOL_INPUTS = (torch.rand(1, 1, 8, 8),) _MHA_INPUT = (torch.rand(1, 8, 32),) +_RUNTIME_INPUTS = ( + torch.rand(1, 8, 16), + torch.rand(1, 8, 16), + torch.rand(1, 8, 8) + 0.5, +) def _to_edge(model, inputs): @@ -186,7 +198,7 @@ def test_fold_is_bit_exact(): def test_sharedqspec_consumer_scale_is_untouched(): - # Adrian's case: a constant scale feeding a SharedQspec pool. The fold + # A constant scale feeding a SharedQspec pool (RFC #19299). The fold # rewrites no scale, so the pool's shared in/out qparams are identical before # and after -- the shared cluster can never be corrupted -- and the divide # still folds. @@ -198,7 +210,7 @@ def test_sharedqspec_consumer_scale_is_untouched(): pool_before = _find(annotated.graph_module, "max_pool2d") pool_after = _find(folded.graph_module, "max_pool2d") # Pin the premise: max_pool2d is a SharedQspec op (input scale == output - # scale). Without this the test would silently stop guarding Adrian's case. + # scale). Without this the test would silently stop guarding that case. assert ( get_input_qparams(pool_before)[0].scale == get_output_qparams(pool_before)[0].scale @@ -255,3 +267,29 @@ def test_constant_first_operand_does_not_crash(): ) assert not result.modified assert _count(program.graph_module, "aten.mul.Tensor") == 1 + + +def test_skips_on_dtype_or_range_mismatch(): + # A quantize whose dtype/clamp range (args 3..5) differs from the dequantize + # makes the removed requantize a real conversion, not a no-op (e.g. the int16 + # activation path). The scale/zero-point relation still holds, so only the + # args[3:6] guard prevents the fold. + program = _run(_to_edge(_AttnDiv(), _ATTN_INPUTS), []) + _perturb_softmax_quantize(program, 5, lambda _: torch.int16) + result = FoldScaleIntoQuantizePass(exported_program=program).call( + program.graph_module + ) + assert not result.modified + assert _count(program.graph_module, "aten.div.Tensor") == 1 + + +def test_skips_non_param_divisor(): + # The divisor is a runtime graph input, not a param/buffer, so is_param_node + # is False: the scale is genuinely non-constant and must not fold -- and the + # pass must not crash reading a param tensor for a non-param node. + program = _run(_to_edge(_RuntimeDiv(), _RUNTIME_INPUTS), []) + result = FoldScaleIntoQuantizePass(exported_program=program).call( + program.graph_module + ) + assert not result.modified + assert _count(program.graph_module, "aten.div.Tensor") == 1