spark_runtime/weights/gguf/
expert_lru_evict.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2// provenance-id: 526f6e616c6420522e205374657369616b
3
4//! [`ExpertLru`]'s victim selection: strict LRU, or a random slot among the
5//! oldest few percent (the cure for the cyclic replay of a working set a
6//! little larger than the cache). Split from `expert_lru.rs` (500-LoC cap).
7
8use anyhow::{Result, bail};
9
10use super::expert_lru::{ExpertLru, NONE};
11
12impl ExpertLru {
13    /// Move `i` to the LRU end without unmapping it (a prediction that was
14    /// not used: the next victim, but still a hit if it is asked for).
15    pub(super) fn demote(&mut self, i: u32) {
16        self.unlink(i);
17        let m = &mut self.meta[i as usize];
18        m.next = NONE;
19        m.prev = self.tail;
20        if self.tail != NONE {
21            self.meta[self.tail as usize].next = i;
22        }
23        self.tail = i;
24        if self.head == NONE {
25            self.head = i;
26        }
27    }
28
29    /// Strict LRU or a random victim among the oldest `pct`% (see `take_victim`).
30    pub fn set_evict_random_pct(&mut self, pct: usize) {
31        self.evict_pct = pct.min(100);
32    }
33
34    /// xorshift64*, seeded once: the same run picks the same victims.
35    fn rng_next(&mut self) -> u64 {
36        let mut x = self.rng;
37        x ^= x >> 12;
38        x ^= x << 25;
39        x ^= x >> 27;
40        self.rng = x;
41        x.wrapping_mul(0x2545_F491_4F6C_DD1D)
42    }
43
44    /// Take a victim among the `evict_pct`% least recently used slots that
45    /// are not pinned in this epoch and not being filled, dropping whatever
46    /// it held. Strict LRU (`evict_pct` = 0, or fewer than 100/pct slots)
47    /// takes the oldest; otherwise a random one of the oldest few percent.
48    ///
49    /// The randomness is the cure for the cyclic case: a request whose
50    /// working set is a few experts larger than the cache re-reads it in the
51    /// same order, and strict LRU evicts exactly the expert needed next at
52    /// every miss (MinHeap at 8,380 slots: 8,418 distinct experts, 5.5
53    /// misses a step on r2/r3, 09-19). A victim drawn from the oldest 5%
54    /// breaks the lockstep: 1.1 misses a step in the replay, Volvo (fits)
55    /// unchanged at 0. The policy never touches the math.
56    pub(super) fn take_victim(&mut self) -> Result<u32> {
57        let want = (self.n_slots * self.evict_pct / 100).max(1);
58        let mut cands: [u32; 8] = [NONE; 8];
59        let mut n = 0usize;
60        // reservoir-sample up to eight of the oldest `want` eligible slots
61        let mut seen = 0usize;
62        let mut i = self.tail;
63        while i != NONE && seen < want {
64            let m = &self.meta[i as usize];
65            if m.epoch != self.epoch && m.ticket.is_none() {
66                if n < cands.len() {
67                    cands[n] = i;
68                    n += 1;
69                } else {
70                    let j = (self.rng_next() % (seen as u64 + 1)) as usize;
71                    if j < cands.len() {
72                        cands[j] = i;
73                    }
74                }
75                seen += 1;
76            }
77            i = self.meta[i as usize].prev;
78        }
79        if n > 0 {
80            let pick = if n == 1 {
81                cands[0]
82            } else {
83                cands[(self.rng_next() % n as u64) as usize]
84            };
85            if let Some(k) = self.meta[pick as usize].key.take() {
86                self.map.remove(&k);
87                self.slot_changes.push((k.0, k.1, -1));
88                self.stats.evictions += 1;
89            }
90            return Ok(pick);
91        }
92        bail!(
93            "expert cache of {} slots is smaller than one token's working set ({} pinned)",
94            self.n_slots,
95            self.map.len()
96        )
97    }
98}
spark_runtime/weights/gguf/expert_lru_evict.rs

spark_runtime/weights/gguf/
expert_lru_evict.rs
