zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-10-22 02:58:16 +08:00
Code Issues Actions 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

17
mlx/backend/metal/kernels/sdpa_vector.h
View File

@@ -6,6 +6,7 @@ using namespace metal;
constant bool has_mask [[function_constant(20)]];
constant bool query_transposed [[function_constant(21)]];
constant bool do_causal [[function_constant(22)]];
template <typename T, int D, int V = D>
[[kernel]] void sdpa_vector(
@@ -77,7 +78,13 @@ template <typename T, int D, int V = D>
// For each key
for (int i = simd_gid; i < N; i += BN) {
if (!has_mask || mask[0]) {
bool use_key = true;
if (do_causal) {
use_key = i <= (N - int(tpg.y) + int(q_seq_idx));
} else if (has_mask) {
use_key = mask[0];
}
if (use_key) {
// Read the key
for (int j = 0; j < qk_per_thread; j++) {
k[j] = keys[j];
@@ -218,7 +225,13 @@ template <typename T, int D, int V = D>
// For each key
for (int i = block_idx * BN + simd_gid; i < N; i += blocks * BN) {
if (!has_mask || mask[0]) {
bool use_key = true;
if (do_causal) {
use_key = i <= (N - int(tpg.y) + int(q_seq_idx));
} else if (has_mask) {
use_key = mask[0];
}
if (use_key) {
// Read the key
for (int i = 0; i < qk_per_thread; i++) {
k[i] = keys[i];

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);
}
}