libmir-metal 0.1.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
121
122
123
124

use super::{
    Array, DecoderCache, Result, Stream, dense_swiglu::DenseSwiGluModel,
    hybrid_linear_moe::HybridLinearMoeModel, hybrid_moe::HybridMoeModel,
};

#[derive(Debug)]
pub enum DecoderModel {
    HybridMoe(HybridMoeModel),
    HybridLinearMoe(HybridLinearMoeModel),
    DenseSwiGlu(DenseSwiGluModel),
}

impl DecoderModel {
    pub(crate) fn prefers_packed_decode(&self, stream: &Stream) -> bool {
        match self {
            Self::DenseSwiGlu(_) => stream.config().batch.dense_decode.packed(),
            Self::HybridMoe(_) | Self::HybridLinearMoe(_) => true,
        }
    }

    pub(crate) fn new_cache(&self) -> Result<DecoderCache> {
        match self {
            Self::HybridMoe(model) => model.new_cache(),
            Self::HybridLinearMoe(model) => model.new_cache(),
            Self::DenseSwiGlu(model) => model.new_cache(),
        }
    }

    pub(crate) fn forward_decode(
        &self,
        token_ids: &Array,
        cache: &mut DecoderCache,
        position: i32,
        stream: &Stream,
    ) -> Result<Array> {
        match self {
            Self::HybridMoe(model) => model.forward_decode(token_ids, cache, position, stream),
            Self::HybridLinearMoe(model) => {
                model.forward_decode(token_ids, cache, position, stream)
            },
            Self::DenseSwiGlu(model) => model.forward_decode(token_ids, cache, position, stream),
        }
    }

    pub(crate) fn forward_packed_decode(
        &self,
        token_ids: &Array,
        caches: &mut [&mut DecoderCache],
        positions: &[i32],
        stream: &Stream,
    ) -> Result<Array> {
        if caches.len() != positions.len() {
            return Err(super::Error::InvalidModel(
                "packed cache and position row counts differ".into(),
            ));
        }
        match self {
            Self::DenseSwiGlu(model) => {
                model.forward_packed_decode(token_ids, caches, positions, stream)
            },
            Self::HybridMoe(model) => {
                model.forward_packed_decode(token_ids, caches, positions, stream)
            },
            Self::HybridLinearMoe(model) => {
                model.forward_packed_decode(token_ids, caches, positions, stream)
            },
        }
    }

    pub(crate) fn forward_greedy_decode(
        &self,
        token_ids: &Array,
        cache: &mut DecoderCache,
        position: i32,
        stream: &Stream,
    ) -> Result<Array> {
        match self {
            Self::HybridMoe(model) => {
                model.forward_greedy_decode(token_ids, cache, position, stream)
            },
            Self::HybridLinearMoe(model) => {
                model.forward_decode(token_ids, cache, position, stream)
            },
            Self::DenseSwiGlu(model) => model.forward_decode(token_ids, cache, position, stream),
        }
    }

    pub(crate) fn forward_prefill(
        &self,
        token_ids: &Array,
        cache: &mut DecoderCache,
        position: i32,
        stream: &Stream,
    ) -> Result<Array> {
        match self {
            Self::HybridMoe(model) => model.forward_prefill(token_ids, cache, position, stream),
            Self::HybridLinearMoe(model) => {
                model.forward_prefill(token_ids, cache, position, stream)
            },
            Self::DenseSwiGlu(model) => model.forward_prefill(token_ids, cache, position, stream),
        }
    }

    pub(crate) fn fusion_summary(&self) -> (usize, usize, usize, usize) {
        match self {
            Self::HybridMoe(model) => model.fusion_summary(),
            Self::HybridLinearMoe(model) => model.fusion_summary(),
            Self::DenseSwiGlu(model) => {
                let (attention, gate_up) = model.fusion_summary();
                (attention, 0, gate_up, 0)
            },
        }
    }

    pub(crate) fn expert_fusion_summary(&self) -> String {
        match self {
            Self::HybridMoe(model) => model.expert_fusion_summary(),
            Self::HybridLinearMoe(model) => model.expert_fusion_summary(),
            Self::DenseSwiGlu(_) => {
                "expert gate/up fusion is not applicable to dense SwiGLU".into()
            },
        }
    }
}