aten.bmm, aten.mul, aten._softmax doesn't get quantized/replaced with Cortex-M kernels in Transformer model

🐛 Describe the bug

Description:
When quantizing a transformer model with CortexMQuantizer and lowering to the Cortex-M backend, aten.bmm/aten.mul/aten._softmax are not quantized(These are mentioned in CMSIS-NN Supported Operators list).

import argparse
import logging

import torch
import torch.nn as nn
from torch.export import export,default_decompositions

from executorch.exir import EdgeCompileConfig, ExecutorchBackendConfig, to_edge
from executorch.backends.cortex_m.quantizer.quantizer import CortexMQuantizer
from executorch.backends.cortex_m.passes.cortex_m_pass_manager import CortexMPassManager

# Post-training static INT8 quantization (PT2E flow).
from torchao.quantization.pt2e.quantize_pt2e import prepare_pt2e, convert_pt2e

logging.basicConfig(level=logging.INFO, format="%(message)s")
logger = logging.getLogger("simple_transformer")

_SDPA_DECOMP_TABLE = {
    op: decomp
    for op, decomp in default_decompositions().items()
    if "scaled_dot_product" or "_safe_softmax" in str(op) in str(op)
}

def make_calibration_inputs(example_inputs, num_batches: int = 8):
    """Yield representative inputs for static PTQ calibration.

    TODO: Replace the synthetic tensors with REAL representative inference
    inputs. Static quantization freezes activation scales from whatever is
    observed here, so synthetic data yields meaningless ranges and poor
    accuracy.
    """
    for _ in range(num_batches):
        yield tuple(torch.randn_like(t) for t in example_inputs)


def quantize_static_int8(model, example_inputs, calib_batches: int = 8):
    """Apply post-training static INT8 quantization via the PT2E flow."""
    model.eval()
    # prepared_graph = export(model, example_inputs).run_decompositions(_SDPA_DECOMP_TABLE)
    prepared_graph = export(model, example_inputs).module()
    from executorch.backends.arm._passes import ConstantFoldingPass
    prepared_graph = ConstantFoldingPass().call(prepared_graph).graph_module  # Pre-computes and removes the call_function op sitting between the linear nodes and their weights

    quantizer = CortexMQuantizer()
    prepared = prepare_pt2e(prepared_graph, quantizer)

    with torch.no_grad():
        for sample in make_calibration_inputs(example_inputs, calib_batches):
            prepared(*sample)

    return convert_pt2e(prepared)


def main(args):
    S, B, D = args.seq_len, 1, args.d_model

    # Minimal one-layer encoder-decoder transformer.
    model = nn.Transformer(
        d_model=D,
        nhead=args.nhead,
        num_encoder_layers=args.num_layers,
        num_decoder_layers=args.num_layers,
        dim_feedforward=args.dim_feedforward,
        dropout=0.0,
        batch_first=False,
    )
    model.eval()

    example_input = (torch.randn(S, B, D), torch.randn(S, B, D))

    print("Applying post-training static INT8 quantization")
    quantized_model = quantize_static_int8(model, example_input, args.calib_batches)
    quantized_exported_program = export(quantized_model, example_input)

    config = EdgeCompileConfig(
        preserve_ops=[torch.ops.aten.linear.default],
        _check_ir_validity=False,
    )
    edge_program_manager = to_edge(quantized_exported_program, compile_config=config)

    pass_manager = CortexMPassManager(edge_program_manager.exported_program())
    edge_program_manager._edge_programs["forward"] = pass_manager.transform()

    et_program = edge_program_manager.to_executorch(
        config=ExecutorchBackendConfig(extract_delegate_segments=False)
    )
    for op in et_program.executorch_program.execution_plan[0].operators:
        print(op.name)

    pte_path = "simple_transformer_int8_CortexM.pte"
    with open(pte_path, "wb") as f:
        f.write(et_program.buffer)
    print(f"Wrote {pte_path}")

    print(model)
    src = torch.rand((S, B, D))
    tgt = torch.rand((S, B, D))
    out = model(src, tgt)
    print("input", src.size())
    print("output", out.size())


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Export a minimal quantized encoder-decoder transformer to .pte for the Cortex-M backend"
    )
    parser.add_argument("--d_model", type=int, help="d_model", default=32)
    parser.add_argument("--nhead", type=int, help="nhead", default=2)
    parser.add_argument("--num_layers", type=int, help="number of encoder and decoder layers", default=1)
    parser.add_argument("--dim_feedforward", type=int, help="dim_feedforward", default=64)
    parser.add_argument("--seq_len", type=int, help="sequence length", default=8)
    parser.add_argument("--calib_batches", type=int, help="number of calibration batches for static PTQ", default=8)

    args = parser.parse_args()
    main(args)

Additional context:
Running decomposition before quantizing the model does quantize aten::_softmax and aten::bmm, aten::mul is the one that stays unquantized. (Separately, I also ran into the decomposed .pte not opening in Netron , unrelated to this issue, but flagging in case it's a known/expected
limitation for graphs with this many ops.)

Versions

Versions

PyTorch version: 2.13.0+cpu

OS: Ubuntu 24.04.4 LTS (x86_64)
Python version: 3.11.15 (main, Mar 11 2026, 17:20:07) [GCC 14.3.0] (64-bit runtime)

CPU:
Architecture: x86_64

Versions of relevant libraries:
[pip3] executorch==1.4.0+3dd7ccd
[pip3] flake8==6.1.0
[pip3] flake8-breakpoint==1.1.0
[pip3] flake8-bugbear==24.4.26
[pip3] flake8-comprehensions==3.14.0
[pip3] flake8-plugin-utils==1.3.3
[pip3] flake8-pyi==23.5.0
[pip3] mypy==1.14.1
[pip3] mypy_extensions==1.1.0
[pip3] numpy==2.4.4
[pip3] pytorch_tokenizers==1.4.1
[pip3] torch==2.13.0+cpu
[pip3] torchao==0.18.0
[pip3] torchaudio==2.11.0+cpu
[pip3] torchdata==0.11.0
[pip3] torchsr==1.0.4
[pip3] torchtune==0.0.0
[pip3] torchvision==0.28.0
[pip3] triton==3.7.1
[conda] executorch 1.4.0+3dd7ccd pypi_0 pypi
[conda] numpy 2.4.4 pypi_0 pypi
[conda] pytorch-tokenizers 1.4.1 pypi_0 pypi
[conda] torch 2.13.0+cpu pypi_0 pypi
[conda] torchao 0.18.0 pypi_0 pypi
[conda] torchaudio 2.11.0+cpu pypi_0 pypi
[conda] torchdata 0.11.0 pypi_0 pypi
[conda] torchfix 0.6.0 pypi_0 pypi
[conda] torchsr 1.0.4 pypi_0 pypi
[conda] torchtune 0.0.0 pypi_0 pypi
[conda] torchvision 0.28.0 pypi_0 pypi
[conda] triton 3.7.1 pypi_0 pypi

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