mlx_rs/
fast.rs

1//! Fast implementations of commonly used multi-op functions.
2
3use std::ffi::CStr;
4
5use crate::error::Result;
6use crate::utils::guard::Guarded;
7use crate::utils::IntoOption;
8use crate::{Array, Stream};
9use mlx_internal_macros::generate_macro;
10
11/// Optimized implementation of `NN.RoPE`.
12#[allow(clippy::too_many_arguments)]
13pub fn rope<'a>(
14    array: impl AsRef<Array>,
15    dimensions: i32,
16    traditional: bool,
17    base: impl Into<Option<f32>>,
18    scale: f32,
19    offset: i32,
20    freqs: impl Into<Option<&'a Array>>,
21) -> Result<Array> {
22    let stream = Stream::thread_local_or_default();
23    let base = base.into();
24    let base = mlx_sys::mlx_optional_float {
25        value: base.unwrap_or(0.0),
26        has_value: base.is_some(),
27    };
28    let freqs = freqs.into();
29    Array::try_from_op(|res| unsafe {
30        mlx_sys::mlx_fast_rope(
31            res,
32            array.as_ref().as_ptr(),
33            dimensions,
34            traditional,
35            base,
36            scale,
37            offset,
38            freqs
39                .map(|a| a.as_ptr())
40                .unwrap_or(mlx_sys::mlx_array_new()),
41            stream.as_ref().as_ptr(),
42        )
43    })
44}
45
46/// Compatibility shim for [`rope`].
47#[allow(clippy::too_many_arguments)]
48#[generate_macro(customize(forwarding_shim = true, root = "$crate::fast"))]
49#[deprecated(
50    since = "0.26.0",
51    note = "use `with_stream` or `with_device` around `rope`"
52)]
53pub fn rope_device<'a>(
54    #[named] array: impl AsRef<Array>,
55    #[named] dimensions: i32,
56    #[named] traditional: bool,
57    #[optional] base: impl Into<Option<f32>>,
58    #[named] scale: f32,
59    #[named] offset: i32,
60    #[optional] freqs: impl Into<Option<&'a Array>>,
61    #[optional] stream: impl AsRef<Stream>,
62) -> Result<Array> {
63    crate::with_stream(stream.as_ref(), || {
64        rope(array, dimensions, traditional, base, scale, offset, freqs)
65    })
66}
67
68/// Optimized implementation of `NN.RoPE` with dynamic (array) offset.
69///
70/// This variant allows specifying the offset as an array, enabling different
71/// offsets for different positions in the input.
72///
73/// # Params
74///
75/// - `array`: Input array
76/// - `dimensions`: The feature dimensions to apply rope to
77/// - `traditional`: If true, uses the traditional rope implementation
78/// - `base`: The base used to compute angular frequency for each dimension
79/// - `scale`: The scale to apply to the positions
80/// - `offset`: An array of position offsets
81/// - `freqs`: Optional precomputed frequencies
82/// - `stream`: Stream to evaluate on
83#[allow(clippy::too_many_arguments)]
84pub fn rope_dynamic<'a>(
85    array: impl AsRef<Array>,
86    dimensions: i32,
87    traditional: bool,
88    base: impl Into<Option<f32>>,
89    scale: f32,
90    offset: impl AsRef<Array>,
91    freqs: impl Into<Option<&'a Array>>,
92) -> Result<Array> {
93    let stream = Stream::thread_local_or_default();
94    let base = base.into();
95    let base = mlx_sys::mlx_optional_float {
96        value: base.unwrap_or(0.0),
97        has_value: base.is_some(),
98    };
99    let freqs = freqs.into();
100    Array::try_from_op(|res| unsafe {
101        mlx_sys::mlx_fast_rope_dynamic(
102            res,
103            array.as_ref().as_ptr(),
104            dimensions,
105            traditional,
106            base,
107            scale,
108            offset.as_ref().as_ptr(),
109            freqs
110                .map(|a| a.as_ptr())
111                .unwrap_or(mlx_sys::mlx_array_new()),
112            stream.as_ref().as_ptr(),
113        )
114    })
115}
116
117/// Compatibility shim for [`rope_dynamic`].
118#[allow(clippy::too_many_arguments)]
119#[generate_macro(customize(forwarding_shim = true, root = "$crate::fast"))]
120#[deprecated(
121    since = "0.26.0",
122    note = "use `with_stream` or `with_device` around `rope_dynamic`"
123)]
124pub fn rope_dynamic_device<'a>(
125    #[named] array: impl AsRef<Array>,
126    #[named] dimensions: i32,
127    #[named] traditional: bool,
128    #[optional] base: impl Into<Option<f32>>,
129    #[named] scale: f32,
130    #[named] offset: impl AsRef<Array>,
131    #[optional] freqs: impl Into<Option<&'a Array>>,
132    #[optional] stream: impl AsRef<Stream>,
133) -> Result<Array> {
134    crate::with_stream(stream.as_ref(), || {
135        rope_dynamic(array, dimensions, traditional, base, scale, offset, freqs)
136    })
137}
138
139const DEFAULT_MASK_MODE: &CStr = c"";
140const CAUSAL_MASK_MODE: &CStr = c"causal";
141
142/// Mask modes for scaled dot product attention.
143#[derive(Debug)]
144pub enum ScaledDotProductAttentionMask<'a> {
145    /// A single mask array
146    Array(&'a Array),
147
148    /// Causal masking (no explicit mask array needed)
149    Causal,
150}
151
152impl<'a> From<&'a Array> for ScaledDotProductAttentionMask<'a> {
153    fn from(mask: &'a Array) -> Self {
154        ScaledDotProductAttentionMask::Array(mask)
155    }
156}
157
158impl<'a> IntoOption<ScaledDotProductAttentionMask<'a>> for &'a Array {
159    fn into_option(self) -> Option<ScaledDotProductAttentionMask<'a>> {
160        Some(ScaledDotProductAttentionMask::Array(self))
161    }
162}
163
164impl ScaledDotProductAttentionMask<'_> {
165    fn as_mode_and_mask(&self) -> (&'static CStr, mlx_sys::mlx_array) {
166        match self {
167            ScaledDotProductAttentionMask::Array(mask) => (DEFAULT_MASK_MODE, mask.as_ptr()),
168            ScaledDotProductAttentionMask::Causal => {
169                (CAUSAL_MASK_MODE, unsafe { mlx_sys::mlx_array_new() })
170            }
171        }
172    }
173}
174
175/// A fast implementation of multi-head attention: `O = softmax(Q @ K.T, dim=-1) @ V`
176///
177/// Supports [Multi-Head Attention](https://arxiv.org/abs/1706.03762), [Grouped Query Attention](https://arxiv.org/abs/2305.13245), and [Multi-Query Attention](https://arxiv.org/abs/1911.02150).
178///
179/// This function will dispatch to an optimized Metal kernel when the query sequence length is 1. It handles other cases with regular MLX operations.
180///
181/// > Note: The softmax operation is performed in float32 precision regardless of input precision (float16 or float32).
182///
183/// > Note: For Grouped Query Attention and Multi-Query Attention, the input arrays for `key` and `value` should not be pre-tiled to match the `query` array.
184pub fn scaled_dot_product_attention<'a>(
185    queries: impl AsRef<Array>,
186    keys: impl AsRef<Array>,
187    values: impl AsRef<Array>,
188    scale: f32,
189    mask: impl IntoOption<ScaledDotProductAttentionMask<'a>>,
190    sinks: impl Into<Option<&'a Array>>,
191) -> Result<Array> {
192    let stream = Stream::thread_local_or_default();
193    let (mask_mode, mask_arr) = mask.into_option().map_or_else(
194        || (DEFAULT_MASK_MODE, unsafe { mlx_sys::mlx_array_new() }),
195        |m| m.as_mode_and_mask(),
196    );
197
198    Array::try_from_op(|res| unsafe {
199        mlx_sys::mlx_fast_scaled_dot_product_attention(
200            res,
201            queries.as_ref().as_ptr(),
202            keys.as_ref().as_ptr(),
203            values.as_ref().as_ptr(),
204            scale,
205            mask_mode.as_ptr(),
206            mask_arr,
207            sinks
208                .into()
209                .map(|a| a.as_ptr())
210                .unwrap_or(mlx_sys::mlx_array_new()),
211            false,
212            stream.as_ref().as_ptr(),
213        )
214    })
215}
216
217/// Compatibility shim for [`scaled_dot_product_attention`].
218#[generate_macro(customize(forwarding_shim = true, root = "$crate::fast"))]
219#[deprecated(
220    since = "0.26.0",
221    note = "use `with_stream` or `with_device` around `scaled_dot_product_attention`"
222)]
223pub fn scaled_dot_product_attention_device<'a>(
224    queries: impl AsRef<Array>,
225    keys: impl AsRef<Array>,
226    values: impl AsRef<Array>,
227    scale: f32,
228    #[optional] mask: impl IntoOption<ScaledDotProductAttentionMask<'a>>,
229    #[optional] sinks: impl Into<Option<&'a Array>>,
230    #[optional] stream: impl AsRef<Stream>,
231) -> Result<Array> {
232    crate::with_stream(stream.as_ref(), || {
233        scaled_dot_product_attention(queries, keys, values, scale, mask, sinks)
234    })
235}
236
237/// Root Mean Square normalization (RMS norm).
238///
239/// The normalization is with respect to the last axis of the input `x`.
240///
241/// # Params
242///
243/// - x: input array
244/// - weight: An optional multiplicative weight. When present, it must be one-dimensional with the
245///   same size as the last axis of `x`. When absent, only normalization is applied.
246/// - eps: A small additive constant for numerical stability
247pub fn rms_norm(x: impl AsRef<Array>, weight: Option<&Array>, eps: f32) -> Result<Array> {
248    let stream = Stream::thread_local_or_default();
249    let empty_weight = weight
250        .is_none()
251        .then(|| unsafe { Array::from_ptr(mlx_sys::mlx_array_new()) });
252    let weight = weight
253        .map(Array::as_ptr)
254        .unwrap_or_else(|| empty_weight.as_ref().unwrap().as_ptr());
255    Array::try_from_op(|res| unsafe {
256        mlx_sys::mlx_fast_rms_norm(
257            res,
258            x.as_ref().as_ptr(),
259            weight,
260            eps,
261            stream.as_ref().as_ptr(),
262        )
263    })
264}
265
266/// Compatibility shim for [`rms_norm`].
267#[generate_macro(customize(forwarding_shim = true, root = "$crate::fast"))]
268#[deprecated(
269    since = "0.26.0",
270    note = "use `with_stream` or `with_device` around `rms_norm`"
271)]
272pub fn rms_norm_device(
273    x: impl AsRef<Array>,
274    weight: impl AsRef<Array>,
275    eps: f32,
276    #[optional] stream: impl AsRef<Stream>,
277) -> Result<Array> {
278    crate::with_stream(stream.as_ref(), || rms_norm(x, Some(weight.as_ref()), eps))
279}
280
281/// Layer normalization.
282///
283/// The normalization is with respect to the last axis of the input `x`.
284///
285/// # Params
286///
287/// - x: input array
288/// - weight: A multiplicative weight to scale the result by. The `weight` should be one-dimensional
289///   with the same size as the last axis of `x`.  If not given no scaling will occur.
290/// - bias: An additive offset to be added to the result. The `bias` should be one-dimensional
291///   with the same size as the last axis of `x`.  It not given no offset will occur.
292/// - eps: A small additive constant for numerical stability
293/// - stream: stream or device to evaluate on
294pub fn layer_norm<'a>(
295    x: impl AsRef<Array>,
296    weight: impl Into<Option<&'a Array>>,
297    bias: impl Into<Option<&'a Array>>,
298    eps: f32,
299) -> Result<Array> {
300    let stream = Stream::thread_local_or_default();
301    Array::try_from_op(|res| unsafe {
302        mlx_sys::mlx_fast_layer_norm(
303            res,
304            x.as_ref().as_ptr(),
305            weight
306                .into()
307                .map(|a| a.as_ptr())
308                .unwrap_or(mlx_sys::mlx_array_new()),
309            bias.into()
310                .map(|a| a.as_ptr())
311                .unwrap_or(mlx_sys::mlx_array_new()),
312            eps,
313            stream.as_ref().as_ptr(),
314        )
315    })
316}
317
318/// Compatibility shim for [`layer_norm`].
319#[generate_macro(customize(forwarding_shim = true, root = "$crate::fast"))]
320#[deprecated(
321    since = "0.26.0",
322    note = "use `with_stream` or `with_device` around `layer_norm`"
323)]
324pub fn layer_norm_device<'a>(
325    #[named] x: impl AsRef<Array>,
326    #[optional] weight: impl Into<Option<&'a Array>>,
327    #[optional] bias: impl Into<Option<&'a Array>>,
328    #[named] eps: f32,
329    #[optional] stream: impl AsRef<Stream>,
330) -> Result<Array> {
331    crate::with_stream(stream.as_ref(), || layer_norm(x, weight, bias, eps))
332}
333
334#[cfg(test)]
335mod tests {
336    use super::*;
337    use crate::{
338        ops::indexing::{ArrayIndexOp, IndexOp},
339        random::normal,
340    };
341    use float_eq::assert_float_eq;
342    use pretty_assertions::assert_eq;
343
344    #[test]
345    fn test_rope() {
346        crate::random::seed(71).unwrap();
347        let a = crate::random::uniform::<_, f32>(0.0, 1.0, &[2, 8, 16], None).unwrap();
348        assert_eq!(a.shape(), [2, 8, 16]);
349        assert_eq!(a.dtype(), crate::Dtype::Float32);
350
351        let result = rope(a, 8, false, 10000., 1.0, 0, None).unwrap();
352        assert_eq!(result.shape(), [2, 8, 16]);
353        assert_eq!(result.dtype(), crate::Dtype::Float32);
354        assert_float_eq!(
355            result.mean(None).unwrap().item_exact::<f32>(),
356            0.456_253_77,
357            abs <= 0.009_125_075
358        );
359        assert_float_eq!(
360            result.sum(None).unwrap().item_exact::<f32>(),
361            116.800_964,
362            abs <= 2.336_019_3
363        );
364    }
365
366    // Test adapted from Python test_fast.py/test_rope - the Python test accepts both
367    // int offset and array offset, which in C/Rust are separate functions
368    #[test]
369    fn test_rope_dynamic() {
370        crate::random::seed(71).unwrap();
371        let a = crate::random::uniform::<_, f32>(0.0, 1.0, &[2, 8, 16], None).unwrap();
372        assert_eq!(a.shape(), [2, 8, 16]);
373        assert_eq!(a.dtype(), crate::Dtype::Float32);
374
375        // Test with array offset - should produce similar results to int offset of 3
376        let offset = crate::Array::from_int(3);
377        let result = rope_dynamic(&a, 8, false, 10000., 1.0, &offset, None).unwrap();
378        assert_eq!(result.shape(), [2, 8, 16]);
379        assert_eq!(result.dtype(), crate::Dtype::Float32);
380
381        // Compare with regular rope using int offset=3
382        let result_int_offset = rope(&a, 8, false, 10000., 1.0, 3, None).unwrap();
383        assert_eq!(result_int_offset.shape(), [2, 8, 16]);
384
385        // The results should be close
386        let diff = &result - &result_int_offset;
387        let max_diff = diff.abs().unwrap().max(None).unwrap().item_exact::<f32>();
388        assert!(max_diff < 1e-5, "Max difference was {}", max_diff);
389    }
390
391    #[test]
392    fn test_rms_norm() {
393        crate::random::seed(103).unwrap();
394        let a = crate::random::uniform::<_, f32>(0.0, 1.0, &[2, 8, 16], None).unwrap();
395        assert_eq!(a.shape(), [2, 8, 16]);
396        assert_eq!(a.dtype(), crate::Dtype::Float32);
397
398        let weight = Array::ones::<f32>(&[16]).unwrap();
399        let result = rms_norm(a, Some(&weight), 1e-5).unwrap();
400        assert_eq!(result.shape(), [2, 8, 16]);
401        assert_eq!(result.dtype(), crate::Dtype::Float32);
402        assert_float_eq!(
403            result.mean(None).unwrap().item_exact::<f32>(),
404            0.872_938_75,
405            abs <= 0.017_458_774
406        );
407        assert_float_eq!(
408            result.sum(None).unwrap().item_exact::<f32>(),
409            223.472_32,
410            abs <= 4.469_446
411        );
412
413        let ones = Array::ones::<f32>(&[2, 4]).unwrap();
414        let normalized = rms_norm(&ones, None, 0.0).unwrap();
415        assert_eq!(normalized.as_slice::<f32>(), &[1.0; 8]);
416    }
417
418    #[test]
419    pub fn test_layer_norm_affine() {
420        crate::random::seed(635).unwrap();
421        let a = crate::random::uniform::<_, f32>(0.0, 1.0, &[2, 8, 16], None).unwrap();
422        assert_eq!(a.shape(), [2, 8, 16]);
423        assert_eq!(a.dtype(), crate::Dtype::Float32);
424
425        let weight = Array::ones::<f32>(&[16]).unwrap();
426        let bias = Array::zeros::<f32>(&[16]).unwrap();
427        let result = layer_norm(a, &weight, &bias, 1e-5).unwrap();
428        let result = result.index((ArrayIndexOp::Ellipsis, 0));
429        assert_eq!(result.shape(), [2, 8]);
430        assert_eq!(result.dtype(), crate::Dtype::Float32);
431        assert_float_eq!(
432            result.mean(None).unwrap().item_exact::<f32>(),
433            0.290_990_38,
434            abs <= 0.005_819_807_8
435        );
436        assert_float_eq!(
437            result.sum(None).unwrap().item_exact::<f32>(),
438            4.655_846,
439            abs <= 0.093_116_924
440        );
441    }
442
443    #[test]
444    #[allow(non_snake_case)]
445    fn test_fast_sdpa() {
446        // This test just makes sure that `scaled_dot_product_attention` is callable
447        // in the various cases, based on the Python test `test_fast_sdpa`.
448
449        let Dk = 64;
450        let scale = 1.0 / (Dk as f32).sqrt();
451        for seq_len in [63, 129, 400] {
452            for dtype in [crate::Dtype::Float32, crate::Dtype::Float16] {
453                let B = 2;
454                let H = 24;
455                let q = normal::<f32>(&[B, H, seq_len, Dk], None, None, None)
456                    .unwrap()
457                    .as_dtype(dtype)
458                    .unwrap();
459                let k = normal::<f32>(&[B, H, seq_len, Dk], None, None, None)
460                    .unwrap()
461                    .as_dtype(dtype)
462                    .unwrap();
463                let v = normal::<f32>(&[B, H, seq_len, Dk], None, None, None)
464                    .unwrap()
465                    .as_dtype(dtype)
466                    .unwrap();
467
468                let result = scaled_dot_product_attention(q, k, v, scale, None, None).unwrap();
469                assert_eq!(result.shape(), [B, H, seq_len, Dk]);
470                assert_eq!(result.dtype(), dtype);
471                result.eval().unwrap();
472            }
473        }
474    }
475
476    // Test adapted from Python test `test_fast_sdpa.py/test_sdpa_attention_sinks`
477    #[test]
478    fn test_fast_sdpa_with_sinks() {
479        let b = 2;
480        let n_q = 8;
481        let t_q = 128;
482        let t_kv = 128;
483        let d = 64;
484
485        let q = normal::<f32>(&[b, n_q, t_q, d], None, None, None).unwrap();
486        let k = normal::<f32>(&[b, n_q, t_kv, d], None, None, None).unwrap();
487        let v = normal::<f32>(&[b, n_q, t_kv, d], None, None, None).unwrap();
488        let scale = (d as f32).powf(-0.5);
489
490        // Test with sinks parameter
491        let sinks = normal::<f32>(&[n_q], None, None, None).unwrap() * 10.0;
492
493        let result = scaled_dot_product_attention(&q, &k, &v, scale, None, &sinks).unwrap();
494        assert_eq!(result.shape(), &[b, n_q, t_q, d]);
495    }
496}
mlx_rs/fast.rs

mlx_rs/
fast.rs
