zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-10-19 16:48:10 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fix causal mask sdpa vec (#2053)

Browse Source 
* fix sdpa vector causal mask

* test
This commit is contained in:
Awni Hannun
2025-04-08 09:11:23 -07:00
committed by GitHub
parent 08a1bf3f10
commit 00794c42bc
2 changed files with 6 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -172,6 +172,7 @@ void sdpa_vector(
std::string hash_name = kname;
hash_name += has_mask ? "_mask" : "_nomask";
hash_name += query_transposed ? "_qt" : "_qnt";
hash_name += do_causal ? "_c" : "_nc";
// Get the kernel
auto& compute_encoder = d.get_command_encoder(s.index);
@@ -268,6 +269,7 @@ void sdpa_vector_2pass(
std::string hash_name = kname;
hash_name += has_mask ? "_mask" : "_nomask";
hash_name += query_transposed ? "_qt" : "_qnt";
hash_name += do_causal ? "_c" : "_nc";
// Get the kernel
auto& compute_encoder = d.get_command_encoder(s.index);