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
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120

use models::weights::RoutedDecoderLayerBindings;

use super::{config::ClampedRoutedConfig, projection::BoundLinear};
use crate::engine::{
    Array, KvCache, ModelTensors, NormWeight, PagedContextMode, Result, Stream,
    paged_attention_min_context,
};

#[derive(Debug)]
pub(super) struct ClampedRoutedAttention {
    norm: NormWeight,
    query: BoundLinear,
    key: BoundLinear,
    value: BoundLinear,
    output: BoundLinear,
    sinks: Array,
    frequencies: Array,
    config: ClampedRoutedConfig,
}

impl ClampedRoutedAttention {
    pub fn load(
        tensors: &ModelTensors,
        bindings: RoutedDecoderLayerBindings<'_>,
        config: ClampedRoutedConfig,
        stream: &Stream,
    ) -> Result<Self> {
        Ok(Self {
            norm: NormWeight::load_name(tensors, &bindings.input_norm.source)?,
            query: BoundLinear::load_binding(tensors, bindings.query, stream)?,
            key: BoundLinear::load_binding(tensors, bindings.key, stream)?,
            value: BoundLinear::load_binding(tensors, bindings.value, stream)?,
            output: BoundLinear::load_binding(tensors, bindings.attention_output, stream)?,
            sinks: tensors.get(&bindings.attention_sinks.source)?,
            frequencies: Array::yarn_rope_frequencies(
                config.head_dim,
                config.rope_base,
                config.rope_factor,
                config.beta_fast,
                config.beta_slow,
                config.original_context,
                stream,
            )?,
            config,
        })
    }

    pub fn forward(
        &self,
        input: &Array,
        cache: &mut KvCache,
        position: i32,
        causal: bool,
        stream: &Stream,
    ) -> Result<Array> {
        let sequence = *input.shape()?.get(1).ok_or_else(|| {
            crate::engine::Error::InvalidModel(
                "clamped-routed attention input has no sequence axis".into(),
            )
        })?;
        let hidden = self.norm.apply(input, self.config.epsilon, stream)?;
        let queries = self
            .query
            .forward(&hidden, stream)?
            .reshape(&[1, sequence, self.config.heads, self.config.head_dim], stream)?;
        let keys = self
            .key
            .forward(&hidden, stream)?
            .reshape(&[1, sequence, self.config.kv_heads, self.config.head_dim], stream)?;
        let values = self
            .value
            .forward(&hidden, stream)?
            .reshape(&[1, sequence, self.config.kv_heads, self.config.head_dim], stream)?;
        let queries = self.rope(&queries, position, stream)?;
        let keys = self.rope(&keys, position, stream)?;
        let values = values.transpose(&[0, 2, 1, 3], stream)?;
        let context = cache.update_for_attention_mode(
            &keys,
            &values,
            stream,
            paged_attention_min_context(stream),
            PagedContextMode::View,
        )?;
        let attended = match context.mask.as_ref() {
            Some(mask) => queries.masked_scaled_dot_product_attention_with_sinks(
                &context.keys,
                &context.values,
                self.config.scale,
                mask,
                &self.sinks,
                stream,
            )?,
            None => queries.scaled_dot_product_attention_with_sinks(
                &context.keys,
                &context.values,
                self.config.scale,
                causal,
                &self.sinks,
                stream,
            )?,
        };
        let output = attended
            .transpose(&[0, 2, 1, 3], stream)?
            .reshape(&[1, sequence, self.config.heads * self.config.head_dim], stream)?;
        self.output.forward(&output, stream)
    }

    fn rope(&self, input: &Array, position: i32, stream: &Stream) -> Result<Array> {
        input
            .transpose(&[0, 2, 1, 3], stream)?
            .rope_with_frequencies(
                self.config.head_dim,
                false,
                &self.frequencies,
                position,
                stream,
            )?
            .multiply_scalar(self.config.rope_concentration, stream)
    }
}