mx.fast.scaled_dot_product_attention produces incorrect attention scores (MLX 0.31.2)

[coldbugs]

Environment

• MLX 0.31.2
• macOS, Apple M4
• Python 3.12

Summary

mx.fast.scaled_dot_product_attention returns mathematically incorrect results compared to manual attention, causing complete model failure (0% accuracy). The manual fallback produces correct results (99.5% accuracy), matching a Candle/Rust implementation with identical weights.

Model config

• seq_len=82 (81 Sudoku tokens + 1 puzzle embedding), hidden=512, heads=8, head_dim=64
• 4 H-layers + 4 L-layers, H_cycles=2, L_cycles=2
• RoPE (base=10000) with pre-computed cos/sin tables shape [82, 64]
• No GQA (num_heads == num_kv_heads)
• Weights: sapientinc/HRM-checkpoint-sudoku-extreme

Behavior

MPS fused kernel: model produces garbage output, 0% accuracy on Sudoku given cells. All 16 ACT steps run (q-values never converge).

Manual attention: model produces correct Sudoku solutions, 99.5% accuracy (214/215 given cells across 10 puzzles), halts in 2-7 steps.

Reproduction

The fix was replacing the MPS fused call with the manual fallback that already exists in the code:

# Broken:
return mx.fast.scaled_dot_product_attention(query, key, value, scale=scale)

# Fixed:
q = mx.transpose(query, (0, 2, 1, 3))
k = mx.transpose(key, (0, 2, 1, 3))
v = mx.transpose(value, (0, 2, 1, 3))
scores = q @ mx.transpose(k, (0, 1, 3, 2)) * scale
attn_weights = mx.softmax(scores, axis=-1)
out = attn_weights @ v
return mx.transpose(out, (0, 2, 1, 3))

Full investigation log with side-by-side comparison against Candle/Rust available on request.

Workaround

Use manual attention unconditionally — the MPS fused kernel produces wrong attention weights for this model configuration, cascading into zero accuracy.

[zcbenz]

Is this a 0.31.2 only issue?

[coldbugs]

Confirmed on MLX 0.31.2 / macOS / Apple M4.

Minimal reproduction with deterministic inputs (no random seed):

import mlx.core as mx
import numpy as np

B, S, H, D = 2, 82, 8, 64
scale = float(D) ** -0.5

def arr(shape):
    n = int(np.prod(shape))
    return mx.array(np.linspace(-0.5, 0.5, n, endpoint=False, dtype=np.float32).reshape(shape))

q = arr((B, S, H, D))
k = arr((B, S, H, D))
v = arr((B, S, H, D))

# manual
qt = mx.transpose(q, (0, 2, 1, 3))
kt = mx.transpose(k, (0, 2, 1, 3))
vt = mx.transpose(v, (0, 2, 1, 3))
manual = mx.transpose(mx.softmax(qt @ mx.transpose(kt, (0, 1, 3, 2)) * scale) @ vt, (0, 2, 1, 3))

# fused
fused = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale)

diff = np.abs(np.array(manual) - np.array(fused))
print(f"max_diff={np.max(diff):.6e}")

Output:

max_diff=4.969254e-01

Worst element: manual=-0.000419, fused=-0.497344 — ~0.5 divergence, three orders of magnitude off. The manual attention matches the Candle/Rust reference implementation exactly.

[coldbugs]

Is this a 0.31.2 only issue?

It also impact 0.31.0 and 0.31.1

[coldbugs]

Root cause analysis from DeepSeek AI:

The bug is a layout mismatch between the Python API and the Metal kernel in mlx/backend/metal/scaled_dot_product_attention.cpp.

Python API sends Q/K/V in [B, L, H, D] layout (batch, seq_len, heads, head_dim):

# HRM shape: [2, 82, 8, 64]
mx.fast.scaled_dot_product_attention(q, k, v, scale=scale)

Metal kernel (steel_attention.h) expects [B, H, L, D] and indexes with strides accordingly:

Q += tidl.z * params->Q_strides[0] + // Batch
     tidl.y * params->Q_strides[1] + // Head
     tidl.x * BQ * params->Q_strides[2]; // Sequence

