[WIP] Quantized SDPA

[CC-Yeh]

Proposed changes

Add Metal quantized SDPA vector kernels based on #1515

With M4, L=32768, H = 32, D = 128, Lq=1 :

Precision	SDPA (ms)	Quant SDPA (ms)	Ops-Based (ms)	Quant Ops-Based (ms)
mxfp4	98.63080	15.32626	43.72120	24.71464
mxfp8	97.37316	18.71779	42.89932	46.47875

TODO:

• Support Affine and NVFP4
• Adapt Faster two pass sdpa #3023 once it's merged.
• Tuning on Larger devices (Need help).
• Cleanup

What improve performance:

• Removed thread storage k, v to reduce register pressure (was waiting on synchronization).
• Fused computation with dequantization
• Tuned reading size ('uint16_t'/'uin32_t') for quantized k/v
• Manual unroll better than clang loop optimizer

Checklist

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

[awni]

The numbers seem quite good.. a little too good to be true 😅

What's the difference between SDPA and Attention in the benchmark? Also what's the query sequence length used for the benchmark?

[CC-Yeh]

The numbers seem quite good.. a little too good to be true 😅

Totally agree, must be missing something 🤔

What's the difference between SDPA and Attention in the benchmark? Also what's the query sequence length used for the benchmark?

Attention is a simple reference implementation built from matmul + softmax + matmul (Maybe too naive?).
SDPA uses mx.fast.scaled_dot_product_attention, which hits the sdpa_vector_2pass kernels when Lq ≤ 8 (this case).

The query sequence length here is 1 (q.shape = (1, 32, 1, 128)), so this benchmark is measuring the single-token decode case, where one new token attends to a long KV cache (L = 32768).

[CC-Yeh]

@awni
Fixed some bugs in dequantizing 8bit and benchmark(unneccessary dequantization steps).
Now the numbers make more sense 😃

[awni]

So if I’m understanding correctly the fused implementation is slower in the quantized case than the unfused ops-based one?

[CC-Yeh]

Fused SDPA is faster: MXFP4 15.33 ms vs 24.71 ms, and MXFP8 26.09 ms vs 46.48 ms to decode a single query.

[awni]

Very nice!!

[awni]

Why not nvfp4?

[CC-Yeh]

It’s on the way! I just wanted to make sure the PR structure was okay first.

[awni]

Why not affine?