llama_cpp_2/context/session.rs
1//! utilities for working with session files
2
3use crate::context::LlamaContext;
4use crate::token::LlamaToken;
5use std::ffi::{CString, NulError};
6use std::fmt::{Debug, Formatter};
7use std::path::{Path, PathBuf};
8
9/// Flags for state sequence operations.
10///
11/// These flags control what parts of the state are included when saving/restoring
12/// sequence state.
13#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14pub struct LlamaStateSeqFlags(pub(crate) llama_cpp_sys_2::llama_state_seq_flags);
15
16impl LlamaStateSeqFlags {
17 /// Work only with partial states, such as SWA KV cache or recurrent cache (e.g. Mamba).
18 ///
19 /// This flag is useful when you only want to save/restore the recurrent state
20 /// without affecting the KV cache.
21 pub const PARTIAL_ONLY: LlamaStateSeqFlags = LlamaStateSeqFlags(1);
22
23 /// Keep the copied data on device (GPU) memory rather than host.
24 ///
25 /// Getting the state for a `seq_id` with this flag invalidates all
26 /// prior states obtained for that `seq_id` with this flag set.
27 pub const ON_DEVICE: LlamaStateSeqFlags = LlamaStateSeqFlags(2);
28
29 /// Create an empty flags set.
30 pub const fn empty() -> LlamaStateSeqFlags {
31 LlamaStateSeqFlags(0)
32 }
33
34 /// Get the raw flags value.
35 pub const fn bits(&self) -> u32 {
36 self.0
37 }
38
39 /// Construct from raw bits.
40 pub const fn from_bits(bits: u32) -> LlamaStateSeqFlags {
41 LlamaStateSeqFlags(bits)
42 }
43
44 /// Check if a flag is set.
45 pub const fn contains(&self, other: LlamaStateSeqFlags) -> bool {
46 (self.0 & other.0) != 0
47 }
48}
49
50impl Default for LlamaStateSeqFlags {
51 fn default() -> Self {
52 Self::empty()
53 }
54}
55
56impl std::ops::BitOr for LlamaStateSeqFlags {
57 type Output = Self;
58 fn bitor(self, other: Self) -> Self {
59 Self(self.0 | other.0)
60 }
61}
62
63/// Failed to save a sequence state file
64#[derive(Debug, Eq, PartialEq, thiserror::Error)]
65pub enum SaveSeqStateError {
66 /// llama.cpp failed to save the sequence state file
67 #[error("Failed to save sequence state file")]
68 FailedToSave,
69
70 /// null byte in string
71 #[error("null byte in string {0}")]
72 NullError(#[from] NulError),
73
74 /// failed to convert path to str
75 #[error("failed to convert path {0} to str")]
76 PathToStrError(PathBuf),
77}
78
79/// Failed to load a sequence state file
80#[derive(Debug, Eq, PartialEq, thiserror::Error)]
81pub enum LoadSeqStateError {
82 /// llama.cpp failed to load the sequence state file
83 #[error("Failed to load sequence state file")]
84 FailedToLoad,
85
86 /// null byte in string
87 #[error("null byte in string {0}")]
88 NullError(#[from] NulError),
89
90 /// failed to convert path to str
91 #[error("failed to convert path {0} to str")]
92 PathToStrError(PathBuf),
93
94 /// Insufficient max length
95 #[error("max_length is not large enough to hold {n_out} (was {max_tokens})")]
96 InsufficientMaxLength {
97 /// The length of the loaded sequence
98 n_out: usize,
99 /// The maximum length
100 max_tokens: usize,
101 },
102}
103
104/// Failed to save a Session file
105#[derive(Debug, Eq, PartialEq, thiserror::Error)]
106pub enum SaveSessionError {
107 /// llama.cpp failed to save the session file
108 #[error("Failed to save session file")]
109 FailedToSave,
110
111 /// null byte in string
112 #[error("null byte in string {0}")]
113 NullError(#[from] NulError),
114
115 /// failed to convert path to str
116 #[error("failed to convert path {0} to str")]
117 PathToStrError(PathBuf),
118}
119
120/// Failed to load a Session file
121#[derive(Debug, Eq, PartialEq, thiserror::Error)]
122pub enum LoadSessionError {
123 /// llama.cpp failed to load the session file
124 #[error("Failed to load session file")]
125 FailedToLoad,
126
127 /// null byte in string
128 #[error("null byte in string {0}")]
129 NullError(#[from] NulError),
130
131 /// failed to convert path to str
132 #[error("failed to convert path {0} to str")]
133 PathToStrError(PathBuf),
134
135 /// Insufficient max length
136 #[error("max_length is not large enough to hold {n_out} (was {max_tokens})")]
137 InsufficientMaxLength {
138 /// The length of the session file
139 n_out: usize,
140 /// The maximum length
141 max_tokens: usize,
142 },
143}
144
145impl LlamaContext<'_> {
146 /// Save the current session to a file.
147 ///
148 /// # Parameters
149 ///
150 /// * `path_session` - The file to save to.
151 /// * `tokens` - The tokens to associate the session with. This should be a prefix of a sequence of tokens that the context has processed, so that the relevant KV caches are already filled.
152 ///
153 /// # Errors
154 ///
155 /// Fails if the path is not a valid utf8, is not a valid c string, or llama.cpp fails to save the session file.
156 #[deprecated(since = "0.1.136", note = "Use `state_save_file` instead")]
157 pub fn save_session_file(
158 &self,
159 path_session: impl AsRef<Path>,
160 tokens: &[LlamaToken],
161 ) -> Result<(), SaveSessionError> {
162 let path = path_session.as_ref();
163 let path = path
164 .to_str()
165 .ok_or_else(|| SaveSessionError::PathToStrError(path.to_path_buf()))?;
166
167 let cstr = CString::new(path)?;
168
169 if unsafe {
170 llama_cpp_sys_2::llama_save_session_file(
171 self.synchronizable_ptr(),
172 cstr.as_ptr(),
173 tokens.as_ptr().cast::<llama_cpp_sys_2::llama_token>(),
174 tokens.len(),
175 )
176 } {
177 Ok(())
178 } else {
179 Err(SaveSessionError::FailedToSave)
180 }
181 }
182 /// Load a session file into the current context.
183 ///
184 /// You still need to pass the returned tokens to the context for inference to work. What this function buys you is that the KV caches are already filled with the relevant data.
185 ///
186 /// # Parameters
187 ///
188 /// * `path_session` - The file to load from. It must be a session file from a compatible context, otherwise the function will error.
189 /// * `max_tokens` - The maximum token length of the loaded session. If the session was saved with a longer length, the function will error.
190 ///
191 /// # Errors
192 ///
193 /// Fails if the path is not a valid utf8, is not a valid c string, or llama.cpp fails to load the session file. (e.g. the file does not exist, is not a session file, etc.)
194 #[deprecated(since = "0.1.136", note = "Use `state_load_file` instead")]
195 pub fn load_session_file(
196 &mut self,
197 path_session: impl AsRef<Path>,
198 max_tokens: usize,
199 ) -> Result<Vec<LlamaToken>, LoadSessionError> {
200 let path = path_session.as_ref();
201 let path = path
202 .to_str()
203 .ok_or(LoadSessionError::PathToStrError(path.to_path_buf()))?;
204
205 let cstr = CString::new(path)?;
206 let mut tokens: Vec<LlamaToken> = Vec::with_capacity(max_tokens);
207 let mut n_out = 0;
208
209 // SAFETY: cast is valid as LlamaToken is repr(transparent)
210 let tokens_out = tokens.as_mut_ptr().cast::<llama_cpp_sys_2::llama_token>();
211
212 let load_session_success = unsafe {
213 llama_cpp_sys_2::llama_load_session_file(
214 self.context.as_mut_ptr(),
215 cstr.as_ptr(),
216 tokens_out,
217 max_tokens,
218 &mut n_out,
219 )
220 };
221 if load_session_success {
222 if n_out > max_tokens {
223 return Err(LoadSessionError::InsufficientMaxLength { n_out, max_tokens });
224 }
225 // SAFETY: we checked that n_out <= max_tokens and llama.cpp promises that n_out tokens will be written
226 unsafe {
227 tokens.set_len(n_out);
228 }
229 Ok(tokens)
230 } else {
231 Err(LoadSessionError::FailedToLoad)
232 }
233 }
234
235 /// Save the full state to a file.
236 ///
237 /// This is the non-deprecated replacement for [`save_session_file`](Self::save_session_file).
238 ///
239 /// # Parameters
240 ///
241 /// * `path_session` - The file to save to.
242 /// * `tokens` - The tokens to associate the state with. This should be a prefix of a sequence
243 /// of tokens that the context has processed, so that the relevant KV caches are already filled.
244 ///
245 /// # Errors
246 ///
247 /// Fails if the path is not a valid utf8, is not a valid c string, or llama.cpp fails to save
248 /// the state file.
249 pub fn state_save_file(
250 &self,
251 path_session: impl AsRef<Path>,
252 tokens: &[LlamaToken],
253 ) -> Result<(), SaveSessionError> {
254 let path = path_session.as_ref();
255 let path = path
256 .to_str()
257 .ok_or_else(|| SaveSessionError::PathToStrError(path.to_path_buf()))?;
258
259 let cstr = CString::new(path)?;
260
261 if unsafe {
262 llama_cpp_sys_2::llama_state_save_file(
263 self.synchronizable_ptr(),
264 cstr.as_ptr(),
265 tokens.as_ptr().cast::<llama_cpp_sys_2::llama_token>(),
266 tokens.len(),
267 )
268 } {
269 Ok(())
270 } else {
271 Err(SaveSessionError::FailedToSave)
272 }
273 }
274
275 /// Load a state file into the current context.
276 ///
277 /// This is the non-deprecated replacement for [`load_session_file`](Self::load_session_file).
278 ///
279 /// You still need to pass the returned tokens to the context for inference to work. What this
280 /// function buys you is that the KV caches are already filled with the relevant data.
281 ///
282 /// # Parameters
283 ///
284 /// * `path_session` - The file to load from. It must be a state file from a compatible context,
285 /// otherwise the function will error.
286 /// * `max_tokens` - The maximum token length of the loaded state. If the state was saved with a
287 /// longer length, the function will error.
288 ///
289 /// # Errors
290 ///
291 /// Fails if the path is not a valid utf8, is not a valid c string, or llama.cpp fails to load
292 /// the state file.
293 pub fn state_load_file(
294 &mut self,
295 path_session: impl AsRef<Path>,
296 max_tokens: usize,
297 ) -> Result<Vec<LlamaToken>, LoadSessionError> {
298 let path = path_session.as_ref();
299 let path = path
300 .to_str()
301 .ok_or(LoadSessionError::PathToStrError(path.to_path_buf()))?;
302
303 let cstr = CString::new(path)?;
304 let mut tokens: Vec<LlamaToken> = Vec::with_capacity(max_tokens);
305 let mut n_out = 0;
306
307 // SAFETY: cast is valid as LlamaToken is repr(transparent)
308 let tokens_out = tokens.as_mut_ptr().cast::<llama_cpp_sys_2::llama_token>();
309
310 let success = unsafe {
311 llama_cpp_sys_2::llama_state_load_file(
312 self.context.as_mut_ptr(),
313 cstr.as_ptr(),
314 tokens_out,
315 max_tokens,
316 &mut n_out,
317 )
318 };
319 if success {
320 if n_out > max_tokens {
321 return Err(LoadSessionError::InsufficientMaxLength { n_out, max_tokens });
322 }
323 // SAFETY: we checked that n_out <= max_tokens and llama.cpp promises that n_out tokens will be written
324 unsafe {
325 tokens.set_len(n_out);
326 }
327 Ok(tokens)
328 } else {
329 Err(LoadSessionError::FailedToLoad)
330 }
331 }
332
333 /// Save state for a single sequence to a file.
334 ///
335 /// This enables saving state for individual sequences, which is useful for multi-sequence
336 /// inference scenarios.
337 ///
338 /// # Parameters
339 ///
340 /// * `filepath` - The file to save to.
341 /// * `seq_id` - The sequence ID whose state to save.
342 /// * `tokens` - The tokens to associate with the saved state.
343 ///
344 /// # Errors
345 ///
346 /// Fails if the path is not a valid utf8, is not a valid c string, or llama.cpp fails to save
347 /// the sequence state file.
348 ///
349 /// # Returns
350 ///
351 /// The number of bytes written on success.
352 pub fn state_seq_save_file(
353 &self,
354 filepath: impl AsRef<Path>,
355 seq_id: i32,
356 tokens: &[LlamaToken],
357 ) -> Result<usize, SaveSeqStateError> {
358 let path = filepath.as_ref();
359 let path = path
360 .to_str()
361 .ok_or_else(|| SaveSeqStateError::PathToStrError(path.to_path_buf()))?;
362
363 let cstr = CString::new(path)?;
364
365 let bytes_written = unsafe {
366 llama_cpp_sys_2::llama_state_seq_save_file(
367 self.synchronizable_ptr(),
368 cstr.as_ptr(),
369 seq_id,
370 tokens.as_ptr().cast::<llama_cpp_sys_2::llama_token>(),
371 tokens.len(),
372 )
373 };
374
375 if bytes_written == 0 {
376 Err(SaveSeqStateError::FailedToSave)
377 } else {
378 Ok(bytes_written)
379 }
380 }
381
382 /// Load state for a single sequence from a file.
383 ///
384 /// This enables loading state for individual sequences, which is useful for multi-sequence
385 /// inference scenarios.
386 ///
387 /// # Parameters
388 ///
389 /// * `filepath` - The file to load from.
390 /// * `dest_seq_id` - The destination sequence ID to load the state into.
391 /// * `max_tokens` - The maximum number of tokens to read.
392 ///
393 /// # Errors
394 ///
395 /// Fails if the path is not a valid utf8, is not a valid c string, or llama.cpp fails to load
396 /// the sequence state file.
397 ///
398 /// # Returns
399 ///
400 /// A tuple of `(tokens, bytes_read)` on success.
401 pub fn state_seq_load_file(
402 &mut self,
403 filepath: impl AsRef<Path>,
404 dest_seq_id: i32,
405 max_tokens: usize,
406 ) -> Result<(Vec<LlamaToken>, usize), LoadSeqStateError> {
407 let path = filepath.as_ref();
408 let path = path
409 .to_str()
410 .ok_or(LoadSeqStateError::PathToStrError(path.to_path_buf()))?;
411
412 let cstr = CString::new(path)?;
413 let mut tokens: Vec<LlamaToken> = Vec::with_capacity(max_tokens);
414 let mut n_out = 0;
415
416 // SAFETY: cast is valid as LlamaToken is repr(transparent)
417 let tokens_out = tokens.as_mut_ptr().cast::<llama_cpp_sys_2::llama_token>();
418
419 let bytes_read = unsafe {
420 llama_cpp_sys_2::llama_state_seq_load_file(
421 self.context.as_mut_ptr(),
422 cstr.as_ptr(),
423 dest_seq_id,
424 tokens_out,
425 max_tokens,
426 &mut n_out,
427 )
428 };
429
430 if bytes_read == 0 {
431 return Err(LoadSeqStateError::FailedToLoad);
432 }
433
434 if n_out > max_tokens {
435 return Err(LoadSeqStateError::InsufficientMaxLength { n_out, max_tokens });
436 }
437
438 // SAFETY: we checked that n_out <= max_tokens and llama.cpp promises that n_out tokens will be written
439 unsafe {
440 tokens.set_len(n_out);
441 }
442
443 Ok((tokens, bytes_read))
444 }
445
446 /// Returns the maximum size in bytes of the state (rng, logits, embedding
447 /// and `kv_cache`) - will often be smaller after compacting tokens
448 #[must_use]
449 pub fn get_state_size(&self) -> usize {
450 unsafe { llama_cpp_sys_2::llama_get_state_size(self.synchronizable_ptr()) }
451 }
452
453 /// Copies the state to the specified destination address.
454 ///
455 /// Returns the number of bytes copied
456 ///
457 /// # Safety
458 ///
459 /// Destination needs to have allocated enough memory.
460 pub unsafe fn copy_state_data(&self, dest: *mut u8) -> usize {
461 unsafe { llama_cpp_sys_2::llama_copy_state_data(self.synchronizable_ptr(), dest) }
462 }
463
464 /// Set the state reading from the specified address
465 /// Returns the number of bytes read
466 ///
467 /// # Safety
468 ///
469 /// help wanted: not entirely sure what the safety requirements are here.
470 pub unsafe fn set_state_data(&mut self, src: &[u8]) -> usize {
471 unsafe { llama_cpp_sys_2::llama_set_state_data(self.context.as_mut_ptr(), src.as_ptr()) }
472 }
473
474 /// Get the size of the state for a single sequence with optional flags.
475 ///
476 /// This is the extended version that supports flags for partial state operations.
477 ///
478 /// # Parameters
479 ///
480 /// * `seq_id` - The sequence ID to get the state size for.
481 /// * `flags` - Optional flags (e.g., [`LlamaStateSeqFlags::PARTIAL_ONLY`]).
482 ///
483 /// # Returns
484 ///
485 /// The size in bytes needed to store the sequence state.
486 #[must_use]
487 pub fn state_seq_get_size_ext(&self, seq_id: i32, flags: LlamaStateSeqFlags) -> usize {
488 unsafe {
489 llama_cpp_sys_2::llama_state_seq_get_size_ext(
490 self.synchronizable_ptr(),
491 seq_id,
492 flags.0,
493 )
494 }
495 }
496
497 /// Copy the state of a single sequence into the specified buffer with optional flags.
498 ///
499 /// This is the extended version that supports flags for partial state operations.
500 ///
501 /// # Parameters
502 ///
503 /// * `dest` - Destination buffer to copy state into.
504 /// * `seq_id` - The sequence ID to get the state for.
505 /// * `flags` - Optional flags (e.g., [`LlamaStateSeqFlags::PARTIAL_ONLY`]).
506 ///
507 /// # Safety
508 ///
509 /// Destination needs to have allocated enough memory.
510 ///
511 /// # Returns
512 ///
513 /// The number of bytes copied.
514 pub unsafe fn state_seq_get_data_ext(
515 &self,
516 dest: *mut u8,
517 seq_id: i32,
518 flags: LlamaStateSeqFlags,
519 ) -> usize {
520 unsafe {
521 llama_cpp_sys_2::llama_state_seq_get_data_ext(
522 self.synchronizable_ptr(),
523 dest,
524 usize::MAX,
525 seq_id,
526 flags.0,
527 )
528 }
529 }
530
531 /// Set the state for a single sequence from the specified buffer with optional flags.
532 ///
533 /// This is the extended version that supports flags for partial state operations.
534 /// Useful for restoring only the recurrent/partial state without affecting the KV cache.
535 ///
536 /// # Parameters
537 ///
538 /// * `src` - Source buffer containing the state data.
539 /// * `dest_seq_id` - The destination sequence ID to load the state into.
540 /// * `flags` - Optional flags (e.g., [`LlamaStateSeqFlags::PARTIAL_ONLY`]).
541 ///
542 /// # Safety
543 ///
544 /// The source buffer must contain valid state data.
545 ///
546 /// # Returns
547 ///
548 /// Positive on success, zero on failure.
549 pub unsafe fn state_seq_set_data_ext(
550 &mut self,
551 src: &[u8],
552 dest_seq_id: i32,
553 flags: LlamaStateSeqFlags,
554 ) -> bool {
555 unsafe {
556 llama_cpp_sys_2::llama_state_seq_set_data_ext(
557 self.context.as_mut_ptr(),
558 src.as_ptr(),
559 src.len(),
560 dest_seq_id,
561 flags.0,
562 ) > 0
563 }
564 }
565
566 /// Serialize sequence `seq_id`'s state into an opaque [`SeqState`].
567 ///
568 /// Enables save/restore of context state at arbitrary points in a
569 /// sequence. This is particularly useful on architectures where
570 /// `clear_kv_cache_seq` cannot roll back partial state (Mamba, RWKV,
571 /// Gated Delta Networks, and other recurrent / hybrid-recurrent
572 /// models): pair with [`LlamaStateSeqFlags::PARTIAL_ONLY`] to save
573 /// just the running recurrent and SWA state, then restore it via
574 /// [`Self::state_seq_set`] to effectively "rewind" the sequence.
575 ///
576 /// The returned [`SeqState`] is opaque — its bytes cannot be
577 /// inspected or forged from safe code, so [`Self::state_seq_set`]
578 /// only ever sees data produced by this method.
579 ///
580 /// Wraps `llama_state_seq_get_data_ext`.
581 ///
582 /// # Errors
583 ///
584 /// Returns [`StateSeqError::SizeMismatch`] if llama.cpp writes a
585 /// different number of bytes than the reported state size.
586 pub fn state_seq_get(
587 &self,
588 seq_id: i32,
589 flags: LlamaStateSeqFlags,
590 ) -> Result<SeqState, crate::StateSeqError> {
591 let size = unsafe {
592 llama_cpp_sys_2::llama_state_seq_get_size_ext(
593 self.synchronizable_ptr(),
594 seq_id,
595 flags.0,
596 )
597 };
598 let mut bytes = vec![0u8; size];
599 let n = unsafe {
600 llama_cpp_sys_2::llama_state_seq_get_data_ext(
601 self.synchronizable_ptr(),
602 bytes.as_mut_ptr(),
603 size,
604 seq_id,
605 flags.0,
606 )
607 };
608 if n != size {
609 return Err(crate::StateSeqError::SizeMismatch {
610 expected: size,
611 actual: n,
612 });
613 }
614 Ok(SeqState { bytes, flags })
615 }
616
617 /// Restore sequence state previously captured by [`Self::state_seq_get`]
618 /// into `seq_id`.
619 ///
620 /// Cross-sequence restore (`seq_id` different from the sequence the
621 /// state was captured from) is supported — llama.cpp treats the
622 /// destination sequence id independently of the source.
623 ///
624 /// Wraps `llama_state_seq_set_data_ext`.
625 ///
626 /// # Errors
627 ///
628 /// Returns [`StateSeqError::SizeMismatch`] if llama.cpp reads a
629 /// different number of bytes than the state buffer contains — this
630 /// covers shape mismatches (different `n_ctx`, `n_layer`, quantization,
631 /// etc.) that llama.cpp's own deserializer detects and aborts on.
632 pub fn state_seq_set(
633 &mut self,
634 state: &SeqState,
635 seq_id: i32,
636 ) -> Result<(), crate::StateSeqError> {
637 let n = unsafe {
638 llama_cpp_sys_2::llama_state_seq_set_data_ext(
639 self.context.as_mut_ptr(),
640 state.bytes.as_ptr(),
641 state.bytes.len(),
642 seq_id,
643 state.flags.0,
644 )
645 };
646 if n != state.bytes.len() {
647 return Err(crate::StateSeqError::SizeMismatch {
648 expected: state.bytes.len(),
649 actual: n,
650 });
651 }
652 Ok(())
653 }
654}
655
656/// Opaque, immutable snapshot of a single sequence's state.
657///
658/// Produced by [`LlamaContext::state_seq_get`] and consumed by
659/// [`LlamaContext::state_seq_set`]. Bytes cannot be constructed, forged,
660/// or mutated from safe code — the only way to obtain a `SeqState` is
661/// via a get call, which guarantees the payload came from llama.cpp
662/// itself. Combined with llama.cpp's own shape validation on the
663/// deserialize path, this closes the "arbitrary bytes into C parser"
664/// unsoundness that motivates the unsafe raw setters.
665///
666/// The state carries the flag set it was captured with; restore uses
667/// those same flags automatically so the byte layout always matches.
668#[derive(Clone)]
669pub struct SeqState {
670 bytes: Vec<u8>,
671 flags: LlamaStateSeqFlags,
672}
673
674impl SeqState {
675 /// The flag set that was used to capture this state.
676 #[must_use]
677 pub fn flags(&self) -> LlamaStateSeqFlags {
678 self.flags
679 }
680
681 /// Size in bytes of the serialized state.
682 #[must_use]
683 pub fn byte_len(&self) -> usize {
684 self.bytes.len()
685 }
686}
687
688impl Debug for SeqState {
689 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
690 f.debug_struct("SeqState")
691 .field("byte_len", &self.bytes.len())
692 .field("flags", &self.flags)
693 .finish()
694 }
695}