The dispatch code in sdpa_full_self_attention_nax reads the shape dimensions as if the input is already [B, H, L, D]:

int B = q.shape(0);   // = 2  ✓
int H = q.shape(1);   // = 82 ✗ should be 8
int D = q.shape(3);   // = 64 ✓
int qL = q.shape(2);  // = 8  ✗ should be 82

Everything downstream — tile counts, alignment checks, memory strides — is computed from these swapped dimensions.

Concrete impact for [2, 82, 8, 64]:

• Kernel believes there are 82 heads of sequence length 8 (instead of 8 heads × length 82)
• Head stride = H*D = 512 when it should be L*D = 5248
• Real heads 8–81 read garbage memory; every real sequence is truncated to 8 tokens

This produces the observed ~0.5 max absolute divergence and cascades into zero model accuracy.

The fix would be either:

1. Transpose Q/K/V to [B, H, L, D] before passing to the Metal kernel
2. Or swap the shape-index interpretations in the dispatch code

[zcbenz]

Can you try setting environment variable MLX_ENABLE_TF32=0?

[coldbugs]

(.venv-vllm-metal) ➜ HRM MLX_ENABLE_TF32=0 python3 tests/test_mlx_sdpa_bug.py
MLX 0.31.2
shape=[2,82,8,64] max_diff=4.969254e-01 BUG — manual=-0.000419 fused=-0.497344

