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

use models::weights::GatedSoftmaxBindings;

use super::{GatedFullAttention, GatedFullAttentionConfig};
use crate::engine::{
    ModelTensors, NormWeight, QuantizedLinear, Result, Stream,
    binding::{BoundLinear, adjusted_norm},
};

impl GatedFullAttention {
    pub fn load_with_norm_shift(
        tensors: &ModelTensors,
        prefix: &str,
        config: GatedFullAttentionConfig,
        group_size: i32,
        norm_shift: f32,
        stream: &Stream,
    ) -> Result<Self> {
        Ok(Self {
            config,
            query: linear(tensors, prefix, "q_proj", group_size)?,
            key: linear(tensors, prefix, "k_proj", group_size)?,
            value: linear(tensors, prefix, "v_proj", group_size)?,
            output: linear(tensors, prefix, "o_proj", group_size)?,
            query_norm: NormWeight::load_adjusted(
                tensors,
                &format!("{prefix}.q_norm"),
                norm_shift,
                stream,
            )?,
            key_norm: NormWeight::load_adjusted(
                tensors,
                &format!("{prefix}.k_norm"),
                norm_shift,
                stream,
            )?,
        })
    }

    pub fn load_bindings(
        tensors: &ModelTensors,
        bindings: GatedSoftmaxBindings<'_>,
        config: GatedFullAttentionConfig,
        norm_shift: f32,
        stream: &Stream,
    ) -> Result<Self> {
        Ok(Self {
            config,
            query: BoundLinear::load(tensors, bindings.query, stream)?,
            key: BoundLinear::load(tensors, bindings.key, stream)?,
            value: BoundLinear::load(tensors, bindings.value, stream)?,
            output: BoundLinear::load(tensors, bindings.output, stream)?,
            query_norm: adjusted_norm(tensors, bindings.query_norm, norm_shift, stream)?,
            key_norm: adjusted_norm(tensors, bindings.key_norm, norm_shift, stream)?,
        })
    }
}

fn linear(
    tensors: &ModelTensors,
    prefix: &str,
    name: &str,
    group_size: i32,
) -> Result<BoundLinear> {
    QuantizedLinear::load(tensors, &format!("{prefix}.{name}"), group_size).map(BoundLinear::Affine)
}