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
119
120
121
122
123
124
125
126
127
128

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

mod hybrid;

#[derive(Debug)]
pub struct DecoderCache {
    storage: CacheStorage,
}

#[derive(Debug)]
enum CacheStorage {
    Attention(Vec<KvCache>),
    HybridLinear(Vec<HybridLinearLayerCache>),
}

impl DecoderCache {
    pub fn new(cache_windows: &[Option<usize>], step: usize) -> Result<Self> {
        Self::new_with_format(cache_windows, step, KvPageFormat::Native, 16)
    }

    pub(crate) fn new_with_format(
        cache_windows: &[Option<usize>],
        step: usize,
        format: KvPageFormat,
        page_size: usize,
    ) -> Result<Self> {
        let caches = cache_windows
            .iter()
            .map(|window| {
                if window.is_none() {
                    KvCache::new_paged_with_format(step, page_size, format)
                } else {
                    KvCache::new_with_window(step, *window)
                }
            })
            .collect::<Result<Vec<_>>>()?;
        Ok(Self { storage: CacheStorage::Attention(caches) })
    }

    pub fn new_hybrid_linear(mixers: &[MixerLowering], step: usize) -> Result<Self> {
        Self::new_hybrid_linear_with_format(mixers, step, KvPageFormat::Native, 16)
    }

    pub(crate) fn new_hybrid_linear_with_format(
        mixers: &[MixerLowering],
        step: usize,
        format: KvPageFormat,
        page_size: usize,
    ) -> Result<Self> {
        let layers = hybrid::new(mixers, step, format, page_size)?;
        Ok(Self {
            storage: CacheStorage::HybridLinear(layers),
        })
    }

    pub(crate) fn attention_caches_mut(&mut self) -> Result<&mut [KvCache]> {
        match &mut self.storage {
            CacheStorage::Attention(caches) => Ok(caches),
            CacheStorage::HybridLinear(_) => {
                Err(Error::InvalidModel("expected attention-only cache".into()))
            },
        }
    }

    pub(crate) fn gated_delta_state(&mut self, index: usize) -> Result<&mut GatedDeltaState> {
        match &mut self.storage {
            CacheStorage::HybridLinear(layers) => hybrid::gated_delta_state(layers, index),
            CacheStorage::Attention(_) => {
                Err(Error::InvalidModel("expected hybrid linear cache".into()))
            },
        }
    }

    pub(crate) fn full_attention_cache(&mut self, index: usize) -> Result<&mut KvCache> {
        match &mut self.storage {
            CacheStorage::HybridLinear(layers) => hybrid::full_attention_cache(layers, index),
            CacheStorage::Attention(_) => {
                Err(Error::InvalidModel("expected hybrid linear cache".into()))
            },
        }
    }

    pub fn reset(&mut self) -> Result<()> {
        match &mut self.storage {
            CacheStorage::Attention(caches) => caches.iter_mut().try_for_each(KvCache::reset),
            CacheStorage::HybridLinear(layers) => hybrid::reset(layers),
        }
    }

    pub fn reserve(&mut self, tokens: usize) -> Result<()> {
        match &mut self.storage {
            CacheStorage::Attention(caches) => {
                caches.iter_mut().try_for_each(|cache| cache.reserve(tokens))
            },
            CacheStorage::HybridLinear(layers) => hybrid::reserve(layers, tokens),
        }
    }

    pub fn cached_tokens(&self) -> Result<usize> {
        match &self.storage {
            CacheStorage::Attention(caches) => {
                caches.first().map_or_else(|| Ok(0), KvCache::offset)
            },
            CacheStorage::HybridLinear(layers) => hybrid::offset(layers),
        }
    }

    pub fn snapshot_at(&self, offset: usize) -> Result<Self> {
        match &self.storage {
            CacheStorage::Attention(caches) => {
                let caches = caches
                    .iter()
                    .map(|cache| cache.snapshot_at(offset))
                    .collect::<Result<Vec<_>>>()?;
                Ok(Self { storage: CacheStorage::Attention(caches) })
            },
            CacheStorage::HybridLinear(layers) => {
                let layers = hybrid::snapshot_at(layers, offset)?;
                Ok(Self {
                    storage: CacheStorage::HybridLinear(layers),
                })
            },
        }
    }
}