VERDICT: BUG present on MLX array([[[[-0.5, -0.499988, -0.499976, ..., -0.499274, -0.499262, -0.49925],
[-0.499238, -0.499226, -0.499214, ..., -0.498511, -0.498499, -0.498488],
[-0.498476, -0.498464, -0.498452, ..., -0.497749, -0.497737, -0.497725],
...,
[-0.496189, -0.496177, -0.496165, ..., -0.495463, -0.495451, -0.495439],
[-0.495427, -0.495415, -0.495403, ..., -0.4947, -0.494688, -0.494677],
[-0.494665, -0.494653, -0.494641, ..., -0.493938, -0.493926, -0.493914]],
[[-0.493902, -0.493891, -0.493879, ..., -0.493176, -0.493164, -0.493152],
[-0.49314, -0.493128, -0.493116, ..., -0.492414, -0.492402, -0.49239],
[-0.492378, -0.492366, -0.492354, ..., -0.491652, -0.49164, -0.491628],
...,
[-0.490091, -0.49008, -0.490068, ..., -0.489365, -0.489353, -0.489341],
[-0.489329, -0.489317, -0.489305, ..., -0.488603, -0.488591, -0.488579],
[-0.488567, -0.488555, -0.488543, ..., -0.487841, -0.487829, -0.487817]],
[[-0.487805, -0.487793, -0.487781, ..., -0.487078, -0.487067, -0.487055],
[-0.487043, -0.487031, -0.487019, ..., -0.486316, -0.486304, -0.486292],
[-0.486281, -0.486269, -0.486257, ..., -0.485554, -0.485542, -0.48553],
...,
[-0.483994, -0.483982, -0.48397, ..., -0.483267, -0.483256, -0.483244],
[-0.483232, -0.48322, -0.483208, ..., -0.482505, -0.482493, -0.482481],
[-0.482469, -0.482458, -0.482446, ..., -0.481743, -0.481731, -0.481719]],
...,
[[-0.0182927, -0.0182808, -0.0182689, ..., -0.0175662, -0.0175543, -0.0175424],
[-0.0175305, -0.0175186, -0.0175067, ..., -0.016804, -0.0167921, -0.0167802],
[-0.0167683, -0.0167564, -0.0167445, ..., -0.0160418, -0.0160299, -0.016018],
...,
[-0.0144817, -0.0144698, -0.0144579, ..., -0.0137552, -0.0137433, -0.0137314],
[-0.0137195, -0.0137076, -0.0136957, ..., -0.012993, -0.0129811, -0.0129692],
[-0.0129573, -0.0129454, -0.0129335, ..., -0.0122308, -0.0122189, -0.012207]],
[[-0.0121951, -0.0121832, -0.0121713, ..., -0.0114687, -0.0114567, -0.0114448],
[-0.0114329, -0.011421, -0.0114091, ..., -0.0107065, -0.0106946, -0.0106826],
[-0.0106707, -0.0106588, -0.0106469, ..., -0.00994426, -0.00993236, -0.00992045],
...,
[-0.00838415, -0.00837224, -0.00836033, ..., -0.00765768, -0.00764577, -0.00763386],
[-0.00762195, -0.00761004, -0.00759813, ..., -0.00689548, -0.00688357, -0.00687167],
[-0.00685976, -0.00684785, -0.00683594, ..., -0.00613329, -0.00612138, -0.00610947]],
[[-0.00609756, -0.00608565, -0.00607374, ..., -0.00537109, -0.00535918, -0.00534728],
[-0.00533537, -0.00532346, -0.00531155, ..., -0.0046089, -0.00459699, -0.00458508],
[-0.00457317, -0.00456126, -0.00454935, ..., -0.0038467, -0.00383479, -0.00382288],
...,
[-0.00228659, -0.00227468, -0.00226277, ..., -0.00156012, -0.00154821, -0.0015363],
[-0.00152439, -0.00151248, -0.00150057, ..., -0.000797923, -0.000786014, -0.000774104],
[-0.000762195, -0.000750286, -0.000738377, ..., -3.57279e-05, -2.38186e-05, -1.19093e-05]]],
[[[0, 1.19093e-05, 2.38186e-05, ..., 0.000726467, 0.000738377, 0.000750286],
[0.000762195, 0.000774104, 0.000786014, ..., 0.00148866, 0.00150057, 0.00151248],
[0.00152439, 0.0015363, 0.00154821, ..., 0.00225086, 0.00226277, 0.00227468],
...,
[0.00381098, 0.00382288, 0.00383479, ..., 0.00453744, 0.00454935, 0.00456126],
[0.00457317, 0.00458508, 0.00459699, ..., 0.00529964, 0.00531155, 0.00532346],
[0.00533537, 0.00534728, 0.00535918, ..., 0.00606183, 0.00607374, 0.00608565]],
[[0.00609756, 0.00610947, 0.00612138, ..., 0.00682403, 0.00683594, 0.00684785],
[0.00685976, 0.00687167, 0.00688357, ..., 0.00758622, 0.00759813, 0.00761004],
[0.00762195, 0.00763386, 0.00764577, ..., 0.00834842, 0.00836033, 0.00837224],
...,
[0.00990854, 0.00992045, 0.00993236, ..., 0.010635, 0.0106469, 0.0106588],
[0.0106707, 0.0106826, 0.0106946, ..., 0.0113972, 0.0114091, 0.011421],
[0.0114329, 0.0114448, 0.0114567, ..., 0.0121594, 0.0121713, 0.0121832]],
[[0.0121951, 0.012207, 0.0122189, ..., 0.0129216, 0.0129335, 0.0129454],
[0.0129573, 0.0129692, 0.0129811, ..., 0.0136838, 0.0136957, 0.0137076],
[0.0137195, 0.0137314, 0.0137433, ..., 0.014446, 0.0144579, 0.0144698],
...,
[0.0160061, 0.016018, 0.0160299, ..., 0.0167326, 0.0167445, 0.0167564],
[0.0167683, 0.0167802, 0.0167921, ..., 0.0174948, 0.0175067, 0.0175186],
[0.0175305, 0.0175424, 0.0175543, ..., 0.018257, 0.0182689, 0.0182808]],
...,
[[0.481707, 0.481719, 0.481731, ..., 0.482434, 0.482446, 0.482458],
[0.482469, 0.482481, 0.482493, ..., 0.483196, 0.483208, 0.48322],
[0.483232, 0.483244, 0.483256, ..., 0.483958, 0.48397, 0.483982],
...,
[0.485518, 0.48553, 0.485542, ..., 0.486245, 0.486257, 0.486269],
[0.486281, 0.486292, 0.486304, ..., 0.487007, 0.487019, 0.487031],
[0.487043, 0.487055, 0.487067, ..., 0.487769, 0.487781, 0.487793]],
[[0.487805, 0.487817, 0.487829, ..., 0.488531, 0.488543, 0.488555],
[0.488567, 0.488579, 0.488591, ..., 0.489294, 0.489305, 0.489317],
[0.489329, 0.489341, 0.489353, ..., 0.490056, 0.490068, 0.49008],
...,
[0.491616, 0.491628, 0.49164, ..., 0.492342, 0.492354, 0.492366],
[0.492378, 0.49239, 0.492402, ..., 0.493105, 0.493116, 0.493128],
[0.49314, 0.493152, 0.493164, ..., 0.493867, 0.493879, 0.493891]],
[[0.493902, 0.493914, 0.493926, ..., 0.494629, 0.494641, 0.494653],
[0.494665, 0.494677, 0.494688, ..., 0.495391, 0.495403, 0.495415],
[0.495427, 0.495439, 0.495451, ..., 0.496153, 0.496165, 0.496177],
...,
[0.497713, 0.497725, 0.497737, ..., 0.49844, 0.498452, 0.498464],
[0.498476, 0.498488, 0.498499, ..., 0.499202, 0.499214, 0.499226],
[0.499238, 0.49925, 0.499262, ..., 0.499964, 0.499976, 0.499988]]]], dtype=float32)
Known affected: 0.31.2

