Metal shaders for memory efficient self attention on large sequences

[bpkeene]

Implements metal shaders for:

o = mx.fast.scaled_dot_product_attention(queries, keys, values, scale=scale, mask=None)

Supports fp16, fp32 dtypes; flexible hidden dimension, currently templated for 64 and 128.

Causal masking for prompt encoding not yet implemented, this shader is focused at present on full self-attention.

Context

We're continuing the work started with #735 and following up with memory savings for the self-attention use case common in diffusion transformer workflows such as Stable Diffusion 3. In particular, we observe memory savings for the 8-b use case of > 5GB for float32. Current shaders are implemented using Steel-like GEMM primitives in MLX style, with potential for performance tuning to improve latency while retaining the memory savings.

Supported

• Supports mx.float16 and mx.float32 dtypes
• Supports head_dim=64,128 (covers most 7b+ LLMs)
• MHA supported (no MQA, no GQA)
• Unsupported use cases still go through MLX primitives under the hood.
• No backward pass implementation (inference-only kernel)

Checklist

Put an x in the boxes that apply.

• 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)

[bpkeene]

Marking as draft, currently working through some numerical issues via separate workflow, and will add CPU side bindings + dispatch, test & docs - sharing a current status and will update this PR

[awni]

These changed to compute_encoder.set_input_array (for inputs) and compute_encoder.set_output_array for outputs of the shader.

[awni]

This changes to auto& compute_encoder ...

[bpkeene]

Hi folks,

Attaching some graphs for measured latency on M3 Max and some estimated memory savings per attention block (empirically observed at several data points, graph here obtained via formulas)

[bpkeene]

Some room for improvement on larger sequences re: latency, with a divergence after ~2300 sequence length, though the memory savings exceeds 1GB ~2k, and is approaching 5GB at 4250 sequence length (SD3 8B use case).

All measurements were with batch size 2, heads = 38, hidden dim = 64, and float32 on M3 Max / 48GB.

[awni]

@bpkeene left a few minor comments. Could you address? Once updated we can run the tests and get this merged.

[bpkeene]

Updated with the requested changes, thank you for the prompt review!