ik_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)]
9#[allow(clippy::module_name_repetitions)]
10pub enum KvCacheConversionError {
11 /// Sequence id conversion to i32 failed
12 #[error("Provided sequence id is too large for a i32")]
13 SeqIdTooLarge(#[source] TryFromIntError),
14 /// Position 0 conversion to i32 failed
15 #[error("Provided start position is too large for a i32")]
16 P0TooLarge(#[source] TryFromIntError),
17 /// Position 1 conversion to i32 failed
18 #[error("Provided end position is too large for a i32")]
19 P1TooLarge(#[source] TryFromIntError),
20 /// Partial sequence couldn't be removed
21 #[error("Couldn't remove partial sequence")]
22 PartialSequenceRemovalFailed(i32, i32),
23}
24
25impl LlamaContext<'_> {
26 /// Copy the cache from one sequence to another.
27 ///
28 /// # Parameters
29 ///
30 /// * `src` - The sequence id to copy the cache from.
31 /// * `dest` - The sequence id to copy the cache to.
32 /// * `size` - The size of the cache to copy.
33 pub fn copy_cache(&mut self, src: i32, dest: i32, size: i32) {
34 unsafe {
35 ik_llama_cpp_sys::llama_kv_cache_seq_cp(self.context.as_ptr(), src, dest, 0, size);
36 }
37 }
38
39 /// Copy the cache from one sequence to another.
40 ///
41 /// # Returns
42 /// A `Result` indicating whether the operation was successful.
43 ///
44 /// # Parameters
45 /// * `src` - The sequence id to copy the cache from.
46 /// * `dest` - The sequence id to copy the cache to.
47 /// * `p0` - The start position of the cache to clear. If `None`, the entire cache is copied up to `p1`.
48 /// * `p1` - The end position of the cache to clear. If `None`, the entire cache is copied starting from `p0`.
49 ///
50 /// # Errors
51 /// If either position exceeds [`i32::MAX`].
52 pub fn copy_kv_cache_seq(
53 &mut self,
54 src: i32,
55 dest: i32,
56 p0: Option<u32>,
57 p1: Option<u32>,
58 ) -> Result<(), KvCacheConversionError> {
59 let p0 = p0
60 .map_or(Ok(-1), i32::try_from)
61 .map_err(KvCacheConversionError::P0TooLarge)?;
62 let p1 = p1
63 .map_or(Ok(-1), i32::try_from)
64 .map_err(KvCacheConversionError::P1TooLarge)?;
65 unsafe {
66 ik_llama_cpp_sys::llama_kv_cache_seq_cp(self.context.as_ptr(), src, dest, p0, p1);
67 }
68 Ok(())
69 }
70
71 /// Clear the kv cache for the given sequence within the specified range `[p0, p1)`
72 /// Returns `false` only when partial sequence removals fail. Full sequence removals always succeed.
73 ///
74 /// # Returns
75 /// A `Result` indicating whether the operation was successful. If the sequence id or
76 /// either position exceeds the maximum i32 value, no removal is attempted and an `Err` is returned.
77 ///
78 /// # Parameters
79 /// * `src` - The sequence id to clear the cache for. If `None`, matches all sequences
80 /// * `p0` - The start position of the cache to clear. If `None`, the entire cache is cleared up to `p1`.
81 /// * `p1` - The end position of the cache to clear. If `None`, the entire cache is cleared from `p0`.
82 ///
83 /// # Errors
84 /// If the sequence id or either position exceeds [`i32::MAX`].
85 pub fn clear_kv_cache_seq(
86 &mut self,
87 src: Option<u32>,
88 p0: Option<u32>,
89 p1: Option<u32>,
90 ) -> Result<bool, KvCacheConversionError> {
91 let src = src
92 .map_or(Ok(-1), i32::try_from)
93 .map_err(KvCacheConversionError::SeqIdTooLarge)?;
94 let p0 = p0
95 .map_or(Ok(-1), i32::try_from)
96 .map_err(KvCacheConversionError::P0TooLarge)?;
97 let p1 = p1
98 .map_or(Ok(-1), i32::try_from)
99 .map_err(KvCacheConversionError::P1TooLarge)?;
100 Ok(unsafe { ik_llama_cpp_sys::llama_kv_cache_seq_rm(self.context.as_ptr(), src, p0, p1) })
101 }
102
103 /// Clear the KV cache
104 pub fn clear_kv_cache(&mut self) {
105 unsafe { ik_llama_cpp_sys::llama_kv_cache_clear(self.context.as_ptr()) }
106 }
107
108 /// Removes all tokens that do not belong to the specified sequence
109 ///
110 /// # Parameters
111 ///
112 /// * `seq_id` - The sequence id to keep
113 pub fn llama_kv_cache_seq_keep(&mut self, seq_id: i32) {
114 unsafe { ik_llama_cpp_sys::llama_kv_cache_seq_keep(self.context.as_ptr(), seq_id) }
115 }
116
117 #[allow(clippy::doc_markdown)]
118 /// Copy all tokens that belong to the specified sequence to another sequence
119 /// If the KV cache is RoPEd, the KV data is updated accordingly:
120 /// - lazily on next [`LlamaContext::decode`]
121 /// - explicitly with [`Self::kv_cache_update`]
122 ///
123 /// # Returns
124 /// A `Result` indicating whether the operation was successful.
125 ///
126 /// # Parameters
127 ///
128 /// * `seq_id_src` - The sequence id to copy from. If negative, matches any sequence.
129 /// * `seq_id_dst` - The sequence id to copy to
130 /// * `p0` - The start position of the cache to copy from. If `None`, the entire cache is updated up to `p1`.
131 /// * `p1` - The end position of the cache to copy from. If `None`, the entire cache is updated starting from `p0`.
132 ///
133 /// # Errors
134 /// If either position exceeds [`i32::MAX`].
135 pub fn kv_cache_seq_cp(
136 &mut self,
137 seq_id_src: i32,
138 seq_id_dst: i32,
139 p0: Option<u32>,
140 p1: Option<u32>,
141 ) -> Result<(), KvCacheConversionError> {
142 let p0 = p0
143 .map_or(Ok(-1), i32::try_from)
144 .map_err(KvCacheConversionError::P0TooLarge)?;
145 let p1 = p1
146 .map_or(Ok(-1), i32::try_from)
147 .map_err(KvCacheConversionError::P1TooLarge)?;
148 unsafe {
149 ik_llama_cpp_sys::llama_kv_cache_seq_cp(
150 self.context.as_ptr(),
151 seq_id_src,
152 seq_id_dst,
153 p0,
154 p1,
155 );
156 }
157 Ok(())
158 }
159
160 #[allow(clippy::doc_markdown)]
161 /// Adds relative position "delta" to all tokens that belong to the specified sequence and have positions in `[p0, p1)`
162 /// If the KV cache is RoPEd, the KV data is updated accordingly:
163 /// - lazily on next [`LlamaContext::decode`]
164 /// - explicitly with [`Self::kv_cache_update`]
165 ///
166 /// # Returns
167 /// A `Result` indicating whether the operation was successful.
168 ///
169 /// # Parameters
170 ///
171 /// * `seq_id` - The sequence id to update
172 /// * `p0` - The start position of the cache to update. If `None`, the entire cache is updated up to `p1`.
173 /// * `p1` - The end position of the cache to update. If `None`, the entire cache is updated starting from `p0`.
174 /// * `delta` - The relative position to add to the tokens
175 ///
176 /// # Errors
177 /// If either position exceeds [`i32::MAX`].
178 pub fn kv_cache_seq_add(
179 &mut self,
180 seq_id: i32,
181 p0: Option<u32>,
182 p1: Option<u32>,
183 delta: i32,
184 ) -> Result<(), KvCacheConversionError> {
185 let p0 = p0
186 .map_or(Ok(-1), i32::try_from)
187 .map_err(KvCacheConversionError::P0TooLarge)?;
188 let p1 = p1
189 .map_or(Ok(-1), i32::try_from)
190 .map_err(KvCacheConversionError::P1TooLarge)?;
191 unsafe {
192 ik_llama_cpp_sys::llama_kv_cache_seq_add(self.context.as_ptr(), seq_id, p0, p1, delta);
193 }
194 Ok(())
195 }
196
197 #[allow(clippy::doc_markdown)]
198 /// Removes all tokens that belong to the specified sequence and have positions in [p0, p1)
199 /// If the KV cache is RoPEd, the KV data is updated accordingly:
200 /// - lazily on next [`LlamaContext::decode`]
201 /// - explicitly with [`Self::kv_cache_update`]
202 ///
203 /// # Returns
204 /// A `Result` indicating whether the operation was successful.
205 ///
206 /// # Parameters
207 ///
208 /// * `seq_id` - The sequence id to update
209 /// * `p0` - The start position of the cache to update. If `None`, the entire cache is updated up to `p1`.
210 /// * `p1` - The end position of the cache to update. If `None`, the entire cache is updated starting from `p0`.
211 ///
212 /// # Errors
213 /// If either position exceeds [`i32::MAX`].
214 pub fn kv_cache_seq_rm(
215 &mut self,
216 seq_id: i32,
217 p0: Option<u32>,
218 p1: Option<u32>,
219 ) -> Result<(), KvCacheConversionError> {
220 let p0 = p0
221 .map_or(Ok(-1), i32::try_from)
222 .map_err(KvCacheConversionError::P0TooLarge)?;
223 let p1 = p1
224 .map_or(Ok(-1), i32::try_from)
225 .map_err(KvCacheConversionError::P1TooLarge)?;
226 if !unsafe {
227 ik_llama_cpp_sys::llama_kv_cache_seq_rm(self.context.as_ptr(), seq_id, p0, p1)
228 } {
229 return Err(KvCacheConversionError::PartialSequenceRemovalFailed(p0, p1));
230 }
231 Ok(())
232 }
233
234 /// Integer division of the positions by factor of `d > 1`
235 /// If the KV cache is `RoPEd`, the KV data is updated accordingly:
236 /// - lazily on next [`LlamaContext::decode`]
237 /// - explicitly with [`Self::kv_cache_update`]
238 ///
239 /// # Returns
240 /// A `Result` indicating whether the operation was successful.
241 ///
242 /// # Parameters
243 ///
244 /// * `seq_id` - The sequence id to update
245 /// * `p0` - The start position of the cache to update. If `None`, the entire cache is updated up to `p1`.
246 /// * `p1` - The end position of the cache to update. If `None`, the entire cache is updated starting from `p0`.
247 /// * `d` - The factor to divide the positions by
248 ///
249 /// # Errors
250 /// If either position exceeds [`i32::MAX`].
251 pub fn kv_cache_seq_div(
252 &mut self,
253 seq_id: i32,
254 p0: Option<u32>,
255 p1: Option<u32>,
256 d: NonZeroU8,
257 ) -> Result<(), KvCacheConversionError> {
258 let p0 = p0
259 .map_or(Ok(-1), i32::try_from)
260 .map_err(KvCacheConversionError::P0TooLarge)?;
261 let p1 = p1
262 .map_or(Ok(-1), i32::try_from)
263 .map_err(KvCacheConversionError::P1TooLarge)?;
264 let d = c_int::from(d.get());
265 unsafe {
266 ik_llama_cpp_sys::llama_kv_cache_seq_div(self.context.as_ptr(), seq_id, p0, p1, d);
267 }
268 Ok(())
269 }
270
271 /// Returns the smallest position present in the KV cache for the specified sequence
272 ///
273 /// # Parameters
274 ///
275 /// * `seq_id` - The sequence id to get the min position for
276 #[must_use]
277 pub fn kv_cache_seq_pos_min(&self, seq_id: i32) -> i32 {
278 unsafe { ik_llama_cpp_sys::llama_kv_cache_seq_pos_min(self.context.as_ptr(), seq_id) }
279 }
280
281 /// Returns the largest position present in the KV cache for the specified sequence
282 ///
283 /// # Parameters
284 ///
285 /// * `seq_id` - The sequence id to get the max position for
286 #[must_use]
287 pub fn kv_cache_seq_pos_max(&self, seq_id: i32) -> i32 {
288 unsafe { ik_llama_cpp_sys::llama_kv_cache_seq_pos_max(self.context.as_ptr(), seq_id) }
289 }
290
291 /// Defragment the KV cache
292 /// This will be applied:
293 /// - lazily on next [`LlamaContext::decode`]
294 /// - explicitly with [`Self::kv_cache_update`]
295 pub fn kv_cache_defrag(&mut self) {
296 unsafe { ik_llama_cpp_sys::llama_kv_cache_defrag(self.context.as_ptr()) }
297 }
298
299 /// Apply the KV cache updates (such as K-shifts, defragmentation, etc.)
300 ///
301 /// # Returns
302 /// The status returned by the underlying `llama_kv_cache_update`. Positive return
303 /// values are not a fatal error, but rather a warning:
304 /// - `0` - success
305 /// - `1` - context overflow in a model where k-shift is not supported
306 pub fn kv_cache_update(&mut self) -> i32 {
307 unsafe { ik_llama_cpp_sys::llama_kv_cache_update(self.context.as_ptr()) }
308 }
309}