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
108
109
110
111
112
113
114
115
116
117
118

use crate::engine::{
    Error, GatedDeltaState, KvCache, KvPageFormat, Result, lowering::MixerLowering,
};

#[derive(Debug)]
pub(super) enum HybridLinearLayerCache {
    Linear(GatedDeltaState),
    Full(KvCache),
}

pub(super) fn new(
    mixers: &[MixerLowering],
    step: usize,
    format: KvPageFormat,
    page_size: usize,
) -> Result<Vec<HybridLinearLayerCache>> {
    mixers
        .iter()
        .map(|mixer| match mixer {
            MixerLowering::Linear => GatedDeltaState::new().map(HybridLinearLayerCache::Linear),
            MixerLowering::Softmax { window: None, .. } => {
                KvCache::new_paged_with_format(step, page_size, format)
                    .map(HybridLinearLayerCache::Full)
            },
            MixerLowering::Softmax { window: Some(window), .. } => {
                KvCache::new_with_window(step, Some(*window)).map(HybridLinearLayerCache::Full)
            },
        })
        .collect()
}

pub(super) fn reset(layers: &mut [HybridLinearLayerCache]) -> Result<()> {
    for layer in layers {
        match layer {
            HybridLinearLayerCache::Linear(state) => state.reset()?,
            HybridLinearLayerCache::Full(cache) => cache.reset()?,
        }
    }
    Ok(())
}

pub(super) fn reserve(layers: &mut [HybridLinearLayerCache], tokens: usize) -> Result<()> {
    for layer in layers {
        if let HybridLinearLayerCache::Full(cache) = layer {
            cache.reserve(tokens)?;
        }
    }
    Ok(())
}

pub(super) fn gated_delta_state(
    layers: &mut [HybridLinearLayerCache],
    index: usize,
) -> Result<&mut GatedDeltaState> {
    match layers.get_mut(index) {
        Some(HybridLinearLayerCache::Linear(state)) => Ok(state),
        Some(HybridLinearLayerCache::Full(_)) => {
            Err(Error::InvalidModel(format!("layer {index} requires a full attention cache")))
        },
        None => Err(Error::InvalidModel(format!("missing hybrid cache layer {index}"))),
    }
}

pub(super) fn full_attention_cache(
    layers: &mut [HybridLinearLayerCache],
    index: usize,
) -> Result<&mut KvCache> {
    match layers.get_mut(index) {
        Some(HybridLinearLayerCache::Full(cache)) => Ok(cache),
        Some(HybridLinearLayerCache::Linear(_)) => {
            Err(Error::InvalidModel(format!("layer {index} requires a linear attention state")))
        },
        None => Err(Error::InvalidModel(format!("missing hybrid cache layer {index}"))),
    }
}

pub(super) fn offset(layers: &[HybridLinearLayerCache]) -> Result<usize> {
    let Some(first) = layers.first() else {
        return Ok(0);
    };
    let offset = layer_offset(first)?;
    for layer in layers {
        if layer_offset(layer)? != offset {
            return Err(Error::InvalidModel("hybrid cache layers have diverging offsets".into()));
        }
    }
    Ok(offset)
}

pub(super) fn snapshot_at(
    layers: &[HybridLinearLayerCache],
    requested_offset: usize,
) -> Result<Vec<HybridLinearLayerCache>> {
    if offset(layers)? != requested_offset {
        return Err(Error::InvalidModel(
            "hybrid linear cache can only snapshot its current recurrent state".into(),
        ));
    }
    layers.iter().map(snapshot).collect()
}

fn layer_offset(layer: &HybridLinearLayerCache) -> Result<usize> {
    match layer {
        HybridLinearLayerCache::Linear(state) => state.offset(),
        HybridLinearLayerCache::Full(cache) => cache.offset(),
    }
}

fn snapshot(layer: &HybridLinearLayerCache) -> Result<HybridLinearLayerCache> {
    match layer {
        HybridLinearLayerCache::Linear(state) => {
            state.snapshot().map(HybridLinearLayerCache::Linear)
        },
        HybridLinearLayerCache::Full(cache) => {
            cache.snapshot_at(cache.offset()?).map(HybridLinearLayerCache::Full)
        },
    }
}