diff --git a/mlx/backend/cuda/scaled_dot_product_attention.cpp b/mlx/backend/cuda/scaled_dot_product_attention.cpp index ca411e91c6..286d500e2c 100644 --- a/mlx/backend/cuda/scaled_dot_product_attention.cpp +++ b/mlx/backend/cuda/scaled_dot_product_attention.cpp @@ -549,6 +549,33 @@ void sdpa_vector( namespace fast { +namespace { + +std::tuple 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() { diff --git a/mlx/backend/metal/kernels/scaled_dot_product_attention.metal b/mlx/backend/metal/kernels/scaled_dot_product_attention.metal index 44ccb834e9..84486d62c3 100644 --- a/mlx/backend/metal/kernels/scaled_dot_product_attention.metal +++ b/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) diff --git a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.h b/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.h index 0d9628e834..29fa7ba396 100644 --- a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.h +++ b/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( &Vs[Vs_offset + kk * LDV_tgp + dd]); - if constexpr (BD == 128) { + if constexpr (BD >= 128) { simdgroup_barrier(mem_flags::mem_none); } diff --git a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.metal b/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.metal index 4bb9ff5873..fbd84004f0 100644 --- a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.metal +++ b/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) \ diff --git a/mlx/backend/metal/scaled_dot_product_attention.cpp b/mlx/backend/metal/scaled_dot_product_attention.cpp index bb8ad808b8..acff685790 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); } @@ -174,9 +175,8 @@ void sdpa_full_self_attention_metal( bool do_causal_, const std::optional& mask, const std::optional& 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 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() { diff --git a/mlx/backend/no_gpu/primitives.cpp b/mlx/backend/no_gpu/primitives.cpp index b7d7a19467..7f60d0d83a 100644 --- a/mlx/backend/no_gpu/primitives.cpp +++ b/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) diff --git a/mlx/fast.cpp b/mlx/fast.cpp index a668fe9abd..df2beebd84 100644 --- a/mlx/fast.cpp +++ b/mlx/fast.cpp @@ -618,7 +618,8 @@ array scaled_dot_product_attention( const std::string& mask_mode /* = "" */, std::optional mask_arr /* = {} */, const std::optional& 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( - 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(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 { 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..61a392e418 100644 --- a/mlx/fast_primitives.h +++ b/mlx/fast_primitives.h @@ -212,12 +212,14 @@ class ScaledDotProductAttention : public Custom { float scale, bool do_causal, bool has_sinks, - bool output_logsumexp) + bool output_logsumexp, + bool force_fused) : Custom(stream, std::move(fallback)), scale_(scale), do_causal_(do_causal), has_sinks_(has_sinks), - output_logsumexp_(output_logsumexp) {} + output_logsumexp_(output_logsumexp), + force_fused_(force_fused) {} static bool use_fallback( const array& q, @@ -228,6 +230,7 @@ class ScaledDotProductAttention : public Custom { bool do_causal, bool is_training, bool output_logsumexp, + bool force_fused, Stream s); static bool supports_bool_mask(); @@ -251,7 +254,12 @@ class ScaledDotProductAttention : public Custom { DEFINE_INPUT_OUTPUT_SHAPE() auto state() const { return std::make_tuple( - nullptr, scale_, do_causal_, has_sinks_, output_logsumexp_); + nullptr, + scale_, + do_causal_, + has_sinks_, + output_logsumexp_, + force_fused_); } private: @@ -259,6 +267,7 @@ class ScaledDotProductAttention : public Custom { bool do_causal_; bool has_sinks_; bool output_logsumexp_; + bool force_fused_; }; class ScaledDotProductAttentionVJP : public Custom { diff --git a/python/src/fast.cpp b/python/src/fast.cpp index e59357bc33..0a50dc79cd 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,11 @@ 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``, use a fused kernel + regardless of the builtin heuristics and raise error when no + fused kernel is available. For certain configurations this would + result in slower kernel getting used but can reduce memory + consumption. Default: ``False``. Returns: array: The output array. diff --git a/python/tests/test_fast_sdpa.py b/python/tests/test_fast_sdpa.py index a5fdb0fbb7..997fa1028c 100644 --- a/python/tests/test_fast_sdpa.py +++ b/python/tests/test_fast_sdpa.py @@ -722,6 +722,100 @@ def test_grad(slow, fast, args): ).sum() test_grad(loss_slow, loss_fast, [q, k, v]) + @unittest.skipIf(not mx.metal.is_available(), "Metal kernel path only") + def test_sdpa_force_fused_metal(self): + if mx.default_device() != mx.gpu: + self.skipTest("requires GPU") + + def make_qkv(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 + + # Full attention kernel. + for D, qL, mask in product((192, 256), (9, 16), (None, "causal")): + with self.subTest(head_dim=D, qL=qL, mask=mask): + q, k, v = make_qkv(qL, 512, D, 8, 4) + scale = D**-0.5 + 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, atol=1e-3, rtol=1e-3)) + + # Vector attention kernel. + for D in (192, 256): + with self.subTest(head_dim=D): + q, k, v = make_qkv(4, 16385, D, 4, 2) + scale = D**-0.5 + ref = mlx_ref_attn(q, k, v, scale=scale) + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, force_fused=True + ) + self.assertTrue(mx.allclose(ref, out, atol=1e-3, rtol=1e-3)) + + # No full attention fused kernels. + with self.assertRaisesRegex(ValueError, "supports head dims"): + q, k, v = make_qkv(16, 512, 512) + mx.fast.scaled_dot_product_attention( + q, k, v, scale=512**-0.5, force_fused=True + ) + with self.assertRaisesRegex( + ValueError, "query sequence to be no longer than the key sequence" + ): + q, k, v = make_qkv(32, 16, 64) + mx.fast.scaled_dot_product_attention( + q, + k, + v, + scale=64**-0.5, + mask="causal", + force_fused=True, + ) + + # No vector attention fused kernels. + with self.assertRaisesRegex(ValueError, "supports head dims"): + q, k, v = make_qkv(1, 128, 72) + mx.fast.scaled_dot_product_attention( + q, k, v, scale=72**-0.5, force_fused=True + ) + with self.assertRaisesRegex(ValueError, "GQA factor to be at most 32"): + q, k, v = make_qkv(8, 128, 64, qH=8, kH=1) + mx.fast.scaled_dot_product_attention( + q, k, v, scale=64**-0.5, force_fused=True + ) + + # No CPU fused kernel. + with mx.stream(mx.cpu): + q, k, v = make_qkv(8, 128, 8) + with self.assertRaisesRegex(ValueError, "require a GPU"): + mx.fast.scaled_dot_product_attention( + q, k, v, scale=64**-0.5, force_fused=True + ) + + @unittest.skipIf(not mx.cuda.is_available(), "CUDA kernel path only") + def test_sdpa_force_fused_cuda(self): + if mx.default_device() != mx.gpu: + self.skipTest("requires GPU") + + def make_qkv(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 + + # Vector attention kernel. + for D in (64, 96, 128): + with self.subTest(head_dim=D): + q, k, v = make_qkv(3, 128, D, 4, 2) + scale = D**-0.5 + ref = mlx_ref_attn(q, k, v, scale=scale) + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, force_fused=True + ) + self.assertTrue(mx.allclose(ref, out, atol=1e-3, rtol=1e-3)) + def test_sdpa_sliced(self): N = 8 D = 64