From c184262d29a84b5d9935cd45ec71fd9bc3c90338 Mon Sep 17 00:00:00 2001 From: Angelos Katharopoulos Date: Thu, 12 Feb 2026 10:46:21 -0800 Subject: [PATCH] Fix donation in sdpa vector (#3121) --- mlx/backend/metal/copy.cpp | 3 +-- mlx/backend/metal/normalization.cpp | 6 ++---- mlx/backend/metal/scaled_dot_product_attention.cpp | 6 +++++- mlx/backend/metal/scan.cpp | 3 +-- 4 files changed, 9 insertions(+), 9 deletions(-) diff --git a/mlx/backend/metal/copy.cpp b/mlx/backend/metal/copy.cpp index 6b791289..24c1506c 100644 --- a/mlx/backend/metal/copy.cpp +++ b/mlx/backend/metal/copy.cpp @@ -109,12 +109,11 @@ void copy_gpu_inplace( auto& compute_encoder = d.get_command_encoder(s.index); compute_encoder.set_compute_pipeline_state(kernel); - bool donate_in = in.data_shared_ptr() == nullptr; inp_offset *= size_of(in.dtype()); out_offset *= size_of(out.dtype()); - compute_encoder.set_input_array(donate_in ? out : in, 0, inp_offset); + compute_encoder.set_input_array(in, 0, inp_offset); compute_encoder.set_output_array(out, 1, out_offset); auto thread_group_size = kernel->maxTotalThreadsPerThreadgroup(); diff --git a/mlx/backend/metal/normalization.cpp b/mlx/backend/metal/normalization.cpp index da0160b2..1dc0338f 100644 --- a/mlx/backend/metal/normalization.cpp +++ b/mlx/backend/metal/normalization.cpp @@ -82,8 +82,7 @@ void RMSNorm::eval_gpu( uint32_t w_stride = (w.ndim() == 1) ? w.strides()[0] : 0; compute_encoder.set_compute_pipeline_state(kernel); - compute_encoder.set_input_array( - x.data_shared_ptr() == nullptr ? out : x, 0); + compute_encoder.set_input_array(x, 0); compute_encoder.set_input_array(w, 1); compute_encoder.set_output_array(out, 2); compute_encoder.set_bytes(eps_, 3); @@ -287,8 +286,7 @@ void LayerNorm::eval_gpu( uint32_t w_stride = (w.ndim() == 1) ? w.strides()[0] : 0; uint32_t b_stride = (b.ndim() == 1) ? b.strides()[0] : 0; compute_encoder.set_compute_pipeline_state(kernel); - compute_encoder.set_input_array( - x.data_shared_ptr() == nullptr ? out : x, 0); + compute_encoder.set_input_array(x, 0); compute_encoder.set_input_array(w, 1); compute_encoder.set_input_array(b, 2); compute_encoder.set_output_array(out, 3); diff --git a/mlx/backend/metal/scaled_dot_product_attention.cpp b/mlx/backend/metal/scaled_dot_product_attention.cpp index c990cda3..37e554f1 100644 --- a/mlx/backend/metal/scaled_dot_product_attention.cpp +++ b/mlx/backend/metal/scaled_dot_product_attention.cpp @@ -711,7 +711,8 @@ void ScaledDotProductAttention::eval_gpu( return (strides[0] == strides[1] * shape[1]); }; - const auto& q = copy_unless(q_copy_unless, q_pre); + bool q_copied = !q_copy_unless(q_pre); + array q = (q_copied) ? contiguous_copy_gpu(q_pre, s) : q_pre; const auto& k = copy_unless(kv_copy_unless, k_pre); const auto& v = copy_unless(kv_copy_unless, v_pre); @@ -719,6 +720,9 @@ void ScaledDotProductAttention::eval_gpu( if (q.is_donatable() && q.flags().row_contiguous && q.size() == o.size()) { o.copy_shared_buffer(q); } else { + if (q_copied) { + copies.push_back(q); + } o.set_data(allocator::malloc(o.nbytes())); } diff --git a/mlx/backend/metal/scan.cpp b/mlx/backend/metal/scan.cpp index b35de64c..b48ec41c 100644 --- a/mlx/backend/metal/scan.cpp +++ b/mlx/backend/metal/scan.cpp @@ -91,8 +91,7 @@ void scan_gpu_inplace( } else { auto& compute_encoder = d.get_command_encoder(s.index); compute_encoder.set_compute_pipeline_state(kernel); - compute_encoder.set_input_array( - in.data_shared_ptr() == nullptr ? out : in, 0); + compute_encoder.set_input_array(in, 0); compute_encoder.set_output_array(out, 1); size_t size = in.shape(axis); size_t stride = in.strides()[axis];