zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-07-15 04:51:13 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

SDPA fixes

Browse Source
This commit is contained in:
paramthakkar123 2025-04-19 08:15:56 +05:30
parent 5f04c0f818
commit 8ac0a20250
2 changed files with 2 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
mlx/fast.cpp
View File

@ -734,8 +734,7 @@ array scaled_dot_product_attention(
const bool sdpa_vector_supported_head_dim =
query_head_dim == value_head_dim &&
(query_head_dim == 64 || query_head_dim == 96 || query_head_dim == 128 ||
query_head_dim == 256);
(query_head_dim % 32 == 0 || query_head_dim % 2 == 0);
const bool sdpa_full_supported_head_dim = query_head_dim == value_head_dim &&
(query_head_dim == 64 || query_head_dim == 80 || query_head_dim == 128);

2
mlx/primitives.cpp
View File

@ -5108,4 +5108,4 @@ bool Hadamard::is_equivalent(const Primitive& other) const {
return scale_ == h_other.scale_;
}
} // namespace mlx::core
} // namespace mlx::core