libmir-metal 0.3.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
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146

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

#[derive(Debug)]
pub struct FusedGateUp {
    arrays: QuantizedArrays,
    input_width: usize,
    gate_width: usize,
    up_width: usize,
    group_size: i32,
    bits: i32,
}

#[derive(Debug)]
pub struct GateUpOutput {
    pub gate: Array,
    pub up: Array,
}

impl FusedGateUp {
    pub(crate) fn new(
        gate: &QuantizedArrays,
        up: &QuantizedArrays,
        group_size: i32,
        bits: i32,
        stream: &Stream,
    ) -> Result<Self> {
        let gate_shape = dimensions(&gate.weight)?;
        let up_shape = dimensions(&up.weight)?;
        if gate_shape.len() != 2 || up_shape.len() != 2 || gate_shape[1] != up_shape[1] {
            return Err(Error::InvalidQuantization(
                "fused gate/up weights are incompatible".into(),
            ));
        }
        let input_width = logical_input_width(gate, group_size)?;
        if logical_input_width(up, group_size)? != input_width {
            return Err(Error::InvalidQuantization(
                "fused gate/up logical input widths are incompatible".into(),
            ));
        }
        Ok(Self {
            arrays: concatenate(gate, up, 0, group_size, bits, stream)?,
            input_width,
            gate_width: gate_shape[0],
            up_width: up_shape[0],
            group_size,
            bits,
        })
    }

    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<GateUpOutput> {
        let output = input.quantized_matmul(&self.arrays, true, stream)?;
        let (gate, up) = split_last(&output, self.gate_width, stream)?;
        Ok(GateUpOutput { gate, up })
    }

    pub(crate) fn forward_pair(&self, input: &Array, stream: &Stream) -> Result<(Array, Array)> {
        let output = self.forward(input, stream)?;
        Ok((output.gate, output.up))
    }

    pub(crate) const fn tuning_geometry(&self) -> (usize, usize, usize, i32, i32) {
        (self.input_width, self.gate_width, self.up_width, self.group_size, self.bits)
    }
}

pub(super) fn concatenate(
    first: &QuantizedArrays,
    second: &QuantizedArrays,
    axis: i32,
    group_size: i32,
    bits: i32,
    stream: &Stream,
) -> Result<QuantizedArrays> {
    let graph = stream.native().graph();
    QuantizedArrays::new(
        Array::from_native(
            graph.concatenate(&[first.weight.native(), second.weight.native()], axis)?,
        )?,
        Array::from_native(
            graph.concatenate(&[first.scales.native(), second.scales.native()], axis)?,
        )?,
        Array::from_native(
            graph.concatenate(&[first.biases.native(), second.biases.native()], axis)?,
        )?,
        group_size,
        bits,
    )
}

pub(super) fn split_last(input: &Array, width: usize, stream: &Stream) -> Result<(Array, Array)> {
    let shape = input.native().shape()?;
    let rank = shape.dimensions().len();
    let total = *shape.dimensions().last().ok_or(Error::ShapeOverflow)?;
    let mut start = vec![0; rank];
    let mut stop = shape.dimensions().to_vec();
    stop[rank - 1] = width;
    let graph = stream.native().graph();
    let first = Array::from_native(graph.slice(input.native(), &start, &stop)?)?;
    start[rank - 1] = width;
    stop[rank - 1] = total;
    Ok((first, Array::from_native(graph.slice(input.native(), &start, &stop)?)?))
}

pub(super) fn split_interleaved_last(
    input: &Array,
    width: usize,
    stream: &Stream,
) -> Result<(Array, Array)> {
    let mut shape = input.shape()?;
    let last = shape.len().checked_sub(1).ok_or(Error::ShapeOverflow)?;
    if usize::try_from(shape[last])? != width.checked_mul(2).ok_or(Error::ShapeOverflow)? {
        return Err(Error::InvalidModel("interleaved gate/up width differs".into()));
    }
    shape[last] = i32::try_from(width)?;
    shape.push(2);
    let paired = input.reshape(&shape, stream)?;
    let mut start = vec![0; shape.len()];
    let mut stop = shape
        .iter()
        .map(|value| Ok(usize::try_from(*value)?))
        .collect::<Result<Vec<_>>>()?;
    stop[last + 1] = 1;
    let graph = stream.native().graph();
    let gate = Array::from_native(graph.slice(paired.native(), &start, &stop)?)?
        .squeeze_axis(-1, stream)?;
    start[last + 1] = 1;
    stop[last + 1] = 2;
    let up = Array::from_native(graph.slice(paired.native(), &start, &stop)?)?
        .squeeze_axis(-1, stream)?;
    Ok((gate, up))
}

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

fn logical_input_width(arrays: &QuantizedArrays, group_size: i32) -> Result<usize> {
    let groups = dimensions(&arrays.scales)?.last().copied().ok_or(Error::ShapeOverflow)?;
    groups.checked_mul(usize::try_from(group_size)?).ok_or(Error::ShapeOverflow)
}