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

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

#[derive(Debug)]
pub struct FusedKeyValue {
    pair: FusedGateUp,
}

impl FusedKeyValue {
    pub(crate) fn new(
        key: &QuantizedArrays,
        value: &QuantizedArrays,
        group_size: i32,
        bits: i32,
        stream: &Stream,
    ) -> Result<Self> {
        Ok(Self {
            pair: FusedGateUp::new(key, value, group_size, bits, stream)?,
        })
    }

    pub(crate) fn warm(&self) -> Result<()> {
        self.pair.warm()
    }

    pub(crate) fn forward(&self, input: &Array, stream: &Stream) -> Result<(Array, Array)> {
        self.pair.forward_pair(input, stream)
    }
}