@@ -28,7 +28,8 @@ impl LlamaContext<'_> {
2828 /// * `dest` - The sequence id to copy the cache to.
2929 /// * `size` - The size of the cache to copy.
3030 pub fn copy_cache ( & mut self , src : i32 , dest : i32 , size : i32 ) {
31- unsafe { llama_cpp_sys_2:: llama_kv_self_seq_cp ( self . context . as_ptr ( ) , src, dest, 0 , size) }
31+ let mem = unsafe { llama_cpp_sys_2:: llama_get_memory ( self . context . as_ptr ( ) ) } ;
32+ unsafe { llama_cpp_sys_2:: llama_memory_seq_cp ( mem, src, dest, 0 , size) }
3233 }
3334
3435 /// Copy the cache from one sequence to another.
@@ -57,9 +58,8 @@ impl LlamaContext<'_> {
5758 let p1 = p1
5859 . map_or ( Ok ( -1 ) , i32:: try_from)
5960 . map_err ( KvCacheConversionError :: P1TooLarge ) ?;
60- unsafe {
61- llama_cpp_sys_2:: llama_kv_self_seq_cp ( self . context . as_ptr ( ) , src, dest, p0, p1) ;
62- }
61+ let mem = unsafe { llama_cpp_sys_2:: llama_get_memory ( self . context . as_ptr ( ) ) } ;
62+ unsafe { llama_cpp_sys_2:: llama_memory_seq_cp ( mem, src, dest, p0, p1) } ;
6363 Ok ( ( ) )
6464 }
6565
@@ -92,18 +92,16 @@ impl LlamaContext<'_> {
9292 let p1 = p1
9393 . map_or ( Ok ( -1 ) , i32:: try_from)
9494 . map_err ( KvCacheConversionError :: P1TooLarge ) ?;
95- Ok ( unsafe { llama_cpp_sys_2:: llama_kv_self_seq_rm ( self . context . as_ptr ( ) , src, p0, p1) } )
95+ let mem = unsafe { llama_cpp_sys_2:: llama_get_memory ( self . context . as_ptr ( ) ) } ;
96+ Ok ( unsafe { llama_cpp_sys_2:: llama_memory_seq_rm ( mem, src, p0, p1) } )
9697 }
9798
98- /// Returns the number of used KV cells (i.e. have at least one sequence assigned to them)
99- #[ must_use]
100- pub fn get_kv_cache_used_cells ( & self ) -> i32 {
101- unsafe { llama_cpp_sys_2:: llama_kv_self_used_cells ( self . context . as_ptr ( ) ) }
102- }
10399
104100 /// Clear the KV cache
105101 pub fn clear_kv_cache ( & mut self ) {
106- unsafe { llama_cpp_sys_2:: llama_kv_self_clear ( self . context . as_ptr ( ) ) }
102+ let mem = unsafe { llama_cpp_sys_2:: llama_get_memory ( self . context . as_ptr ( ) ) } ;
103+ // clear both metadata and data buffers to match previous semantics
104+ unsafe { llama_cpp_sys_2:: llama_memory_clear ( mem, true ) }
107105 }
108106
109107 /// Removes all tokens that do not belong to the specified sequence
@@ -112,7 +110,8 @@ impl LlamaContext<'_> {
112110 ///
113111 /// * `seq_id` - The sequence id to keep
114112 pub fn llama_kv_cache_seq_keep ( & mut self , seq_id : i32 ) {
115- unsafe { llama_cpp_sys_2:: llama_kv_self_seq_keep ( self . context . as_ptr ( ) , seq_id) }
113+ let mem = unsafe { llama_cpp_sys_2:: llama_get_memory ( self . context . as_ptr ( ) ) } ;
114+ unsafe { llama_cpp_sys_2:: llama_memory_seq_keep ( mem, seq_id) }
116115 }
117116
118117 #[ allow( clippy:: doc_markdown) ]
@@ -146,9 +145,8 @@ impl LlamaContext<'_> {
146145 let p1 = p1
147146 . map_or ( Ok ( -1 ) , i32:: try_from)
148147 . map_err ( KvCacheConversionError :: P1TooLarge ) ?;
149- unsafe {
150- llama_cpp_sys_2:: llama_kv_self_seq_add ( self . context . as_ptr ( ) , seq_id, p0, p1, delta) ;
151- }
148+ let mem = unsafe { llama_cpp_sys_2:: llama_get_memory ( self . context . as_ptr ( ) ) } ;
149+ unsafe { llama_cpp_sys_2:: llama_memory_seq_add ( mem, seq_id, p0, p1, delta) } ;
152150 Ok ( ( ) )
153151 }
154152
@@ -183,7 +181,8 @@ impl LlamaContext<'_> {
183181 . map_or ( Ok ( -1 ) , i32:: try_from)
184182 . map_err ( KvCacheConversionError :: P1TooLarge ) ?;
185183 let d = c_int:: from ( d. get ( ) ) ;
186- unsafe { llama_cpp_sys_2:: llama_kv_self_seq_div ( self . context . as_ptr ( ) , seq_id, p0, p1, d) }
184+ let mem = unsafe { llama_cpp_sys_2:: llama_get_memory ( self . context . as_ptr ( ) ) } ;
185+ unsafe { llama_cpp_sys_2:: llama_memory_seq_div ( mem, seq_id, p0, p1, d) }
187186 Ok ( ( ) )
188187 }
189188
@@ -194,19 +193,7 @@ impl LlamaContext<'_> {
194193 /// * `seq_id` - The sequence id to get the max position for
195194 #[ must_use]
196195 pub fn kv_cache_seq_pos_max ( & self , seq_id : i32 ) -> i32 {
197- unsafe { llama_cpp_sys_2:: llama_kv_self_seq_pos_max ( self . context . as_ptr ( ) , seq_id) }
198- }
199-
200- /// Defragment the KV cache
201- /// This will be applied:
202- /// - lazily on next [`LlamaContext::decode`]
203- /// - explicitly with [`Self::kv_cache_update`]
204- pub fn kv_cache_defrag ( & mut self ) {
205- unsafe { llama_cpp_sys_2:: llama_kv_self_defrag ( self . context . as_ptr ( ) ) }
206- }
207-
208- /// Apply the KV cache updates (such as K-shifts, defragmentation, etc.)
209- pub fn kv_cache_update ( & mut self ) {
210- unsafe { llama_cpp_sys_2:: llama_kv_self_update ( self . context . as_ptr ( ) ) }
196+ let mem = unsafe { llama_cpp_sys_2:: llama_get_memory ( self . context . as_ptr ( ) ) } ;
197+ unsafe { llama_cpp_sys_2:: llama_memory_seq_pos_max ( mem, seq_id) }
211198 }
212199}
0 commit comments