babichjacob_llama_cpp_2/context/kv_cache.rs
1//! utilities for working with the kv cache
2
3use crate::context::LlamaContext;
4use std::ffi::c_int;
5use std::num::{NonZeroU8, TryFromIntError};
6
7/// Errors that can occur when attempting to prepare values for the kv cache
8#[derive(Debug, Eq, PartialEq, thiserror::Error)]
9pub enum KvCacheConversionError {
10 /// Sequence id conversion to i32 failed
11 #[error("Provided sequence id is too large for a i32")]
12 SeqIdTooLarge(#[source] TryFromIntError),
13 /// Position 0 conversion to i32 failed
14 #[error("Provided start position is too large for a i32")]
15 P0TooLarge(#[source] TryFromIntError),
16 /// Position 1 conversion to i32 failed
17 #[error("Provided end position is too large for a i32")]
18 P1TooLarge(#[source] TryFromIntError),
19}
20
21impl LlamaContext<'_> {
22 /// Copy the cache from one sequence to another.
23 ///
24 /// # Parameters
25 ///
26 /// * `src` - The sequence id to copy the cache from.
27 /// * `dest` - The sequence id to copy the cache to.
28 /// * `size` - The size of the cache to copy.
29 pub fn copy_cache(&mut self, src: i32, dest: i32, size: i32) {
30 unsafe { llama_cpp_sys_2::llama_kv_cache_seq_cp(self.context.as_ptr(), src, dest, 0, size) }
31 }
32
33 /// Copy the cache from one sequence to another.
34 ///
35 /// # Returns
36 /// A `Result` indicating whether the operation was successful. If the either position exceeds
37 /// the maximum i32 value, no copy is attempted and an `Err` is returned.
38 ///
39 /// # Parameters
40 ///
41 /// * `src` - The sequence id to copy the cache from.
42 /// * `dest` - The sequence id to copy the cache to.
43 /// * `p0` - The start position of the cache to clear. If `None`, the entire cache is copied up to `p1`.
44 /// * `p1` - The end position of the cache to clear. If `None`, the entire cache is copied starting from `p0`.
45 pub fn copy_kv_cache_seq(
46 &mut self,
47 src: i32,
48 dest: i32,
49 p0: Option<u32>,
50 p1: Option<u32>,
51 ) -> Result<(), KvCacheConversionError> {
52 let p0 = p0
53 .map_or(Ok(-1), i32::try_from)
54 .map_err(|e| KvCacheConversionError::P0TooLarge(e))?;
55 let p1 = p1
56 .map_or(Ok(-1), i32::try_from)
57 .map_err(|e| KvCacheConversionError::P1TooLarge(e))?;
58 unsafe {
59 llama_cpp_sys_2::llama_kv_cache_seq_cp(self.context.as_ptr(), src, dest, p0, p1);
60 }
61 Ok(())
62 }
63
64 /// Clear the kv cache for the given sequence within the specified range `[p0, p1)`
65 /// Returns `false` only when partial sequence removals fail. Full sequence removals always succeed.
66 ///
67 /// # Returns
68 /// A `Result` indicating whether the operation was successful. If the sequence id or
69 /// either position exceeds the maximum i32 value, no removal is attempted and an `Err` is returned.
70 ///
71 /// # Parameters
72 ///
73 /// * `src` - The sequence id to clear the cache for. If `None`, matches all sequences
74 /// * `p0` - The start position of the cache to clear. If `None`, the entire cache is cleared up to `p1`.
75 /// * `p1` - The end position of the cache to clear. If `None`, the entire cache is cleared from `p0`.
76 pub fn clear_kv_cache_seq(
77 &mut self,
78 src: Option<u32>,
79 p0: Option<u32>,
80 p1: Option<u32>,
81 ) -> Result<bool, KvCacheConversionError> {
82 let src = src
83 .map_or(Ok(-1), i32::try_from)
84 .map_err(|e| KvCacheConversionError::SeqIdTooLarge(e))?;
85 let p0 = p0
86 .map_or(Ok(-1), i32::try_from)
87 .map_err(|e| KvCacheConversionError::P0TooLarge(e))?;
88 let p1 = p1
89 .map_or(Ok(-1), i32::try_from)
90 .map_err(|e| KvCacheConversionError::P1TooLarge(e))?;
91 Ok(unsafe { llama_cpp_sys_2::llama_kv_cache_seq_rm(self.context.as_ptr(), src, p0, p1) })
92 }
93
94 /// Returns the number of used KV cells (i.e. have at least one sequence assigned to them)
95 #[must_use]
96 pub fn get_kv_cache_used_cells(&self) -> i32 {
97 unsafe { llama_cpp_sys_2::llama_get_kv_cache_used_cells(self.context.as_ptr()) }
98 }
99
100 /// Clear the KV cache
101 pub fn clear_kv_cache(&mut self) {
102 unsafe { llama_cpp_sys_2::llama_kv_cache_clear(self.context.as_ptr()) }
103 }
104
105 /// Removes all tokens that do not belong to the specified sequence
106 ///
107 /// # Parameters
108 ///
109 /// * `seq_id` - The sequence id to keep
110 pub fn llama_kv_cache_seq_keep(&mut self, seq_id: i32) {
111 unsafe { llama_cpp_sys_2::llama_kv_cache_seq_keep(self.context.as_ptr(), seq_id) }
112 }
113
114 #[allow(clippy::doc_markdown)]
115 /// Adds relative position "delta" to all tokens that belong to the specified sequence and have positions in `[p0, p1)`
116 /// If the KV cache is RoPEd, the KV data is updated accordingly:
117 /// - lazily on next [`LlamaContext::decode`]
118 /// - explicitly with [`Self::kv_cache_update`]
119 ///
120 /// # Returns
121 /// A `Result` indicating whether the operation was successful. If either position
122 /// exceeds the maximum i32 value, no update is attempted and an `Err` is returned.
123 ///
124 /// # Parameters
125 ///
126 /// * `seq_id` - The sequence id to update
127 /// * `p0` - The start position of the cache to update. If `None`, the entire cache is updated up to `p1`.
128 /// * `p1` - The end position of the cache to update. If `None`, the entire cache is updated starting from `p0`.
129 /// * `delta` - The relative position to add to the tokens
130 pub fn kv_cache_seq_add(
131 &mut self,
132 seq_id: i32,
133 p0: Option<u32>,
134 p1: Option<u32>,
135 delta: i32,
136 ) -> Result<(), KvCacheConversionError> {
137 let p0 = p0
138 .map_or(Ok(-1), i32::try_from)
139 .map_err(|e| KvCacheConversionError::P0TooLarge(e))?;
140 let p1 = p1
141 .map_or(Ok(-1), i32::try_from)
142 .map_err(|e| KvCacheConversionError::P1TooLarge(e))?;
143 unsafe {
144 llama_cpp_sys_2::llama_kv_cache_seq_add(self.context.as_ptr(), seq_id, p0, p1, delta);
145 }
146 Ok(())
147 }
148
149 /// Integer division of the positions by factor of `d > 1`
150 /// If the KV cache is `RoPEd`, the KV data is updated accordingly:
151 /// - lazily on next [`LlamaContext::decode`]
152 /// - explicitly with [`Self::kv_cache_update`]
153 ///
154 /// # Returns
155 /// A `Result` indicating whether the operation was successful. If either position
156 /// exceeds the maximum i32 value, no update is attempted and an `Err` is returned.
157 ///
158 /// # Parameters
159 ///
160 /// * `seq_id` - The sequence id to update
161 /// * `p0` - The start position of the cache to update. If `None`, the entire cache is updated up to `p1`.
162 /// * `p1` - The end position of the cache to update. If `None`, the entire cache is updated starting from `p0`.
163 /// * `d` - The factor to divide the positions by
164 pub fn kv_cache_seq_div(
165 &mut self,
166 seq_id: i32,
167 p0: Option<u32>,
168 p1: Option<u32>,
169 d: NonZeroU8,
170 ) -> Result<(), KvCacheConversionError> {
171 let p0 = p0
172 .map_or(Ok(-1), i32::try_from)
173 .map_err(|e| KvCacheConversionError::P0TooLarge(e))?;
174 let p1 = p1
175 .map_or(Ok(-1), i32::try_from)
176 .map_err(|e| KvCacheConversionError::P1TooLarge(e))?;
177 let d = c_int::from(d.get());
178 unsafe { llama_cpp_sys_2::llama_kv_cache_seq_div(self.context.as_ptr(), seq_id, p0, p1, d) }
179 Ok(())
180 }
181
182 /// Returns the largest position present in the KV cache for the specified sequence
183 ///
184 /// # Parameters
185 ///
186 /// * `seq_id` - The sequence id to get the max position for
187 #[must_use]
188 pub fn kv_cache_seq_pos_max(&self, seq_id: i32) -> i32 {
189 unsafe { llama_cpp_sys_2::llama_kv_cache_seq_pos_max(self.context.as_ptr(), seq_id) }
190 }
191
192 /// Defragment the KV cache
193 /// This will be applied:
194 /// - lazily on next [`LlamaContext::decode`]
195 /// - explicitly with [`Self::kv_cache_update`]
196 pub fn kv_cache_defrag(&mut self) {
197 unsafe { llama_cpp_sys_2::llama_kv_cache_defrag(self.context.as_ptr()) }
198 }
199
200 /// Apply the KV cache updates (such as K-shifts, defragmentation, etc.)
201 pub fn kv_cache_update(&mut self) {
202 unsafe { llama_cpp_sys_2::llama_kv_cache_update(self.context.as_ptr()) }
203 }
204
205 /// Returns the number of tokens in the KV cache (slow, use only for debug)
206 /// If a KV cell has multiple sequences assigned to it, it will be counted multiple times
207 #[must_use]
208 pub fn get_kv_cache_token_count(&self) -> i32 {
209 unsafe { llama_cpp_sys_2::llama_get_kv_cache_token_count(self.context.as_ptr()) }
210 }
211
212 /// Create an empty KV cache view. (use only for debugging purposes)
213 ///
214 /// # Parameters
215 ///
216 /// * `n_max_seq` - Maximum number of sequences that can exist in a cell. It's not an error
217 /// if there are more sequences in a cell than this value, however they will
218 /// not be visible in the view `cells_sequences`.
219 #[must_use]
220 pub fn new_kv_cache_view(&self, n_max_seq: i32) -> KVCacheView {
221 let view =
222 unsafe { llama_cpp_sys_2::llama_kv_cache_view_init(self.context.as_ptr(), n_max_seq) };
223 KVCacheView { view, ctx: self }
224 }
225}
226
227/// Information associated with an individual cell in the KV cache view.
228#[derive(Debug)]
229pub struct KVCacheViewCell {
230 /// The position for this cell. Takes KV cache shifts into account.
231 /// May be negative if the cell is not populated.
232 pub pos: llama_cpp_sys_2::llama_pos,
233}
234
235/// An updateable view of the KV cache. (use only for debugging purposes)
236#[derive(Debug)]
237pub struct KVCacheView<'a> {
238 ctx: &'a LlamaContext<'a>,
239 view: llama_cpp_sys_2::llama_kv_cache_view,
240}
241
242impl<'a> KVCacheView<'a> {
243 /// Update the KV cache view structure with the current state of the KV cache. (use only for debugging purposes)
244 pub fn update(&mut self) {
245 unsafe {
246 llama_cpp_sys_2::llama_kv_cache_view_update(self.ctx.context.as_ptr(), &mut self.view);
247 }
248 }
249
250 /// Number of KV cache cells. This will be the same as the context size.
251 #[must_use]
252 pub fn n_cells(&self) -> i32 {
253 self.view.n_cells
254 }
255
256 /// Number of tokens in the cache. For example, if there are two populated
257 /// cells, the first with 1 sequence id in it and the second with 2 sequence
258 /// ids then you'll have 3 tokens.
259 #[must_use]
260 pub fn token_count(&self) -> i32 {
261 self.view.token_count
262 }
263
264 /// Number of populated cache cells.
265 #[must_use]
266 pub fn used_cells(&self) -> i32 {
267 self.view.used_cells
268 }
269
270 /// Maximum contiguous empty slots in the cache.
271 #[must_use]
272 pub fn max_contiguous(&self) -> i32 {
273 self.view.max_contiguous
274 }
275
276 /// Index to the start of the `max_contiguous` slot range. Can be negative
277 /// when cache is full.
278 #[must_use]
279 pub fn max_contiguous_idx(&self) -> i32 {
280 self.view.max_contiguous_idx
281 }
282
283 /// Information for individual cells.
284 ///
285 /// # Panics
286 ///
287 /// - if `n_cells` does not fit into usize.
288 pub fn cells(&self) -> impl Iterator<Item = KVCacheViewCell> {
289 unsafe {
290 std::slice::from_raw_parts(
291 self.view.cells,
292 usize::try_from(self.view.n_cells).expect("failed to fit n_cells into usize"),
293 )
294 }
295 .iter()
296 .map(|&cell| KVCacheViewCell { pos: cell.pos })
297 }
298
299 /// The sequences for each cell. There will be `n_max_seq` items per cell.
300 ///
301 /// # Panics
302 ///
303 /// - if `n_cells * n_max_seq` does not fit into usize.
304 /// - if `n_max_seq` does not fit into usize.
305 pub fn cells_sequences(&self) -> impl Iterator<Item = &[llama_cpp_sys_2::llama_seq_id]> {
306 unsafe {
307 std::slice::from_raw_parts(
308 self.view.cells_sequences,
309 usize::try_from(self.view.n_cells * self.view.n_seq_max)
310 .expect("failed to fit n_cells * n_max_seq into usize"),
311 )
312 }
313 .chunks(usize::try_from(self.view.n_seq_max).expect("failed to fit n_max_seq into usize"))
314 }
315}
316
317impl<'a> Drop for KVCacheView<'a> {
318 fn drop(&mut self) {
319 unsafe {
320 llama_cpp_sys_2::llama_kv_cache_view_free(&mut self.view);
321 }
322 }
323}