[Performance] CPU MultiHeadAttention with attention_bias is ~4x slower than the unfused ops #33873

"""ONNX Runtime CPU: `MultiHeadAttention` with an `attention_bias` runs slower than the MatMul/Softmax ops it replaces.

One decode step at Qwen3-0.6B's sizes: batch 8, one new token over 129 keys, 16 heads, head size 128, float32.
"""

import statistics
import time

import numpy as np
import onnxruntime as ort
from onnx import TensorProto, helper, numpy_helper

BATCH, NEW, TOTAL, HEADS, HEAD_SIZE = 8, 1, 129, 16, 128
F32 = TensorProto.FLOAT


def inputs():
    return [
        helper.make_tensor_value_info("query", F32, [BATCH, NEW, HEADS * HEAD_SIZE]),
        helper.make_tensor_value_info("key", F32, [BATCH, TOTAL, HEADS * HEAD_SIZE]),
        helper.make_tensor_value_info("value", F32, [BATCH, TOTAL, HEADS * HEAD_SIZE]),
        helper.make_tensor_value_info("bias", F32, [BATCH, 1, NEW, TOTAL]),
    ]


def mha() -> bytes:
    node = helper.make_node("MultiHeadAttention", ["query", "key", "value", "", "", "bias"], ["output"], domain="com.microsoft", num_heads=HEADS)
    return _model([node], [])


def unfused() -> bytes:
    # The same attention as plain ops: split heads, scaled QK^T + bias, Softmax, times V, merge heads.
    heads = numpy_helper.from_array(np.array([BATCH, -1, HEADS, HEAD_SIZE], np.int64), "heads")
    merge = numpy_helper.from_array(np.array([BATCH, -1, HEADS * HEAD_SIZE], np.int64), "merge")
    scale = numpy_helper.from_array(np.array(HEAD_SIZE**-0.5, np.float32), "scale")
    nodes = [
        helper.make_node("Reshape", ["query", "heads"], ["q4"]),
        helper.make_node("Transpose", ["q4"], ["q"], perm=[0, 2, 1, 3]),
        helper.make_node("Reshape", ["key", "heads"], ["k4"]),
        helper.make_node("Transpose", ["k4"], ["kt"], perm=[0, 2, 3, 1]),
        helper.make_node("Reshape", ["value", "heads"], ["v4"]),
        helper.make_node("Transpose", ["v4"], ["v"], perm=[0, 2, 1, 3]),
        helper.make_node("MatMul", ["q", "kt"], ["scores"]),
        helper.make_node("Mul", ["scores", "scale"], ["scaled"]),
        helper.make_node("Add", ["scaled", "bias"], ["biased"]),
        helper.make_node("Softmax", ["biased"], ["weights"], axis=-1),
        helper.make_node("MatMul", ["weights", "v"], ["o4"]),
        helper.make_node("Transpose", ["o4"], ["o"], perm=[0, 2, 1, 3]),
        helper.make_node("Reshape", ["o", "merge"], ["output"]),
    ]
    return _model(nodes, [heads, merge, scale])


def _model(nodes, inits) -> bytes:
    graph = helper.make_graph(nodes, "attention", inputs(), [helper.make_tensor_value_info("output", F32, None)], inits)
    opsets = [helper.make_opsetid("", 20), helper.make_opsetid("com.microsoft", 1)]
    return helper.make_model(graph, opset_imports=opsets, ir_version=10).SerializeToString()


def run(model: bytes, feeds: dict):
    options = ort.SessionOptions()
    options.intra_op_num_threads = 16
    session = ort.InferenceSession(model, options, providers=["CPUExecutionProvider"])
    for _ in range(20):
        session.run(None, feeds)
    times = []
    for _ in range(200):
        start = time.perf_counter()
        output = session.run(None, feeds)[0]
        times.append(time.perf_counter() - start)
    return statistics.median(times) * 1e6, output


rng = np.random.default_rng(0)
feeds = {
    "query": rng.standard_normal((BATCH, NEW, HEADS * HEAD_SIZE)).astype(np.float32),
    "key": rng.standard_normal((BATCH, TOTAL, HEADS * HEAD_SIZE)).astype(np.float32),
    "value": rng.standard_normal((BATCH, TOTAL, HEADS * HEAD_SIZE)).astype(np.float32),
    "bias": np.zeros((BATCH, 1, NEW, TOTAL), np.float32),
}
print("onnxruntime", ort.__version__)
fused_us, fused_out = run(mha(), feeds)
plain_us, plain_out = run(unfused(), feeds)
print(f"  MultiHeadAttention with attention_bias: {fused_us:7.1f} us")
print(f"  the same as MatMul/Softmax ops        : {plain_us:7.1f} us   (max diff {np.abs(fused_out - plain_out).max():.1e})")