NA/M5 addmm / matmul on MPS ~10–20% slower than PyTorch for 1280×1280 BF16 (hurts Nano Chat training)

[Anemll]

Summary
On an Mac M5, mlx.addmm and A @ B are consistently ~1.1–1.2× slower than PyTorch’s addmm / matmul on MPS for 1280×1280 BF16 matrices. This GEMM shape is representative of Nano Chat–style transformer training, so the gap directly reduces end‑to‑end training throughput vs PyTorch.

It sped up after NA support but its still behind since official Neural Accelerator Support

Repro
Gist with script and logs: https://gist.github.com/Anemll/5420800c3d29c7fae18a2b9b10907b14
Key details:
Shape: (1280, 1280)
Dtype: bfloat16
Device: PyTorch mps, MLX Metal
Warmup: 30, Iterations: 1000
Ops tested:
addmm: beta * C + alpha * (A @ B)
matmul: A @ B
Sync:
PyTorch: torch.mps.synchronize()
MLX: mx.eval(result)
Results (MLX 0.31.0)
PyTorch MPS:
addmm: ~0.596 ms
matmul: ~0.549 ms
MLX:
addmm: ~0.651 ms
matmul: ~0.646 ms
Ratios:
addmm: 1.09× slower (MLX / PyTorch)
matmul: 1.18× slower
Motivation (Nano Chat)
Small LLM (“Nano Chat”) training is GEMM‑bound at these sizes; this 10–20% gap in BF16 GEMM on M‑series Macs translates almost directly into slower training steps vs PyTorch MPS. Closing this gap would make MLX more competitive as the default backend for local Nano Chat training.

[jagrit06]

I will get to addressing this issue
In the meanwhile, just so I also have an idea of the real world use case associated, do the tests you're doing with NanoChat also require the added C matrix to be full 1280x1280 ?

[Anemll]

Yes, the real-world use case requires full 1280×1280 C matrices (not bias vectors).

NanoChat model uses model_dim = 1280 (20 layers × 64 aspect ratio). The Muon optimizer's Newton-Schulz iteration operates directly on weight matrices:

Weight Matrix	Shape	NS iteration operates on
Attention Q/K/V/O	(1280, 1280)	1280×1280
MLP up-projection	(1280, 5120)	Transposed → 1280×1280
MLP down-projection	(5120, 1280)	Transposed → 1280×1280

Each Newton-Schulz step (5 steps per optimizer update) does:

A = X @ X.T                                        # (1280,1280) @ (1280,1280)
B = mx.addmm(b * A, A, A, beta=1.0, alpha=c)       # full matrix C, not a bias
X = mx.addmm(a * X, B, X, beta=1.0, alpha=1.0)     # full matrix C, not a bias

So the C argument to addmm is always a full 1280×1280 matrix (result of a prior computation), not a broadcast bias vector.

Update on performance: I just re-ran the benchmarks on MLX 0.30.0 and the gap has narrowed considerably from the ~3x I originally reported:

Shape	PyTorch MPS (ms)	MLX (ms)	Ratio
1280×1280	5.11	6.19	1.21x
1280×320	0.82	1.02	1.24x
5120×1280	15.19	17.02	1.12x

Average: ~1.19xThe addmm fusion itself is actually faster than PyTorch now — the remaining gap seems to come from plain matmul (X @ X.T).