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

use super::{
    Array, Error, QuantizedArrays, Result, Stream,
    fused_gate_up::{concatenate, split_last},
};

#[derive(Debug)]
pub struct FusedAttention {
    arrays: QuantizedArrays,
    query_width: usize,
    key_width: usize,
    has_value: bool,
}

#[derive(Debug)]
pub struct AttentionProjectionOutput {
    pub query: Array,
    pub key: Array,
    pub value: Option<Array>,
}

impl FusedAttention {
    pub(crate) fn new(
        query: &QuantizedArrays,
        key: &QuantizedArrays,
        value: Option<&QuantizedArrays>,
        group_size: i32,
        bits: i32,
        stream: &Stream,
    ) -> Result<Self> {
        let query_shape = dimensions(&query.weight)?;
        let key_shape = dimensions(&key.weight)?;
        validate_pair(&query_shape, &key_shape)?;
        let mut arrays = concatenate(query, key, 0, group_size, bits, stream)?;
        if let Some(value) = value {
            let value_shape = dimensions(&value.weight)?;
            validate_pair(&query_shape, &value_shape)?;
            arrays = concatenate(&arrays, value, 0, group_size, bits, stream)?;
        }
        Ok(Self {
            arrays,
            query_width: query_shape[0],
            key_width: key_shape[0],
            has_value: value.is_some(),
        })
    }

    pub(crate) fn warm(&self) -> Result<()> {
        self.arrays.weight.async_eval()?;
        self.arrays.scales.async_eval()?;
        self.arrays.biases.async_eval()
    }

    pub(crate) fn forward(
        &self,
        input: &Array,
        stream: &Stream,
    ) -> Result<AttentionProjectionOutput> {
        let output =
            input.quantized_matmul(&self.arrays, true, stream)?.astype_like(input, stream)?;
        let (query, remainder) = split_last(&output, self.query_width, stream)?;
        let (key, value) = split_last(&remainder, self.key_width, stream)?;
        Ok(AttentionProjectionOutput {
            query,
            key,
            value: self.has_value.then_some(value),
        })
    }
}

fn validate_pair(first: &[usize], second: &[usize]) -> Result<()> {
    if first.len() != 2 || second.len() != 2 || first[1] != second[1] {
        return Err(Error::InvalidQuantization("fused attention weights are incompatible".into()));
    }
    Ok(())
}

fn dimensions(array: &Array) -> Result<Vec<usize>> {
    Ok(array.native().shape()?.dimensions().to_vec())
}