diff --git a/mlx/backend/cuda/scaled_dot_product_attention.cpp b/mlx/backend/cuda/scaled_dot_product_attention.cpp index ca411e91c6..37aae3701f 100644 --- a/mlx/backend/cuda/scaled_dot_product_attention.cpp +++ b/mlx/backend/cuda/scaled_dot_product_attention.cpp @@ -8,6 +8,8 @@ #include +#include + 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; } diff --git a/mlx/backend/metal/scaled_dot_product_attention.cpp b/mlx/backend/metal/scaled_dot_product_attention.cpp index cb4af8523e..cb444ca21c 100644 --- a/mlx/backend/metal/scaled_dot_product_attention.cpp +++ b/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() { diff --git a/mlx/backend/no_gpu/primitives.cpp b/mlx/backend/no_gpu/primitives.cpp index b7d7a19467..733a3d2789 100644 --- a/mlx/backend/no_gpu/primitives.cpp +++ b/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; } diff --git a/mlx/fast.cpp b/mlx/fast.cpp index a668fe9abd..402293947e 100644 --- a/mlx/fast.cpp +++ b/mlx/fast.cpp @@ -618,6 +618,7 @@ array scaled_dot_product_attention( const std::string& mask_mode /* = "" */, std::optional mask_arr /* = {} */, const std::optional& 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, diff --git a/mlx/fast.h b/mlx/fast.h index 934fadc2b7..c5f664df79 100644 --- a/mlx/fast.h +++ b/mlx/fast.h @@ -53,6 +53,7 @@ MLX_API array scaled_dot_product_attention( const std::string& mask_mode = "", std::optional mask_arr = {}, const std::optional& sinks = {}, + bool force_fused = false, StreamOrDevice s = {}); using TemplateArg = std::variant; diff --git a/mlx/fast_primitives.h b/mlx/fast_primitives.h index 0d2f861045..da7e55f8c8 100644 --- a/mlx/fast_primitives.h +++ b/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& inputs, std::vector& outputs) override { throw std::runtime_error("NYI"); diff --git a/python/src/fast.cpp b/python/src/fast.cpp index e59357bc33..c2b8b66583 100644 --- a/python/src/fast.cpp +++ b/python/src/fast.cpp @@ -234,6 +234,7 @@ void init_fast(nb::module_& parent_module) { const float scale, const std::variant& mask, const std::optional& sinks, + bool force_fused, mx::StreamOrDevice s) { bool has_mask = !std::holds_alternative(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(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. diff --git a/python/tests/test_fast_sdpa.py b/python/tests/test_fast_sdpa.py index 2418997bc8..8bc29b8b1a 100644 --- a/python/tests/test_fast_sdpa.py +++ b/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)