Repository navigation
Metal SDPA: no fused path for head_dim=512 (vector or full kernel), and the fallback is silent #3885
Description
Summary
mx.fast.scaled_dot_product_attention has no fused path for head_dim = 512 — neither the vector (decode) kernel nor the full (prefill) kernel accepts it — and the fallback to the composite eager path is silent, with no way to query or log it. Models with d=512 attention (e.g. Gemma 4's global layers) run every such attention op, at every sequence length, on the unfused path with materialized scores.
This complements #3658 / #3293, which cover extending the full kernel to d=192/256. This issue is about (a) d=512 having no coverage anywhere, including the vector kernel where decode lives, and (b) fallback discoverability.
Current routing (use_fallback, scaled_dot_product_attention.cpp#L621-L633)
- vector (qL ≤ 8):
query_head_dim ∈ {64, 96, 128, 256}or(192, 128)— no 512 - full (qL > 8):
query_head_dim ∈ {64, 80, 128}— no 192/256 (tracked in Allow force using fused kernel in scaled_dot_product_attention #3658), no 512
Gemma 4 31B (dense) has a 5:1 hybrid layout: 50 sliding layers (n_q=32, n_kv=16, d=256) and 10 global layers (n_q=32, n_kv=4, d=512, 25% partial RoPE). The sliding layers get the vector kernel at decode; the global layers never fuse, and at prefill no layer of this model fuses at all.
Measured impact (M4 Max 128 GB, bf16, Gemma-4-31B QAT Q4; MLX pin 51b2768, routing re-verified against 0c537a4)
Decode (qL=1), isolated mx.fast.scaled_dot_product_attention calls, additive mask:
| shape | seqK=1024 | seqK=4096 |
|---|---|---|
| d=512, 32 q / 4 kv heads (fallback) | 168 µs | 260 µs |
| d=256, 32 q / 8 kv heads (same KV bytes, vector kernel) | 146 µs | 184 µs |
≈ +15% @1k and +41% @4k per call for the same bytes, plus ~8 extra dispatches per call. Small per token (10 layers ≈ +0.8 ms at 4K), but it is pure routing loss.
Prefill is where it bites: the fallback materializes [B, n_kv, n_rep, qL, kL] scores. At L=4096 a single d=512 global layer runs 84 ms and materializes ~1 GB of transient scores; at the context lengths in #3658 this is the difference between running and OOM. Chunking bounds the transient but not the arithmetic (see #3658's analysis — same conclusion from our side).
Requests
- Vector-kernel d=512 (decode is the common case for these layers). If register pressure at d=512 is the blocker, a split-D variant (as the existing
(192,128)combo suggests is feasible) would cover it. - Full-kernel d≥192/256: already asked in Allow force using fused kernel in scaled_dot_product_attention #3658 / fix: add head_dim=256 to fused SDPA full attention kernel #3293 — just noting d=512 exists in the wild too (Gemma 4 global layers), so parameterizing rather than enumerating dims may be worth it.
- Fallback discoverability: today the only ways to notice the composite path are a Metal System Trace or reading
use_fallback. Either of these would have saved us days:- an env var (e.g.
MLX_FAST_LOG_FALLBACK=1) that logs once per distinct shape when amx.fast.*op takes its fallback, and/or - a queryable predicate (
mx.fast.sdpa_is_fused(q_shape, k_shape, ...)or a documented support matrix in the docs).
- an env var (e.g.
Context: found while building an MLX-based Gemma-4 inference engine and chasing why prefill attention time didn't match flash expectations; happy to provide traces or run candidate kernels on M4 Max.
Activity
hxu296 commentedon Jul 21, 2026
This is comparing d=512 fallback vs. d=256 fused? Maybe I'm not following this, but is this comparison apple-to-apple?
apocryphx commentedon Jul 22, 2026
In plain terms: Gemma 4's global layers use an attention head size (d=512) that MLX's fast attention kernels don't support, so every such call silently falls back to a slower generic path. There is no fast d=512 path to compare against directly, so the table instead compares against a supported shape that reads exactly the same number of KV-cache bytes — and at decode time, attention speed is set almost entirely by how many bytes you read. The gap between the two rows therefore measures the cost of the missing kernel, not a difference in the work being done.
So to your question, @hxu296: correct — it isn't the same shape, and it can't be, because no fused d=512 path exists; that absence is this issue. The second row is a control normalized on bytes, not geometry.
The reasoning: decode attention (qL=1) at these sizes is KV-bandwidth-bound. The two rows move identical KV traffic — d·n_kv = 512·4 = 256·8, so K+V reads are byte-for-byte the same (33.5 MB at seqK=4096, bf16); Q/O differ but at qL=1 they're KB-scale against MB-scale KV. Arithmetic intensity differs (the d=512 row does 4 MACs per KV element vs 2), but both sit far below the M4 Max roofline balance point (~56 FLOP/byte at bf16), so the extra MACs ride along with the same bytes rather than adding wall time. Under that normalization, the same bytes arriving 15–41% later is routing loss, not workload difference.
To check it isn't pin-specific, I re-ran this morning against the current wheel (
mlx==0.29.3, routing identical for these shapes: vector = {64, 96, 128, 256}, full = {64, 80, 128} — scaled_dot_product_attention.cpp#L395-L400), one fresh process per row, Python-side timing (absolutes carry per-call dispatch overhead; the ratios are the point):+19% @1k, +32% @4k — same pattern. Bench script below if you want to reproduce.
bench (fresh process per row)
hxu296 commentedon Jul 22, 2026
Ah yes, the kv bytes are the same given that you doubled the kv heads for the d=256 config. But it'd be nice to add a template instantiation for the fused kernel in a custom build and compare d=512 fused vs. unfused (you could route on an env var to test the two variants in the same build), so we could see whether fused will be beneficial for this machine. My worry is that d=512 fused would have too much register pressure to be beneficial, but it requires some testing to know
apocryphx commentedon Jul 23, 2026
Ran your experiment on the M4 Max — naive
(512, 512)instantiation plus env-gated routing, so one build serves both paths for the same shape:Same shape (d=512, 32q/4kv, qL=1, bf16, additive mask), same static lib, min of 5 fresh processes × 500 evals each (C++ harness), M4 Max 128 GB:
Observations:
qk_per_thread16 vs 8). Pressure costs, but doesn't erase the win.(192, 128)combo suggests the machinery tolerates that kind of specialization.Honest framing on magnitude: for Gemma-4-31B decode (10 global layers) this saves ~0.1–1 ms/token depending on context — real but modest; the larger costs in this issue remain the prefill full-kernel gap and fallback discoverability. But it establishes that d=512 vector coverage is trivial to add, numerically correct, and never worse than the fallback. Happy to turn it into a PR (with or without the env gate) if that's useful.
hxu296 commentedon Jul 23, 2026
This looks encouraging. I think this has the potential to turn into a PR if a wider sweep demonstrates no regression vs. unfused. The ideal case is for d=512 fused to beat d=512 unfused on nearly all shape classes, so the kernel selection logic can stay simple.
If the kernel selection logic is very complex, then we need to trade off perf vs. more complex gating. The
MLX_SDPA_VECTOR_512should stay as a test-time knob and be removed if we turn this into a PR.Some axes of the sweep I can think of:
You already demonstrated fused is better for a sweep across axis 4 when fixing other axes. It'd be great if you could sweep over all axes for the full landscape
apocryphx commentedon Jul 24, 2026
Ran the full sweep — twice, as it turned out: once with background load on the machine, once quiet. The aggregate runtime was identical (ratio 0.994) and the win/loss regions came out the same in both runs, so the landscape below (quiet run) is robust, not a scheduling artifact. Methodology as before: qL=1 decode, bf16, B=1, min of 3 fresh processes per cell, 500/300/100 evals by depth, M4 Max 128 GB. Speedup = fallback/fused; bold ≥ 1.05×, italic < 0.95×:
105 cells: 50 wins (≥1.05×), 39 ties, 16 losses (<0.95×), zero failures — including the small-N cells that exercise the 1-pass kernel at 1024-thread groups, so no register cliff anywhere.
The losses aren't scattered; they're two mechanisms:
nkv × blocksthreadgroups, so one KV head starves the grid while carrying d=512's register load. Kernel-side fixable later (scaleblockswith 1/nkv).Which makes the selection logic land nicely for your "no regression, stay simple" bar:
nkv >= 2 && gqa >= 8→ zero regressions. Worst covered cell 0.97×; everything else ties or wins up to 1.31×. One condition, and it covers the d=512 models that actually exist (Gemma 4 global layers: 32q/4kv, g8). Everything outside keeps today's fallback, so nothing changes for other shapes.nkv >= 2 && gqa >= 4widens coverage (g4 wins up to 1.19×) at the cost of one 0.91× cell (64q/16kv @16k).I'd PR the first gate (env knob removed, condition folded into
sdpa_vector_supported_head_dim), with the MQAblocksscaling as a possible follow-up that could widen it later. Raw µs in the fold.raw µs (fallback/fused per cell, quiet run)
apocryphx commentedon Jul 24, 2026
Ran the same sweep on an M1 Max 64 GB — and it turned into a correctness find, not just a perf datapoint. Short version: the d=512 vector kernels are not launchable within pipeline limits on M1-class GPUs, and the 1-pass path fails silently.
The register cliff is real, and it bites twice
On M1 Max (
applegpu_g13s) the compiler register-limits both d=512 pipelines to 832 threads per threadgroup (probed viamaxTotalThreadsPerThreadgroup; every other head dim gets 1024):Two different failure modes fall out:
L >= 1024on this device): threadgroup is32 × gqa × qL, so gqa=32 requests 1024 > 832 andcheck_kernel_threadgroup_sizethrows. Visible, at least.L < 1024): dispatches a flat 1024 with no size check at all. The launch is invalid, the GPU drops it, and the kernel returns all zeros — for every d=512 shape, not just gqa=32.MTL_DEBUG_LAYER=1flags it immediately (threadsPerThreadgroup ... must be <= 832). My first sweep's short-L columns showed 1.4–2× "wins" that were a no-op kernel returning garbage quickly.The M4 numbers above are unaffected — its pipelines genuinely allow 1024 — but the missing check in
sdpa_vector's dispatch would ship silent wrong results to every M1/M2-class user of a d=512 model.Two fixes tried
[[max_total_threads_per_threadgroup(1024)]]on both kernels: makes every launch legal, numerics correct — and spills registers on every d=512 launch. 90/105 cells regress (worst 0.37×). Dead on arrival.BD×BDtranspose buffer, and each simd writingBD/BNoutput slices. Numerically verified at all depths, validation clean — patch available. But it doesn't rescue perf on M1: short-L fused still loses to the composite fallback nearly everywhere (0.5–1.0×). The M4's short-L wins just don't transfer.Corrected M1 Max landscape
Same methodology (qL=1 decode, bf16, B=1, min of 3 fresh processes, 500/300/100 evals by depth). Short-L columns use the BN=16 kernel (the only correct fused option here);
L >= 1024columns are the unmodified 2-pass. Speedup = fallback/fused, bold ≥ 1.05×, italic < 0.95×:105 cells: 23 wins, 24 ties, 52 losses, 6 failures (gqa=32 2-pass). The deep-GQA wins are the same region as the M4 — up to 1.43× at 16K — but shallow contexts flip to fallback territory on this GPU.
Implications for the selection gate
The
nkv >= 2 && gqa >= 8gate is not safe cross-device: on M1 it admits 64q/2kv (g32), which crashes, and it routes short-L shapes that regress. Two ways to slice it:nkv >= 2 && gqa >= 8 && gqa <= 16 && N >= 4096. On this M1 data: worst cell 0.99×, wins to 1.43×. On the M4 data above: worst 0.97×, wins to 1.29×. Zero regressions on both machines, at the cost of the M4's shallow-context wins.kernel->maxTotalThreadsPerThreadgroup()against the required threadgroup size at selection time. 1024-capable devices (M4) keep the full gate; register-limited devices either fall back or restrict to the deep-GQA region. This is also the only robust guard against future register-limited targets.Either way, two things seem non-negotiable: the 1-pass dispatch needs the missing
check_kernel_threadgroup_size(silent zeros must become at worst a loud error), and gqa=32 needs to be excluded or the 2-pass threadgroup shape reworked before unconditional routing is on the table.Raw µs in the fold.
raw µs (fallback/fused per cell, M1 Max)
apocryphx commentedon Jul 24, 2026
Validated the M1 findings from the M4 side, with the BN=16 kernel applied:
Why M4 never hit the bug — pipeline probe on the same metallib:
The 832 cap is M1-generational; the M4's original short-L numbers were genuine computation, not dropped launches. Numerics of the BN=16 variant on M4: clean (bf16-rounding agreement with the fallback across MHA/g4/g8/g32, both 1-pass lengths).
But BN=16 costs the M4 its short-L wins: re-running the 256/512 columns under the M1-safe kernel flips them from mostly-wins to 7 wins / 13 ties / 22 losses (of 42). So both devices now point the same way: the 1-pass region isn't worth routing at d=512, on any hardware.
That converges to a PR spec with no device-dependent logic:
(512, 512)instantiations; 1-pass carries the BN=16 shape for D >= 512 as correctness insurance (withgroup_dimsto match), though the gate below never routes it;check_kernel_threadgroup_sizeadded tosdpa_vector's dispatch — silent zeros become a loud error on any future register-limited target, for every head dim;query_head_dim == 512 && n_kv >= 2 && gqa >= 8 && gqa <= 16 && kL >= 4096, no env var. On the M4 data: worst covered cell 0.97×, wins to 1.31×. On the M1 data: worst 0.99×, wins to 1.43×.kL >= 4096only ever selects the unmodified 2-pass kernel on both device classes.Deep-context is also where d=512 decode actually happens in practice (Gemma 4 chat sessions sit at 4K+ KV depth routinely), so the gate gives up little. If this spec looks right, I'll open the PR.
apocryphx commentedon Jul 24, 2026
Ran the blocks-heuristic probe on an M4 iPad Pro (16 GB) to chase the deep-context losses there — result: mostly not the heuristic.
First, the full iPad landscape (same grid and methodology; in-process with cache clears between arms — iOS has no processes; BN=16 kernel throughout): 77 wins / 20 ties / 8 losses of 105 — by far the greenest device, exactly as the bandwidth math predicts (120 GB/s: the fallback's extra traffic is expensive, the fused kernel's occupancy tax is cheap). Short contexts win 1.1–1.4× with the same BN=16 kernel that loses short-L on both Macs. Pipeline probe: all d=512 pipelines report maxThreads=1024 — the 832 cap is confirmed M1-generational, on iPadOS too.
The 8 losses cluster at 16K depth with few KV heads.
MTLDevice.architecture.namehere isapplegpu_g16g— ends in'g', a third device class that takes the else-branch of the 2-passblocksheuristic (64 blocks), which no Mac exercises. But sweepingMLX_SDPA_BLOCKSover the losing cells (plus two winning controls) moves almost nothing — only nkv=2@16K improves (0.74× → 0.92× at 512):The loss severity instead tracks threadgroup size = 32·gqa: g32 → 0.40×, g16 → ~0.75×, g8 → ~0.9×, g4 and below → ties. With d=512's register footprint, large threadgroups leave one resident group per core — on a 10-core GPU there is no latency hiding left at depth. Kernel-occupancy property, not host tuning.
Revised proposal — still no new selection machinery, and each piece unambiguous on all three devices tested (M4 Max / M1 Max / M4 iPad):
check_kernel_threadgroup_sizeinsdpa_vector's dispatch, unconditional — the silent-zeros fix; protects every head dim on every register-limited target.blocks:(devc == 'd' || devc == 's') && n_kv >= 2 && 8 <= gqa <= 16 && kL >= 4096. Zero regressions where it routes (worst 0.97× on M4 Max, 0.99× on M1 Max, wins to 1.43×); phone/tablet-class devices keep today's fallback everywhere.Left deliberately on the table: the iPad's broad shallow-context wins can't be routed device-blind (that region is exactly where BN=16 loses on both Macs) — per-device-class tuning could pick them up in a follow-up. The spec above takes only what's unambiguous. Happy to open the PR with this shape.
raw µs (fallback/fused per cell, M4 iPad Pro, iPadOS 26.5.2)
apocryphx commentedon Jul 24, 2026
One correction and one open validation gap, before anyone takes the gate at face value:
Correction: M4 Max reports
applegpu_g16s— so the architecture suffix is'g'= base chips (and iPads/iPhones),'s'= Pro and Max,'d'= Ultra. Two implications: base-M-series Macs (e.g. MacBook Air M4, 120 GB/s — bandwidth-identical to the iPad tested above) are correctly excluded by the's'/'d'condition, which is a nice accident; but the condition includes Pro-class chips (M1/M4 Pro), which none of our three devices represent, anddevccannot distinguish Pro from Max.Why Pro is probably fine, and why "probably" isn't enough: Pro chips take the same
's'blocks branch validated on both Maxes (256 blocks at 16K, not the 64 the iPad's'g'branch picks), with ~2× the iPad's cores and bandwidth — and the only in-gate iPad losses were at kL=16384 (kL=4096 was green on all three devices, 0.97×–1.39×). Interpolation says Max-like. But the iPad result is exactly why interpolation shouldn't be trusted here: the occupancy behavior at 32·gqa-sized threadgroups didn't extrapolate from the Maxes either.If anyone has an M1 Pro or M4 Pro handy, the 30-line bench from the earlier comment run over the five gate cells at kL ∈ {4096, 16384} would close this. Failing that, the PR could ship the gate at
kL >= 4096capped below 16K (green on every device tested, including the iPad) and lift the 16K+ range for's'-class in a follow-up once Pro numbers exist. Ultra ('d') is likewise untested — in the direction of more headroom, but stated for the record.