mirtal 0.1.0

Native, explicit-stream Rust API for Apple Metal
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61

use cxx::SharedPtr;
use mirtal_sys::ffi;

use crate::{Array, Error, Graph, Result};

#[derive(Debug, Default, Clone, Copy)]
/// Mask policy for scaled dot-product attention.
pub enum AttentionMask<'array> {
    #[default]
    /// Applies no attention mask.
    None,
    /// Applies MLX's causal attention mask.
    Causal,
    /// Uses a device-resident tensor as the attention mask.
    Array(&'array Array),
}

#[derive(Debug, Clone, Copy)]
/// Options for scaled dot-product attention.
pub struct ScaledDotProductAttention<'array> {
    /// Multiplier applied to query-key scores.
    pub scale: f32,
    /// Mask policy applied before softmax.
    pub mask: AttentionMask<'array>,
    /// Optional device-resident attention-sink values.
    pub sinks: Option<&'array Array>,
}

impl Graph<'_> {
    /// Computes scaled dot-product attention on this graph's stream.
    pub fn scaled_dot_product_attention(
        self,
        queries: &Array,
        keys: &Array,
        values: &Array,
        options: ScaledDotProductAttention<'_>,
    ) -> Result<Array> {
        if !options.scale.is_finite() || options.scale <= 0.0 {
            return Err(Error::InvalidAttention("scale must be finite and positive".into()));
        }
        let (mask_kind, mask) = match options.mask {
            AttentionMask::None => (0, SharedPtr::null()),
            AttentionMask::Causal => (1, SharedPtr::null()),
            AttentionMask::Array(mask) => (2, mask.raw_clone()),
        };
        let sinks = options.sinks.map_or_else(SharedPtr::null, Array::raw_clone);
        let native = ffi::NativeAttentionOptions { scale: options.scale, mask_kind };
        Array::from_raw(
            ffi::sdpa(
                queries.native()?,
                keys.native()?,
                values.native()?,
                &native,
                mask,
                sinks,
                self.native()?,
            )?,
            "scaled dot-product attention",
        )
    }
}