L-electron-Rare/mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[CUDA] Use cuDNN SDPA for decoding when using fixed-size KV cache (#3113)

Browse Source
This commit is contained in:
Cheng
2026-02-10 09:15:45 +09:00
committed by GitHub
parent 5e018de4e5
commit 54bb3eea42
3 changed files with 156 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mlx/backend/cuda/scaled_dot_product_attention.cpp
+124 -22
View File
@@ -57,6 +57,58 @@ void malloc_with_same_layout(
{true, false, false});
}
bool use_cudnn_for_decoding(
const array& q,
const array& k,
const array& v,
bool has_arr_mask) {
if (q.shape(2) != 1) {
return false;
}
if (has_arr_mask) {
return false;
}
// The cuDNN SDPA is faster than vector kernel but for small sequence the
// overhead would kill the advantage.
constexpr int kv_cache_step = 256; // number is from mlx-lm
if (k.shape(2) < kv_cache_step) {
return false;
}
// When called during graph building the strides is not available, and we
// rely on |supports_sdpa_vector| to decide whether to use fast sdpa since
// we can fallback to |sdpa_vector|.
if ((k.status() != array::evaluated) || (v.status() != array::evaluated)) {
return false;
}
// Check if k/v are slices from fixed-size kv cache.
auto is_slice = [](const array& kv) {
// Get pre-sliced sequence length from strides, and check if the buffer
// belongs to a contiguous kv cache.
int64_t T_kv = kv.strides(1) / kv.strides(2);
if (kv.size() / kv.shape(2) * T_kv != kv.buffer_size() / kv.itemsize()) {
return false;
}
// It is possible to use heuristic to check slices, but for now just make
// mlx-lm work.
return T_kv % kv_cache_step == 0;
};
return is_slice(k) && is_slice(v);
}
// Get original kv from slices, i.e. undo keys[..., :offset, :]
array unslice_kv(const array& kv) {
Shape shape = kv.shape();
shape[2] = /* T_kv */ kv.strides(1) / kv.strides(2);
array copy(shape, kv.dtype(), nullptr, {});
copy.copy_shared_buffer(
kv,
make_contiguous_strides(shape),
{true, true, false},
/* data_size */ kv.buffer_size() / kv.itemsize(),
/* offset */ -kv.offset());
return copy;
}
constexpr int QKV_NDIM = 4;
struct SDPACacheKey {
@@ -81,7 +133,8 @@ inline BytesKey<SDPACacheKey> build_sdpa_cache_key(
const array& v,
bool do_causal,
const std::optional<array>& mask_arr,
bool output_logsumexp = true) {
bool decoding = false,
bool output_logsumexp = false) {
BytesKey<SDPACacheKey> cache_key;
cache_key.pod = {
encoder.device().cuda_device(),
@@ -101,12 +154,19 @@ inline BytesKey<SDPACacheKey> build_sdpa_cache_key(
cache_key.pod.mask_shape = vector_key<QKV_NDIM>(mask_arr->shape());
cache_key.pod.mask_strides = vector_key<QKV_NDIM>(mask_arr->strides());
}
if (decoding) {
int64_t T_kv = k.strides(1) / k.strides(2);
cache_key.pod.k_shape[2] = T_kv;
cache_key.pod.v_shape[2] = T_kv;
cache_key.pod.k_strides.fill(0);
cache_key.pod.v_strides.fill(0);
}
return cache_key;
}
auto& sdpa_cache() {
static LRUBytesKeyCache<SDPACacheKey, DnnGraph> cache(
"MLX_CUDA_SDPA_CACHE_SIZE", /* default_capacity */ 64);
"MLX_CUDA_SDPA_CACHE_SIZE", /* default_capacity */ 256);
return cache;
}
@@ -122,6 +182,8 @@ enum UIDS {
V,
SCALE,
BIAS,
SEQ_LEN_Q,
SEQ_LEN_KV,
O,
STATS,
// Backward graph:
@@ -138,6 +200,8 @@ DnnGraph build_sdpa_graph(
const array& v,
bool do_causal,
const std::optional<array>& mask_arr,
const std::optional<array>& seq_len_q,
const std::optional<array>& seq_len_kv,
bool output_logsumexp,
const array& o,
const std::optional<array>& stats) {
@@ -157,6 +221,11 @@ DnnGraph build_sdpa_graph(
if (mask_arr) {
options.set_bias(graph.tensor("BIAS", BIAS, *mask_arr));
}
if (seq_len_q && seq_len_kv) {
options.set_padding_mask(true);
options.set_seq_len_q(graph.tensor("SEQ_LEN_Q", SEQ_LEN_Q, *seq_len_q));
options.set_seq_len_kv(graph.tensor("SEQ_LEN_KV", SEQ_LEN_KV, *seq_len_kv));
}
auto [o_, stats_] = graph.sdpa(q_, k_, v_, options);
graph.tensor(o_, O, o)->set_output(true);
@@ -222,9 +291,10 @@ bool supports_sdpa_cudnn(
const array& q,
const array& k,
const array& v,
bool has_arr_mask,
bool do_causal,
Stream s) {
static bool enabled = env::get_var("MLX_CUDA_USE_CUDNN_SPDA", 1);
static bool enabled = env::get_var("MLX_CUDA_USE_CUDNN_SDPA", 1);
if (!enabled) {
return false;
}
@@ -234,8 +304,8 @@ bool supports_sdpa_cudnn(
return false;
}
// Only use cuDNN for prefilling (T_q > 1) and training (T_q == T_kv).
if ((q.shape(2) == 1) && (q.shape(2) != k.shape(2))) {
// Only use cuDNN for decoding when k/v are slices from fixed-size kv cache.
if ((q.shape(2) == 1) && !use_cudnn_for_decoding(q, k, v, has_arr_mask)) {
return false;
}
@@ -256,8 +326,8 @@ bool supports_sdpa_cudnn(
void sdpa_cudnn(
const array& q,
const array& k,
const array& v,
array k,
array v,
float scale,
array& o,
std::optional<array>& stats,
@@ -270,6 +340,24 @@ void sdpa_cudnn(
malloc_with_same_layout(encoder, o, q);
// For decoding, unslice k/v and apply padding mask.
std::optional<array> seq_len_q;
std::optional<array> seq_len_kv;
bool decoding = use_cudnn_for_decoding(q, k, v, mask_arr.has_value());
if (decoding) {
int B = q.shape(0);
std::vector<int> seq_len_q_vec(B, q.shape(2));
std::vector<int> seq_len_kv_vec(B, k.shape(2));
seq_len_q = array(seq_len_q_vec.begin(), {B, 1, 1, 1});
seq_len_kv = array(seq_len_kv_vec.begin(), {B, 1, 1, 1});
encoder.add_temporary(*seq_len_q);
encoder.add_temporary(*seq_len_kv);
k = unslice_kv(k);
v = unslice_kv(v);
encoder.add_temporary(k);
encoder.add_temporary(v);
}
encoder.set_input_array(q);
encoder.set_input_array(k);
encoder.set_input_array(v);
@@ -277,6 +365,10 @@ void sdpa_cudnn(
if (mask_arr) {
encoder.set_input_array(*mask_arr);
}
if (seq_len_q && seq_len_kv) {
encoder.set_input_array(*seq_len_q);
encoder.set_input_array(*seq_len_kv);
}
if (output_logsumexp) {
stats->set_data(cu::malloc_async(stats->nbytes(), encoder));
encoder.set_output_array(*stats);
@@ -284,11 +376,21 @@ void sdpa_cudnn(
// Search cache.
auto cache_key = build_sdpa_cache_key(
encoder, q, k, v, do_causal, mask_arr, output_logsumexp);
encoder, q, k, v, do_causal, mask_arr, decoding, output_logsumexp);
auto it = sdpa_cache().find(cache_key);
if (it == sdpa_cache().end()) {
auto graph = build_sdpa_graph(
handle, q, k, v, do_causal, mask_arr, output_logsumexp, o, stats);
handle,
q,
k,
v,
do_causal,
mask_arr,
seq_len_q,
seq_len_kv,
output_logsumexp,
o,
stats);
it = sdpa_cache().emplace(cache_key, std::move(graph)).first;
}
auto& graph = it->second;
@@ -302,6 +404,10 @@ void sdpa_cudnn(
if (mask_arr) {
variant_pack[BIAS] = gpu_ptr<void>(*mask_arr);
}
if (seq_len_q && seq_len_kv) {
variant_pack[SEQ_LEN_Q] = gpu_ptr<void>(*seq_len_q);
variant_pack[SEQ_LEN_KV] = gpu_ptr<void>(*seq_len_kv);
}
if (output_logsumexp) {
variant_pack[STATS] = gpu_ptr<void>(*stats);
}
@@ -376,9 +482,7 @@ bool supports_sdpa_vector(
const array& q,
const array& k,
const array& v,
bool has_mask,
bool has_arr_mask,
bool do_causal,
bool output_logsumexp);
void sdpa_vector(
const array& q,
@@ -406,9 +510,8 @@ bool ScaledDotProductAttention::use_fallback(
return true;
}
return !supports_sdpa_vector(
q, k, v, has_mask, has_arr_mask, do_causal, output_logsumexp) &&
!supports_sdpa_cudnn(q, k, v, do_causal, s);
return !supports_sdpa_cudnn(q, k, v, has_arr_mask, do_causal, s) &&
!supports_sdpa_vector(q, k, v, has_arr_mask, output_logsumexp);
}
bool ScaledDotProductAttention::supports_bool_mask() {
@@ -438,14 +541,7 @@ void ScaledDotProductAttention::eval_gpu(
stats = outputs[1];
}
if (supports_sdpa_vector(
q, k, v, has_mask, has_arr_mask, do_causal_, output_logsumexp_)) {
if (has_sinks_) {
sdpa_vector(q, k, v, scale_, out, do_causal_, inputs.back(), s);
} else {
sdpa_vector(q, k, v, scale_, out, do_causal_, std::nullopt, s);
}
} else {
if (supports_sdpa_cudnn(q, k, v, has_arr_mask, do_causal_, s)) {
sdpa_cudnn(
q,
k,
@@ -457,6 +553,12 @@ void ScaledDotProductAttention::eval_gpu(
mask_arr,
output_logsumexp_,
s);
} else {
if (has_sinks_) {
sdpa_vector(q, k, v, scale_, out, do_causal_, inputs.back(), s);
} else {
sdpa_vector(q, k, v, scale_, out, do_causal_, std::nullopt, s);
}
}
}
mlx/backend/cuda/scaled_dot_product_attention.cu
-2
View File
@@ -665,9 +665,7 @@ bool supports_sdpa_vector(
const array& q,
const array& k,
const array& v,
bool has_mask,
bool has_arr_mask,
bool do_causal,
bool output_logsumexp) {
if (output_logsumexp) {
return false;
python/tests/test_fast_sdpa.py
+32
View File
@@ -605,6 +605,38 @@ class TestFastSDPA(mlx_tests.MLXTestCase):
).sum()
test_grad(loss_slow, loss_fast, [q, k, v])
def test_sdpa_sliced(self):
N = 8
D = 64
scale = D**-0.5
for B, T_q, T_kv, offset, mask in product(
(1, 2, 4),
(1, 8),
(256, 512),
(8, 9, 64, 79),
(None, "causal"),
):
with self.subTest(B=B, T_q=T_q, T_kv=T_kv, offset=offset, mask=mask):
q = mx.random.normal((B, N, T_q, D), mx.float16)
k = mx.random.normal((B, N, T_kv, D), mx.float16)
v = mx.random.normal((B, N, T_kv, D), mx.float16)
k = k[..., :offset, :]
v = v[..., :offset, :]
ref = mlx_ref_attn(q, k, v, scale=scale, mask=mask)
for i in range(2):
out = mx.fast.scaled_dot_product_attention(
q, k, v, scale=scale, mask=mask
)
if B == 1:
tolerance = {"rtol": 1e-3, "atol": 1e-3}
else:
tolerance = {"rtol": 1e-2, "atol": 1e-2}
self.assertTrue(mx.allclose(ref, out, **tolerance))
if __name__ == "__main__":
mlx_tests.MLXTestRunner(failfast=True)