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

use models::weights::TextTensorLayout;

use super::{LayerConfig, attention, weights::LayerWeights};
use crate::engine::{Array, ModelTensors, Result, Stream};

#[derive(Debug)]
pub(super) struct TextEmbeddingLayer {
    config: LayerConfig,
    weights: LayerWeights,
}

impl TextEmbeddingLayer {
    pub fn load(
        tensors: &ModelTensors,
        config: LayerConfig,
        layout: &TextTensorLayout,
        stream: &Stream,
    ) -> Result<Self> {
        Ok(Self {
            config,
            weights: LayerWeights::load(tensors, config, layout, stream)?,
        })
    }

    pub fn forward(&self, input: &Array, stream: &Stream) -> Result<Array> {
        let normalized = self.weights.input_norm.apply(input, self.config.rms_norm_eps, stream)?;
        let attention =
            attention::forward(&normalized, &self.weights.attention, self.config, stream)?;
        let hidden = input.add(&attention, stream)?;
        let normalized =
            self.weights
                .post_attention_norm
                .apply(&hidden, self.config.rms_norm_eps, stream)?;
        let gate = self.weights.gate.forward(&normalized, stream)?;
        let up = self.weights.up.forward(&normalized, stream)?;
        let activated = gate.silu_mul(&up, stream)?;
        hidden.add(&self.weights.down.forward(&activated, stream)?, stream)
    }
}