feat(attention): add Flash Attention VJP for vector path (L≤8)

Implement fused backward pass for scaled_dot_product_attention on
short sequences (L≤8) using the vector kernel approach. This eliminates
the O(N²) memory requirement of unfused attention by recomputing the
attention matrix on-the-fly during backpropagation.

Key changes:
- Add sdpa_vector_vjp.h with GPU kernels for computing dQ, dK, dV
- Extend forward pass to output logsumexp (LSE) when needed for VJP
- Add comprehensive Python tests for gradient correctness
- Fix CUDA cuDNN backward to handle masks via set_bias() (removes
  unnecessary fallback)

Performance (M3 Max, L≤8):
- 1.1-1.4x faster than unfused attention for backward pass
- Memory: O(N) instead of O(N²) for attention matrix

The STEEL VJP for longer sequences (L>8) will be added in a follow-up PR.

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,12 @@ 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 */) {
   // 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,7 @@ 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)
 if(MLX_METAL_VERSION GREATER_EQUAL 320)
   build_kernel(fence)
 endif()

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

@@ -81,8 +81,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;
@@ -91,7 +93,9 @@ 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;
   }
 
@@ -118,13 +122,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;
@@ -156,7 +161,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
@@ -247,15 +252,18 @@ template <typename T, int D, int V = D>
   sums += o_offset * blocks + block_idx;
   maxs += o_offset * blocks + block_idx;
 
-  // Read the query
+  // 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];
   }
 
   U max_score = Limits<U>::finite_min;
   U sum_exp_score = 0;
   if (has_sinks && block_idx == 0) {
-    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;
   }
 
@@ -278,13 +286,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;