Add TurboQuant KV cache compression with native Metal SDPA kernel

[arozanov]

Proposed changes

Adds TurboQuant (arXiv 2504.19874) as a native Metal SDPA kernel for KV cache compression.

• New QuantizationMode::TurboQuant
• sdpa_vector_turbo Metal kernel: reads 3-bit packed K indices with codebook dequant
• Pre-rotated query optimization: no WHT butterfly in attention inner loop
• TurboQuantSDPA primitive with full eval_gpu dispatch
• Python API: mx.fast.turboquant_sdpa()

Benchmarks (M4 Pro 48GB, 28 query heads, 4 KV heads, D=128):

Context	Apple SDPA	TurboQuant	Speedup
1K	0.161ms	0.108ms	1.5x
4K	0.164ms	0.109ms	1.5x
8K	0.209ms	0.106ms	2.0x
16K	0.511ms	0.104ms	4.9x

TurboQuant reads 3-bit packed data (4.6x less memory bandwidth than fp16). Kernel time constant at ~0.1ms regardless of context length.

Related: standalone package at https://github.com/arozanov/turboquant-mlx and mlx-lm PR at ml-explore/mlx-lm#1067

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)

[zcbenz]

Thanks for the PR and I think it is nicely written.

But I think we need to have a generic quantized sdpa kernel first and then add things like turbo quant as part of the API, rather than a specialized turbo quant kernel directly. #3026 is a nice step but as you can see we don't quite have enough maintainers to review.

On turbo quant itself, we don't just add builtin kernels for every new hot thing, users should use custom extensions or custom kernels instead.

[arozanov]

Thanks for the PR and I think it is nicely written.

But I think we need to have a generic quantized sdpa kernel first and then add things like turbo quant as part of the API, rather than a specialized turbo quant kernel directly. #3026 is a nice step but as you can see we don't quite have enough maintainers to review.

On turbo quant itself, we don't just add builtin kernels for every new hot thing, users should use custom extensions or custom kernels instead.

Makes sense. Happy to help with #3026 if useful - the packed-integer codepath here (3-bit indices + codebook lookup) could serve as a reference for one of the quantization modes a generic API would need.
I'll package this as a custom Metal kernel extension via mx.fast.metal_kernel() in the meantime.

[deceptech-packet-ninja]

Bug: N reads wrong dimension of k_norms

While testing this kernel locally (built from source), I found that line 450 in scaled_dot_product_attention.cpp:

int N = k_norms.shape(1); // kv sequence length

reads shape(1) which is the head dimension (H_kv), not the sequence length (T). For k_norms of shape (B, H_kv, T), this should be:

int N = k_norms.shape(2); // kv sequence length

With shape(1), the kernel only processes H_kv keys (e.g., 8 instead of 500), causing incorrect attention output. After changing to shape(2), the kernel produces exact match (1.000 cosine similarity) against the reference mx.fast.scaled_dot_product_attention.

Kernel benchmarks after the fix (M-series, Llama-3-8B config: B=1, H_q=32, H_kv=8, D=128):

Context	Standard SDPA	TurboQuant SDPA	Speedup
256	0.43ms	0.44ms	1.0x
512	0.49ms	0.22ms	2.2x
1024	0.52ms	0.20ms	2.6x
4096	1.91ms	0.30ms	6.3x

The kernel is impressively fast — essentially constant time regardless of context length.

[arozanov]

Bug: N reads wrong dimension of k_norms

While testing this kernel locally (built from source), I found that line 450 in scaled_dot_product_attention.cpp:

int N = k_norms.shape(1); // kv sequence length

reads shape(1) which is the head dimension (H_kv), not the sequence length (T). For k_norms of shape (B, H_kv, T), this should be:

int N = k_norms.shape(2); // kv sequence length

With shape(1), the kernel only processes H_kv keys (e.g., 8 instead of 500), causing incorrect attention output. After changing to shape(2), the kernel produces exact match (1.000 cosine similarity) against the reference mx.fast.scaled_dot_product_attention.

Kernel benchmarks after the fix (M-series, Llama-3-8B config: B=1, H_q=32, H_kv=8, D=128):

Context Standard SDPA TurboQuant SDPA Speedup
256 0.43ms 0.44ms 1.0x
512 0.49ms 0.22ms 2.2x
1024 0.52ms 0.20ms 2.6x
4096 1.91ms 0.30ms 6.3x
The kernel is impressively fast — essentially constant time regardless of context length.

Good catch, thanks! Fixed: k_norms.shape(1) → k_norms.shape(2). The shape is (B, H_kv, T) so dimension 2 is the sequence length.