float32 matmul loses precision to float16 when the right operand has a broadcast batch dimension (also corrupts batch-1 SDPA fallback)

[SirsiMaster]

Summary

mx.matmul on float32 inputs silently loses precision to float16 when the right operand has a broadcast (stride-0) batch dimension. Multiplying by an exact 1.0 changes the values:

import mlx.core as mx
mx.random.seed(9)
v    = mx.random.normal((1, 1, 1, 512)).astype(mx.float32)
ones = mx.ones((1, 16, 1, 1)).astype(mx.float32)
out  = ones @ v            # K = 1: out must equal v exactly
print(mx.abs(out - v).max().item())   # 1.9e-03  (should be 0.0)

Every output element lands on the fp16 grid. Two confirmations:

# fp16 ulp at 1.0 is 2^-10
b = mx.full((1,1,1,512), 1 + 2**-10).astype(mx.float32)   # fp16-representable
print(mx.abs(ones @ b - (1 + 2**-10)).max().item())        # 0.0
b = mx.full((1,1,1,512), 1 + 2**-12).astype(mx.float32)   # NOT representable
print(mx.abs(ones @ b - (1 + 2**-12)).max().item())        # 2.441406e-04 == 2^-12 exactly

Shape dependence

shape	broadcast?	max err
[1,1] @ [1,512]	–	0.0
[1,1,1,1] @ [1,1,1,512]	no	0.0
[1,16,1,1] @ [1,1,1,512]	yes	1.9e-03
[1,16,1,1] @ [1,16,1,512] (materialized)	no	0.0
[2,16,1,1] @ [1,1,1,512]	yes	1.9e-03

Materializing the broadcast (mx.contiguous(mx.broadcast_to(...))) restores exactness, so the loss is specific to the stride-0 batched path.

Real-world impact: fast.scaled_dot_product_attention at batch 1

For head dims outside the fused kernels' support list, SDPA runs the unfused fallback, whose softmax @ values matmul hits this path via the GQA expand_dims broadcast. At L=1 the output must equal V exactly (softmax of a single score is 1); it does not:

mx.random.seed(9)
q = mx.random.normal((1, 16, 1, 512)).astype(mx.float32)
k = mx.random.normal((1, 1, 1, 512)).astype(mx.float32)
v = mx.random.normal((1, 1, 1, 512)).astype(mx.float32)
out = mx.fast.scaled_dot_product_attention(q, k, v, scale=1.0, mask="causal")
print(mx.abs(out - v).max().item())   # ~1.9e-03; 20/20 random seeds deviate

At B >= 2 the same call is exact (different matmul dispatch). Batch 1 is the single-user interactive case, so every batch-1 decode step through the fallback carries fp16-scale noise on f32 activations — we found it because two independent attention implementations disagreed only at batch 1, and the L=1 ground-truth case (out must equal V, no arithmetic involved) showed which side was wrong.

Environment

• MLX 0.32.0 (pip), Python 3.14
• Apple M5 Max, macOS 26 (Darwin 25.5.0)

Happy to test patches. Related report from the same project: #3949 (rope batch rows, now fixed in 0.32.0 as far as we can measure — thank you!).

[SirsiMaster]

Correction and a much sharper repro: the trigger is M ≥ 2, not the broadcast

My original report blamed a broadcast (stride-0) batch dimension. That was a red herring — my repro happened to use 16 output rows, and the row count was doing the work. The actual condition is far broader:

import mlx.core as mx
mx.random.seed(7)
b = mx.random.normal((1, 1, 1, 512)).astype(mx.float32)   # K = 1, N = 512

for M in (1, 2, 4, 8, 64):
    one = mx.ones((1, 1, M, 1)).astype(mx.float32)
    print(M, mx.abs((one @ b) - b).max().item())
# 1  -> 0.0            (exact)
# 2  -> 1.7095e-03
# 4  -> 1.7095e-03
# 8  -> 1.7095e-03
# 64 -> 1.7095e-03

K = 1 and the left operand is all ones, so the product must equal b exactly. It does at M = 1 and does not at M ≥ 2.

No broadcast is required. Plain 2-D, fully contiguous, no batch dims at all:

b2 = b.reshape(1, 512)
mx.abs((mx.ones((1, 1)) @ b2) - b2).max().item()   # 0.0
mx.abs((mx.ones((8, 1)) @ b2) - b2).max().item()   # 1.7095e-03

Transposing does not matter either — q @ k.swapaxes(-1,-2) and q @ mx.contiguous(k.swapaxes(-1,-2)) agree to 0.0e+00.

It is the fp16 grid, confirmed by choosing values on and off it (ulp at 1.0 is 2⁻¹⁰):

for val in (1 + 2**-10, 1 + 2**-12):
    bb = mx.full((1, 1, 1, 512), val).astype(mx.float32)
    print(val, mx.abs((mx.ones((1, 1, 8, 1)) @ bb) - val).max().item())
# 1+2^-10 (representable in fp16) -> 0.0
# 1+2^-12 (not representable)     -> 2.441406e-04  == 2^-12 exactly

So: float32 matmul appears to round through float16 whenever it dispatches to the matrix-matrix kernel; the matrix-vector path (M = 1) is exact. If this is a deliberate reduced-precision GEMM for f32 inputs, it would be worth documenting, because it is not what float32 implies and there is no opt-out that I can find.

Downstream effect we measured

fast.scaled_dot_product_attention inherits this for any head_dim outside its fused list (64/80/128), where it runs the unfused fallback and every prefill call has M = sequence length. Against a float64 CPU reference, gemma-4 shapes:

head_dim	T	MLX fallback error
256	8	9.6e-07
256	23	3.2e-03
256	64	1.0e-02
512	8	5.9e-03
512	64	1.7e-02

The error grows with sequence length, which is consistent with more fp16-rounded terms entering each softmax-weighted sum. Our own Metal kernel over the same inputs holds ~1e-6 at every one of those shapes, which is how we noticed.

Environment

MLX 0.32.0 (pip), Python 3.14, Apple M5 Max, macOS 26.

Apologies for the imprecise first diagnosis — happy to test any patch.

[Sarves1911]

Hi! I’d like to investigate this issue. I’ll first reproduce it on the current main branch and add a focused regression test before proposing a fix. Please let me know if someone is already working on it or if there is a preferred direction for the fix.

[XXXXRT666]

try MLX_ENABLE_TF32=0