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 let mem = unsafe { llama_cpp_sys_2::llama_get_memory(self.context.as_ptr()) };
35 unsafe { llama_cpp_sys_2::llama_memory_seq_cp(mem, src, dest, 0, size) }
36 }
37
38 /// Copy the cache from one sequence to another.
39 ///
40 /// # Returns
41 /// A `Result` indicating whether the operation was successful.
42 ///
43 /// # Parameters
44 /// * `src` - The sequence id to copy the cache from.
45 /// * `dest` - The sequence id to copy the cache to.
46 /// * `p0` - The start position of the cache to clear. If `None`, the entire cache is copied up to `p1`.
47 /// * `p1` - The end position of the cache to clear. If `None`, the entire cache is copied starting from `p0`.
48 ///
49 /// # Errors
50 /// If either position exceeds [`i32::MAX`].
51 pub fn copy_kv_cache_seq(
52 &mut self,
53 src: i32,
54 dest: i32,
55 p0: Option<u32>,
56 p1: Option<u32>,
57 ) -> Result<(), KvCacheConversionError> {
58 let p0 = p0
59 .map_or(Ok(-1), i32::try_from)
60 .map_err(KvCacheConversionError::P0TooLarge)?;
61 let p1 = p1
62 .map_or(Ok(-1), i32::try_from)
63 .map_err(KvCacheConversionError::P1TooLarge)?;
64 let mem = unsafe { llama_cpp_sys_2::llama_get_memory(self.context.as_ptr()) };
65 unsafe { llama_cpp_sys_2::llama_memory_seq_cp(mem, src, dest, p0, p1) };
66 Ok(())
67 }
68
69 /// Clear the kv cache for the given sequence within the specified range `[p0, p1)`
70 /// Returns `false` only when partial sequence removals fail. Full sequence removals always succeed.
71 ///
72 /// # Returns
73 /// A `Result` indicating whether the operation was successful. If the sequence id or
74 /// either position exceeds the maximum i32 value, no removal is attempted and an `Err` is returned.
75 ///
76 /// # Parameters
77 /// * `src` - The sequence id to clear the cache for. If `None`, matches all sequences
78 /// * `p0` - The start position of the cache to clear. If `None`, the entire cache is cleared up to `p1`.
79 /// * `p1` - The end position of the cache to clear. If `None`, the entire cache is cleared from `p0`.
80 ///
81 /// # Errors
82 /// If the sequence id or either position exceeds [`i32::MAX`].
83 pub fn clear_kv_cache_seq(
84 &mut self,
85 src: Option<u32>,
86 p0: Option<u32>,
87 p1: Option<u32>,
88 ) -> Result<bool, KvCacheConversionError> {
89 let src = src
90 .map_or(Ok(-1), i32::try_from)
91 .map_err(KvCacheConversionError::SeqIdTooLarge)?;
92 let p0 = p0
93 .map_or(Ok(-1), i32::try_from)
94 .map_err(KvCacheConversionError::P0TooLarge)?;
95 let p1 = p1
96 .map_or(Ok(-1), i32::try_from)
97 .map_err(KvCacheConversionError::P1TooLarge)?;
98 let mem = unsafe { llama_cpp_sys_2::llama_get_memory(self.context.as_ptr()) };
99 Ok(unsafe { llama_cpp_sys_2::llama_memory_seq_rm(mem, src, p0, p1) })
100 }
101
102 /// Clear the KV cache
103 pub fn clear_kv_cache(&mut self) {
104 let mem = unsafe { llama_cpp_sys_2::llama_get_memory(self.context.as_ptr()) };
105 // clear both metadata and data buffers to match previous semantics
106 unsafe { llama_cpp_sys_2::llama_memory_clear(mem, true) }
107 }
108
109 /// Removes all tokens that do not belong to the specified sequence
110 ///
111 /// # Parameters
112 ///
113 /// * `seq_id` - The sequence id to keep
114 pub fn llama_kv_cache_seq_keep(&mut self, seq_id: i32) {
115 let mem = unsafe { llama_cpp_sys_2::llama_get_memory(self.context.as_ptr()) };
116 unsafe { llama_cpp_sys_2::llama_memory_seq_keep(mem, seq_id) }
117 }
118
119 #[allow(clippy::doc_markdown)]
120 /// Copy all tokens that belong to the specified sequence to another sequence
121 /// If the KV cache is RoPEd, the KV data is updated accordingly:
122 /// - lazily on next [`LlamaContext::decode`]
123 /// - explicitly with [`Self::kv_cache_update`]
124 ///
125 /// # Returns
126 /// A `Result` indicating whether the operation was successful.
127 ///
128 /// # Parameters
129 ///
130 /// * `seq_id_src` - The sequence id to copy from. If negative, matches any sequence.
131 /// * `seq_id_dst` - The sequence id to copy to
132 /// * `p0` - The start position of the cache to copy from. If `None`, the entire cache is updated up to `p1`.
133 /// * `p1` - The end position of the cache to copy from. If `None`, the entire cache is updated starting from `p0`.
134 ///
135 /// # Errors
136 /// If either position exceeds [`i32::MAX`].
137 pub fn kv_cache_seq_cp(
138 &mut self,
139 seq_id_src: i32,
140 seq_id_dst: i32,
141 p0: Option<u32>,
142 p1: Option<u32>,
143 ) -> Result<(), KvCacheConversionError> {
144 let p0 = p0
145 .map_or(Ok(-1), i32::try_from)
146 .map_err(KvCacheConversionError::P0TooLarge)?;
147 let p1 = p1
148 .map_or(Ok(-1), i32::try_from)
149 .map_err(KvCacheConversionError::P1TooLarge)?;
150 let mem = unsafe { llama_cpp_sys_2::llama_get_memory(self.context.as_ptr()) };
151 unsafe { llama_cpp_sys_2::llama_memory_seq_cp(mem, seq_id_src, seq_id_dst, p0, p1) };
152 Ok(())
153 }
154
155 #[allow(clippy::doc_markdown)]
156 /// Adds relative position "delta" to all tokens that belong to the specified sequence and have positions in `[p0, p1)`
157 /// If the KV cache is RoPEd, the KV data is updated accordingly:
158 /// - lazily on next [`LlamaContext::decode`]
159 /// - explicitly with [`Self::kv_cache_update`]
160 ///
161 /// # Returns
162 /// A `Result` indicating whether the operation was successful.
163 ///
164 /// # Parameters
165 ///
166 /// * `seq_id` - The sequence id to update
167 /// * `p0` - The start position of the cache to update. If `None`, the entire cache is updated up to `p1`.
168 /// * `p1` - The end position of the cache to update. If `None`, the entire cache is updated starting from `p0`.
169 /// * `delta` - The relative position to add to the tokens
170 ///
171 /// # Errors
172 /// If either position exceeds [`i32::MAX`].
173 pub fn kv_cache_seq_add(
174 &mut self,
175 seq_id: i32,
176 p0: Option<u32>,
177 p1: Option<u32>,
178 delta: i32,
179 ) -> Result<(), KvCacheConversionError> {
180 let p0 = p0
181 .map_or(Ok(-1), i32::try_from)
182 .map_err(KvCacheConversionError::P0TooLarge)?;
183 let p1 = p1
184 .map_or(Ok(-1), i32::try_from)
185 .map_err(KvCacheConversionError::P1TooLarge)?;
186 let mem = unsafe { llama_cpp_sys_2::llama_get_memory(self.context.as_ptr()) };
187 unsafe { llama_cpp_sys_2::llama_memory_seq_add(mem, seq_id, p0, p1, delta) };
188 Ok(())
189 }
190
191 #[allow(clippy::doc_markdown)]
192 /// Removes all tokens that belong to the specified sequence and have positions in [p0, p1)
193 /// If the KV cache is RoPEd, the KV data is updated accordingly:
194 /// - lazily on next [`LlamaContext::decode`]
195 /// - explicitly with [`Self::kv_cache_update`]
196 ///
197 /// # Returns
198 /// A `Result` indicating whether the operation was successful.
199 ///
200 /// # Parameters
201 ///
202 /// * `seq_id` - The sequence id to update
203 /// * `p0` - The start position of the cache to update. If `None`, the entire cache is updated up to `p1`.
204 /// * `p1` - The end position of the cache to update. If `None`, the entire cache is updated starting from `p0`.
205 ///
206 /// # Errors
207 /// If either position exceeds [`i32::MAX`].
208 pub fn kv_cache_seq_rm(
209 &mut self,
210 seq_id: i32,
211 p0: Option<u32>,
212 p1: Option<u32>,
213 ) -> Result<(), KvCacheConversionError> {
214 let p0 = p0
215 .map_or(Ok(-1), i32::try_from)
216 .map_err(KvCacheConversionError::P0TooLarge)?;
217 let p1 = p1
218 .map_or(Ok(-1), i32::try_from)
219 .map_err(KvCacheConversionError::P1TooLarge)?;
220 let mem = unsafe { llama_cpp_sys_2::llama_get_memory(self.context.as_ptr()) };
221 if !unsafe { llama_cpp_sys_2::llama_memory_seq_rm(mem, seq_id, p0, p1) } {
222 return Err(KvCacheConversionError::PartialSequenceRemovalFailed(p0, p1));
223 }
224 Ok(())
225 }
226
227 /// Integer division of the positions by factor of `d > 1`
228 /// If the KV cache is `RoPEd`, the KV data is updated accordingly:
229 /// - lazily on next [`LlamaContext::decode`]
230 /// - explicitly with [`Self::kv_cache_update`]
231 ///
232 /// # Returns
233 /// A `Result` indicating whether the operation was successful.
234 ///
235 /// # Parameters
236 ///
237 /// * `seq_id` - The sequence id to update
238 /// * `p0` - The start position of the cache to update. If `None`, the entire cache is updated up to `p1`.
239 /// * `p1` - The end position of the cache to update. If `None`, the entire cache is updated starting from `p0`.
240 /// * `d` - The factor to divide the positions by
241 ///
242 /// # Errors
243 /// If either position exceeds [`i32::MAX`].
244 pub fn kv_cache_seq_div(
245 &mut self,
246 seq_id: i32,
247 p0: Option<u32>,
248 p1: Option<u32>,
249 d: NonZeroU8,
250 ) -> Result<(), KvCacheConversionError> {
251 let p0 = p0
252 .map_or(Ok(-1), i32::try_from)
253 .map_err(KvCacheConversionError::P0TooLarge)?;
254 let p1 = p1
255 .map_or(Ok(-1), i32::try_from)
256 .map_err(KvCacheConversionError::P1TooLarge)?;
257 let d = c_int::from(d.get());
258 let mem = unsafe { llama_cpp_sys_2::llama_get_memory(self.context.as_ptr()) };
259 unsafe { llama_cpp_sys_2::llama_memory_seq_div(mem, seq_id, p0, p1, d) }
260 Ok(())
261 }
262
263 /// Returns the smallest position present in the KV cache for the specified sequence
264 ///
265 /// # Parameters
266 ///
267 /// * `seq_id` - The sequence id to get the min position for
268 #[must_use]
269 pub fn kv_cache_seq_pos_min(&self, seq_id: i32) -> i32 {
270 let mem = unsafe { llama_cpp_sys_2::llama_get_memory(self.context.as_ptr()) };
271 unsafe { llama_cpp_sys_2::llama_memory_seq_pos_min(mem, seq_id) }
272 }
273
274 /// Returns the largest position present in the KV cache for the specified sequence
275 ///
276 /// # Parameters
277 ///
278 /// * `seq_id` - The sequence id to get the max position for
279 #[must_use]
280 pub fn kv_cache_seq_pos_max(&self, seq_id: i32) -> i32 {
281 let mem = unsafe { llama_cpp_sys_2::llama_get_memory(self.context.as_ptr()) };
282 unsafe { llama_cpp_sys_2::llama_memory_seq_pos_max(mem, seq_id) }
283 }
284}