[Metal] Add Flash Attention VJP for training

[Brooooooklyn]

Summary

Implements fused backward pass (VJP) for scaled_dot_product_attention on Metal GPU. This enables efficient gradient computation during training without falling back to unfused (decomposed) attention operations.

Changes

New Files

• mlx/backend/metal/kernels/sdpa_vector_vjp.h - Vector VJP kernel for short sequences (L ≤ 8)
• mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_vjp_dq.h - STEEL dQ gradient kernel
• mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_vjp_dkv.h - STEEL dK/dV gradient kernel

Modified Files

• mlx/backend/metal/scaled_dot_product_attention.cpp - VJP dispatch logic (+840 lines)
• mlx/fast.cpp / mlx/fast_primitives.h - Logsumexp caching, VJP routing
• python/tests/test_fast_sdpa.py - Comprehensive VJP tests (+220 lines)

Implementation Notes

Uses a two-kernel approach to avoid atomic operations:

1. dQ kernel (steel_attention_vjp_dq.h):

   • Computes query gradients via outer loop over KV blocks
   • Uses log2 domain for numerical stability
   • Proper clamping to prevent overflow (exp2 arg clamped to [-88, 0])

2. dK/dV kernel (steel_attention_vjp_dkv.h):

   • Uses K-row ownership model where each simdgroup owns exclusive rows
   • Eliminates race conditions in GQA where multiple query heads share KV
   • No atomic operations needed

3. Vector VJP (sdpa_vector_vjp.h):

   • Optimized path for short sequences (L ≤ 8)
   • Uses float32 accumulators for half/bfloat16 precision
   • Shared memory reduction for efficiency

Key Features

• Float32 accumulators for half/bfloat16 precision
• Logsumexp caching from forward pass for VJP reuse
• Proper GQA (grouped query attention) support
• Causal mask support

Limitations

• Falls back to unfused attention for mask/sinks gradients (per existing design)
• Requires logsumexp from forward pass (training mode only)
• Head dimension D=256 not supported in vector VJP (32KB threadgroup memory limit)

Test Plan

• Existing test_sdpa_grad passes
• New comprehensive VJP tests added:
  • test_sdpa_grad_vector_path - short sequences (L=1,4,7,8)
  • test_sdpa_grad_steel_path - longer sequences (L=16,32,128,256)
  • test_sdpa_grad_head_dims - head dimensions (D=32,64,96,128)
  • test_sdpa_grad_gqa - GQA configurations (4:1, 8:1, 16:1, MHA)
  • test_sdpa_grad_dtypes - float16, bfloat16, float32
  • test_sdpa_grad_edge_cases - L=1, non-power-of-2, large batch, qL≠kvL

All 21 SDPA tests pass (1 skipped for unrelated disabled feature).

[Brooooooklyn]

Notes: I'm working on https://github.com/mlx-node/mlx-node and trying to port some features in trl.
This pull request was generated by Claude Code. I am trying to reduce the computation and memory usage of GRPO training by utilizing the full flash attention feature.

[Copilot]

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

[Copilot]

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

[Brooooooklyn]

@awni @zcbenz Do you have any interesting review this PR?

[zcbenz]

Can you share some benchmarking numbers?

[Brooooooklyn]

Benchmark Summary for PR

┌──────────────────────────┬────────────────────┬──────────────────────────────────────────┐                                                                                              
│         Category         │   Speedup Range    │                  Notes                   │                                                                                              
├──────────────────────────┼────────────────────┼──────────────────────────────────────────┤                                                                                              
│ Forward Only             │ 1.0x - 1.8x faster │ Consistently faster across all configs   │                                                                                              
├──────────────────────────┼────────────────────┼──────────────────────────────────────────┤                                                                                              
│ Backward Only (float16)  │ 0.8x - 1.3x        │ Mixed, mostly break-even to faster       │                                                                                              
├──────────────────────────┼────────────────────┼──────────────────────────────────────────┤                                                                                              
│ Backward Only (bfloat16) │ 0.95x - 2.0x       │ Generally faster, especially short STEEL │                                                                                              
├──────────────────────────┼────────────────────┼──────────────────────────────────────────┤                                                                                              
│ Full VJP (fwd+bwd)       │ 0.8x - 1.8x        │ Vector path fastest, STEEL break-even    │                                                                                              
└──────────────────────────┴────────────────────┴──────────────────────────────────────────┘                                                                                              

Key observations:

• Vector path (L≤8): Generally 1.1x-1.8x faster for VJP
• STEEL path (L>8): Mostly break-even (~1.0x) for backward, which is expected since the main benefit is memory reduction
• Long sequences (L≥1024): VJP is ~1.0x (break-even on speed, but saves memory by not materializing attention matrix)

[zcbenz]

