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

use models::semantic::{
    ActivationSpec, DecoderLayerSpec, FeedForwardSpec, MixerSpec, NormalizationKind,
    NormalizationSpec,
};

use super::super::{Error, Result};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NormalizationLowering {
    Rms,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MixerLowering {
    Softmax { sinks: bool, window: Option<usize> },
    Linear,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FeedForwardLowering {
    Dense,
    DenseAndRouted,
    SharedRouted,
    ClampedRouted,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct LayerLowering {
    pub index: usize,
    pub input_norm: NormalizationLowering,
    pub post_attention_norm: NormalizationLowering,
    pub mixer: MixerLowering,
    pub feed_forward: FeedForwardLowering,
}

pub fn lower_layer(layer: &DecoderLayerSpec) -> Result<LayerLowering> {
    Ok(LayerLowering {
        index: layer.index,
        input_norm: lower_normalization(layer.index, "input", layer.input_norm)?,
        post_attention_norm: lower_normalization(
            layer.index,
            "post-attention",
            layer.post_attention_norm,
        )?,
        mixer: lower_mixer(&layer.mixer),
        feed_forward: lower_feed_forward(layer.index, &layer.feed_forward)?,
    })
}

fn lower_normalization(
    layer: usize,
    position: &str,
    normalization: NormalizationSpec,
) -> Result<NormalizationLowering> {
    match normalization.kind {
        NormalizationKind::Rms => Ok(NormalizationLowering::Rms),
        NormalizationKind::Layer => Err(Error::InvalidModel(format!(
            "Metal cannot lower layer {layer} {position} layer normalization"
        ))),
    }
}

const fn lower_mixer(mixer: &MixerSpec) -> MixerLowering {
    match mixer {
        MixerSpec::SoftmaxAttention(attention) => MixerLowering::Softmax {
            sinks: attention.sinks,
            window: attention.window,
        },
        MixerSpec::LinearAttention(_) => MixerLowering::Linear,
    }
}

fn lower_feed_forward(layer: usize, feed_forward: &FeedForwardSpec) -> Result<FeedForwardLowering> {
    match feed_forward {
        FeedForwardSpec::Dense {
            activation: ActivationSpec::SwiGlu { clamp: None, up_shift: 0.0, .. },
            ..
        } => Ok(FeedForwardLowering::Dense),
        FeedForwardSpec::DenseAndRouted {
            dense_activation: ActivationSpec::GeluTanh,
            routed,
            ..
        } if routed.activation == ActivationSpec::GeluTanh => {
            Ok(FeedForwardLowering::DenseAndRouted)
        },
        FeedForwardSpec::Routed { shared: Some(_), .. } => Ok(FeedForwardLowering::SharedRouted),
        FeedForwardSpec::Routed { routed, shared: None }
            if matches!(
                routed.activation,
                ActivationSpec::SwiGlu { clamp: Some(_), up_shift, .. } if up_shift != 0.0
            ) =>
        {
            Ok(FeedForwardLowering::ClampedRouted)
        },
        _ => Err(Error::InvalidModel(format!(
            "Metal cannot lower layer {layer} feed-forward activation composition"
        ))),
    }
}