libmir-metal 0.2.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

use models::weights::DenseDecoderLayerBindings;

use super::config::DenseSwiGluLayerConfig;
use crate::engine::{
    Array, ModelTensors, NormWeight, QuantizedLinear, Result, Stream, binding::affine_linear,
};

#[derive(Debug)]
pub(super) struct AttentionWeights {
    pub(super) query: QuantizedLinear,
    pub(super) key: QuantizedLinear,
    pub(super) value: QuantizedLinear,
    pub(super) output: QuantizedLinear,
    pub(super) query_norm: Option<NormWeight>,
    pub(super) key_norm: Option<NormWeight>,
    pub(super) rope_frequencies: Option<Array>,
}

#[derive(Debug)]
pub(super) struct MlpWeights {
    pub(super) gate: QuantizedLinear,
    pub(super) up: QuantizedLinear,
    pub(super) down: QuantizedLinear,
}

#[derive(Debug)]
pub(super) struct DenseSwiGluWeights {
    pub(super) input_norm: NormWeight,
    pub(super) post_attention_norm: NormWeight,
    pub(super) attention: AttentionWeights,
    pub(super) mlp: MlpWeights,
}

impl DenseSwiGluWeights {
    pub(super) fn load_bindings(
        tensors: &ModelTensors,
        bindings: DenseDecoderLayerBindings<'_>,
        config: DenseSwiGluLayerConfig,
        stream: &Stream,
    ) -> Result<Self> {
        let norm = |binding: Option<&models::weights::TensorBinding>| {
            binding
                .map(|binding| NormWeight::load_name(tensors, &binding.source))
                .transpose()
        };
        Ok(Self {
            input_norm: NormWeight::load_name(tensors, &bindings.input_norm.source)?,
            post_attention_norm: NormWeight::load_name(
                tensors,
                &bindings.post_attention_norm.source,
            )?,
            attention: AttentionWeights {
                query: affine_linear(tensors, bindings.attention.query)?,
                key: affine_linear(tensors, bindings.attention.key)?,
                value: affine_linear(tensors, bindings.attention.value)?,
                output: affine_linear(tensors, bindings.attention.output)?,
                query_norm: norm(bindings.attention.query_norm)?,
                key_norm: norm(bindings.attention.key_norm)?,
                rope_frequencies: rope_frequencies(config, stream)?,
            },
            mlp: MlpWeights {
                gate: affine_linear(tensors, bindings.gate)?,
                up: affine_linear(tensors, bindings.up)?,
                down: affine_linear(tensors, bindings.down)?,
            },
        })
    }
}

fn rope_frequencies(config: DenseSwiGluLayerConfig, stream: &Stream) -> Result<Option<Array>> {
    config
        .rope_scaling
        .map(|scaling| {
            let (factor, low_frequency_factor, high_frequency_factor, original_context_len) =
                scaling.piecewise_frequency().ok_or_else(|| {
                    crate::engine::Error::InvalidModel(
                        "dense SwiGLU does not support this RoPE scaling mode".into(),
                    )
                })?;
            Array::piecewise_rope_frequencies(
                config.head_dim,
                config.rope_base,
                factor.to_string().parse()?,
                low_frequency_factor.to_string().parse()?,
                high_frequency_factor.to_string().parse()?,
                i32::try_from(original_context_len)?,
                stream,
            )
        })
        .transpose()
}