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

use super::{LayerConfig, weights::AttentionWeights};
use crate::engine::{Array, Result, RopeOptions, Stream};

pub(super) fn forward(
    input: &Array,
    weights: &AttentionWeights,
    config: LayerConfig,
    stream: &Stream,
) -> Result<Array> {
    let sequence = input.shape()?.get(1).copied().ok_or_else(|| {
        crate::engine::Error::InvalidModel("embedding attention input has no sequence axis".into())
    })?;
    let queries = weights
        .query
        .forward(input, stream)?
        .reshape(&[1, sequence, config.heads, config.head_dim], stream)?;
    let keys = weights
        .key
        .forward(input, stream)?
        .reshape(&[1, sequence, config.kv_heads, config.head_dim], stream)?;
    let values = weights
        .value
        .forward(input, stream)?
        .reshape(&[1, sequence, config.kv_heads, config.head_dim], stream)?;
    let queries = normalize(queries, weights.query_norm.as_ref(), config, stream)?;
    let keys = normalize(keys, weights.key_norm.as_ref(), config, stream)?;
    let queries = rotate(&queries, config, stream)?;
    let keys = rotate(&keys, config, stream)?;
    let values = values.transpose(&[0, 2, 1, 3], stream)?;
    let output = queries.scaled_dot_product_attention(
        &keys,
        &values,
        config.attention_scale,
        true,
        stream,
    )?;
    weights.output.forward(
        &output
            .transpose(&[0, 2, 1, 3], stream)?
            .reshape(&[1, sequence, config.query_width], stream)?,
        stream,
    )
}

fn normalize(
    input: Array,
    weight: Option<&crate::engine::NormWeight>,
    config: LayerConfig,
    stream: &Stream,
) -> Result<Array> {
    match weight {
        Some(weight) => weight.apply(&input, config.rms_norm_eps, stream),
        None => Ok(input),
    }
}

fn rotate(input: &Array, config: LayerConfig, stream: &Stream) -> Result<Array> {
    input.transpose(&[0, 2, 1, 3], stream)?.rope(
        RopeOptions {
            dimensions: config.head_dim,
            traditional: false,
            base: Some(config.rope_base),
            scale: 1.0,
            offset: 0,
        },
        stream,
    )
}