mx.fast.scaled_dot_product_attention: 2.3–2.8× latency jump from L=8 to L=12 queries (kernel-routing gap), flat until L≥48

[pierre427]

While benchmarking speculative-decoding verify shapes we found a sharp cliff in
fp16 mx.fast.scaled_dot_product_attention when the number of query rows
crosses from 8 to 12. Latency then stays flat until L≈48 and only recovers
efficiency at L≥64 — so the L∈[12, 48] window pays ~L=48 cost regardless of L.

GQA 32 q-heads / 8 kv-heads, head_dim 128, fp16, no mask, M5 Max:

S (kv len)	L=8	L=12	L=16	L=32	L=48	L=64
16384	0.813 ms	1.906 ms	1.903 ms	1.904 ms	1.904 ms	1.588 ms
32768	1.319 ms	3.633 ms	3.646 ms	3.643 ms	3.630 ms	2.992 ms

• L=8→12 is a 2.34× (16k) / 2.75× (32k) jump for 1.5× the work.
• Per-query-row cost at L=12 is ~1.6–1.9× the L=8 cost and worse than L=1.
• The flat plateau L=12..48 suggests the vector kernel's small-L path caps at
  L≤8 and everything above falls to a path tuned for much larger L (steel /
  full attention), with nothing covering the 12–48 window.

Why it matters: multi-token verify steps in speculative decoding (MTP,
prompt-lookup / n-gram, draft-model) live exactly in L≈4–16, and batched short
continuations hit 12–48. The window between "decode-shaped" and "prefill-shaped"
is becoming the hot path for spec-decode workloads.

Observed on the affine-quantized side too: the decomposed quantized attention
(qmm→softmax→qmm) beats fp16 SDPA at L=16 at these S (0.35–0.48× its
latency), which is only possible because of this cliff.

Repro (self-contained):

import time, mlx.core as mx
NQ, NKV, D = 32, 8, 128
for S in (16384, 32768):
    k = mx.random.normal((1, NKV, S, D)).astype(mx.float16)
    v = mx.random.normal((1, NKV, S, D)).astype(mx.float16)
    mx.eval(k, v)
    for L in (1, 2, 4, 8, 12, 16, 24, 32, 48, 64):
        q = mx.random.normal((1, NQ, L, D)).astype(mx.float16); mx.eval(q)
        f = lambda: mx.fast.scaled_dot_product_attention(q, k, v, scale=D**-0.5)
        for _ in range(10): mx.eval(f())
        mx.synchronize(); t0 = time.perf_counter()
        for _ in range(40): mx.eval(f())
        mx.synchronize()
        print(S, L, f"{(time.perf_counter()-t0)/40*1e3:.3f} ms")

Happy to run additional shapes (head dims, GQA ratios, masks) if useful.

[katlun-lgtm]

Repro reproduces the routing here — this is a tile-selection gap, not a kernel-perf regression. Two stacked effects, and they line up with all four of your observations:

1. Vector→full-attn cliff at qL≤8. use_fallback in backend/metal/scaled_dot_product_attention.cpp gates the fast vector path on query_sequence_length <= 8 && (query_sequence_length * gqa_factor) <= 32. With GQA 32/8, gqa_factor=4, so 4*qL ≤ 32 ⇒ the wall is exactly at qL=8. It's the vector kernel's threadgroup budget — group_dims(32, gqa_factor, qL) needs gqa*qL ≤ 32. Anything above 8 leaves the decode-shaped path.

2. Only a bq=64 tile exists for head_dim=128 on the nax path (the plateau). On M5 Max is_nax_available() is true, so sdpa_full_self_attention_nax runs with bq=64 hardcoded, and steel_attention_nax.metal only instantiates bq=64 for D=128. NQ = ceil(qL/64) = 1 for every qL∈[9,64], so L=12/16/32/48 all process a single 64-row query tile → the flat plateau.

