diff --git a/mlx/backend/cuda/scaled_dot_product_attention.cpp b/mlx/backend/cuda/scaled_dot_product_attention.cpp index 286d500e2c..e6c8abf122 100644 --- a/mlx/backend/cuda/scaled_dot_product_attention.cpp +++ b/mlx/backend/cuda/scaled_dot_product_attention.cpp @@ -1,4 +1,4 @@ -// Copyright © 2025 Apple Inc. +// Copyright © 2025-2026 Apple Inc. #include "mlx/backend/cuda/cudnn_utils.h" #include "mlx/backend/cuda/device.h" @@ -583,6 +583,7 @@ bool ScaledDotProductAttention::use_fallback( bool has_mask, bool has_arr_mask, bool do_causal, + bool has_sinks, bool is_training, bool output_logsumexp, bool force_fused, diff --git a/mlx/backend/metal/kernels/CMakeLists.txt b/mlx/backend/metal/kernels/CMakeLists.txt index 00f90ac862..987d55e885 100644 --- a/mlx/backend/metal/kernels/CMakeLists.txt +++ b/mlx/backend/metal/kernels/CMakeLists.txt @@ -54,7 +54,7 @@ build_kernel(layer_norm) build_kernel(random) build_kernel(rms_norm) build_kernel(rope) -build_kernel(scaled_dot_product_attention sdpa_vector.h) +build_kernel(scaled_dot_product_attention sdpa_blocked.h sdpa_vector.h) build_kernel(gated_delta_update gated_delta_update.h) if(MLX_METAL_VERSION GREATER_EQUAL 320) build_kernel(fence) diff --git a/mlx/backend/metal/kernels/scaled_dot_product_attention.metal b/mlx/backend/metal/kernels/scaled_dot_product_attention.metal index 68e2a6c06a..f445119256 100644 --- a/mlx/backend/metal/kernels/scaled_dot_product_attention.metal +++ b/mlx/backend/metal/kernels/scaled_dot_product_attention.metal @@ -2,10 +2,28 @@ // clang-format off #include "mlx/backend/metal/kernels/utils.h" +#include "mlx/backend/metal/kernels/sdpa_blocked.h" #include "mlx/backend/metal/kernels/sdpa_vector.h" using namespace metal; +instantiate_kernel( + "sdpa_blocked_scale_copy_bfloat16", + sdpa_blocked_scale_copy, + bfloat16_t) +instantiate_kernel( + "sdpa_blocked_scale_copy_float16", + sdpa_blocked_scale_copy, + float16_t) +instantiate_kernel( + "sdpa_blocked_causal_mask_bfloat16", + sdpa_blocked_causal_mask, + bfloat16_t) +instantiate_kernel( + "sdpa_blocked_causal_mask_float16", + sdpa_blocked_causal_mask, + float16_t) + // SDPA vector instantiations #define instantiate_sdpa_vector_aggregation(type, value_dim) \ instantiate_kernel( \ diff --git a/mlx/backend/metal/kernels/sdpa_blocked.h b/mlx/backend/metal/kernels/sdpa_blocked.h new file mode 100644 index 0000000000..be4e900c67 --- /dev/null +++ b/mlx/backend/metal/kernels/sdpa_blocked.h @@ -0,0 +1,45 @@ +// Copyright © 2026 Apple Inc. + +#pragma once + +#include + +#include "mlx/backend/metal/kernels/defines.h" + +using namespace metal; + +template +[[kernel]] void sdpa_blocked_scale_copy( + const device T* q [[buffer(0)]], + device T* q_block [[buffer(1)]], + constant float& scale [[buffer(2)]], + constant int64_t* q_strides [[buffer(3)]], + constant int& heads [[buffer(4)]], + constant int& rows [[buffer(5)]], + constant int& head_dim [[buffer(6)]], + constant int& q_offset [[buffer(7)]], + uint3 elem [[thread_position_in_grid]]) { + int d = elem.x; + int row = elem.y; + int head = elem.z % heads; + int batch = elem.z / heads; + int64_t q_index = int64_t(batch) * q_strides[0] + + int64_t(head) * q_strides[1] + int64_t(q_offset + row) * q_strides[2] + d; + size_t q_block_index = (size_t(elem.z) * rows + row) * head_dim + d; + q_block[q_block_index] = q[q_index] * static_cast(scale); +} + +template +[[kernel]] void sdpa_blocked_causal_mask( + device T* scores [[buffer(0)]], + constant int& tile_rows [[buffer(1)]], + constant int& columns [[buffer(2)]], + uint2 elem [[thread_position_in_grid]]) { + int column = elem.x; + int score_row = elem.y; + int row = score_row % tile_rows; + if (column > row) { + size_t index = size_t(score_row) * columns + columns - tile_rows + column; + scores[index] = Limits::finite_min; + } +} diff --git a/mlx/backend/metal/scaled_dot_product_attention.cpp b/mlx/backend/metal/scaled_dot_product_attention.cpp index 6ece8c43fb..cd1f12a532 100644 --- a/mlx/backend/metal/scaled_dot_product_attention.cpp +++ b/mlx/backend/metal/scaled_dot_product_attention.cpp @@ -1,4 +1,5 @@ -// Copyright © 2024 Apple Inc. +// Copyright © 2024-2026 Apple Inc. +#include #include #include "mlx/backend/common/compiled.h" @@ -7,6 +8,7 @@ #include "mlx/backend/metal/kernels.h" #include "mlx/backend/metal/kernels/defines.h" #include "mlx/backend/metal/kernels/steel/attn/params.h" +#include "mlx/backend/metal/matmul.h" #include "mlx/backend/metal/utils.h" #include "mlx/fast_primitives.h" #include "mlx/utils.h" @@ -15,6 +17,189 @@ namespace mlx::core::fast { namespace { +bool use_blocked_d512( + const array& q, + const array& k, + const array& v, + bool has_arr_mask, + bool do_causal, + bool has_sinks, + Stream s) { + return s.device == Device::gpu && !metal::is_nax_available() && + (q.dtype() == float16 || q.dtype() == bfloat16) && q.shape(-1) == 512 && + v.shape(-1) == 512 && q.shape(2) >= 1024 && q.shape(2) <= k.shape(2) && + do_causal && !has_arr_mask && !has_sinks; +} + +void sdpa_blocked_softmax(const Stream& s, metal::Device& d, array& scores) { + constexpr int looped_limit = 4096; + constexpr int simd_size = 32; + + int axis_size = scores.shape().back(); + int n_rows = scores.size() / axis_size; + std::string kernel_name = + axis_size > looped_limit ? "looped_softmax_" : "block_softmax_"; + if (scores.dtype() != float32) { + kernel_name += "precise_"; + } + kernel_name += type_to_name(scores); + + auto kernel = get_softmax_kernel(d, kernel_name, true, scores); + auto& compute_encoder = metal::get_command_encoder(s); + + size_t threadgroup_size; + if (axis_size <= looped_limit) { + size_t reads = (axis_size + SOFTMAX_N_READS - 1) / SOFTMAX_N_READS; + size_t simds = (reads + simd_size - 1) / simd_size; + threadgroup_size = simd_size * simds; + assert(threadgroup_size <= kernel->maxTotalThreadsPerThreadgroup()); + } else { + threadgroup_size = kernel->maxTotalThreadsPerThreadgroup(); + } + + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.set_input_array(scores, 0); + compute_encoder.set_output_array(scores, 1); + compute_encoder.set_bytes(axis_size, 2); + compute_encoder.dispatch_threads( + MTL::Size(n_rows * threadgroup_size, 1, 1), + MTL::Size(threadgroup_size, 1, 1)); +} + +void sdpa_full_self_attention_blocked( + const Stream& s, + metal::Device& d, + const array& q, + const array& k, + const array& v, + float scale, + array& o) { + // This block size balances score-buffer memory and GEMM efficiency. + constexpr int block_rows = 256; + + int B = q.shape(0); + int H = q.shape(1); + int qL = q.shape(2); + int kL = k.shape(2); + int D = q.shape(3); + int kvH = k.shape(1); + int gqa = H / kvH; + + array score_buffer({B, H, block_rows, kL}, q.dtype(), nullptr, {}); + score_buffer.set_data(allocator::malloc(score_buffer.nbytes())); + array q_buffer({B, H, block_rows, D}, q.dtype(), nullptr, {}); + q_buffer.set_data(allocator::malloc(q_buffer.nbytes())); + array out_buffer({B, H, block_rows, D}, o.dtype(), nullptr, {}); + out_buffer.set_data(allocator::malloc(out_buffer.nbytes())); + + auto& compute_encoder = metal::get_command_encoder(s); + auto scale_kernel = + d.get_kernel("sdpa_blocked_scale_copy_" + type_to_name(q)); + auto mask_kernel = + d.get_kernel("sdpa_blocked_causal_mask_" + type_to_name(q)); + Shape batch_shape = {B, kvH}; + Strides k_batch_strides = {k.strides(0), k.strides(1)}; + Strides v_batch_strides = {v.strides(0), v.strides(1)}; + + std::vector copies; + for (int q_offset = 0; q_offset < qL; q_offset += block_rows) { + int rows = std::min(block_rows, qL - q_offset); + // Skip columns that are masked for every row in this query block. + int columns = std::min(kL, kL - qL + q_offset + rows); + + array q_block({B, H, rows, D}, q.dtype(), nullptr, {}); + q_block.copy_shared_buffer( + q_buffer, q_block.strides(), q_block.flags(), q_block.size()); + compute_encoder.set_compute_pipeline_state(scale_kernel); + compute_encoder.set_input_array(q, 0); + compute_encoder.set_output_array(q_block, 1); + compute_encoder.set_bytes(scale, 2); + compute_encoder.set_vector_bytes(q.strides(), 3); + compute_encoder.set_bytes(H, 4); + compute_encoder.set_bytes(rows, 5); + compute_encoder.set_bytes(D, 6); + compute_encoder.set_bytes(q_offset, 7); + auto scale_group_dims = get_block_dims(D, rows, B * H); + compute_encoder.dispatch_threads( + MTL::Size(D, rows, B * H), scale_group_dims); + + array scores({B, H, rows, columns}, q.dtype(), nullptr, {}); + scores.copy_shared_buffer( + score_buffer, scores.strides(), scores.flags(), scores.size()); + + Strides q_batch_strides = {int64_t(H) * rows * D, int64_t(gqa) * rows * D}; + steel_matmul( + /* const Stream& s = */ s, + /* metal::Device& d = */ d, + /* const array& a = */ q_block, + /* const array& b = */ k, + /* array& out = */ scores, + /* int M = */ gqa * rows, + /* int N = */ columns, + /* int K = */ D, + /* int batch_size_out = */ B * kvH, + /* int lda = */ D, + /* int ldb = */ k.strides(2), + /* bool transpose_a = */ false, + /* bool transpose_b = */ true, + /* std::vector& copies = */ copies, + /* Shape batch_shape = */ batch_shape, + /* Strides A_batch_stride = */ q_batch_strides, + /* Strides B_batch_stride = */ k_batch_strides); + + int total_rows = B * H * rows; + compute_encoder.set_compute_pipeline_state(mask_kernel); + compute_encoder.set_output_array(scores, 0); + compute_encoder.set_bytes(rows, 1); + compute_encoder.set_bytes(columns, 2); + auto mask_group_dims = get_block_dims(rows, total_rows, 1); + compute_encoder.dispatch_threads( + MTL::Size(rows, total_rows, 1), mask_group_dims); + + sdpa_blocked_softmax(s, d, scores); + + array out_block({B, H, rows, D}, o.dtype(), nullptr, {}); + out_block.copy_shared_buffer( + out_buffer, out_block.strides(), out_block.flags(), out_block.size()); + + Strides scores_batch_strides = { + int64_t(H) * rows * columns, int64_t(gqa) * rows * columns}; + steel_matmul( + /* const Stream& s = */ s, + /* metal::Device& d = */ d, + /* const array& a = */ scores, + /* const array& b = */ v, + /* array& out = */ out_block, + /* int M = */ gqa * rows, + /* int N = */ D, + /* int K = */ columns, + /* int batch_size_out = */ B * kvH, + /* int lda = */ columns, + /* int ldb = */ v.strides(2), + /* bool transpose_a = */ false, + /* bool transpose_b = */ false, + /* std::vector& copies = */ copies, + /* Shape batch_shape = */ batch_shape, + /* Strides A_batch_stride = */ scores_batch_strides, + /* Strides B_batch_stride = */ v_batch_strides); + copy_gpu_inplace( + /* const array& in = */ out_block, + /* array& out = */ o, + /* const Shape& data_shape = */ out_block.shape(), + /* const Strides& i_strides = */ out_block.strides(), + /* const Strides& o_strides = */ o.strides(), + /* int64_t i_offset = */ 0, + /* int64_t o_offset = */ int64_t(q_offset) * o.strides(2), + /* CopyType ctype = */ CopyType::GeneralGeneral, + /* const Stream& s = */ s); + } + + compute_encoder.add_temporaries(std::move(copies)); + compute_encoder.add_temporary(std::move(score_buffer)); + compute_encoder.add_temporary(std::move(q_buffer)); + compute_encoder.add_temporary(std::move(out_buffer)); +} + void sdpa_full_self_attention_nax( const Stream& s, metal::Device& d, @@ -823,6 +1008,7 @@ bool ScaledDotProductAttention::use_fallback( bool has_mask, bool has_arr_mask, bool do_causal, + bool has_sinks, bool is_training, bool output_logsumexp, bool force_fused, @@ -845,6 +1031,9 @@ bool ScaledDotProductAttention::use_fallback( // forward and backward. return true; } + if (use_blocked_d512(q, k, v, has_arr_mask, do_causal, has_sinks, s)) { + return false; + } if (!has_fused) { return true; } @@ -1012,8 +1201,13 @@ void ScaledDotProductAttention::eval_gpu( ? std::optional{copy_unless(is_matrix_contiguous, inputs[3])} : std::nullopt; - sdpa_full_self_attention_metal( - s, d, q, k, v, scale_, o, do_causal_, mask, sinks); + if (use_blocked_d512( + q, k, v, has_arr_mask, do_causal_, sinks.has_value(), s)) { + sdpa_full_self_attention_blocked(s, d, q, k, v, scale_, o); + } else { + sdpa_full_self_attention_metal( + s, d, q, k, v, scale_, o, do_causal_, mask, sinks); + } } metal::get_command_encoder(s).add_temporaries(std::move(copies)); diff --git a/mlx/backend/no_gpu/primitives.cpp b/mlx/backend/no_gpu/primitives.cpp index c37515f9ef..d9226e3586 100644 --- a/mlx/backend/no_gpu/primitives.cpp +++ b/mlx/backend/no_gpu/primitives.cpp @@ -1,4 +1,4 @@ -// Copyright © 2023-2024 Apple Inc. +// Copyright © 2023-2026 Apple Inc. #include "mlx/primitives.h" #include "mlx/distributed/primitives.h" @@ -30,6 +30,7 @@ bool fast::ScaledDotProductAttention::use_fallback( bool has_mask, bool has_arr_mask, bool do_causal, + bool has_sinks, bool is_training, bool output_logsumexp, bool force_fused, diff --git a/mlx/fast.cpp b/mlx/fast.cpp index add78fa7c1..4d68511b28 100644 --- a/mlx/fast.cpp +++ b/mlx/fast.cpp @@ -1,4 +1,4 @@ -// Copyright © 2023-2024 Apple Inc. +// Copyright © 2023-2026 Apple Inc. #include #include @@ -934,6 +934,7 @@ array scaled_dot_product_attention( has_mask, has_arr_mask, do_causal, + has_sinks, is_training, output_logsumexp, force_fused, diff --git a/mlx/fast_primitives.h b/mlx/fast_primitives.h index cef13bb0ff..181aac342b 100644 --- a/mlx/fast_primitives.h +++ b/mlx/fast_primitives.h @@ -1,4 +1,4 @@ -// Copyright © 2024 Apple Inc. +// Copyright © 2024-2026 Apple Inc. #include #include @@ -289,6 +289,7 @@ class ScaledDotProductAttention : public Custom { bool has_mask, bool has_arr_mask, bool do_causal, + bool has_sinks, bool is_training, bool output_logsumexp, bool force_fused, diff --git a/python/tests/test_fast_sdpa.py b/python/tests/test_fast_sdpa.py index fefe7325b3..1907ca126b 100644 --- a/python/tests/test_fast_sdpa.py +++ b/python/tests/test_fast_sdpa.py @@ -929,6 +929,53 @@ 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_d512_default_metal(self): + if mx.default_device() != mx.gpu: + self.skipTest("requires GPU") + mx.random.seed(0) + + cases = ( + (1, 4, 2, 1024, 1057, 1152), + (2, 2, 1, 1025, 1025, None), + (1, 2, 1, 1024, 4103, 4128), + ) + for dtype, (B, qH, kH, qL, kL, cache_len) in product( + (mx.float16, mx.bfloat16), cases + ): + with self.subTest(dtype=dtype, qL=qL, kL=kL): + q = (0.5 * mx.random.normal((B, qH, qL, 512))).astype(dtype) + kv_len = cache_len or kL + k_cache = (0.5 * mx.random.normal((B, kH, kv_len, 512))).astype(dtype) + v_cache = (0.5 * mx.random.normal((B, kH, kv_len, 512))).astype(dtype) + k = k_cache[..., :kL, :] + v = v_cache[..., :kL, :] + scale = 512**-0.5 + ref = mlx_ref_attn(q, k, v, scale=scale, mask="causal") + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, mask="causal" + ) + tol = 1e-2 if dtype == mx.bfloat16 else 1e-3 + self.assertTrue(mx.allclose(ref, out, atol=tol, rtol=tol)) + + q = mx.random.normal((1, 2, 1024, 512), mx.float16) + k = mx.random.normal((1, 1, 1057, 512), mx.float16) + v = mx.random.normal((1, 1, 1057, 512), mx.float16) + scale = 512**-0.5 + mask = mx.random.uniform(shape=(1, 1, 1024, 1057)) > 0.2 + sinks = mx.random.normal((2,), mx.float16) + fallback_cases = ( + (v, mask, None), + (v, "causal", sinks), + (v[..., :128], "causal", None), + ) + for v_arg, mask_arg, sinks_arg in fallback_cases: + ref = mlx_ref_attn(q, k, v_arg, scale=scale, mask=mask_arg, sinks=sinks_arg) + out = mx.fast.scaled_dot_product_attention( + q, k, v_arg, scale=scale, mask=mask_arg, sinks=sinks_arg + ) + self.assertTrue(mx.allclose(ref, out, atol=1e-3, rtol=1e-3)) + @unittest.skipIf(not mx.metal.is_available(), "Metal kernel path only") def test_sdpa_force_fused_metal(self): if mx.default_device() != mx.gpu: