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
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130

use models::weights::RoutedExpertBindings;

use super::{
    super::{
        expert_tuning,
        fused_gate_up::{split_interleaved_last, split_last},
    },
    Array, BoundLinear, Error, FusedExpertGateUp, ModelTensors, Result, Stream,
};

#[derive(Debug)]
pub(super) enum RoutedGateUp {
    Separate {
        gate: BoundLinear,
        up: BoundLinear,
        fused: Option<Box<FusedExpertGateUp>>,
    },
    Fused {
        projection: BoundLinear,
        width: usize,
        interleaved: bool,
    },
}

impl RoutedGateUp {
    pub(super) fn load(
        tensors: &ModelTensors,
        bindings: RoutedExpertBindings<'_>,
        stream: &Stream,
    ) -> Result<(Self, BoundLinear)> {
        match bindings {
            RoutedExpertBindings::SeparateGateUp { gate, up, down } => Ok((
                Self::Separate {
                    gate: BoundLinear::load(tensors, gate, stream)?,
                    up: BoundLinear::load(tensors, up, stream)?,
                    fused: None,
                },
                BoundLinear::load(tensors, down, stream)?,
            )),
            RoutedExpertBindings::InterleavedGateUp { gate_up, down } => {
                let output = gate_up.shape.get(1).copied().ok_or(Error::ShapeOverflow)?;
                if !output.is_multiple_of(2) {
                    return Err(Error::InvalidModel(
                        "fused expert gate/up width must be even".into(),
                    ));
                }
                Ok((
                    Self::Fused {
                        projection: BoundLinear::load(tensors, gate_up, stream)?,
                        width: output / 2,
                        interleaved: gate_up.transforms.contains(
                            &models::weights::BindingTransform::FusedGateUp { interleaved: true },
                        ),
                    },
                    BoundLinear::load(tensors, down, stream)?,
                ))
            },
            RoutedExpertBindings::Individual { .. } => Err(Error::InvalidModel(
                "shared routed Metal execution requires stacked expert tensors".into(),
            )),
        }
    }

    pub(super) fn enable(&mut self, stream: &Stream) -> Result<bool> {
        match self {
            Self::Separate { gate, up, fused } => {
                if fused.is_none() {
                    *fused = gate.fuse_expert_gate_up(up, stream)?.map(Box::new);
                }
                fused.as_deref().map_or(Ok(()), FusedExpertGateUp::warm)?;
                Ok(fused.is_some())
            },
            Self::Fused { .. } => Ok(true),
        }
    }

    pub(super) fn fused_bytes(&self) -> Result<Option<usize>> {
        match self {
            Self::Separate { gate, up, .. } => gate.fused_expert_gate_up_bytes(up),
            Self::Fused { .. } => Ok(Some(0)),
        }
    }

    pub(super) const fn is_fused(&self) -> bool {
        match self {
            Self::Separate { fused, .. } => fused.is_some(),
            Self::Fused { .. } => true,
        }
    }

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

    pub(super) fn gather(
        &self,
        input: &Array,
        indices: &Array,
        sorted: bool,
        stream: &Stream,
    ) -> Result<(Array, Array)> {
        match self {
            Self::Separate { gate, up, fused } => {
                let fused = (!sorted).then_some(fused.as_deref()).flatten();
                fused.map_or_else(
                    || {
                        Ok((
                            gate.gather(input, indices, sorted, stream)?,
                            up.gather(input, indices, sorted, stream)?,
                        ))
                    },
                    |projection| {
                        expert_tuning::forward(gate, up, projection, input, indices, stream)
                    },
                )
            },
            Self::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)
                }
            },
        }
    }
}