diff --git a/mlx/backend/cuda/scaled_dot_product_attention.cpp b/mlx/backend/cuda/scaled_dot_product_attention.cpp index 54700bdc..1cec0aa9 100644 --- a/mlx/backend/cuda/scaled_dot_product_attention.cpp +++ b/mlx/backend/cuda/scaled_dot_product_attention.cpp @@ -57,6 +57,58 @@ void malloc_with_same_layout( {true, false, false}); } +bool use_cudnn_for_decoding( + const array& q, + const array& k, + const array& v, + bool has_arr_mask) { + if (q.shape(2) != 1) { + return false; + } + if (has_arr_mask) { + return false; + } + // The cuDNN SDPA is faster than vector kernel but for small sequence the + // overhead would kill the advantage. + constexpr int kv_cache_step = 256; // number is from mlx-lm + if (k.shape(2) < kv_cache_step) { + return false; + } + // When called during graph building the strides is not available, and we + // rely on |supports_sdpa_vector| to decide whether to use fast sdpa since + // we can fallback to |sdpa_vector|. + if ((k.status() != array::evaluated) || (v.status() != array::evaluated)) { + return false; + } + // Check if k/v are slices from fixed-size kv cache. + auto is_slice = [](const array& kv) { + // Get pre-sliced sequence length from strides, and check if the buffer + // belongs to a contiguous kv cache. + int64_t T_kv = kv.strides(1) / kv.strides(2); + if (kv.size() / kv.shape(2) * T_kv != kv.buffer_size() / kv.itemsize()) { + return false; + } + // It is possible to use heuristic to check slices, but for now just make + // mlx-lm work. + return T_kv % kv_cache_step == 0; + }; + return is_slice(k) && is_slice(v); +} + +// Get original kv from slices, i.e. undo keys[..., :offset, :] +array unslice_kv(const array& kv) { + Shape shape = kv.shape(); + shape[2] = /* T_kv */ kv.strides(1) / kv.strides(2); + array copy(shape, kv.dtype(), nullptr, {}); + copy.copy_shared_buffer( + kv, + make_contiguous_strides(shape), + {true, true, false}, + /* data_size */ kv.buffer_size() / kv.itemsize(), + /* offset */ -kv.offset()); + return copy; +} + constexpr int QKV_NDIM = 4; struct SDPACacheKey { @@ -81,7 +133,8 @@ inline BytesKey build_sdpa_cache_key( const array& v, bool do_causal, const std::optional& mask_arr, - bool output_logsumexp = true) { + bool decoding = false, + bool output_logsumexp = false) { BytesKey cache_key; cache_key.pod = { encoder.device().cuda_device(), @@ -101,12 +154,19 @@ inline BytesKey build_sdpa_cache_key( cache_key.pod.mask_shape = vector_key(mask_arr->shape()); cache_key.pod.mask_strides = vector_key(mask_arr->strides()); } + if (decoding) { + int64_t T_kv = k.strides(1) / k.strides(2); + cache_key.pod.k_shape[2] = T_kv; + cache_key.pod.v_shape[2] = T_kv; + cache_key.pod.k_strides.fill(0); + cache_key.pod.v_strides.fill(0); + } return cache_key; } auto& sdpa_cache() { static LRUBytesKeyCache cache( - "MLX_CUDA_SDPA_CACHE_SIZE", /* default_capacity */ 64); + "MLX_CUDA_SDPA_CACHE_SIZE", /* default_capacity */ 256); return cache; } @@ -122,6 +182,8 @@ enum UIDS { V, SCALE, BIAS, + SEQ_LEN_Q, + SEQ_LEN_KV, O, STATS, // Backward graph: @@ -138,6 +200,8 @@ DnnGraph build_sdpa_graph( const array& v, bool do_causal, const std::optional& mask_arr, + const std::optional& seq_len_q, + const std::optional& seq_len_kv, bool output_logsumexp, const array& o, const std::optional& stats) { @@ -157,6 +221,11 @@ DnnGraph build_sdpa_graph( if (mask_arr) { options.set_bias(graph.tensor("BIAS", BIAS, *mask_arr)); } + if (seq_len_q && seq_len_kv) { + options.set_padding_mask(true); + options.set_seq_len_q(graph.tensor("SEQ_LEN_Q", SEQ_LEN_Q, *seq_len_q)); + options.set_seq_len_kv(graph.tensor("SEQ_LEN_KV", SEQ_LEN_KV, *seq_len_kv)); + } auto [o_, stats_] = graph.sdpa(q_, k_, v_, options); graph.tensor(o_, O, o)->set_output(true); @@ -222,9 +291,10 @@ bool supports_sdpa_cudnn( const array& q, const array& k, const array& v, + bool has_arr_mask, bool do_causal, Stream s) { - static bool enabled = env::get_var("MLX_CUDA_USE_CUDNN_SPDA", 1); + static bool enabled = env::get_var("MLX_CUDA_USE_CUDNN_SDPA", 1); if (!enabled) { return false; } @@ -234,8 +304,8 @@ bool supports_sdpa_cudnn( return false; } - // Only use cuDNN for prefilling (T_q > 1) and training (T_q == T_kv). - if ((q.shape(2) == 1) && (q.shape(2) != k.shape(2))) { + // Only use cuDNN for decoding when k/v are slices from fixed-size kv cache. + if ((q.shape(2) == 1) && !use_cudnn_for_decoding(q, k, v, has_arr_mask)) { return false; } @@ -256,8 +326,8 @@ bool supports_sdpa_cudnn( void sdpa_cudnn( const array& q, - const array& k, - const array& v, + array k, + array v, float scale, array& o, std::optional& stats, @@ -270,6 +340,24 @@ void sdpa_cudnn( malloc_with_same_layout(encoder, o, q); + // For decoding, unslice k/v and apply padding mask. + std::optional seq_len_q; + std::optional seq_len_kv; + bool decoding = use_cudnn_for_decoding(q, k, v, mask_arr.has_value()); + if (decoding) { + int B = q.shape(0); + std::vector seq_len_q_vec(B, q.shape(2)); + std::vector seq_len_kv_vec(B, k.shape(2)); + seq_len_q = array(seq_len_q_vec.begin(), {B, 1, 1, 1}); + seq_len_kv = array(seq_len_kv_vec.begin(), {B, 1, 1, 1}); + encoder.add_temporary(*seq_len_q); + encoder.add_temporary(*seq_len_kv); + k = unslice_kv(k); + v = unslice_kv(v); + encoder.add_temporary(k); + encoder.add_temporary(v); + } + encoder.set_input_array(q); encoder.set_input_array(k); encoder.set_input_array(v); @@ -277,6 +365,10 @@ void sdpa_cudnn( if (mask_arr) { encoder.set_input_array(*mask_arr); } + if (seq_len_q && seq_len_kv) { + encoder.set_input_array(*seq_len_q); + encoder.set_input_array(*seq_len_kv); + } if (output_logsumexp) { stats->set_data(cu::malloc_async(stats->nbytes(), encoder)); encoder.set_output_array(*stats); @@ -284,11 +376,21 @@ void sdpa_cudnn( // Search cache. auto cache_key = build_sdpa_cache_key( - encoder, q, k, v, do_causal, mask_arr, output_logsumexp); + encoder, q, k, v, do_causal, mask_arr, decoding, output_logsumexp); auto it = sdpa_cache().find(cache_key); if (it == sdpa_cache().end()) { auto graph = build_sdpa_graph( - handle, q, k, v, do_causal, mask_arr, output_logsumexp, o, stats); + handle, + q, + k, + v, + do_causal, + mask_arr, + seq_len_q, + seq_len_kv, + output_logsumexp, + o, + stats); it = sdpa_cache().emplace(cache_key, std::move(graph)).first; } auto& graph = it->second; @@ -302,6 +404,10 @@ void sdpa_cudnn( if (mask_arr) { variant_pack[BIAS] = gpu_ptr(*mask_arr); } + if (seq_len_q && seq_len_kv) { + variant_pack[SEQ_LEN_Q] = gpu_ptr(*seq_len_q); + variant_pack[SEQ_LEN_KV] = gpu_ptr(*seq_len_kv); + } if (output_logsumexp) { variant_pack[STATS] = gpu_ptr(*stats); } @@ -376,9 +482,7 @@ bool supports_sdpa_vector( const array& q, const array& k, const array& v, - bool has_mask, bool has_arr_mask, - bool do_causal, bool output_logsumexp); void sdpa_vector( const array& q, @@ -406,9 +510,8 @@ bool ScaledDotProductAttention::use_fallback( return true; } - return !supports_sdpa_vector( - q, k, v, has_mask, has_arr_mask, do_causal, output_logsumexp) && - !supports_sdpa_cudnn(q, k, v, do_causal, s); + return !supports_sdpa_cudnn(q, k, v, has_arr_mask, do_causal, s) && + !supports_sdpa_vector(q, k, v, has_arr_mask, output_logsumexp); } bool ScaledDotProductAttention::supports_bool_mask() { @@ -438,14 +541,7 @@ void ScaledDotProductAttention::eval_gpu( stats = outputs[1]; } - if (supports_sdpa_vector( - q, k, v, has_mask, has_arr_mask, do_causal_, output_logsumexp_)) { - if (has_sinks_) { - sdpa_vector(q, k, v, scale_, out, do_causal_, inputs.back(), s); - } else { - sdpa_vector(q, k, v, scale_, out, do_causal_, std::nullopt, s); - } - } else { + if (supports_sdpa_cudnn(q, k, v, has_arr_mask, do_causal_, s)) { sdpa_cudnn( q, k, @@ -457,6 +553,12 @@ void ScaledDotProductAttention::eval_gpu( mask_arr, output_logsumexp_, s); + } else { + if (has_sinks_) { + sdpa_vector(q, k, v, scale_, out, do_causal_, inputs.back(), s); + } else { + sdpa_vector(q, k, v, scale_, out, do_causal_, std::nullopt, s); + } } } diff --git a/mlx/backend/cuda/scaled_dot_product_attention.cu b/mlx/backend/cuda/scaled_dot_product_attention.cu index 111cfee1..df8df3e3 100644 --- a/mlx/backend/cuda/scaled_dot_product_attention.cu +++ b/mlx/backend/cuda/scaled_dot_product_attention.cu @@ -665,9 +665,7 @@ bool supports_sdpa_vector( const array& q, const array& k, const array& v, - bool has_mask, bool has_arr_mask, - bool do_causal, bool output_logsumexp) { if (output_logsumexp) { return false; diff --git a/python/tests/test_fast_sdpa.py b/python/tests/test_fast_sdpa.py index 401e53d7..1a3bd167 100644 --- a/python/tests/test_fast_sdpa.py +++ b/python/tests/test_fast_sdpa.py @@ -605,6 +605,38 @@ class TestFastSDPA(mlx_tests.MLXTestCase): ).sum() test_grad(loss_slow, loss_fast, [q, k, v]) + def test_sdpa_sliced(self): + N = 8 + D = 64 + scale = D**-0.5 + + for B, T_q, T_kv, offset, mask in product( + (1, 2, 4), + (1, 8), + (256, 512), + (8, 9, 64, 79), + (None, "causal"), + ): + with self.subTest(B=B, T_q=T_q, T_kv=T_kv, offset=offset, mask=mask): + q = mx.random.normal((B, N, T_q, D), mx.float16) + k = mx.random.normal((B, N, T_kv, D), mx.float16) + v = mx.random.normal((B, N, T_kv, D), mx.float16) + + k = k[..., :offset, :] + v = v[..., :offset, :] + + ref = mlx_ref_attn(q, k, v, scale=scale, mask=mask) + + for i in range(2): + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, mask=mask + ) + if B == 1: + tolerance = {"rtol": 1e-3, "atol": 1e-3} + else: + tolerance = {"rtol": 1e-2, "atol": 1e-2} + self.assertTrue(mx.allclose(ref, out, **tolerance)) + if __name__ == "__main__": mlx_tests.MLXTestRunner(failfast=True)