metrale_telemetry/
prefix_cache.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: The [`PrefixCache`] trait (KV block reuse across requests that
4//! share a prompt prefix, plus the SSM snapshot index) and the prefix-cache
5//! hit and miss counters.
6//!
7//! Implementations: [`NoPrefixCaching`] (caching disabled) and the
8//! `metrale-cache` crate's `RadixTree`.
9//!
10//! Owner: telemetry.
11//! Invariants: none beyond the types.
12
13use std::sync::atomic::Ordering;
14
15mod no_caching;
16mod tier_evict;
17pub use no_caching::NoPrefixCaching;
18pub use tier_evict::TierEvict;
19
20pub fn record_cache_hit(matched_tokens: usize) {
21    let m = crate::run_metrics::metrics();
22    m.cache_hits.fetch_add(1, Ordering::Relaxed);
23    m.cache_hit_tokens
24        .fetch_add(matched_tokens as u64, Ordering::Relaxed);
25}
26
27pub fn record_cache_miss() {
28    crate::run_metrics::metrics()
29        .cache_misses
30        .fetch_add(1, Ordering::Relaxed);
31}
32
33pub fn cache_hit_count() -> u64 {
34    crate::run_metrics::metrics()
35        .cache_hits
36        .load(Ordering::Relaxed)
37}
38pub fn cache_miss_count() -> u64 {
39    crate::run_metrics::metrics()
40        .cache_misses
41        .load(Ordering::Relaxed)
42}
43pub fn cache_hit_tokens_total() -> u64 {
44    crate::run_metrics::metrics()
45        .cache_hit_tokens
46        .load(Ordering::Relaxed)
47}
48
49/// 2026-09-26: What [`PrefixCache::evict`] removed.
50#[derive(Debug, Clone, Default)]
51pub struct EvictedBlocks {
52    /// 2026-09-26: KV blocks the cache no longer stores; the caller returns the
53    /// cache's reference on each (`PagedKvCache::return_evicted_block`).
54    pub physical: Vec<u32>,
55    /// 2026-09-26: The disk-block ids the evicted blocks carried; the caller
56    /// `dec_disk_ref`s each. Empty without `--high-speed-swap`.
57    pub disk_block_ids: Vec<u32>,
58}
59
60/// 2026-09-26: What an insert started or stopped storing, so the caller can
61/// take and release the matching references.
62///
63/// The cache holds one KV reference on each block it stores, taken when it
64/// starts storing it and handed back on eviction. Where a node already exists
65/// for a chunk, the insert keeps the node's block, which can differ from the
66/// inserting sequence's block, so the caller references the blocks listed
67/// here rather than its own block table.
68#[derive(Debug, Clone, Default)]
69pub struct InsertAcquired {
70    /// 2026-09-26: Disk-block ids the cache newly references; the caller
71    /// `inc_disk_ref`s each.
72    pub disk_block_ids: Vec<u32>,
73    /// 2026-09-26: KV blocks the cache started storing (new nodes, and a new
74    /// partial-suffix slot); the caller `inc_ref`s each once.
75    pub blocks: Vec<u32>,
76    /// 2026-09-26: KV blocks the cache stopped storing: a partial-suffix slot
77    /// that was replaced. The caller `dec_ref`s each once.
78    pub released_blocks: Vec<u32>,
79}
80
81impl EvictedBlocks {
82    pub fn is_empty(&self) -> bool {
83        self.physical.is_empty()
84    }
85
86    pub fn len(&self) -> usize {
87        self.physical.len()
88    }
89}
90
91/// 2026-09-26: Result of looking up a token sequence in the prefix cache.
92#[derive(Debug, Clone)]
93pub struct PrefixMatch {
94    /// 2026-09-26: KV blocks to reuse, in token order.
95    pub matched_blocks: Vec<u32>,
96    /// 2026-09-26: `--high-speed-swap` disk-block ids parallel to
97    /// `matched_blocks`, `u32::MAX` where a block has none; empty when no
98    /// matched block has one.
99    pub matched_disk_block_ids: Vec<u32>,
100    /// 2026-09-26: Tokens matched. Block-aligned unless partial-tail sharing
101    /// (`METRALE_PREFIX_SUBBLOCK`) matched into a block.
102    pub matched_tokens: usize,
103    /// 2026-09-26: The resident SSM snapshot slot of the deepest snapshot
104    /// within the match, if that snapshot is resident.
105    pub ssm_snapshot: Option<usize>,
106    /// 2026-09-26: Tokens covered by `ssm_snapshot`, at most `matched_tokens`;
107    /// SSM state for the tokens between is not covered.
108    pub ssm_snapshot_tokens: usize,
109    /// 2026-09-26: The tier key (prefix hash) when the deepest snapshot within
110    /// the match is spilled; `ssm_snapshot` is then `None`. The caller faults
111    /// it in and calls [`PrefixCache::promote_snapshot`].
112    pub ssm_snapshot_tier_key: Option<u64>,
113    /// 2026-09-26: Tokens covered by `ssm_snapshot_tier_key`.
114    pub ssm_snapshot_tier_tokens: usize,
115    /// 2026-09-26: Whether the matched snapshot is a session tail
116    /// ([`PrefixCache::insert_tail_snapshot`]).
117    pub ssm_snapshot_is_tail: bool,
118}
119
120impl PrefixMatch {
121    /// 2026-09-26: Empty match (no cached prefix found).
122    pub fn empty() -> Self {
123        Self {
124            matched_blocks: Vec::new(),
125            matched_disk_block_ids: Vec::new(),
126            matched_tokens: 0,
127            ssm_snapshot: None,
128            ssm_snapshot_tokens: 0,
129            ssm_snapshot_tier_key: None,
130            ssm_snapshot_tier_tokens: 0,
131            ssm_snapshot_is_tail: false,
132        }
133    }
134
135    /// 2026-09-26: Whether any prefix was matched.
136    pub fn is_empty(&self) -> bool {
137        self.matched_tokens == 0
138    }
139}
140
141/// 2026-09-26: A prefix cache. Every method takes `&self`, so an
142/// implementation keeps its state behind interior mutability.
143pub trait PrefixCache: Send + Sync {
144    /// 2026-09-26: Whether this cache stores blocks and holds references.
145    /// `NoPrefixCaching` answers `false`, and callers then skip the cache's
146    /// bookkeeping.
147    fn is_active(&self) -> bool {
148        true
149    }
150
151    /// 2026-09-26: Look up `tokens` and return the cached KV blocks of the
152    /// longest matching prefix, taking a reference on each full matched block
153    /// so it survives eviction until [`PrefixCache::release`].
154    ///
155    /// `session_hash` scopes session-tail snapshots; `adapter_id` keys the
156    /// cache, so only blocks and snapshots computed under the same adapter
157    /// match (`0` is the base model).
158    fn lookup(
159        &self,
160        tokens: &[u32],
161        block_size: usize,
162        session_hash: u64,
163        adapter_id: u64,
164    ) -> PrefixMatch;
165
166    /// 2026-09-26: The number of tokens `lookup` would match, without taking
167    /// references, touching LRU state or counting a hit or miss. Keyed by
168    /// `adapter_id` like `lookup`. The default answers 0.
169    fn peek_matched_tokens(&self, _tokens: &[u32], _block_size: usize, _adapter_id: u64) -> usize {
170        0
171    }
172
173    /// 2026-09-30: Take a reference on each full block of the prefix `lookup` would match,
174    /// without the snapshot lookup or a hit or miss count, so eviction skips that prefix until
175    /// [`PrefixCache::release_matched`] with the returned token count. The default pins nothing
176    /// and answers 0.
177    fn pin_prefix(&self, _tokens: &[u32], _block_size: usize, _adapter_id: u64) -> usize {
178        0
179    }
180
181    /// 2026-09-26: Insert a completed prefill's blocks.
182    ///
183    /// `block_table[i]` is the block for tokens
184    /// `[i*block_size .. (i+1)*block_size]`. `disk_block_ids` is empty or
185    /// parallel to `block_table` (`--high-speed-swap`). `matched_tokens` is
186    /// the prefix the sequence already holds through `lookup`; the sequence
187    /// gets a reference on every node past it, so that its `release` leaves
188    /// the cache's own reference in place. The caller applies the returned
189    /// [`InsertAcquired`].
190    fn insert(
191        &self,
192        tokens: &[u32],
193        block_table: &[u32],
194        disk_block_ids: &[u32],
195        block_size: usize,
196        matched_tokens: usize,
197        adapter_id: u64,
198    ) -> InsertAcquired;
199
200    /// 2026-09-26: `insert`, then register `snapshot_id` (a snapshot-pool
201    /// slot) for the whole of `tokens`, tagged with `session_hash`. Returns
202    /// the displaced resident slot for the caller to free, if any, and the
203    /// insert's [`InsertAcquired`].
204    #[allow(clippy::too_many_arguments)]
205    fn insert_with_snapshot(
206        &self,
207        tokens: &[u32],
208        block_table: &[u32],
209        disk_block_ids: &[u32],
210        block_size: usize,
211        snapshot_id: usize,
212        session_hash: u64,
213        matched_tokens: usize,
214        adapter_id: u64,
215    ) -> (Option<usize>, InsertAcquired);
216
217    /// 2026-09-26: Register `snapshot_id` for `tokens`, the sequence up to the
218    /// snapshot point, tagged with `session_hash`. Returns the displaced
219    /// resident slot for the caller to free, if any.
220    #[allow(clippy::too_many_arguments)]
221    fn insert_intermediate_snapshot(
222        &self,
223        tokens: &[u32],
224        block_table: &[u32],
225        disk_block_ids: &[u32],
226        block_size: usize,
227        snapshot_id: usize,
228        session_hash: u64,
229        matched_tokens: usize,
230        adapter_id: u64,
231    ) -> Option<usize>;
232
233    /// 2026-09-26: Register the session's tail snapshot in the index only. For
234    /// a nonzero `session_hash` it replaces the session's previous tail and
235    /// tail sibling. Returns the displaced snapshot ids to free.
236    fn insert_tail_snapshot(
237        &self,
238        tokens: &[u32],
239        snapshot_id: usize,
240        session_hash: u64,
241        adapter_id: u64,
242    ) -> Vec<usize>;
243
244    /// 2026-09-26: Register the tail's sibling snapshot in the index. Call it
245    /// after `insert_tail_snapshot`, which removes the session's previous
246    /// sibling. Returns a displaced snapshot id to free, if any.
247    fn insert_tail_sibling_snapshot(
248        &self,
249        tokens: &[u32],
250        snapshot_id: usize,
251        session_hash: u64,
252        adapter_id: u64,
253    ) -> Option<usize>;
254
255    /// 2026-09-26: Drop the sequence's reference on each full block of
256    /// `tokens` the cache holds, when the sequence finishes. `adapter_id` must
257    /// be the one used at `lookup` and `insert`.
258    fn release(&self, tokens: &[u32], block_size: usize, adapter_id: u64);
259
260    /// 2026-09-26: Drop the references one `lookup` took: the full blocks of
261    /// the first `matched_tokens` tokens and no more, for a caller that rolls
262    /// back a lookup before any sequence owns its blocks.
263    fn release_matched(
264        &self,
265        tokens: &[u32],
266        block_size: usize,
267        matched_tokens: usize,
268        adapter_id: u64,
269    );
270
271    /// 2026-09-26: Evict least-recently-used leaves that no sequence holds
272    /// until at least `num_blocks` blocks are freed or none is left. A leaf's
273    /// partial-suffix block goes with it, so one more block than asked can
274    /// come back. The caller applies the [`EvictedBlocks`].
275    fn evict(&self, num_blocks: usize) -> EvictedBlocks;
276
277    /// 2026-09-26: Evict one SSM snapshot from the index; returns its slot for
278    /// the caller to free.
279    fn evict_snapshot_lru(&self) -> Option<usize>;
280
281    /// 2026-09-26: Pick a resident snapshot to evict: one at least
282    /// `min_tokens` deep is marked spilled and stays findable, a shallower one
283    /// is removed ([`TierEvict`]). `min_tokens == 0` always spills. `None`
284    /// when nothing is resident; the default answers `None`.
285    fn evict_snapshot_to_tier(&self, min_tokens: usize) -> Option<TierEvict> {
286        let _ = min_tokens;
287        None
288    }
289
290    /// 2026-09-26: After the caller faulted a spilled snapshot into
291    /// `new_slot`, make its entry resident there. `false` if the key is
292    /// unknown; the default answers `false`.
293    fn promote_snapshot(&self, key: u64, new_slot: usize) -> bool {
294        let _ = (key, new_slot);
295        false
296    }
297
298    /// 2026-09-26: Remove the spilled entry for `key`, for a caller whose
299    /// stored bytes for it are gone, so later lookups stop returning it. A
300    /// resident entry is left alone, since its slot is live. Returns whether
301    /// an entry was removed; the default answers `false`.
302    fn forget_snapshot_tier_key(&self, key: u64) -> bool {
303        let _ = key;
304        false
305    }
306
307    /// 2026-09-26: Number of SSM snapshots in the snapshot index.
308    fn snapshot_count(&self) -> usize;
309
310    /// 2026-09-26: `(entries, cached_blocks)`. `RadixTree` stores one block
311    /// per entry and answers its node count for both.
312    fn stats(&self) -> (usize, usize);
313}
314
315#[cfg(test)]
316mod tests {
317    use super::*;
318
319    #[test]
320    fn test_prefix_match_empty() {
321        let m = PrefixMatch::empty();
322        assert!(m.is_empty());
323        assert_eq!(m.matched_tokens, 0);
324        assert!(m.matched_blocks.is_empty());
325    }
326
327    #[test]
328    fn test_no_prefix_caching_is_noop() {
329        let cache = NoPrefixCaching;
330        let tokens = vec![1, 2, 3, 4, 5, 6, 7, 8];
331        let block_table = vec![0, 1];
332        let disk_block_ids: Vec<u32> = vec![];
333
334        let m = cache.lookup(&tokens, 4, 0, 0);
335        assert!(m.is_empty());
336
337        let new_acq = cache.insert(&tokens, &block_table, &disk_block_ids, 4, 0, 0);
338        assert!(new_acq.disk_block_ids.is_empty());
339        assert!(new_acq.blocks.is_empty());
340        cache.release(&tokens, 4, 0);
341
342        let evicted = cache.evict(10);
343        assert!(evicted.is_empty());
344
345        assert_eq!(cache.evict_snapshot_lru(), None);
346        assert_eq!(cache.snapshot_count(), 0);
347
348        assert_eq!(cache.stats(), (0, 0));
349    }
350}
metrale_telemetry/prefix_cache.rs

metrale_telemetry/
prefix_cache.rs
