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

mlx_rs/
fast.rs