I ran the benchmark script on a M4 air and it seems that fused VJP is slower for most cases:

$ python benchmarks/python/sdpa_vjp_bench.py 
SDPA VJP Benchmark - dtype=float16
=====================================================================================

[Forward + Backward (VJP)]
  B  H_q  H_kv      L    D |    unfused      fused  speedup     path
-------------------------------------------------------------------------------------
  2    8     8      1   64 |      0.07ms      0.07ms    1.09x   vector
  2    8     8      4   64 |      0.07ms      0.07ms    1.03x   vector
  2    8     8      8   64 |      0.07ms      0.07ms    1.05x   vector
  2    8     8      8  128 |      0.08ms      0.08ms    0.99x   vector
  2    8     8     32   64 |      0.10ms      0.11ms    0.89x    STEEL
  2    8     8     64   64 |      0.18ms      0.21ms    0.84x    STEEL
  2    8     8    128   64 |      0.45ms      0.54ms    0.82x    STEEL
  2    8     8    128  128 |      0.59ms      0.69ms    0.85x    STEEL
  2    8     8    256  128 |      1.93ms      2.32ms    0.83x    STEEL
  1   32     8    512   64 |     11.45ms     14.91ms    0.77x    STEEL
  1   32     8    512  128 |     14.26ms     17.73ms    0.80x    STEEL
  1   32     8   1024   64 |     50.49ms     64.39ms    0.78x    STEEL
  1   32     8   1024  128 |     56.06ms     69.96ms    0.80x    STEEL
  1   32     8   2048  128 |    244.84ms    299.90ms    0.82x    STEEL
  2   32     8    256   64 |      6.16ms      8.06ms    0.76x    STEEL
  2   32     4    256   64 |      6.23ms      8.02ms    0.78x    STEEL

[Brooooooklyn]

Interesting, slight difference in the results between the M4 Air and my MacBook Pro 16 with the M3 Max and 128GB of RAM.
It appears to be affected by memory bandwidth limitations on the M4 Air; The VJP kernels need to:

• Recompute attention scores
• Load logsumexp from memory
• Compute P = exp(S - LSE)
• Compute gradients

The primary benefit of Flash Attention VJP is memory, not speed:

• Unfused: O(N²) memory for attention matrix
• Fused: O(N) memory - no attention matrix materialization

For long sequences (4K+), memory savings can be critical even if speed is similar or slightly slower.

[zcbenz]

I tested on a M3 Max and the results are still mixed:

$ python benchmarks/python/sdpa_vjp_bench.py 
SDPA VJP Benchmark - dtype=float16
=====================================================================================

[Forward + Backward (VJP)]
  B  H_q  H_kv      L    D |    unfused      fused  speedup     path
-------------------------------------------------------------------------------------
  2    8     8      1   64 |      0.30ms      0.26ms    1.17x   vector
  2    8     8      4   64 |      0.34ms      0.24ms    1.43x   vector
  2    8     8      8   64 |      0.28ms      0.25ms    1.14x   vector
  2    8     8      8  128 |      0.29ms      0.33ms    0.88x   vector
  2    8     8     32   64 |      0.33ms      0.36ms    0.92x    STEEL
  2    8     8     64   64 |      0.43ms      0.28ms    1.55x    STEEL
  2    8     8    128   64 |      0.29ms      0.33ms    0.86x    STEEL
  2    8     8    128  128 |      0.31ms      0.31ms    0.98x    STEEL
  2    8     8    256  128 |      0.39ms      0.40ms    0.97x    STEEL
  1   32     8    512   64 |      1.02ms      1.01ms    1.02x    STEEL
  1   32     8    512  128 |      1.30ms      1.34ms    0.97x    STEEL
  1   32     8   1024   64 |      3.75ms      3.75ms    1.00x    STEEL
  1   32     8   1024  128 |      4.81ms      4.81ms    1.00x    STEEL
  1   32     8   2048  128 |     16.65ms     16.66ms    1.00x    STEEL
  2   32     8    256   64 |      0.63ms      0.66ms    0.95x    STEEL
  2   32     4    256   64 |      0.65ms      0.67ms    0.97x    STEEL

I'm not really familiar with how VJP works for flash attention but generally speaking reduced memory usage should make op faster when it is memory-bound, which is usually the case for mac, so I think fused op being slower likely means something is off.

But anyway the result on vector path looks promising, maybe separate that part into an independent PR first? It would also make reviewing much easier.

[Brooooooklyn]

Done, I only preserve the vector path in this pull request, while I continue to investigate the reasons behind the slow performance of the STEEL path.

[zcbenz]

I'm not familiar with the Metal implementation so we will need other maintainers to review the code.

[Brooooooklyn]

@awni, can you take a look at this PR?

[Brooooooklyn]

Completely rewritten; it is now closed and will reopen as a new one.