libmir-metal 0.1.0

Metal inference backend for libmir
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
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94

use super::PagedAttentionScratch;
use crate::engine::{Array, Result, Stream};

#[derive(Debug, Clone, Copy)]
pub struct PagedAttention<'a> {
    pub key_pages: &'a Array,
    pub value_pages: &'a Array,
    pub page_table: &'a Array,
    pub page_dependency: &'a Array,
    pub page_size: usize,
    pub context_tokens: usize,
}

impl Array {
    pub fn scaled_dot_product_attention(
        &self,
        keys: &Self,
        values: &Self,
        scale: f32,
        causal: bool,
        stream: &Stream,
    ) -> Result<Self> {
        let mask = if causal {
            mirtal::AttentionMask::Causal
        } else {
            mirtal::AttentionMask::None
        };
        Self::from_native(stream.native().graph().scaled_dot_product_attention(
            self.native(),
            keys.native(),
            values.native(),
            mirtal::ScaledDotProductAttention { scale, mask, sinks: None },
        )?)?
        .astype_like(self, stream)
    }

    pub fn masked_scaled_dot_product_attention(
        &self,
        keys: &Self,
        values: &Self,
        scale: f32,
        mask: &Self,
        stream: &Stream,
    ) -> Result<Self> {
        Self::from_native(stream.native().graph().scaled_dot_product_attention(
            self.native(),
            keys.native(),
            values.native(),
            mirtal::ScaledDotProductAttention {
                scale,
                mask: mirtal::AttentionMask::Array(mask.native()),
                sinks: None,
            },
        )?)?
        .astype_like(self, stream)
    }

    pub fn paged_scaled_dot_product_attention(
        &self,
        paged: PagedAttention<'_>,
        scale: f32,
        stream: &Stream,
    ) -> Result<Self> {
        self.paged_scaled_dot_product_attention_with_scratch(
            paged,
            &PagedAttentionScratch::default(),
            scale,
            stream,
        )
    }

    pub(crate) fn paged_scaled_dot_product_attention_with_scratch(
        &self,
        paged: PagedAttention<'_>,
        scratch: &PagedAttentionScratch,
        scale: f32,
        stream: &Stream,
    ) -> Result<Self> {
        Self::from_native(stream.paged_attention(
            [
                self.native(),
                paged.key_pages.native(),
                paged.value_pages.native(),
                paged.page_table.native(),
                paged.page_dependency.native(),
            ],
            scratch,
            paged.page_size,
            paged.context_tokens,
            scale,
        )?)?
        .astype_like(self, stream)
    }
}