libmir-metal 0.3.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

use crate::engine::{
    Array, NormWeight, Result, Stream,
    binding::BoundLinear,
    fused_gate_up::{split_interleaved_last, split_last},
};

mod load;

#[derive(Debug)]
pub(super) struct AttentionWeights {
    pub(super) query: BoundLinear,
    pub(super) key: BoundLinear,
    pub(super) value: Option<BoundLinear>,
    pub(super) output: BoundLinear,
    pub(super) query_norm: NormWeight,
    pub(super) key_norm: NormWeight,
    pub(super) rope_frequencies: Option<Array>,
}

#[derive(Debug)]
pub(super) struct DenseWeights {
    pub(super) gate: BoundLinear,
    pub(super) up: BoundLinear,
    pub(super) down: BoundLinear,
}

#[derive(Debug)]
pub(super) struct RouterWeights {
    pub(super) projection: BoundLinear,
    pub(super) norm_scale: Array,
    pub(super) expert_scale: Array,
}

#[derive(Debug)]
pub(super) struct ExpertWeights {
    pub(super) gate_up: ExpertGateUpWeights,
    pub(super) down: BoundLinear,
}

#[derive(Debug)]
pub(super) enum ExpertGateUpWeights {
    Separate {
        gate: BoundLinear,
        up: BoundLinear,
    },
    Fused {
        projection: BoundLinear,
        width: usize,
        interleaved: bool,
    },
}

impl ExpertWeights {
    pub(super) fn gather_gate_up(
        &self,
        input: &Array,
        indices: &Array,
        sorted: bool,
        stream: &Stream,
    ) -> Result<(Array, Array)> {
        match &self.gate_up {
            ExpertGateUpWeights::Separate { gate, up } => Ok((
                gate.gather(input, indices, sorted, stream)?,
                up.gather(input, indices, sorted, stream)?,
            )),
            ExpertGateUpWeights::Fused { projection, width, interleaved } => {
                let output = projection.gather(input, indices, sorted, stream)?;
                if *interleaved {
                    split_interleaved_last(&output, *width, stream)
                } else {
                    split_last(&output, *width, stream)
                }
            },
        }
    }

    pub(super) const fn separate(&self) -> Option<(&BoundLinear, &BoundLinear)> {
        match &self.gate_up {
            ExpertGateUpWeights::Separate { gate, up } => Some((gate, up)),
            ExpertGateUpWeights::Fused { .. } => None,
        }
    }

    pub(super) fn tuning_format(&self) -> (i32, i32) {
        match &self.gate_up {
            ExpertGateUpWeights::Separate { gate, .. } => gate.tuning_format(),
            ExpertGateUpWeights::Fused { projection, .. } => projection.tuning_format(),
        }
    }
}

#[derive(Debug)]
pub(super) struct LayerWeights {
    pub(super) input_norm: NormWeight,
    pub(super) post_attention_norm: NormWeight,
    pub(super) pre_dense_norm: NormWeight,
    pub(super) post_dense_norm: NormWeight,
    pub(super) pre_expert_norm: NormWeight,
    pub(super) post_expert_norm: NormWeight,
    pub(super) post_feed_forward_norm: NormWeight,
    pub(super) layer_scalar: Array,
    pub(super) attention: AttentionWeights,
    pub(super) dense: DenseWeights,
    pub(super) router: RouterWeights,
    pub(super) experts: ExpertWeights,
}