[Performance] Vector SDPA achieves only ~25% of memory bandwidth for GQA (4 KV heads, head_dim 256) at long context on M1 Ultra

[KaiserCaiser]

Describe the issue

Single-query (decode-time) mx.fast.scaled_dot_product_attention runs far
below memory-bandwidth speed for a GQA shape with few KV heads and large
head_dim at long context. On an M1 Ultra with mlx 0.32.0:

• Shape: q=(1, 32, 1, 256), k=v=(1, 4, 100000, 256), bf16, mask=None
• Measured: 2.2-2.4 ms per call
• KV bytes touched: 2 × 4 × 100000 × 256 × 2B ≈ 0.41 GB → ~170-190 GB/s achieved
• The same machine sustains ~680 GB/s on a plain reduction over a large
  contiguous array (benchmark below), so the kernel achieves ~25-28% of
  practically available bandwidth.

For a 16-full-attention-layer model (Qwen3.6-27B hybrid) this makes decode
attention cost ~39 ms/token at 100k context — about 43% of the total token
time — and is the dominant reason decode throughput trails
bandwidth-proportional expectations at long context.

MLX_SDPA_BLOCKS (new in 0.32.0) was swept at this shape: the default
selection is already optimal (8/16/32/64/128 are all equal or worse), so this
does not look like a block-count tuning issue.

q_len=2 costs 5.39 ms (≈2.2×), so batched-verify speculative decoding at long
context is also gated by the same kernel.

To Reproduce

import time
import mlx.core as mx

B, Hq, Hkv, D, S = 1, 32, 4, 256, 100000

# reference: practically achievable bandwidth on this machine
rows = 8 * 1024**3 // (4 * 65536)
a = mx.random.normal((rows, 65536))
mx.eval(a)
mx.eval(a.sum())
t0 = time.time()
for _ in range(5):
    mx.eval(a.sum())
bw = 8 / ((time.time() - t0) / 5)
print(f"reduction bandwidth: {bw:.0f} GB/s")

k = mx.random.normal((B, Hkv, S, D), dtype=mx.bfloat16)
v = mx.random.normal((B, Hkv, S, D), dtype=mx.bfloat16)
q = mx.random.normal((B, Hq, 1, D), dtype=mx.bfloat16)
mx.eval(k, v, q)
scale = D ** -0.5

mx.eval(mx.fast.scaled_dot_product_attention(q, k, v, scale=scale, mask=None))
t0 = time.time()
n = 50
for _ in range(n):
    mx.eval(mx.fast.scaled_dot_product_attention(q, k, v, scale=scale, mask=None))
dt = (time.time() - t0) / n
kv_gb = 2 * Hkv * S * D * 2 / 1e9
print(f"sdpa q_len=1: {dt*1000:.2f} ms  -> {kv_gb/dt:.0f} GB/s "
      f"({kv_gb/dt/bw*100:.0f}% of reduction bandwidth)")

Output on M1 Ultra (128GB), macOS 26.2, mlx 0.32.0 (pip):

reduction bandwidth: 676 GB/s
sdpa q_len=1: 2.17 ms  -> 189 GB/s (28% of reduction bandwidth)

Expected behavior

Decode-time SDPA at long context is memory-bound; something closer to
50-80% of achievable bandwidth would be expected (llama.cpp's Metal FA
achieves ~75%+ on comparable GQA shapes on M-series).

Happy to run patches/benchmarks on this machine (M1 Ultra 128GB, real
workload = 27B hybrid model at 80-160k context).

---

This investigation (bandwidth baseline, kernel benchmarks, and this report)
was performed by Claude (model: claude-fable-5) via Claude Code, operating
directly on the reporter's hardware.

[KaiserCaiser]

Two follow-ups:

Correction on the model shape. The real model (Qwen3.6-27B) uses
num_attention_heads=24, not 32 — the 32 in my report came from a library
default rather than the checkpoint's config.json. The benchmark above is
self-contained (it measures the 32-head shape it constructs), and the KV
bytes / bandwidth math is unchanged since query-head count doesn't affect KV
size. At the true shape the stock kernel measures essentially the same:

• q=(1, 24, 1, 256), k=v=(1, 4, 100000, 256), bf16: 2.15 ms (~190 GB/s)

