zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-08-09 18:56:39 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Stop matrix copies with new attention kernel (#1639)

Browse Source
This commit is contained in:
Jagrit Digani 2024-12-02 14:12:38 -08:00 committed by GitHub
parent 1445dcaa60
commit 9d40e521d7
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 2 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
mlx/backend/metal/scaled_dot_product_attention.cpp
View File

@ -288,11 +288,9 @@ void ScaledDotProductAttention::eval_gpu(
strides[0] == strides[1] * shape[1]; strides[0] == strides[1] * shape[1];
}; };
// Checks that the last two dims are row contiguous. // Checks that the headdim dimension has stride 1.
auto is_matrix_contiguous = [](const array& arr) { auto is_matrix_contiguous = [](const array& arr) {
auto& strides = arr.strides(); return arr.strides(3) == 1;
auto& shape = arr.shape();
return strides[3] == 1 && strides[2] == shape[3];
}; };
// We are in vector mode ie single query // We are in vector mode ie single query