llama_cpp_4/context/kv_cache.rs
1//! utilities for working with the kv cache
2
3use crate::context::LlamaContext;
4use std::num::{NonZeroU8, TryFromIntError};
5
6/// Errors that can occur when attempting to prepare values for the kv cache
7#[derive(Debug, Eq, PartialEq, thiserror::Error)]
8pub enum KvCacheConversionError {
9 /// Sequence id conversion to i32 failed
10 #[error("Provided sequence id is too large for a i32")]
11 SeqIdTooLarge(#[source] TryFromIntError),
12 /// Position 0 conversion to i32 failed
13 #[error("Provided start position is too large for a i32")]
14 P0TooLarge(#[source] TryFromIntError),
15 /// Position 1 conversion to i32 failed
16 #[error("Provided end position is too large for a i32")]
17 P1TooLarge(#[source] TryFromIntError),
18}
19
20impl LlamaContext<'_> {
21 /// How many backend buffers a state write of this context's memory would
22 /// touch.
23 ///
24 /// `0` for an empty cache. Useful before
25 /// [`state_seq_save_file`](crate::context::LlamaContext::state_seq_save_file)
26 /// and friends: a KV cache spread across several backends is written buffer
27 /// by buffer, so the count is a proxy for how much work a save will be, and
28 /// a cache reporting zero has nothing to save yet.
29 ///
30 /// Wraps `llama_memory_get_buffers` through the `ext_shim`, because
31 /// upstream returns a `std::vector` by value and bindgen cannot express
32 /// that — see the note in `build.rs`.
33 #[must_use]
34 pub fn memory_buffer_count(&self) -> usize {
35 // SAFETY: the memory handle is borrowed from a live context, and a null
36 // `out` asks the shim for the count without writing anything.
37 unsafe {
38 let mem = llama_cpp_sys_4::llama_get_memory(self.context.as_ptr());
39 llama_cpp_sys_4::llama_memory_buffers_collect(mem, std::ptr::null_mut(), 0)
40 }
41 }
42
43 /// Copy the cache from one sequence to another.
44 ///
45 /// # Parameters
46 ///
47 /// * `src` - The sequence id to copy the cache from.
48 /// * `dest` - The sequence id to copy the cache to.
49 /// * `size` - The size of the cache to copy.
50 pub fn copy_cache(&mut self, src: i32, dest: i32, size: i32) {
51 unsafe {
52 let mem = llama_cpp_sys_4::llama_get_memory(self.context.as_ptr());
53 llama_cpp_sys_4::llama_memory_seq_cp(mem, src, dest, 0, size);
54 }
55 }
56
57 /// Copy the cache from one sequence to another.
58 ///
59 /// # Parameters
60 ///
61 /// * `src` - The sequence id to copy the cache from.
62 /// * `dest` - The sequence id to copy the cache to.
63 /// * `p0` - The start position of the cache to clear. If `None`, the entire cache is copied up to `p1`.
64 /// * `p1` - The end position of the cache to clear. If `None`, the entire cache is copied starting from `p0`.
65 ///
66 /// # Errors
67 ///
68 /// Returns [`KvCacheConversionError`] if either position exceeds the maximum `i32` value.
69 pub fn copy_kv_cache_seq(
70 &mut self,
71 src: i32,
72 dest: i32,
73 p0: Option<u32>,
74 p1: Option<u32>,
75 ) -> Result<(), KvCacheConversionError> {
76 let p0 = p0
77 .map_or(Ok(-1), i32::try_from)
78 .map_err(KvCacheConversionError::P0TooLarge)?;
79 let p1 = p1
80 .map_or(Ok(-1), i32::try_from)
81 .map_err(KvCacheConversionError::P1TooLarge)?;
82 unsafe {
83 let mem = llama_cpp_sys_4::llama_get_memory(self.context.as_ptr());
84 llama_cpp_sys_4::llama_memory_seq_cp(mem, src, dest, p0, p1);
85 }
86 Ok(())
87 }
88
89 /// Clear the kv cache for the given sequence within the specified range `[p0, p1)`.
90 ///
91 /// Returns `false` only when partial sequence removals fail. Full sequence removals always succeed.
92 ///
93 /// # Parameters
94 ///
95 /// * `src` - The sequence id to clear the cache for. If `None`, matches all sequences
96 /// * `p0` - The start position of the cache to clear. If `None`, the entire cache is cleared up to `p1`.
97 /// * `p1` - The end position of the cache to clear. If `None`, the entire cache is cleared from `p0`.
98 ///
99 /// # Errors
100 ///
101 /// Returns [`KvCacheConversionError`] if the sequence id or either position exceeds the maximum `i32` value.
102 pub fn clear_kv_cache_seq(
103 &mut self,
104 src: Option<u32>,
105 p0: Option<u32>,
106 p1: Option<u32>,
107 ) -> Result<bool, KvCacheConversionError> {
108 let src = src
109 .map_or(Ok(-1), i32::try_from)
110 .map_err(KvCacheConversionError::SeqIdTooLarge)?;
111 let p0 = p0
112 .map_or(Ok(-1), i32::try_from)
113 .map_err(KvCacheConversionError::P0TooLarge)?;
114 let p1 = p1
115 .map_or(Ok(-1), i32::try_from)
116 .map_err(KvCacheConversionError::P1TooLarge)?;
117 Ok(unsafe {
118 let mem = llama_cpp_sys_4::llama_get_memory(self.context.as_ptr());
119 llama_cpp_sys_4::llama_memory_seq_rm(mem, src, p0, p1)
120 })
121 }
122
123 /// Clear the KV cache
124 pub fn clear_kv_cache(&mut self) {
125 unsafe {
126 let mem = llama_cpp_sys_4::llama_get_memory(self.context.as_ptr());
127 llama_cpp_sys_4::llama_memory_clear(mem, true);
128 }
129 }
130
131 /// Removes all tokens that do not belong to the specified sequence
132 ///
133 /// # Parameters
134 ///
135 /// * `seq_id` - The sequence id to keep
136 pub fn llama_kv_cache_seq_keep(&mut self, seq_id: i32) {
137 unsafe {
138 let mem = llama_cpp_sys_4::llama_get_memory(self.context.as_ptr());
139 llama_cpp_sys_4::llama_memory_seq_keep(mem, seq_id);
140 }
141 }
142
143 #[allow(clippy::doc_markdown)]
144 /// Adds relative position "delta" to all tokens that belong to the specified sequence and have positions in `[p0, p1)`.
145 ///
146 /// # Parameters
147 ///
148 /// * `seq_id` - The sequence id to update
149 /// * `p0` - The start position of the cache to update. If `None`, the entire cache is updated up to `p1`.
150 /// * `p1` - The end position of the cache to update. If `None`, the entire cache is updated starting from `p0`.
151 /// * `delta` - The relative position to add to the tokens
152 ///
153 /// # Errors
154 ///
155 /// Returns [`KvCacheConversionError`] if either position exceeds the maximum `i32` value.
156 pub fn kv_cache_seq_add(
157 &mut self,
158 seq_id: i32,
159 p0: Option<u32>,
160 p1: Option<u32>,
161 delta: i32,
162 ) -> Result<(), KvCacheConversionError> {
163 let p0 = p0
164 .map_or(Ok(-1), i32::try_from)
165 .map_err(KvCacheConversionError::P0TooLarge)?;
166 let p1 = p1
167 .map_or(Ok(-1), i32::try_from)
168 .map_err(KvCacheConversionError::P1TooLarge)?;
169 unsafe {
170 let mem = llama_cpp_sys_4::llama_get_memory(self.context.as_ptr());
171 llama_cpp_sys_4::llama_memory_seq_add(mem, seq_id, p0, p1, delta);
172 }
173 Ok(())
174 }
175
176 /// Integer division of the positions by factor of `d > 1`.
177 ///
178 /// # Parameters
179 ///
180 /// * `seq_id` - The sequence id to update
181 /// * `p0` - The start position of the cache to update. If `None`, the entire cache is updated up to `p1`.
182 /// * `p1` - The end position of the cache to update. If `None`, the entire cache is updated starting from `p0`.
183 /// * `d` - The factor to divide the positions by
184 ///
185 /// # Errors
186 ///
187 /// Returns [`KvCacheConversionError`] if either position exceeds the maximum `i32` value.
188 pub fn kv_cache_seq_div(
189 &mut self,
190 seq_id: i32,
191 p0: Option<u32>,
192 p1: Option<u32>,
193 d: NonZeroU8,
194 ) -> Result<(), KvCacheConversionError> {
195 let p0 = p0
196 .map_or(Ok(-1), i32::try_from)
197 .map_err(KvCacheConversionError::P0TooLarge)?;
198 let p1 = p1
199 .map_or(Ok(-1), i32::try_from)
200 .map_err(KvCacheConversionError::P1TooLarge)?;
201 let d = i32::from(d.get());
202 unsafe {
203 let mem = llama_cpp_sys_4::llama_get_memory(self.context.as_ptr());
204 llama_cpp_sys_4::llama_memory_seq_div(mem, seq_id, p0, p1, d);
205 }
206 Ok(())
207 }
208
209 /// Returns the largest position present in the KV cache for the specified sequence
210 ///
211 /// # Parameters
212 ///
213 /// * `seq_id` - The sequence id to get the max position for
214 #[must_use]
215 pub fn kv_cache_seq_pos_max(&self, seq_id: i32) -> i32 {
216 unsafe {
217 let mem = llama_cpp_sys_4::llama_get_memory(self.context.as_ptr());
218 llama_cpp_sys_4::llama_memory_seq_pos_max(mem, seq_id)
219 }
220 }
221}