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

use std::time::Instant;

use super::{HybridMoeLayer, profile_components};
use crate::engine::{
    Array, KvCache, Result, Stream,
    hybrid_moe::{
        attention::{self, DecodeContext},
        feed_forward,
    },
};

impl HybridMoeLayer {
    pub fn forward_uncached_decode(&self, input: &Array, stream: &Stream) -> Result<Array> {
        self.forward_decode(input, None, 0, false, stream)
    }

    pub fn forward_decode(
        &self,
        input: &Array,
        cache: Option<&mut KvCache>,
        position: i32,
        causal: bool,
        stream: &Stream,
    ) -> Result<Array> {
        let hidden = self.mix(input, cache, position, causal, None, stream)?;
        self.feed_forward(&hidden, causal, stream)
    }

    pub(super) fn mix(
        &self,
        input: &Array,
        cache: Option<&mut KvCache>,
        position: i32,
        causal: bool,
        mask: Option<&Array>,
        stream: &Stream,
    ) -> Result<Array> {
        let profile = !causal && profile_components(stream);
        let started = Instant::now();
        let normalized = self.weights.input_norm.apply(input, self.config.rms_norm_eps, stream)?;
        let attention = attention::forward_decode(
            &normalized,
            &self.weights.attention,
            self.config,
            self.fused_attention.as_ref(),
            self.fused_key_value.as_ref(),
            DecodeContext { cache, position, causal, mask, stream },
        )?;
        let attention =
            self.weights
                .post_attention_norm
                .apply(&attention, self.config.rms_norm_eps, stream)?;
        let hidden = input.add(&attention, stream)?;
        if profile {
            emit_profile(&hidden, stream, self.config.layer_index, "attention", started)?;
        }
        Ok(hidden)
    }

    pub(super) fn feed_forward(
        &self,
        input: &Array,
        causal: bool,
        stream: &Stream,
    ) -> Result<Array> {
        let profile = !causal && profile_components(stream);
        let started = Instant::now();
        let feed_forward = feed_forward::forward(
            input,
            &self.weights,
            self.config,
            self.fused_gate_up.as_ref(),
            self.fused_expert_gate_up.as_ref(),
            stream,
        )?;
        let feed_forward = self.weights.post_feed_forward_norm.apply(
            &feed_forward,
            self.config.rms_norm_eps,
            stream,
        )?;
        let output =
            input.add(&feed_forward, stream)?.multiply(&self.weights.layer_scalar, stream)?;
        if profile {
            emit_profile(&output, stream, self.config.layer_index, "feed_forward", started)?;
        }
        Ok(output)
    }
}

fn emit_profile(
    output: &Array,
    stream: &Stream,
    layer: usize,
    component: &str,
    started: Instant,
) -> Result<()> {
    output.async_eval()?;
    stream.synchronize()?;
    tracing::debug!(layer, component, "MLX hybrid MoE component synchronized");
    tracing::debug!(
        layer,
        component,
        milliseconds = started.elapsed().as_secs_f64() * 1_000.0,
        "MLX hybrid MoE component profile"
    );
    Ok(())
}