cherenkov/
prefix_cache.rs

1//! Bounded LRU checkpoints of complete hybrid-model sequence state.
2
3use crate::{
4    options::Options,
5    qwen4_exp::gpu::{Gpu, MAX_NB, PrefixState},
6    runner::PrefillResume,
7};
8use anyhow::{Result, ensure};
9use std::collections::VecDeque;
10use std::time::{Duration, Instant};
11
12struct Entry {
13    tokens: Vec<u32>,
14    following: Option<u32>,
15    complete: bool,
16    next: u32,
17    drafts: Vec<u32>,
18    logits: Vec<f32>,
19    drafting: bool,
20    state: PrefixState,
21    bytes: usize,
22    touched: Instant,
23}
24
25pub(crate) struct PrefixCache {
26    entries: VecDeque<Entry>,
27    capacity: usize,
28    used: usize,
29    evictions: u64,
30    max_entries: usize,
31    idle: Option<Duration>,
32}
33
34fn matches(tokens: &[u32], following: Option<u32>, complete: bool, request: &[u32]) -> bool {
35    request.starts_with(tokens)
36        && if request.len() == tokens.len() {
37            complete
38        } else {
39            // The last MTP cache row depends on the token AFTER the trunk prefix.
40            // Never reuse it for a different continuation.
41            following.is_none_or(|next| request[tokens.len()] == next)
42        }
43}
44
45impl PrefixCache {
46    pub(crate) fn new(capacity: usize, max_entries: usize, idle_seconds: u64) -> Self {
47        Self {
48            entries: VecDeque::new(),
49            capacity,
50            used: 0,
51            evictions: 0,
52            max_entries,
53            idle: (idle_seconds > 0).then(|| Duration::from_secs(idle_seconds)),
54        }
55    }
56
57    fn evict_oldest(&mut self) {
58        if let Some(entry) = self.entries.pop_front() {
59            self.used -= entry.bytes;
60            self.evictions += 1;
61        }
62    }
63
64    pub(crate) fn expire(&mut self) {
65        if let Some(idle) = self.idle {
66            while self
67                .entries
68                .front()
69                .is_some_and(|e| e.touched.elapsed() >= idle)
70            {
71                self.evict_oldest();
72            }
73        }
74    }
75
76    pub(crate) fn stats(&self) -> (usize, usize, u64) {
77        (self.entries.len(), self.used, self.evictions)
78    }
79
80    fn store(&mut self, gpu: &Gpu<'_>, ids: &[u32], seed: &PrefillResume, drafting: bool) {
81        let pos = gpu.pos;
82        let bytes = gpu.prefix_state_bytes()
83            + pos * 4
84            + std::mem::size_of::<Entry>()
85            + seed.drafts.len() * 4
86            + seed.logits.as_ref().map_or(0, |l| l.len() * 4);
87
88        if bytes > self.capacity || pos == 0 {
89            return;
90        }
91
92        if let Some(i) = self
93            .entries
94            .iter()
95            .position(|e| e.tokens == ids[..pos] && e.drafting == drafting)
96        {
97            self.used -= self.entries.remove(i).unwrap().bytes;
98        }
99
100        // Evict BEFORE copying state: allocation peaks also stay within the cache budget.
101        while self.used + bytes > self.capacity || self.entries.len() >= self.max_entries {
102            self.evict_oldest();
103        }
104
105        let following = drafting.then(|| ids.get(pos).copied().unwrap_or(seed.next));
106
107        self.entries.push_back(Entry {
108            tokens: ids[..pos].to_vec(),
109            following,
110            complete: pos == ids.len(),
111            next: seed.next,
112            drafts: seed.drafts.clone(),
113            logits: seed.logits.clone().unwrap_or_default(),
114            drafting,
115            state: gpu.save_prefix(),
116            bytes,
117            touched: Instant::now(),
118        });
119
120        self.used += bytes;
121    }
122
123    pub(crate) fn begin(
124        &mut self,
125        gpu: &mut Gpu<'_>,
126        ids: &[u32],
127        stable_boundaries: &[usize],
128        options: &Options,
129    ) -> Result<Prefill> {
130        self.expire();
131
132        let drafting = options.effective_drafts() > 0;
133        let mut cached = 0;
134        let mut seed = PrefillResume::default();
135
136        if let Some(i) = self
137            .entries
138            .iter()
139            .enumerate()
140            .filter(|(_, e)| {
141                e.drafting == drafting && matches(&e.tokens, e.following, e.complete, ids)
142            })
143            .max_by_key(|(_, e)| e.tokens.len())
144            .map(|(i, _)| i)
145        {
146            let mut entry = self.entries.remove(i).unwrap();
147            entry.touched = Instant::now();
148
149            gpu.restore_prefix(&entry.state)?;
150
151            cached = entry.tokens.len();
152            seed = PrefillResume {
153                next: entry.next,
154                drafts: entry.drafts.clone(),
155                logits: Some(entry.logits.clone()),
156            };
157
158            self.entries.push_back(entry);
159        }
160
161        let mut boundaries = vec![ids.len()];
162
163        if self.capacity > 0 {
164            boundaries.extend(
165                stable_boundaries
166                    .iter()
167                    .copied()
168                    .filter(|&b| b > cached && b < ids.len()),
169            );
170
171            if ids.len() > 1 {
172                boundaries.push(ids.len() - 1);
173            }
174        }
175
176        boundaries.sort_unstable();
177        boundaries.dedup();
178
179        Ok(Prefill {
180            cached,
181            boundaries,
182            seed,
183        })
184    }
185}
186
187/// The cursor stays with its request; one advance processes at most one chunk.
188pub(crate) struct Prefill {
189    pub cached: usize,
190    boundaries: Vec<usize>,
191    pub seed: PrefillResume,
192}
193
194pub(crate) struct PrefillProgress {
195    pub done: bool,
196    /// Successful engine rows suitable for pacing; excludes boundary tails,
197    /// memory-limited chunks, and the small-row execution path.
198    pub pacing_tokens: Option<usize>,
199}
200
201impl Prefill {
202    pub(crate) fn advance(
203        &mut self,
204        gpu: &mut Gpu<'_>,
205        cache: &mut PrefixCache,
206        ids: &[u32],
207        options: &Options,
208        quantum: usize,
209    ) -> Result<PrefillProgress> {
210        let Some(end) = self.boundaries.iter().copied().find(|&end| end > gpu.pos) else {
211            return Ok(PrefillProgress {
212                done: true,
213                pacing_tokens: None,
214            });
215        };
216        let pf_min = std::env::var("CHERENKOV_PREFILL_MIN")
217            .ok()
218            .and_then(|s| s.parse().ok())
219            .unwrap_or(64);
220        let remaining = end - gpu.pos;
221        let engine = remaining >= pf_min;
222        let capacity = if engine {
223            let fit = gpu.prefill_rows_fit(false)?;
224
225            std::env::var("CHERENKOV_PREFILL_CHUNK")
226                .ok()
227                .and_then(|s| s.parse().ok())
228                .unwrap_or(fit)
229                .min(fit)
230        } else {
231            std::env::var("CHERENKOV_ROWS_MAX")
232                .ok()
233                .and_then(|s| s.parse().ok())
234                .unwrap_or(MAX_NB)
235                .clamp(1, MAX_NB)
236        };
237        let n = chunk_len(remaining, capacity, quantum)?;
238        let pos = gpu.pos;
239        self.seed = prefill_rows(
240            gpu,
241            &ids[pos..pos + n],
242            ids.get(pos + n).copied(),
243            options.effective_drafts(),
244            engine,
245        )?;
246
247        if gpu.pos == end || engine {
248            cache.store(gpu, ids, &self.seed, options.effective_drafts() > 0);
249        }
250
251        gpu.prefill_release();
252
253        Ok(PrefillProgress {
254            done: gpu.pos == ids.len(),
255            pacing_tokens: pacing_sample(remaining, capacity, quantum, engine).then_some(n),
256        })
257    }
258}
259
260fn pacing_sample(remaining: usize, capacity: usize, quantum: usize, engine: bool) -> bool {
261    engine && remaining >= quantum && capacity >= quantum
262}
263
264/// Balance a boundary's chunks without exceeding memory or scheduling limits.
265fn chunk_len(remaining: usize, capacity: usize, quantum: usize) -> Result<usize> {
266    let size = capacity.min(quantum);
267
268    ensure!(
269        remaining > 0 && size > 0,
270        "prefill chunk requires tokens and capacity"
271    );
272
273    Ok(remaining.div_ceil(remaining.div_ceil(size)))
274}
275
276/// Advance one prompt chunk and retain the predictions needed to resume decode.
277fn prefill_rows(
278    gpu: &mut Gpu<'_>,
279    rows: &[u32],
280    following: Option<u32>,
281    drafts: usize,
282    engine: bool,
283) -> Result<PrefillResume> {
284    let last = following.is_none();
285
286    if engine {
287        let (next, draft) = gpu.prefill_chunk(rows, following, false)?;
288        let mut seed = PrefillResume {
289            next,
290            drafts: Vec::new(),
291            logits: Some(gpu.logits().to_vec()),
292        };
293
294        if drafts == 0 || !last {
295            return Ok(seed);
296        }
297
298        seed.drafts.push(draft);
299
300        if drafts >= 2 {
301            seed.drafts.push(gpu.mtp_chain(draft)?);
302        }
303
304        return Ok(seed);
305    }
306
307    let result = gpu.step_rows(rows, false, false)?;
308
309    gpu.commit(rows.len())?;
310
311    let mut seed = PrefillResume {
312        next: result[rows.len() - 1],
313        drafts: Vec::new(),
314        logits: Some(gpu.logits_row(rows.len() - 1).to_vec()),
315    };
316
317    if drafts == 0 {
318        return Ok(seed);
319    }
320
321    let mut next = rows[1..].to_vec();
322
323    next.push(following.unwrap_or(seed.next));
324
325    seed.drafts = gpu.mtp_draft(&next, if last { drafts } else { 1 })?;
326
327    Ok(seed)
328}
329
330#[cfg(test)]
331#[path = "../tests/unit/prefix_cache.rs"]
332mod tests;
cherenkov/prefix_cache.rs

cherenkov/
prefix_cache.rs
