L-electron-Rare/mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix donation in sdpa vector (#3121)

Browse Source
This commit is contained in:
Angelos Katharopoulos
2026-02-12 10:46:21 -08:00
committed by GitHub
parent 72e94c81e1
commit c184262d29
4 changed files with 9 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mlx/backend/metal/copy.cpp
+1 -2
View File
@@ -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();
mlx/backend/metal/normalization.cpp
+2 -4
View File
@@ -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);
mlx/backend/metal/scaled_dot_product_attention.cpp
+5 -1
View File
@@ -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()));
}
mlx/backend/metal/scan.cpp
+1 -2
View File
@@ -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];