[1R053]

We hit the same class of failure independently in our project — wanted to share what we found, because we think this issue (and at least one other open report) reduces to a Metal Toolchain codegen bug rather than an MLX bug.

Diagnosis: metalfe-32023.883 miscompiles NAX kernels for matmuls with M > 8

When MLX is built on macOS 26 without pinning CMAKE_OSX_DEPLOYMENT_TARGET below 26.0, the Metal frontend reports __METAL_VERSION__ = 400 and the CMake gate at mlx/backend/metal/kernels/CMakeLists.txt:157 (MLX_METAL_VERSION ≥ 400 AND MACOS_SDK_VERSION ≥ 26.2) enables the Metal-4 tensor "NAX" GEMM/qmm kernel codepath (files like steel_gemm_fused_nax.metal, quantized_nax.metal). The resulting compiled kernels return silently wrong matmul outputs for any matmul with M > 8 rows, where M is the left-hand activation dimension. Matmuls with M ≤ 8 are bit-exact.

For Transformer inference this means every prefill matmul (attention QK^T, attention output projection, MLP gate/up/down, lm_head — all of which share M = prompt length) is corrupted whenever the prompt is more than 8 tokens long. The error compounds across layers, so the first generated token after prefill is already garbage. Decode with M=1 stays correct, which is why some users see "kind-of-works on short prompts, totally collapses on long prompts" — the @AlexaHLE SDPA report here matches that fingerprint exactly.

Repro (against an MLX built from source on macOS 26 with no deployment-target pin)

import mlx.core as mx
import numpy as np

def check(M, K, N, tag):
    a = np.sin(np.arange(M * K) * 0.013).astype(np.float32).reshape(M, K)
    b = np.cos(np.arange(K * N) * 0.007).astype(np.float32).reshape(K, N)
    mx.eval(c := mx.matmul(mx.array(a), mx.array(b)))
    err = float(np.abs(np.asarray(c) - a @ b).max())
    print(f"{tag}: M={M}  max abs err = {err:.6f}")
    return err

check( 8, 64, 64, "below NAX threshold")  # ~1e-4, correct
check(16, 64, 64, "above NAX threshold")  # ~18, gross divergence
check(64, 64, 64, "above NAX threshold")  # ~30, gross divergence

Confirm the NAX path is actually compiled in:

strings build/mlx/backend/metal/libmlx.a | grep -ciE \
  'steel_gemm_fused_nax_|_qmm_n_nax_|_qmm_t_nax_|_gather_qmm_n_nax_|_gather_qmm_t_nax_'

Non-zero → NAX kernels present → bug will reproduce on the affected toolchain.

Smoking-gun evidence: Apple's own MLX shipping pipeline avoids the NAX path

The mlx.metallib shipped inside Apple's iOS-26.x simulator runtimes — at PrivateFrameworks/GESS.framework/mlx.metallib — contains zero NAX kernel symbols. We checked three different runtime builds (iOS 26.4 / 2026-03-16 + 2026-04-09, iOS 26.5 / 2026-05-07); all three are 2,427,468 bytes and all pass the strings | grep ...nax... check at 0 matches. A build from source on the same host produces a libmlx.a with 7 such symbols and the divergent-matmul behaviour above. Apple's internal MLX shipping pipeline appears to apply the equivalent of -DCMAKE_OSX_DEPLOYMENT_TARGET=14.0 to avoid the NAX path entirely.

