kime_engine/
cache.rs

1//! The answer cache from spec/11-serving.md. A question's answer depends only on its own row of
2//! token ids, its option markers and its type, since every question is its own sequence and the
3//! engine gives the same bits whatever else is in the batch. So one entry is the logits and the
4//! act head output of one row, keyed by the blake3 of what went to the device. The answer JSON is
5//! built from them again on a hit, which keeps the labels the request sent.
6//!
7//! Each cache belongs to one loaded model, so the model, its version and its precision are part
8//! of the key without being hashed. Eviction keeps two generations: new entries go in the young
9//! one, and when it is full the old one is dropped and the young one takes its place. A hit in
10//! the old generation moves the entry back to the young one, so what is used stays in, as with
11//! an LRU, for one hash map insert rather than a linked list.
12
13use std::collections::HashMap;
14use std::sync::atomic::{AtomicU64, Ordering};
15use std::sync::{Mutex, PoisonError};
16
17use kime_tok::layout::CompatSequence;
18
19/// The blake3 hash of one row.
20pub(crate) type Key = [u8; 32];
21
22/// What the device gave for one row.
23#[derive(Debug, Clone, PartialEq)]
24pub(crate) struct Entry {
25    pub(crate) logits: Box<[f32]>,
26    pub(crate) act: [f32; 2],
27}
28
29/// How a request uses the cache, from `kime.cache`.
30#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
31pub enum CacheMode {
32    /// Answers from the cache when it can and keeps what it computes.
33    #[default]
34    Use,
35    /// Neither reads nor writes the cache.
36    Bypass,
37    /// Computes every answer and overwrites the cache with it.
38    Refresh,
39}
40
41impl CacheMode {
42    /// The mode `kime.cache` names: `"use"`, `"bypass"` or `"refresh"`.
43    #[must_use]
44    pub fn parse(s: &str) -> Option<CacheMode> {
45        match s {
46            "use" => Some(CacheMode::Use),
47            "bypass" => Some(CacheMode::Bypass),
48            "refresh" => Some(CacheMode::Refresh),
49            _ => None,
50        }
51    }
52}
53
54/// Hits and misses since the model loaded, and what the cache holds now.
55#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
56pub struct CacheStats {
57    /// The most answers it holds, 0 when it is off.
58    pub capacity: usize,
59    /// Answers it holds.
60    pub entries: usize,
61    /// Questions answered from it.
62    pub hits: u64,
63    /// Questions looked up and not found, then computed.
64    pub misses: u64,
65}
66
67#[derive(Default)]
68struct Generations {
69    young: HashMap<Key, Entry>,
70    old: HashMap<Key, Entry>,
71}
72
73pub(crate) struct AnswerCache {
74    /// The most entries in one generation, half the capacity.
75    half: usize,
76    maps: Mutex<Generations>,
77    hits: AtomicU64,
78    misses: AtomicU64,
79}
80
81impl AnswerCache {
82    /// A cache for about `capacity` answers. It holds between half of that and all of it.
83    pub(crate) fn new(capacity: usize) -> Self {
84        AnswerCache {
85            half: capacity.div_ceil(2).max(1),
86            maps: Mutex::default(),
87            hits: AtomicU64::new(0),
88            misses: AtomicU64::new(0),
89        }
90    }
91
92    pub(crate) fn key(seq: &CompatSequence, qtype: u8) -> Key {
93        let mut h = blake3::Hasher::new();
94        h.update(&[qtype]);
95        h.update(&(seq.ids.len() as u64).to_le_bytes());
96        for id in &seq.ids {
97            h.update(&id.to_le_bytes());
98        }
99        for m in &seq.markers {
100            h.update(&m.to_le_bytes());
101        }
102        *h.finalize().as_bytes()
103    }
104
105    /// The entries for `keys`, in order, with `None` for each miss. The counts go up.
106    pub(crate) fn get(&self, keys: &[Key]) -> Vec<Option<Entry>> {
107        let mut g = self.maps.lock().unwrap_or_else(PoisonError::into_inner);
108        let out: Vec<Option<Entry>> = keys
109            .iter()
110            .map(|k| {
111                if let Some(e) = g.young.get(k) {
112                    return Some(e.clone());
113                }
114                let e = g.old.remove(k)?;
115                self.put(&mut g, *k, e.clone());
116                Some(e)
117            })
118            .collect();
119        drop(g);
120        let hits = out.iter().filter(|e| e.is_some()).count() as u64;
121        self.hits.fetch_add(hits, Ordering::Relaxed);
122        self.misses.fetch_add(keys.len() as u64 - hits, Ordering::Relaxed);
123        out
124    }
125
126    /// The entries for `keys` when every one is there, counted as hits. When one is missing
127    /// nothing is counted, since the caller looks them up again with [`AnswerCache::get`].
128    pub(crate) fn all(&self, keys: &[Key]) -> Option<Vec<Entry>> {
129        let g = self.maps.lock().unwrap_or_else(PoisonError::into_inner);
130        let out = keys
131            .iter()
132            .map(|k| g.young.get(k).or_else(|| g.old.get(k)).cloned())
133            .collect::<Option<Vec<_>>>()?;
134        drop(g);
135        self.hits.fetch_add(keys.len() as u64, Ordering::Relaxed);
136        Some(out)
137    }
138
139    /// Keeps `entries`, replacing any already there.
140    pub(crate) fn insert(&self, entries: impl IntoIterator<Item = (Key, Entry)>) {
141        let mut g = self.maps.lock().unwrap_or_else(PoisonError::into_inner);
142        for (k, e) in entries {
143            g.old.remove(&k);
144            self.put(&mut g, k, e);
145        }
146    }
147
148    fn put(&self, g: &mut Generations, k: Key, e: Entry) {
149        if g.young.len() >= self.half && !g.young.contains_key(&k) {
150            g.old = std::mem::take(&mut g.young);
151        }
152        g.young.insert(k, e);
153    }
154
155    pub(crate) fn stats(&self) -> CacheStats {
156        let g = self.maps.lock().unwrap_or_else(PoisonError::into_inner);
157        CacheStats {
158            capacity: self.half * 2,
159            entries: g.young.len() + g.old.len(),
160            hits: self.hits.load(Ordering::Relaxed),
161            misses: self.misses.load(Ordering::Relaxed),
162        }
163    }
164}
165
166#[cfg(test)]
167mod tests {
168    use super::*;
169
170    fn entry(x: f32) -> Entry {
171        Entry { logits: vec![x].into(), act: [x, 0.0] }
172    }
173
174    #[test]
175    fn keeps_what_is_used_and_counts() {
176        let c = AnswerCache::new(4);
177        let k = |i: u8| [i; 32];
178        c.insert((0..2).map(|i| (k(i), entry(f32::from(i)))));
179        // The young generation is full at 2, so 2 and 3 push 0 and 1 to the old one.
180        c.insert((2..4).map(|i| (k(i), entry(f32::from(i)))));
181        assert_eq!(c.stats().entries, 4);
182        // 0 is used, so it moves back to the young one, which turns the generations over and
183        // drops 1.
184        assert_eq!(c.get(&[k(0)]), vec![Some(entry(0.0))]);
185        c.insert([(k(4), entry(4.0))]);
186        let got = c.get(&[k(0), k(1), k(4), k(9)]);
187        assert_eq!(got, vec![Some(entry(0.0)), None, Some(entry(4.0)), None]);
188        let s = c.stats();
189        assert_eq!((s.hits, s.misses, s.capacity), (3, 2, 4));
190        assert!(s.entries <= 4, "{s:?}");
191    }
192
193    #[test]
194    fn insert_overwrites() {
195        let c = AnswerCache::new(10);
196        c.insert([([1; 32], entry(1.0))]);
197        c.insert([([1; 32], entry(2.0))]);
198        assert_eq!(c.get(&[[1; 32]]), vec![Some(entry(2.0))]);
199        assert_eq!(c.stats().entries, 1);
200    }
201
202    #[test]
203    fn modes() {
204        assert_eq!(CacheMode::parse("bypass"), Some(CacheMode::Bypass));
205        assert_eq!(CacheMode::parse("refresh"), Some(CacheMode::Refresh));
206        assert_eq!(CacheMode::parse("use"), Some(CacheMode::Use));
207        assert_eq!(CacheMode::parse("off"), None);
208    }
209}
kime_engine/cache.rs

kime_engine/
cache.rs
