Add force_fused option to scaled_dot_product_attention

mlx/backend/cuda/scaled_dot_product_attention.cpp

@@ -549,6 +549,33 @@ void sdpa_vector(
 
 namespace fast {
 
+namespace {
+
+std::tuple<bool, std::string> has_fused_kernel(
+    const array& q,
+    const array& k,
+    const array& v,
+    bool has_arr_mask,
+    bool do_causal,
+    bool output_logsumexp,
+    Stream s) {
+  if (s.device != Device::gpu) {
+    return {false, "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)) {
+    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 {false, msg.str()};
+  }
+  return {true, ""};
+}
+
+} // namespace
+
 bool ScaledDotProductAttention::use_fallback(
     const array& q,
     const array& k,
@@ -558,13 +585,21 @@ bool ScaledDotProductAttention::use_fallback(
     bool do_causal,
     bool is_training,
     bool output_logsumexp,
+    bool force_fused,
     Stream s) {
-  if (s.device == Device::cpu) {
-    return true;
+  auto [has_fused, reason] =
+      has_fused_kernel(q, k, v, has_arr_mask, do_causal, output_logsumexp, s);
+  if (force_fused) {
+    if (!has_fused) {
+      std::ostringstream msg;
+      msg << "[scaled_dot_product_attention] force_fused=True but no fused "
+             "kernel is available: "
+          << reason;
+      throw std::invalid_argument(msg.str());
+    }
+    return false;
   }
-
-  return !supports_sdpa_cudnn(q, k, v, has_arr_mask, do_causal, s) &&
-      !supports_sdpa_vector(q, k, v, has_arr_mask, output_logsumexp);
+  return !has_fused;
 }
 
 bool ScaledDotProductAttention::supports_bool_mask() {

mlx/backend/metal/kernels/scaled_dot_product_attention.metal

@@ -33,10 +33,12 @@ using namespace metal;
   instantiate_sdpa_vector(type, 96, 96)          \
   instantiate_sdpa_vector(type, 128, 128)        \
   instantiate_sdpa_vector(type, 192, 128)        \
+  instantiate_sdpa_vector(type, 192, 192)        \
   instantiate_sdpa_vector(type, 256, 256)        \
   instantiate_sdpa_vector_aggregation(type, 64)  \
   instantiate_sdpa_vector_aggregation(type, 96)  \
   instantiate_sdpa_vector_aggregation(type, 128) \
+  instantiate_sdpa_vector_aggregation(type, 192) \
   instantiate_sdpa_vector_aggregation(type, 256)
 
 instantiate_sdpa_vector_heads(float)

mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.h

@@ -428,7 +428,7 @@ template <
       for (short id = 0; id < TD; id++) {
         STEEL_PRAGMA_UNROLL
         for (short ik = 0; ik < TK; ik++) {
-          if constexpr (BD == 128) {
+          if constexpr (BD >= 128) {
             simdgroup_barrier(mem_flags::mem_none);
           }
 
@@ -438,7 +438,7 @@ template <
           Vtile.template load<T, 1, 1, LDV_tgp, 1>(
               &Vs[Vs_offset + kk * LDV_tgp + dd]);
 
-          if constexpr (BD == 128) {
+          if constexpr (BD >= 128) {
             simdgroup_barrier(mem_flags::mem_none);
           }
 

mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.metal

@@ -12,6 +12,8 @@
   attention, dtype, bq, bk, bd, wm, wn, mtype, float)
 
 #define instantiate_attn_shapes_helper(iname, itype, mname, mtype)  \
+    instantiate_attn(iname, itype, 32, 16, 256, 4, 1, mname, mtype) \
+    instantiate_attn(iname, itype, 32, 16, 192, 4, 1, mname, mtype) \
     instantiate_attn(iname, itype, 32, 16, 128, 4, 1, mname, mtype) \
     instantiate_attn(iname, itype, 32, 32,  96, 4, 1, mname, mtype) \
     instantiate_attn(iname, itype, 32, 32,  80, 4, 1, mname, mtype) \

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);
 }
 
@@ -174,9 +175,8 @@ void sdpa_full_self_attention_metal(
     bool do_causal_,
     const std::optional<array>& mask,
     const std::optional<array>& sinks) {
-  // NAX tiles the head dim in units of kU=16 and steps TD by 2, so it needs
-  // a multiple of 32; 72 and 80 take the classic steel kernel.
-  if (metal::is_nax_available() && q.shape(3) != 80 && q.shape(3) != 72 &&
+  if (metal::is_nax_available() &&
+      (q.shape(3) == 64 || q.shape(3) == 96 || q.shape(3) == 128) &&
       (env::enable_tf32() || q.dtype() != float32)) {
     return sdpa_full_self_attention_nax(
         /* const Stream& s = */ s,
@@ -325,6 +325,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);
 }
 
@@ -591,28 +592,23 @@ void sdpa_vector_2pass(
   compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
 }
 
-} // namespace
-
-bool ScaledDotProductAttention::use_fallback(
+std::tuple<bool, std::string> has_fused_kernel(
     const array& q,
     const array& k,
     const array& v,
     bool has_mask,
     bool has_arr_mask,
     bool do_causal,
-    bool is_training,
     bool output_logsumexp,
     Stream s) {
-  if (is_training) {
-    // It's faster for training on Metal to use the unfused SDPA for both
-    // forward and backward.
-    return true;
+  if (s.device != Device::gpu) {
+    return {false, "the fused kernels require a GPU (Metal) stream."};
   }
   if (output_logsumexp) {
-    return true;
-  }
-  if (s.device == Device::cpu) {
-    return true;
+    return {
+        false,
+        "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);
@@ -623,27 +619,103 @@ 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 == 72 || query_head_dim == 80 ||
-       query_head_dim == 96 || query_head_dim == 128);
+  std::ostringstream msg;
+  if (query_sequence_length > 8) {
+    const bool supported_head_dim = query_head_dim == value_head_dim &&
+        (query_head_dim == 64 || query_head_dim == 72 || query_head_dim == 80 ||
+         query_head_dim == 96 || query_head_dim == 128 ||
+         query_head_dim == 192 || query_head_dim == 256);
+    if (!supported_head_dim) {
+      msg << "the full attention kernel supports head dims "
+          << "{64, 72, 80, 96, 128, 192, 256} with matching query/value head "
+          << "dims; got query head dim " << query_head_dim
+          << " and value head dim " << value_head_dim << ".";
+      return {false, 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 {false, msg.str()};
+    }
+  } else {
+    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 == 192 ||
+          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, 192, 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 {false, 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 {false, 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 {false, msg.str()};
+    }
+  }
+  return {true, ""};
+}
 
-  const bool sdpa_full_supported_mask = !has_mask || has_arr_mask ||
-      (query_sequence_length <= key_sequence_length && do_causal);
+} // namespace
 
-  const bool supports_sdpa_full = query_sequence_length > 8 &&
-      sdpa_full_supported_mask && sdpa_full_supported_head_dim;
+bool ScaledDotProductAttention::use_fallback(
+    const array& q,
+    const array& k,
+    const array& v,
+    bool has_mask,
+    bool has_arr_mask,
+    bool do_causal,
+    bool is_training,
+    bool output_logsumexp,
+    bool force_fused,
+    Stream s) {
+  auto [has_fused, reason] = has_fused_kernel(
+      q, k, v, has_mask, has_arr_mask, do_causal, output_logsumexp, s);
+  if (force_fused) {
+    if (!has_fused) {
+      std::ostringstream msg;
+      msg << "[scaled_dot_product_attention] force_fused=True but no fused "
+             "kernel is available: "
+          << reason;
+      throw std::invalid_argument(msg.str());
+    }
+    return false;
+  }
 
-  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;
+  if (is_training) {
+    // It's faster for training on Metal to use the unfused SDPA for both
+    // forward and backward.
+    return true;
+  }
+  if (!has_fused) {
+    return true;
+  }
 
-  return !(supports_sdpa_full || supports_sdpa_vector);
+  // Unfused path is faster for following shapes.
+  const int query_sequence_length = q.shape(2);
+  const int query_head_dim = q.shape(-1);
+  const int value_head_dim = v.shape(-1);
+  if (query_sequence_length > 8) {
+    return query_head_dim == 192 || query_head_dim == 256;
+  } else {
+    return query_head_dim == value_head_dim && query_head_dim == 192;
+  }
 }
 
 bool ScaledDotProductAttention::supports_bool_mask() {

mlx/backend/no_gpu/primitives.cpp

@@ -32,20 +32,26 @@ bool fast::ScaledDotProductAttention::use_fallback(
     bool do_causal,
     bool is_training,
     bool output_logsumexp,
+    bool force_fused,
     Stream s) {
+  if (force_fused) {
+    throw std::invalid_argument(
+        "[scaled_dot_product_attention] force_fused=True but no fused "
+        "kernel is available in CPU backend.");
+  }
   return true;
 }
 
-bool fast::ScaledDotProductAttention::supports_bool_mask() {
-  return false;
-}
-
 bool fast::ScaledDotProductAttentionVJP::use_fallback(
     const array& q,
     Stream s) {
   return true;
 }
 
+bool fast::ScaledDotProductAttention::supports_bool_mask() {
+  return false;
+}
+
 NO_GPU(Abs)
 NO_GPU(Add)
 NO_GPU(AddMM)

mlx/fast.cpp

@@ -618,7 +618,8 @@ array scaled_dot_product_attention(
     const std::string& mask_mode /* = "" */,
     std::optional<array> mask_arr /* = {} */,
     const std::optional<array>& sinks /* = {} */,
-    StreamOrDevice s /* = {}*/) {
+    bool force_fused /* = false */,
+    StreamOrDevice s /* = {} */) {
   for (const auto& tensor : {queries, keys, values}) {
     if (tensor.ndim() != 4) {
       std::ostringstream msg;
@@ -834,6 +835,7 @@ array scaled_dot_product_attention(
           do_causal,
           is_training,
           output_logsumexp,
+          force_fused,
           stream)) {
     if (has_bool_mask && !ScaledDotProductAttention::supports_bool_mask()) {
       // Convert bool mask to additive mask.
@@ -846,7 +848,13 @@ array scaled_dot_product_attention(
     }
     Shape out_shape{q.shape(0), q.shape(1), q.shape(2), v.shape(-1)};
     auto primitive = std::make_shared<ScaledDotProductAttention>(
-        stream, fallback, scale, do_causal, has_sinks, output_logsumexp);
+        stream,
+        fallback,
+        scale,
+        do_causal,
+        has_sinks,
+        output_logsumexp,
+        force_fused);
     if (output_logsumexp) {
       return array::make_arrays(
           {std::move(out_shape), Shape{q.shape(0), q.shape(1), q.shape(2), 1}},
@@ -912,7 +920,8 @@ bool ScaledDotProductAttention::is_equivalent(const Primitive& other) const {
       static_cast<const ScaledDotProductAttention&>(other);
   return scale_ == a_other.scale_ && do_causal_ == a_other.do_causal_ &&
       has_sinks_ == a_other.has_sinks_ &&
-      output_logsumexp_ == a_other.output_logsumexp_;
+      output_logsumexp_ == a_other.output_logsumexp_ &&
+      force_fused_ == a_other.force_fused_;
 }
 
 bool ScaledDotProductAttentionVJP::is_equivalent(const Primitive& other) const {

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>;