PERF: detect scaled-dot-product-attention pattern and lower to mx.fast.scaled_dot_product_attention

Part of #2085.

Current implementation

PyTensor expresses attention as softmax((Q @ K.transpose) * scale) @ V. The MLX backend dispatches each piece (matmul, transpose, mul, softmax, matmul) and lets mx.compile fuse what it can. MLX has a hand-written mx.fast.scaled_dot_product_attention (flash-attention-style) that is faster forward and provides a fused backward kernel — important for transformer training/inference loops.

Reproducer + metrics

import time, statistics, mlx.core as mx, numpy as np

def sync(): mx.synchronize()

def bench(fn, *args, n=30, runs=5):
    sync()
    for _ in range(8): mx.eval(fn(*args))
    sync()
    out = []
    for _ in range(runs):
        sync(); t0 = time.perf_counter()
        for _ in range(n): mx.eval(fn(*args))
        sync()
        out.append((time.perf_counter() - t0) / n)
    return statistics.median(out) * 1e6

B, H, S, hd = 8, 8, 128, 64
rng = np.random.default_rng(0)
Q = mx.array(rng.standard_normal((B, H, S, hd)).astype(np.float32))
K = mx.array(rng.standard_normal((B, H, S, hd)).astype(np.float32))
V = mx.array(rng.standard_normal((B, H, S, hd)).astype(np.float32))
mx.eval(Q, K, V); sync()
scale = 1.0 / np.sqrt(hd).astype(np.float32)

@mx.compile
def manual(q, k, v):
    return mx.softmax((q @ k.transpose(0, 1, 3, 2)) * scale, axis=-1) @ v

@mx.compile
def fast(q, k, v):
    return mx.fast.scaled_dot_product_attention(q, k, v, scale=scale)

print(f"manual: {bench(manual, Q, K, V):.1f} us")
print(f"fast:   {bench(fast,   Q, K, V):.1f} us")

Implementation	Median (us)
Manual (matmul + softmax + matmul)	527
mx.fast.scaled_dot_product_attention	401

→ 1.3× speedup forward.

End-to-end on a transformer block (B=8, S=128, D=512, H=8) with both LayerNorm and SDPA fused (post #2088 + this issue):

	Median (us)
Raw MLX manual block	669
Raw MLX with fast.layer_norm + SDPA	567
PyTensor MLX (current, trust_input)	774

→ Closing the gap to ~570 µs would match raw MLX with all fast ops.

Proposed change

Same shape as #2088. Define an MLXScaledDotProductAttention(scale) Op and a node rewriter.

Files:

• pytensor/link/mlx/ops.py — add MLXScaledDotProductAttention
• pytensor/link/mlx/dispatch/fused.py — add dispatcher
• pytensor/link/mlx/rewriting/fused_attention.py — pattern matcher

Sketch

# pytensor/link/mlx/ops.py
class MLXScaledDotProductAttention(Op):
    __props__ = ("scale",)
    def __init__(self, scale): self.scale = scale
    def make_node(self, q, k, v):
        out = q.type()
        return Apply(self, [q, k, v], [out])
    def perform(self, node, inputs, outputs):
        q, k, v = inputs
        scores = (q @ np.swapaxes(k, -1, -2)) * self.scale
        scores -= scores.max(axis=-1, keepdims=True)
        attn = np.exp(scores) / np.exp(scores).sum(axis=-1, keepdims=True)
        outputs[0][0] = attn @ v


# pytensor/link/mlx/dispatch/fused.py
@mlx_funcify.register(MLXScaledDotProductAttention)
def mlx_funcify_MLXSDPA(op, **kwargs):
    scale = op.scale
    def sdpa(q, k, v):
        return mx.fast.scaled_dot_product_attention(q, k, v, scale=scale)
    return sdpa


# pytensor/link/mlx/rewriting/fused_attention.py
@node_rewriter(tracks=[Dot, BatchedDot])  # tail matmul of (attn @ V)
def fuse_scaled_dot_product_attention(fgraph, node):
    """Detect softmax((Q @ K.T) * scale) @ V → MLXScaledDotProductAttention."""
    matched = match_sdpa_subgraph(node)
    if matched is None:
        return None
    q, k, v, scale = matched
    return [MLXScaledDotProductAttention(scale=scale)(q, k, v)]

Important guards in the matcher

mx.fast.scaled_dot_product_attention does not currently support attention masks or causal masks via the simple call shown above (it does via separate kwargs). The pattern matcher must:

• Reject the rewrite if the softmax input has additional terms (e.g. + mask) added between Q @ K.T * scale and softmax.
• Reject if dropout / scaling other than the constant pre-softmax scale is present.
• Confirm Q, K, V have matching final two dims and same leading batch dims.

A conservative initial version that only matches the bare softmax(Q @ K.T * scale) @ V pattern is fine; mask/causal support can come later by extending MLXScaledDotProductAttention with optional inputs.

Numerics

With fp32 inputs the maximum absolute deviation between manual and mx.fast.scaled_dot_product_attention measured during this analysis was 6.56e-7.

Acceptance criteria

• MLXScaledDotProductAttention Op defined and dispatched.
• Rewriter fires only on the safe pattern (no mask, no extra additions to scores).
• Numerics within 1e-5 of the manual graph.
• Benchmark in tests/link/mlx/ shows ≥ 1.25× speedup at the shape above.

Start by reading the existing MLX fused-operation patterns in pytensor/link/mlx/ops.py, pytensor/link/mlx/dispatch/fused.py, and pytensor/link/mlx/rewriting/fused_attention.py. Implement and test the conservative bare softmax(Q @ K.T * scale) @ V match, confirming that masked or otherwise modified score expressions are rejected. Run the MLX tests under tests/link/mlx/ and verify numerical agreement within 1e-5 and the stated speedup benchmark.