Add some fast Metal MLX SDPA kernels #2584
Open
+1,992
−1
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
* Sketch the sdpa kernel * Add full sdpa kernel, * Add test * Add vectorized kernel for decoding * Update tests * Add some docs * Fix sdpa_vector names * Add softcapping for vectorized sdpa * Add softcapping for full sdpa * Add support for head dim 32, 96, 256 * Add support for head dim 32, 96, 256 * Update docs * Add update notice * Clippy and format
| @@ -7,5 +7,6 @@ | |||
| "candle-pyo3" | |||
| ], | |||
| "python.testing.unittestEnabled": false, | |||
| "python.testing.pytestEnabled": true | |||
| "python.testing.pytestEnabled": true, | |||
| "rust-analyzer.cargo.features": ["metal"] | |||
There was a problem hiding this comment.
Wouldn't that break the editor for non metal users?
| // q = (bs, qhead, seq, hidden) | ||
| // k/v = (bs, kv_head, seq, hidden) | ||
|
|
||
| let _hidden = q_shape[q_shape.len() - 1]; |
There was a problem hiding this comment.
If variables such as _hidden are not planned to be used, maybe remove them?
Comment on lines
+1812
to
+1831
| encoder.set_buffer(0, Some(&q_buffer), q_offset as NSUInteger); | ||
| encoder.set_buffer(1, Some(&k_buffer), k_offset as NSUInteger); | ||
| encoder.set_buffer(2, Some(&v_buffer), v_offset as NSUInteger); | ||
| encoder.set_buffer(3, Some(&output), 0); | ||
|
|
||
| encoder.set_bytes( | ||
| 4, | ||
| std::mem::size_of::<MLXFastAttentionParams>() as u64, | ||
| ¶ms as *const MLXFastAttentionParams as *const c_void, | ||
| ); | ||
| encoder.set_bytes( | ||
| 6, | ||
| (std::mem::size_of::<i32>() * batch_shape.len()) as u64, | ||
| batch_shape.as_ptr() as *const i32 as *const c_void, | ||
| ); | ||
| encoder.set_bytes( | ||
| 7, | ||
| (std::mem::size_of::<usize>() * batch_strides.len()) as u64, | ||
| batch_strides.as_ptr() as *const c_void, | ||
| ); |
There was a problem hiding this comment.
Couldn't the EncoderParam trait or better the set_params! macro be used to simplify this?
Comment on lines
+1920
to
+1949
| encoder.set_buffer(0, Some(&q_buffer), q_offset as NSUInteger); | ||
| encoder.set_buffer(1, Some(&k_buffer), k_offset as NSUInteger); | ||
| encoder.set_buffer(2, Some(&v_buffer), v_offset as NSUInteger); | ||
| encoder.set_buffer(3, Some(&output), 0); | ||
|
|
||
| encoder.set_bytes( | ||
| 4, | ||
| std::mem::size_of::<i32>() as u64, | ||
| &gqa_factor as *const i32 as *const c_void, | ||
| ); | ||
| encoder.set_bytes( | ||
| 5, | ||
| std::mem::size_of::<i32>() as u64, | ||
| &n as *const i32 as *const c_void, | ||
| ); | ||
| encoder.set_bytes( | ||
| 6, | ||
| std::mem::size_of::<usize>() as u64, | ||
| &stride as *const usize as *const c_void, | ||
| ); | ||
| encoder.set_bytes( | ||
| 7, | ||
| std::mem::size_of::<f32>() as u64, | ||
| &alpha as *const f32 as *const c_void, | ||
| ); | ||
| encoder.set_bytes( | ||
| 8, | ||
| std::mem::size_of::<f32>() as u64, | ||
| &softcapping as *const f32 as *const c_void, | ||
| ); |
There was a problem hiding this comment.
Same comment as above, set_params! should be an easy win here.
| candle::bail!("query `n_heads` must be a multiple of `n_kv_heads`"); | ||
| } | ||
|
|
||
| let k_head = k_l.dims()[k_l.dims().len() - 1]; |
There was a problem hiding this comment.
k_l.dim(D::Minus1)? would be simpler and make for a better error message than a panic (the same applies to a bunch of places in this function)
|
|
||
| impl candle::CustomOp3 for Sdpa { | ||
| fn name(&self) -> &'static str { | ||
| "sdpa" |
There was a problem hiding this comment.
Let's call this metal-sdpa instead.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This PR adds some MLX SDPA kernels on Metal.
I can observe about a 26% performance improvement with Llama 3.1 8b @ q4k and @ q8_0 when testing through mistral.rs on my Candle fork. I updated the quantized_llama.rs file here to use the new function.
This PR adds a function
candle_nn::ops::sdpa. The MLX attention kernels don't support masking yet, so the performance gains are only for decoding on Metal. Once/if they do, I'll update them - otherwise we can explore using Flash Attention kernels for Metal from llama.cpp.