3. L=64 recovers because it's tile-aligned. align_Q = (qL % bq)==0 is true only at 64; L=12..48 run the align_Q=false masked partial-tile variant over the full 64 rows, which is slower per tile. That's the 1.90→1.59 ms (16k) / 3.63→2.99 (32k) step down.

4. The decomposed quantized path beating fp16 SDPA at L=16 is a direct consequence — it scales with actual L while fp16 pays the 64-tile cost.

A low-risk direction would be to right-size the query tile on this path: add bq=16 / bq=32 nax instantiations for D=128 and pick bq from qL (qL≤16?16:qL≤32?32:64), which would bring the spec-decode verify window (L≈4–16) down to a 16-row tile instead of 64. I don't have an M5 in front of me to bench that; if it's useful I can validate on M5 Max and put up a PR with before/after numbers.

[katlun-lgtm]

Followed up on my own suggestion — I benched the analogous tile-shrink on an M3 Max (non-nax path, bq=32) before claiming it'd help on the nax path. It doesn't: adding a bq=16/WM=2 tile for D=128 and routing qL≤16 to it gave no speedup (L=12 @ S=16k: 2.43ms → 2.49ms; correctness clean vs an fp32 reference).

The reason is that NQ = ceil(qL/bq) is already 1 for qL=12 at bq=32 (and at bq=64 on the nax path), so shrinking the query tile doesn't change the launch. At these S the kernel is K/V-streaming-bound — the steel full-attn path launches only grid=(1, H, 1) threadgroups, each streaming all S keys serially, and latency scales ~linearly in S (L=16: 2.46→4.63 ms for S=16k→32k). The L=8→12 cliff is really the loss of the 2pass vector path's block-parallel-over-S reduction when routing to the head-parallel-only steel kernel.

So the effective fix isn't the tile table — it's either extending the 2pass block-parallel reduction to qL>8 (blocked today by the gqa*qL ≤ 32 threadgroup cap on the vector kernel) or adding split-K to the steel full-attn kernel so the S reduction parallelizes across threadgroups instead of one per head. Happy to prototype the split-K direction if it's of interest.

[katlun-lgtm]

Prototyped the first of the two directions above — extending the 2-pass vector path's block-parallel reduction to qL > 8 — and it collapses the cliff on M3 Max. (I went this way rather than steel split-K because the 2-pass kernel already has the split-S + logsumexp merge machinery; split-K would duplicate it for the same window.)

S	L=8	L=12	L=16
16384	1.13 → 1.14 ms	2.49 → 1.44 ms (1.73×)	2.50 → 1.96 ms (1.28×)
32768	2.03 → 2.03 ms	4.63 → 2.53 ms (1.83×)	4.65 → 3.56 ms (1.31×)

(fp16, GQA 32/8, D=128, mask=None, 40-iter, baseline/patched interleaved A/B/A/B. Causal is similar: L12 2.58 → 1.52 ms. MHA 32/32: L12 2.52 → 1.75 ms. qL ≤ 8 decode and L ≥ 17 are unchanged.)

Mechanism: each simdgroup in sdpa_vector_2pass_1 now processes up to 4 consecutive query rows serially, reusing the K/V rows already in registers, so the threadgroup constraint relaxes from gqa * qL ≤ 32 to gqa * ceil(qL/4) ≤ 32 and the block-parallel reduction over S survives into the verify window. The original single-row body is kept verbatim behind a function constant, so decode codegen is untouched. Routing sends 8 < qL ≤ 16 to the vector path when kL ≥ 1024. I also tried extending to qL ≤ 32, but L=24–32 was no better than the steel path on M3 Max, so the change stops at 16 — split-K on the steel kernel remains the right tool for that range (and for the nax side generally).

Opened the change as #3838. Caveat: measured on M3 Max (non-nax full-attention baseline). If you can run your repro on the branch on the M5 Max, that would confirm the win transfers to the nax path — the baseline structure it replaces (one query tile per head, serial over S) is the same.