Add force_fused option to scaled_dot_product_attention

mlx/backend/cuda/scaled_dot_product_attention.cpp

@@ -8,6 +8,8 @@
 
 #include <nvtx3/nvtx3.hpp>
 
+#include <sstream>
+
 namespace mlx::core {
 
 namespace {
@@ -567,6 +569,30 @@ bool ScaledDotProductAttention::use_fallback(
       !supports_sdpa_vector(q, k, v, has_arr_mask, output_logsumexp);
 }
 
+std::string ScaledDotProductAttention::fused_unsupported_reason(
+    const array& q,
+    const array& k,
+    const array& v,
+    bool has_mask,
+    bool has_arr_mask,
+    bool do_causal,
+    bool output_logsumexp,
+    Stream s) {
+  if (s.device == Device::cpu) {
+    return "the fused kernels require a GPU stream.";
+  }
+  if (supports_sdpa_cudnn(q, k, v, has_arr_mask, do_causal, s) ||
+      supports_sdpa_vector(q, k, v, has_arr_mask, output_logsumexp)) {
+    return "";
+  }
+  std::ostringstream msg;
+  msg << "neither the cuDNN attention nor the vector attention kernel "
+      << "supports this configuration; got query shape " << q.shape()
+      << ", key shape " << k.shape() << ", value shape " << v.shape()
+      << " with dtype " << q.dtype() << ".";
+  return msg.str();
+}
+
 bool ScaledDotProductAttention::supports_bool_mask() {
   return false;
 }

mlx/backend/metal/scaled_dot_product_attention.cpp

@@ -160,6 +160,7 @@ void sdpa_full_self_attention_nax(
   MTL::Size grid_dims = MTL::Size(NQ, H, B);
   MTL::Size group_dims = MTL::Size(32, wm, wn);
 
+  check_kernel_threadgroup_size(kernel, group_dims, hash_name);
   compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
 }
 
@@ -323,6 +324,7 @@ void sdpa_full_self_attention_metal(
   MTL::Size grid_dims = MTL::Size(NQ, H, B);
   MTL::Size group_dims = MTL::Size(32, wm, wn);
 
+  check_kernel_threadgroup_size(kernel, group_dims, hash_name);
   compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
 }
 
@@ -606,11 +608,26 @@ bool ScaledDotProductAttention::use_fallback(
     // forward and backward.
     return true;
   }
-  if (output_logsumexp) {
-    return true;
-  }
+  return !fused_unsupported_reason(
+              q, k, v, has_mask, has_arr_mask, do_causal, output_logsumexp, s)
+              .empty();
+}
+
+std::string ScaledDotProductAttention::fused_unsupported_reason(
+    const array& q,
+    const array& k,
+    const array& v,
+    bool has_mask,
+    bool has_arr_mask,
+    bool do_causal,
+    bool output_logsumexp,
+    Stream s) {
   if (s.device == Device::cpu) {
-    return true;
+    return "the fused kernels require a GPU stream.";
+  }
+  if (output_logsumexp) {
+    return "the fused forward does not produce the logsumexp required for "
+           "the fused VJP; use default routing when training.";
   }
 
   const int value_head_dim = v.shape(-1);
@@ -621,27 +638,55 @@ bool ScaledDotProductAttention::use_fallback(
   const int num_kv_heads = k.shape(1);
   const int gqa_factor = num_query_heads / num_kv_heads;
 
-  const bool sdpa_vector_supported_head_dim =
-      (query_head_dim == value_head_dim &&
-       (query_head_dim == 64 || query_head_dim == 96 || query_head_dim == 128 ||
-        query_head_dim == 256)) ||
-      (query_head_dim == 192 && value_head_dim == 128);
-  const bool sdpa_full_supported_head_dim = query_head_dim == value_head_dim &&
-      (query_head_dim == 64 || query_head_dim == 80 || query_head_dim == 96 ||
-       query_head_dim == 128);
-
-  const bool sdpa_full_supported_mask = !has_mask || has_arr_mask ||
-      (query_sequence_length <= key_sequence_length && do_causal);
-
-  const bool supports_sdpa_full = query_sequence_length > 8 &&
-      sdpa_full_supported_mask && sdpa_full_supported_head_dim;
-
-  const bool supports_sdpa_vector = (query_sequence_length <= 8) &&
-      (query_sequence_length <= key_sequence_length) &&
-      sdpa_vector_supported_head_dim &&
-      (query_sequence_length * gqa_factor) <= 32;
-
-  return !(supports_sdpa_full || supports_sdpa_vector);
+  std::ostringstream msg;
+  if (query_sequence_length > 8) {
+    // Full attention kernel (query length over 8 routes here)
+    const bool supported_head_dim = query_head_dim == value_head_dim &&
+        (query_head_dim == 64 || query_head_dim == 80 || query_head_dim == 96 ||
+         query_head_dim == 128);
+    if (!supported_head_dim) {
+      msg << "the full attention kernel supports head dims {64, 80, 96, 128} "
+          << "with matching query/value head dims; got query head dim "
+          << query_head_dim << " and value head dim " << value_head_dim << ".";
+      return msg.str();
+    }
+    if (has_mask && !has_arr_mask &&
+        !(query_sequence_length <= key_sequence_length && do_causal)) {
+      msg << "the full attention kernel with a causal mask requires the "
+          << "query sequence to be no longer than the key sequence; got "
+          << "query length " << query_sequence_length << " and key length "
+          << key_sequence_length << ".";
+      return msg.str();
+    }
+  } else {
+    // Vector attention kernel (query length of at most 8 routes here)
+    const bool supported_head_dim =
+        (query_head_dim == value_head_dim &&
+         (query_head_dim == 64 || query_head_dim == 96 ||
+          query_head_dim == 128 || query_head_dim == 256)) ||
+        (query_head_dim == 192 && value_head_dim == 128);
+    if (!supported_head_dim) {
+      msg << "the vector attention kernel supports head dims "
+          << "{64, 96, 128, 256} with matching query/value head dims, or "
+          << "query head dim 192 with value head dim 128; got query head dim "
+          << query_head_dim << " and value head dim " << value_head_dim << ".";
+      return msg.str();
+    }
+    if (query_sequence_length > key_sequence_length) {
+      msg << "the vector attention kernel requires the query sequence to be "
+          << "no longer than the key sequence; got query length "
+          << query_sequence_length << " and key length " << key_sequence_length
+          << ".";
+      return msg.str();
+    }
+    if (query_sequence_length * gqa_factor > 32) {
+      msg << "the vector attention kernel requires the query length times "
+          << "the GQA factor to be at most 32; got query length "
+          << query_sequence_length << " and GQA factor " << gqa_factor << ".";
+      return msg.str();
+    }
+  }
+  return "";
 }
 
 bool ScaledDotProductAttention::supports_bool_mask() {

mlx/backend/no_gpu/primitives.cpp

@@ -36,6 +36,18 @@ bool fast::ScaledDotProductAttention::use_fallback(
   return true;
 }
 
+std::string fast::ScaledDotProductAttention::fused_unsupported_reason(
+    const array& q,
+    const array& k,
+    const array& v,
+    bool has_mask,
+    bool has_arr_mask,
+    bool do_causal,
+    bool output_logsumexp,
+    Stream s) {
+  return "there is no GPU backend in this build of MLX.";
+}
+
 bool fast::ScaledDotProductAttention::supports_bool_mask() {
   return false;
 }

mlx/fast.cpp

@@ -618,6 +618,7 @@ array scaled_dot_product_attention(
     const std::string& mask_mode /* = "" */,
     std::optional<array> mask_arr /* = {} */,
     const std::optional<array>& sinks /* = {} */,
+    bool force_fused /* = false */,
     StreamOrDevice s /* = {}*/) {
   for (const auto& tensor : {queries, keys, values}) {
     if (tensor.ndim() != 4) {
@@ -825,7 +826,18 @@ array scaled_dot_product_attention(
   bool is_training = detail::in_grad_tracing();
   bool has_fast_vjp = !ScaledDotProductAttentionVJP::use_fallback(q, stream);
   bool output_logsumexp = is_training && has_fast_vjp;
-  if (!ScaledDotProductAttention::use_fallback(
+  if (force_fused) {
+    auto reason = ScaledDotProductAttention::fused_unsupported_reason(
+        q, k, v, has_mask, has_arr_mask, do_causal, output_logsumexp, stream);
+    if (!reason.empty()) {
+      std::ostringstream msg;
+      msg << "[scaled_dot_product_attention] force_fused=true but no fused "
+          << "kernel is available: " << reason;
+      throw std::invalid_argument(msg.str());
+    }
+  }
+  if (force_fused ||
+      !ScaledDotProductAttention::use_fallback(
           q,
           k,
           v,

mlx/fast.h

@@ -53,6 +53,7 @@ MLX_API array scaled_dot_product_attention(
     const std::string& mask_mode = "",
     std::optional<array> mask_arr = {},
     const std::optional<array>& sinks = {},
+    bool force_fused = false,
     StreamOrDevice s = {});
 
 using TemplateArg = std::variant<int, bool, Dtype>;

mlx/fast_primitives.h

@@ -231,6 +231,20 @@ class ScaledDotProductAttention : public Custom {
       Stream s);
   static bool supports_bool_mask();
 
+  // Returns an empty string when a fused kernel can handle the given
+  // configuration on `s`, and otherwise a human readable reason why the
+  // fused path is unavailable. Capability only — routing heuristics (such
+  // as preferring the unfused path while training) do not produce a reason.
+  static std::string fused_unsupported_reason(
+      const array& q,
+      const array& k,
+      const array& v,
+      bool has_mask,
+      bool has_arr_mask,
+      bool do_causal,
+      bool output_logsumexp,
+      Stream s);
+
   void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
       override {
     throw std::runtime_error("NYI");

python/src/fast.cpp

@@ -234,6 +234,7 @@ void init_fast(nb::module_& parent_module) {
          const float scale,
          const std::variant<std::monostate, std::string, mx::array>& mask,
          const std::optional<mx::array>& sinks,
+         bool force_fused,
          mx::StreamOrDevice s) {
         bool has_mask = !std::holds_alternative<std::monostate>(mask);
         bool has_str_mask =
@@ -250,16 +251,32 @@ void init_fast(nb::module_& parent_module) {
               throw std::invalid_argument(msg.str());
             }
             return mx::fast::scaled_dot_product_attention(
-                queries, keys, values, scale, mask_str, std::nullopt, sinks, s);
+                queries,
+                keys,
+                values,
+                scale,
+                mask_str,
+                std::nullopt,
+                sinks,
+                force_fused,
+                s);
           } else {
             auto mask_arr = std::get<mx::array>(mask);
             return mx::fast::scaled_dot_product_attention(
-                queries, keys, values, scale, "", mask_arr, sinks, s);
+                queries,
+                keys,
+                values,
+                scale,
+                "",
+                mask_arr,
+                sinks,
+                force_fused,
+                s);
           }
 
         } else {
           return mx::fast::scaled_dot_product_attention(
-              queries, keys, values, scale, "", {}, sinks, s);
+              queries, keys, values, scale, "", {}, sinks, force_fused, s);
         }
       },
       "q"_a,
@@ -269,9 +286,10 @@ void init_fast(nb::module_& parent_module) {
       "scale"_a,
       "mask"_a = nb::none(),
       "sinks"_a = nb::none(),
+      "force_fused"_a = false,
       "stream"_a = nb::none(),
       nb::sig(
-          "def scaled_dot_product_attention(q: array, k: array, v: array, *, scale: float,  mask: None | str | array = None, sinks: array | None = None, stream: StreamOrDevice = None) -> array"),
+          "def scaled_dot_product_attention(q: array, k: array, v: array, *, scale: float,  mask: None | str | array = None, sinks: array | None = None, force_fused: bool = False, stream: StreamOrDevice = None) -> array"),
       R"pbdoc(
         A fast implementation of multi-head attention: ``O = softmax(Q @ K.T, dim=-1) @ V``.
 
@@ -313,6 +331,13 @@ void init_fast(nb::module_& parent_module) {
                last query aligns with the last key.
             sinks (array, optional): An optional array of attention sinks.
                Default: ``None``.
+            force_fused (bool, optional): If ``True``, bypass the routing
+               heuristics and always use the fused kernel, raising an error
+               when no fused kernel supports the given configuration
+               (unsupported head dims, mask constraints, or a CPU stream).
+               Whether the single-query (vector) or the full attention kernel
+               runs still follows the query length as with default routing.
+               Default: ``False``.
 
         Returns:
             array: The output array.

python/tests/test_fast_sdpa.py

@@ -728,5 +728,68 @@ def test_sdpa_sliced(self):
                     self.assertTrue(mx.allclose(ref, out, **tolerance))
 
 
+class TestSDPAForceFused(mlx_tests.MLXTestCase):
+    @unittest.skipIf(not mx.is_available(mx.gpu), "GPU kernel path only")
+    def test_force_fused_matches_reference(self):
+        mx.random.seed(0)
+        # Vector kernel (qL 1) and full kernel (qL 32) shapes, with GQA
+        for B, qL, kL, D, qH, kH in [
+            (1, 1, 128, 64, 8, 8),
+            (1, 1, 128, 128, 8, 2),
+            (1, 32, 128, 64, 8, 8),
+            (1, 32, 128, 128, 8, 2),
+        ]:
+            q = mx.random.normal((B, qH, qL, D), mx.float16)
+            k = mx.random.normal((B, kH, kL, D), mx.float16)
+            v = mx.random.normal((B, kH, kL, D), mx.float16)
+            scale = D**-0.5
+            for mask in (None, "causal"):
+                ref = mlx_ref_attn(q, k, v, scale=scale, mask=mask)
+                out = mx.fast.scaled_dot_product_attention(
+                    q, k, v, scale=scale, mask=mask, force_fused=True
+                )
+                self.assertTrue(mx.allclose(ref, out, rtol=1e-2, atol=1e-2))
+
+    @unittest.skipIf(not mx.metal.is_available(), "Metal support matrix")
+    def test_force_fused_unsupported_raises(self):
+        def make(qL, kL, D, qH=8, kH=8):
+            q = mx.random.normal((1, qH, qL, D), mx.float16)
+            k = mx.random.normal((1, kH, kL, D), mx.float16)
+            v = mx.random.normal((1, kH, kL, D), mx.float16)
+            return q, k, v
+
+        # Head dim unsupported by the full kernel
+        q, k, v = make(32, 128, 72)
+        with self.assertRaises(ValueError):
+            mx.fast.scaled_dot_product_attention(q, k, v, scale=1.0, force_fused=True)
+        # Head dim unsupported by the vector kernel
+        q, k, v = make(1, 128, 72)
+        with self.assertRaises(ValueError):
+            mx.fast.scaled_dot_product_attention(q, k, v, scale=1.0, force_fused=True)
+        # Vector kernel: query length times GQA factor over 32
+        q, k, v = make(8, 128, 64, qH=8, kH=1)
+        with self.assertRaises(ValueError):
+            mx.fast.scaled_dot_product_attention(q, k, v, scale=1.0, force_fused=True)
+        # Full kernel: causal mask with more queries than keys
+        q, k, v = make(32, 16, 64)
+        with self.assertRaises(ValueError):
+            mx.fast.scaled_dot_product_attention(
+                q, k, v, scale=1.0, mask="causal", force_fused=True
+            )
+        # The same shapes still run on the default (fallback) path
+        q, k, v = make(32, 128, 72)
+        mx.eval(mx.fast.scaled_dot_product_attention(q, k, v, scale=1.0))
+
+    def test_force_fused_raises_on_cpu(self):
+        q = mx.random.normal((1, 8, 1, 64))
+        k = mx.random.normal((1, 8, 128, 64))
+        v = mx.random.normal((1, 8, 128, 64))
+        with mx.stream(mx.cpu):
+            with self.assertRaises(ValueError):
+                mx.fast.scaled_dot_product_attention(
+                    q, k, v, scale=1.0, force_fused=True
+                )
+
+
 if __name__ == "__main__":
     mlx_tests.MLXTestRunner(failfast=True)