[Stonesjtu]

Summary

Unroll the generic sdpa_vector_2pass_1 inner key loop by four iterations. Each unrolled iteration loads four K fragments and four V fragments before applying the corresponding score and online-softmax updates. The original scalar loop handles the tail.

This PR does not change the specialized sdpa_vector_2pass_1_gqa kernel.

Motivation

The generic 2-pass kernel previously completed one K-load/score/V-load/online-softmax chain before starting the next key assigned to the SIMD group. Grouping four independent K/V loads is intended to expose more memory-level parallelism to the GPU. TK=4 limits the additional per-thread storage compared with larger unroll factors.

Reproduction

The PR adds benchmarks/python/sdpa_vector_decode_bench.py.

The reported A/B used these exact revisions:

• parent: d63f0f939a1822cb95fe031694ed5d135a495958
• PR: d9a200d

Because the benchmark file does not exist in the parent revision, save it before switching revisions:

git show d9a200d:benchmarks/python/sdpa_vector_decode_bench.py \
  > /tmp/sdpa_vector_decode_bench.py

# Build and run the parent.
git checkout d63f0f939a1822cb95fe031694ed5d135a495958
# Rebuild MLX using the normal development build procedure.
python /tmp/sdpa_vector_decode_bench.py

# Build and run the PR head.
git checkout d9a200d
# Rebuild MLX using the same build configuration.
python benchmarks/python/sdpa_vector_decode_bench.py

My two builds used identical settings:

CMAKE_ARGS='-DMLX_METAL_JIT=OFF -DCMAKE_OSX_DEPLOYMENT_TARGET=15.0'

The script defaults to eight fp16, single-token, no-mask decode shapes. On an M1 Pro, all eight select the generic 2-pass kernel modified by this PR. The script obtains the GPU architecture via mx.device_info(mx.gpu) and labels each requested shape as vector-1pass, generic-2pass, or gqa-specialized according to the current dispatch conditions.

Additional shapes can be requested with repeated arguments:

python benchmarks/python/sdpa_vector_decode_bench.py \
  --shape 32,32,32768,128 \
  --shape 16,4,32768,128

Timing method:

• 10 warmup chains
• 5 rounds × 30 timed samples
• 8 dependent SDPA calls per sample
• reports best sample and median of the five per-round medians, divided by 8
• mx.eval() synchronizes each chain

The latency is host-inclusive operator latency; Python graph construction and dispatch overhead are amortized over the eight calls. logical GB/s counts K/V reads issued by query heads and is not a DRAM counter. It can exceed physical DRAM bandwidth for GQA because shared K/V data can be served from cache.

Results

Apple M1 Pro (applegpu_g13s), fp16, no mask. Values are median-of-round-medians from the included script.

QH	KVH	GQA	context	parent (us)	PR (us)	speedup
32	32	1	4,096	667.5	474.5	1.41x
32	32	1	32,768	5,964.2	3,596.7	1.66x
4	4	1	65,536	968.6	867.4	1.12x
8	4	2	4,096	74.1	69.7	1.06x
8	4	2	32,768	495.5	475.0	1.04x
16	4	4	4,096	109.6	84.9	1.29x
16	4	4	32,768	697.8	469.6	1.49x
32	4	8	4,096	170.9	134.5	1.27x

Correctness

For each shape, the script compares one fused SDPA call against an explicit matmul + fp32 softmax + matmul reference. Relative L∞ error (max(abs(actual - expected)) / max(abs(expected))) ranged from 7.2e-4 to 1.2e-3 on both revisions.

Scope and limitations

• The reported benchmark covers fp16 single-token decode without a mask or attention sinks.
• It intentionally excludes shapes routed to sdpa_vector_2pass_1_gqa, since that kernel is unchanged.
• Results above are from one Apple M1 Pro; the included script is intended to collect equivalent data on other Apple GPU generations.

Test plan

• Full parent and PR builds with identical configuration
• Same benchmark script run on the parent and PR revisions
• Correctness against an explicit fp32-softmax reference
• Generic 2-pass coverage for GQA=1, 2, 4, and GQA=8 below its specialized-path threshold
• Context lengths from 4K to 64K

[zcbenz]

Please provide the benchmark script that we can use to reproduce the claimed performance gain.

[Stonesjtu]

Added the reproducible benchmark and updated the PR description with the exact A/B procedure and results.

[zcbenz]

Thanks a lot for the update, I tried this change on M4 and M5 Max and can confirm it is bringing speedup for most shapes. The only regression I saw is with shape QH=32,KVH=32,GQA=1,context=32,768 on M5 Max, which only runs 0.9x of TK=1.

For follow-up we should tune TK based on hardwares, but for now I think we should merge this since it is a clear win for most cases.