feat(attention): add Flash Attention VJP for Metal GPU

Implement fused backward pass (VJP) for scaled_dot_product_attention
on Metal GPU, enabling efficient training without falling back to
unfused attention.

- **dQ Kernel** (steel_attention_vjp_dq.h): Computes query gradients
  - Outer loop over KV blocks, inner accumulation for dQ
  - Uses log2 domain for numerical stability

- **dK/dV Kernel** (steel_attention_vjp_dkv.h): Computes key/value gradients
  - K-row ownership model eliminates atomic operations
  - Each simdgroup owns exclusive K rows to prevent races

- Optimized path for short sequences (L ≤ 8)
- Uses shared memory for efficient reduction

- Float32 accumulators for half/bfloat16 precision
- Logsumexp caching from forward pass
- Proper GQA (grouped query attention) support
- Causal mask support
- Comprehensive test coverage for all code paths

- No gradient support for mask or attention sinks (falls back to unfused)
- Requires logsumexp from forward pass (training mode only)
- Head dimension D=256 not supported in vector VJP (threadgroup memory)

Co-Authored-By: Claude <noreply@anthropic.com>

Author: Brooooooklyn

mlx/backend/cuda/scaled_dot_product_attention.cpp

@@ -402,7 +402,6 @@ bool ScaledDotProductAttention::use_fallback(
     bool has_mask,
     bool has_arr_mask,
     bool do_causal,
-    bool is_training,
     bool output_logsumexp,
     Stream s) {
   if (s.device == Device::cpu) {
@@ -460,7 +459,16 @@ void ScaledDotProductAttention::eval_gpu(
   }
 }
 
-bool ScaledDotProductAttentionVJP::use_fallback(const array& q, Stream s) {
+bool ScaledDotProductAttentionVJP::use_fallback(
+    const array& q,
+    Stream s,
+    bool has_mask,
+    bool has_sinks,
+    int /* n_kv_heads */) {
+  // Force unfused attention when masks/sinks present
+  if (has_mask || has_sinks) {
+    return true;
+  }
   // The frontend adds a padding mask when sequence length is not a multiple of
   // tile size.
   if (q.shape(2) % 128 != 0) {

mlx/backend/metal/kernels/CMakeLists.txt

@@ -53,7 +53,29 @@ build_kernel(layer_norm)
 build_kernel(random)
 build_kernel(rms_norm)
 build_kernel(rope)
-build_kernel(scaled_dot_product_attention sdpa_vector.h)
+build_kernel(scaled_dot_product_attention sdpa_vector.h sdpa_vector_vjp.h)
+
+set(STEEL_ATTN_HEADERS
+    steel/defines.h
+    steel/utils.h
+    steel/gemm/gemm.h
+    steel/gemm/mma.h
+    steel/gemm/loader.h
+    steel/gemm/transforms.h
+    steel/utils/type_traits.h
+    steel/utils/integral_constant.h
+    steel/attn/attn.h
+    steel/attn/loader.h
+    steel/attn/mma.h
+    steel/attn/params.h
+    steel/attn/transforms.h
+    steel/attn/kernels/steel_attention.h
+    steel/attn/kernels/steel_attention_vjp_dq.h
+    steel/attn/kernels/steel_attention_vjp_dkv.h)
+
+build_kernel(steel/attn/kernels/steel_attention ${STEEL_ATTN_HEADERS})
+build_kernel(steel/attn/kernels/steel_attention_vjp_dq ${STEEL_ATTN_HEADERS})
+build_kernel(steel/attn/kernels/steel_attention_vjp_dkv ${STEEL_ATTN_HEADERS})
 if(MLX_METAL_VERSION GREATER_EQUAL 320)
   build_kernel(fence)
 endif()
@@ -81,22 +103,6 @@ set(STEEL_HEADERS
     steel/utils/type_traits.h
     steel/utils/integral_constant.h)
 
-set(STEEL_ATTN_HEADERS
-    steel/defines.h
-    steel/utils.h
-    steel/gemm/gemm.h
-    steel/gemm/mma.h
-    steel/gemm/loader.h
-    steel/gemm/transforms.h
-    steel/utils/type_traits.h
-    steel/utils/integral_constant.h
-    steel/attn/attn.h
-    steel/attn/loader.h
-    steel/attn/mma.h
-    steel/attn/params.h
-    steel/attn/transforms.h
-    steel/attn/kernels/steel_attention.h)
-
 set(STEEL_NAX_HEADERS
     steel/defines.h
     steel/utils.h

mlx/backend/metal/kernels/scaled_dot_product_attention.metal

@@ -3,6 +3,7 @@
 // clang-format off
 #include "mlx/backend/metal/kernels/utils.h"
 #include "mlx/backend/metal/kernels/sdpa_vector.h"
+#include "mlx/backend/metal/kernels/sdpa_vector_vjp.h"
 
 using namespace metal;
 
@@ -41,4 +42,42 @@ using namespace metal;
 instantiate_sdpa_vector_heads(float)
 instantiate_sdpa_vector_heads(bfloat16_t)
 instantiate_sdpa_vector_heads(float16_t)
+
+// SDPA vector VJP instantiations
+#define instantiate_sdpa_vector_vjp(type, qk_dim, value_dim)    \
+  instantiate_kernel(                                           \
+      "sdpa_vector_vjp_" #type "_" #qk_dim "_" #value_dim,      \
+      sdpa_vector_vjp,                                          \
+      type,                                                     \
+      qk_dim,                                                   \
+      value_dim)
+
+// Note: D=256 exceeds Metal's 32KB threadgroup memory limit for vector VJP kernel
+#define instantiate_sdpa_vector_vjp_heads(type)    \
+  instantiate_sdpa_vector_vjp(type, 64, 64)        \
+  instantiate_sdpa_vector_vjp(type, 96, 96)        \
+  instantiate_sdpa_vector_vjp(type, 128, 128)
+
+instantiate_sdpa_vector_vjp_heads(float)
+instantiate_sdpa_vector_vjp_heads(bfloat16_t)
+instantiate_sdpa_vector_vjp_heads(float16_t)
+
+// SDPA vector VJP accumulate instantiations (for half/bfloat16 with float32 accumulators)
+#define instantiate_sdpa_vector_vjp_accumulate(type, qk_dim, value_dim)    \
+  instantiate_kernel(                                           \
+      "sdpa_vector_vjp_accumulate_" #type "_" #qk_dim "_" #value_dim,      \
+      sdpa_vector_vjp_accumulate,                                          \
+      type,                                                     \
+      qk_dim,                                                   \
+      value_dim)
+
+// Note: D=256 exceeds Metal's 32KB threadgroup memory limit for vector VJP kernel
+#define instantiate_sdpa_vector_vjp_accumulate_heads(type)    \
+  instantiate_sdpa_vector_vjp_accumulate(type, 64, 64)        \
+  instantiate_sdpa_vector_vjp_accumulate(type, 96, 96)        \
+  instantiate_sdpa_vector_vjp_accumulate(type, 128, 128)
+
+// Note: Only instantiate for half/bfloat16 since float32 doesn't need accumulate variant
+instantiate_sdpa_vector_vjp_accumulate_heads(bfloat16_t)
+instantiate_sdpa_vector_vjp_accumulate_heads(float16_t)
     // clang-format on

mlx/backend/metal/kernels/sdpa_vector.h

@@ -80,8 +80,10 @@ template <typename T, int D, int V = D>
   out += o_offset * V + simd_gid * v_per_thread;
 
   // Read the query and 0 the output accumulator
+  // Scale by M_LOG2E_F to match STEEL attention domain (exp2 instead of exp)
+  const U log2e_scale = static_cast<U>(scale * M_LOG2E_F);
   for (int i = 0; i < qk_per_thread; i++) {
-    q[i] = static_cast<U>(scale) * queries[i];
+    q[i] = log2e_scale * queries[i];
   }
   for (int i = 0; i < v_per_thread; i++) {
     o[i] = 0;
@@ -90,7 +92,8 @@ template <typename T, int D, int V = D>
   U max_score = Limits<U>::finite_min;
   U sum_exp_score = 0;
   if (has_sinks && simd_gid == 0) {
-    max_score = static_cast<U>(sinks[q_batch_head_idx % num_q_heads]);
+    // Scale sink by M_LOG2E_F to match log2 domain
+    max_score = static_cast<U>(M_LOG2E_F) * static_cast<U>(sinks[q_batch_head_idx % num_q_heads]);
     sum_exp_score = 1;
   }
 
@@ -117,13 +120,14 @@ template <typename T, int D, int V = D>
       }
       score = simd_sum(score);
       if (float_mask) {
-        score += static_cast<U>(fmask[0]);
+        // Scale float mask by M_LOG2E_F to match log2 domain
+        score += static_cast<U>(M_LOG2E_F) * static_cast<U>(fmask[0]);
       }
 
-      // Update the accumulators
+      // Update the accumulators (using exp2 to match STEEL attention)
       U new_max = max(max_score, score);
-      U factor = fast::exp(max_score - new_max);
-      U exp_score = fast::exp(score - new_max);
+      U factor = fast::exp2(max_score - new_max);
+      U exp_score = fast::exp2(score - new_max);
 
       max_score = new_max;
       sum_exp_score = sum_exp_score * factor + exp_score;
@@ -155,7 +159,7 @@ template <typename T, int D, int V = D>
   threadgroup_barrier(mem_flags::mem_threadgroup);
   max_score = max_scores[simd_lid];
   U new_max = simd_max(max_score);
-  U factor = fast::exp(max_score - new_max);
+  U factor = fast::exp2(max_score - new_max);
   sum_exp_score = simd_sum(sum_exp_scores[simd_lid] * factor);
 
   // Now we need to aggregate all the outputs
@@ -252,8 +256,10 @@ template <typename T, int D, int V = D>
   maxs += o_offset * blocks + block_idx;
 
   // Read the query and 0 the output accumulator
+  // Scale by M_LOG2E_F to match STEEL attention domain (exp2 instead of exp)
+  const U log2e_scale = static_cast<U>(scale * M_LOG2E_F);
   for (int i = 0; i < qk_per_thread; i++) {
-    q[i] = static_cast<U>(scale) * queries[i];
+    q[i] = log2e_scale * queries[i];
   }
   for (int i = 0; i < v_per_thread; i++) {
     o[i] = 0;
@@ -263,7 +269,8 @@ template <typename T, int D, int V = D>
   U sum_exp_score = 0;
   if (has_sinks && block_idx == 0 && simd_gid == 0) {
     int q_head_idx = q_batch_head_idx % num_q_heads;
-    max_score = static_cast<U>(sinks[q_head_idx]);
+    // Scale sink by M_LOG2E_F to match log2 domain
+    max_score = static_cast<U>(M_LOG2E_F) * static_cast<U>(sinks[q_head_idx]);
     sum_exp_score = 1;
   }
 
@@ -291,13 +298,14 @@ template <typename T, int D, int V = D>
       score = simd_sum(score);
 
       if (float_mask) {
-        score += fmask[0];
+        // Scale float mask by M_LOG2E_F to match log2 domain
+        score += static_cast<U>(M_LOG2E_F) * static_cast<U>(fmask[0]);
       }
 
-      // Update the accumulators
+      // Update the accumulators (using exp2 to match STEEL attention)
       U new_max = max(max_score, score);
-      U factor = fast::exp(max_score - new_max);
-      U exp_score = fast::exp(score - new_max);
+      U factor = fast::exp2(max_score - new_max);
+      U exp_score = fast::exp2(score - new_max);
 
       max_score = new_max;
       sum_exp_score = sum_exp_score * factor + exp_score;
@@ -329,7 +337,7 @@ template <typename T, int D, int V = D>
   threadgroup_barrier(mem_flags::mem_threadgroup);
   max_score = (simd_lid < BN) ? max_scores[simd_lid] : -1e9;
   U new_max = simd_max(max_score);
-  U factor = fast::exp(max_score - new_max);
+  U factor = fast::exp2(max_score - new_max);
   sum_exp_score = (simd_lid < BN) ? sum_exp_scores[simd_lid] : 0;
   sum_exp_score = simd_sum(sum_exp_score * factor);
 
@@ -342,7 +350,7 @@ template <typename T, int D, int V = D>
   // Now we need to aggregate all the outputs
   for (int i = 0; i < v_per_thread; i++) {
     outputs[simd_lid * BN + simd_gid] =
-        o[i] * fast::exp(max_scores[simd_gid] - new_max);
+        o[i] * fast::exp2(max_scores[simd_gid] - new_max);
     threadgroup_barrier(mem_flags::mem_threadgroup);
 
     // And write the output
@@ -390,7 +398,7 @@ template <typename T, int D>
   // First everybody reads the max and sum_exp
   U max_score = maxs[simd_lid];
   U new_max = simd_max(max_score);
-  U factor = fast::exp(max_score - new_max);
+  U factor = fast::exp2(max_score - new_max);
   U sum_exp_score = simd_sum(sums[simd_lid] * factor);
 
   // Now read the block into registers and then use shared memory to transpose