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

causal vector sdpa (#2018)

Browse Source 
* causal vector sdpa

* get rid of memory threshold
This commit is contained in:
Awni Hannun
2025-03-28 12:36:13 -07:00
committed by GitHub
parent 98b901ad66
commit 05d7118561
6 changed files with 45 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -138,6 +138,7 @@ void sdpa_vector(
const array& v,
array& out,
float scale,
bool do_causal,
const std::optional<array>& mask) {
// Set the kernel name
std::string kname;
@@ -166,6 +167,7 @@ void sdpa_vector(
metal::MTLFCList func_consts = {
{&has_mask, MTL::DataType::DataTypeBool, 20},
{&query_transposed, MTL::DataType::DataTypeBool, 21},
{&do_causal, MTL::DataType::DataTypeBool, 22},
};
std::string hash_name = kname;
hash_name += has_mask ? "_mask" : "_nomask";
@@ -214,6 +216,7 @@ void sdpa_vector_2pass(
const array& v,
array& out,
float scale,
bool do_causal,
const std::optional<array>& mask) {
// Set the kernel name
std::string kname;
@@ -260,6 +263,7 @@ void sdpa_vector_2pass(
metal::MTLFCList func_consts = {
{&has_mask, MTL::DataType::DataTypeBool, 20},
{&query_transposed, MTL::DataType::DataTypeBool, 21},
{&do_causal, MTL::DataType::DataTypeBool, 22},
};
std::string hash_name = kname;
hash_name += has_mask ? "_mask" : "_nomask";
@@ -401,12 +405,13 @@ void ScaledDotProductAttention::eval_gpu(
// We route to the 2 pass fused attention if
// - The device is large and the sequence length long
// - The sequence length is even longer and we have gqa
bool do_causal = do_causal_ && q.shape(2) > 1;
char devc = d.get_architecture().back();
if ((devc == 'd' && k.shape(2) >= 1024) ||
(k.shape(1) < q.shape(1) && k.shape(2) >= 4096)) {
sdpa_vector_2pass(s, d, q, k, v, o, scale_, mask);
sdpa_vector_2pass(s, d, q, k, v, o, scale_, do_causal, mask);
} else {
sdpa_vector(s, d, q, k, v, o, scale_, mask);
sdpa_vector(s, d, q, k, v, o, scale_, do_causal, mask);
}
}