The gap is closable. A hand-written split-S flash-decoding kernel (via
mx.fast.metal_kernel, ~100 lines MSL) at the true shape reaches 1.41 ms
(~290 GB/s, 1.52×) on the same M1 Ultra, with max|Δ| ≤ 5e-4 vs
mx.fast.scaled_dot_product_attention on random inputs. Structure: one
threadgroup per (kv_head, key-split), cooperative uint4 staging of K/V rows
into threadgroup memory, each simdgroup covering 2 query heads × 1/2 of the
keys (register spills kill it beyond ~2 heads/simdgroup), split partials
merged with a logsumexp pass. In the full model this recovers ~15.7 ms/token
at 100k context (16 attention layers), +19% end-to-end decode.

A load-only variant of the same threadgroup structure sustains ~475 GB/s, so
there is likely further headroom for a native kernel with proper software
pipelining. Happy to share the kernel source or benchmark patches on this
hardware.

(Same setup as the report: investigation by Claude claude-fable-5 via
Claude Code on the reporter's M1 Ultra.)

[hxu296]

one threadgroup per (kv_head, key-split), cooperative uint4 staging of K/V rows into threadgroup memory

You could reduce a lot of shmem traffic by changing how you partitioned work into simdgroups.

Could you try whether the following will help:

• dispatch is the same as yours. One TG per (kv_head, key-split)
• simdgroup further splits work for TG's key-split. Each simdgroup is responsible for attending all n_group query heads for this kv head, and accumulate partial o/denominator/max for this simdgroup's kv partition. Note that this is not a fixed 1/2 kv split, so this simdgroup only needs its partition of partial kv in register, which is reused across all n_group query heads.
• then, as you already did, merge simdgroup partials in a TG-reduction pass.

Note that this way, the entire TG only reads kv once, and it avoids any shmem->register traffic on the kv.

[KaiserCaiser]

Thanks — tried your scheme on the same M1 Ultra (implemented via mx.fast.metal_kernel, so register allocation is compiler-scheduled; caveat below):

• As suggested (each simdgroup: all 6 query heads × a contiguous key sub-partition, KV read device→registers, no threadgroup memory): correctness is fine, but 4.84 ms @ S=100k (~85 GB/s) — 3.4× slower than the staged variant. The signature matches register spilling: ~110+ live floats/lane (6 heads × 512 B accumulators + q).
• Same kernel with fp16 accumulators (halves the persistent per-lane state): jumps 3.3× to 1.47 ms @100k / 1.83 ms @140k, i.e. it lands on the same plateau as my shared-memory staging variant (1.41 / 2.12 ms) — and beats it at 140k. Surprisingly, max|Δ| vs mx.fast.sdpa stays ≤ 1.2e-4 even at 100k on random inputs.

So your premise holds (the shmem round-trip is avoidable), but at metal_kernel level the binding constraint on M1 Ultra becomes the register file: full-precision state for all n_group heads doesn't fit. Both routes plateau at ~280–310 GB/s vs ~475 GB/s for a load-only kernel with the same dispatch — closing that last gap presumably needs software pipelining / register control beyond what metal_kernel exposes. Happy to benchmark a native-kernel patch on this hardware.

(claude-fable-5 via Claude Code, same setup as above.)

[hxu296]

Thanks for testing it out and collecting this data.

I benchmarked the shmem staging variant vs. direct register kv-reuse variant on a M5-class machine. And it didn't throttle on register pressure. This is presumably due to the introduction of GPU dynamic caching since M3, which manages register allocation more intelligently.

Perhaps for pre-M3 class of machines, shmem staging is the more profitable path.

[KaiserCaiser]

That explains the data cleanly — thanks for running the M5 comparison. Dynamic caching (M3+) vs fixed register allocation (M1/M2) accounts for exactly the divergence we measured.

For completeness, one more data point for anyone landing on this thread with pre-M3 hardware: we ended up shipping the register kv-reuse variant with fp16 accumulators in production on the M1 Ultra — it matches the shmem-staging variant at 100k (1.47 vs 1.41 ms) and beats it at 140k (1.83 vs 2.12 ms), with max|Δ| ≤ 1.2e-4 vs mx.fast.sdpa. So the summary across architectures seems to be:

• M3+: register kv-reuse with fp32 accumulators (dynamic caching absorbs the pressure)
• M1/M2: either shmem staging, or register kv-reuse with fp16 accumulators — both plateau at ~280–310 GB/s at metal_kernel level

Either way the stock vector SDPA's ~190 GB/s at this shape leaves a consistent ~1.5× on the table, which is the original issue. Happy to keep benchmarking patches on the M1 Ultra side.

(claude-fable-5 via Claude Code, same setup.)

[hxu296]

Regarding f16 accumulators. It seems flashinfer has a mode to use fp16 for the q@k.T (since the reduction applies to the short head_dim and values are scaled) but still used fp32 for the score@v (since the reduction applies to the kv_dim).