Workaround (and why it's perf-painful)

cmake .. -DCMAKE_OSX_DEPLOYMENT_TARGET=14.0  # no NAX kernels compiled

This restores correctness but costs about 2.4–2.7× on long-prefill throughput compared with what NAX is supposed to deliver. Note also that the ollama project (#16053) found a leak where MLX builds AIR objects with -mmacosx-version-min correctly but the final metal vs metallib link step can still stamp the resulting library with the SDK's macOS 26 target — so the workaround itself can be silently defeated unless the cmake-build dir is invalidated when the deployment target changes. Worth being aware of when reviewing fixes.

Related reports that we believe share this root cause

• Fix nax condition for iphone #3083 — A18 NAX hardware-detection gate (different surface, same kernel family).
• Fix for NAX overflow. #3092 — NAX overflow in ConvTransposed1d (different bug, same kernel family).
• MLX models not working on M5 MacBook Pro lmstudio-ai/lmstudio-bug-tracker#1356 — Unable to load function steel_gemm_fused_nax_* on M5 / macOS 26.2 (likely a different surface of the same broken Apple-side NAX rollout).
• [Metal] MLX fails to build metal library on macOS 26 Tahoe / Metal Toolchain 32023 #3337 — same Metal Toolchain 32023 series, different bug (namespace migration).
• mlx: fix macOS 26 target leakage in v3 metallib ollama/ollama#16053 — deployment-target leakage in the metallib link step.

Suggested actions

1. MLX side: tighten the CMake gate at kernels/CMakeLists.txt:157 to also probe metalfe's version (or just hard-disable NAX on macOS 26 for now). This mirrors what the iOS-sim mlx.metallib is already doing in practice.
2. Apple Toolchain side: this needs a metalfe codegen fix for subtile_matmad_nax. We'd be happy to upstream a standalone Metal/Swift repro for an Apple Feedback Assistant filing if that's useful — please let us know.

Full investigation write-up (with the verified-locally numerics, environment table, and per-kernel detail): happy to share, just ask.

[zcbenz]

The mx.fast.scaled_dot_product_attention does not mean to be bit-level accurate with manual attention scores, mathematically they are the same thing but float computation would give different results when you change the sequence of computations.

You get similar results with PyTorch:

import mlx.core as mx
import numpy as np
import torch
import torch.nn.functional as F

B, S, H, D = 2, 82, 8, 64
scale = float(D) ** -0.5

def arr(shape):
    n = int(np.prod(shape))
    return mx.array(np.linspace(-0.5, 0.5, n, endpoint=False, dtype=np.float32).reshape(shape))

q = arr((B, S, H, D))
k = arr((B, S, H, D))
v = arr((B, S, H, D))

# MLX manual
qt = mx.transpose(q, (0, 2, 1, 3))
kt = mx.transpose(k, (0, 2, 1, 3))
vt = mx.transpose(v, (0, 2, 1, 3))
manual = mx.transpose(mx.softmax(qt @ mx.transpose(kt, (0, 1, 3, 2)) * scale) @ vt, (0, 2, 1, 3))

# MLX fused
fused = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale)

diff_mlx = np.abs(np.array(manual) - np.array(fused))
print(f"MLX manual vs fused max_diff = {np.max(diff_mlx):.6e}")

# PyTorch implementation
q_pt = torch.tensor(np.array(q)).transpose(1, 2)
k_pt = torch.tensor(np.array(k)).transpose(1, 2)
v_pt = torch.tensor(np.array(v)).transpose(1, 2)

pt_fused = F.scaled_dot_product_attention(q_pt, k_pt, v_pt, scale=scale)
pt_fused_np = pt_fused.transpose(1, 2).numpy()

diff_pt = np.abs(np.array(manual) - pt_fused_np)
print(f"MLX manual vs PyTorch fused max_diff = {np.max(diff_pt):.6e}")

MLX manual vs fused max_diff = 4.969254e-01
MLX manual vs PyTorch fused max_diff = 3.308013e-01