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.context.as_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_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.context.as_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_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.context.as_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_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.context.as_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.context.as_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_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(self.context.as_ptr(), seq_id, flags.0)
490        }
491    }
492
493    /// Copy the state of a single sequence into the specified buffer with optional flags.
494    ///
495    /// This is the extended version that supports flags for partial state operations.
496    ///
497    /// # Parameters
498    ///
499    /// * `dest` - Destination buffer to copy state into.
500    /// * `seq_id` - The sequence ID to get the state for.
501    /// * `flags` - Optional flags (e.g., [`LlamaStateSeqFlags::PARTIAL_ONLY`]).
502    ///
503    /// # Safety
504    ///
505    /// Destination needs to have allocated enough memory.
506    ///
507    /// # Returns
508    ///
509    /// The number of bytes copied.
510    pub unsafe fn state_seq_get_data_ext(
511        &self,
512        dest: *mut u8,
513        seq_id: i32,
514        flags: LlamaStateSeqFlags,
515    ) -> usize {
516        unsafe {
517            llama_cpp_sys_2::llama_state_seq_get_data_ext(
518                self.context.as_ptr(),
519                dest,
520                usize::MAX,
521                seq_id,
522                flags.0,
523            )
524        }
525    }
526
527    /// Set the state for a single sequence from the specified buffer with optional flags.
528    ///
529    /// This is the extended version that supports flags for partial state operations.
530    /// Useful for restoring only the recurrent/partial state without affecting the KV cache.
531    ///
532    /// # Parameters
533    ///
534    /// * `src` - Source buffer containing the state data.
535    /// * `dest_seq_id` - The destination sequence ID to load the state into.
536    /// * `flags` - Optional flags (e.g., [`LlamaStateSeqFlags::PARTIAL_ONLY`]).
537    ///
538    /// # Safety
539    ///
540    /// The source buffer must contain valid state data.
541    ///
542    /// # Returns
543    ///
544    /// Positive on success, zero on failure.
545    pub unsafe fn state_seq_set_data_ext(
546        &mut self,
547        src: &[u8],
548        dest_seq_id: i32,
549        flags: LlamaStateSeqFlags,
550    ) -> bool {
551        unsafe {
552            llama_cpp_sys_2::llama_state_seq_set_data_ext(
553                self.context.as_ptr(),
554                src.as_ptr(),
555                src.len(),
556                dest_seq_id,
557                flags.0,
558            ) > 0
559        }
560    }
561
562    /// Serialize sequence `seq_id`'s state into an opaque [`SeqState`].
563    ///
564    /// Enables save/restore of context state at arbitrary points in a
565    /// sequence. This is particularly useful on architectures where
566    /// `clear_kv_cache_seq` cannot roll back partial state (Mamba, RWKV,
567    /// Gated Delta Networks, and other recurrent / hybrid-recurrent
568    /// models): pair with [`LlamaStateSeqFlags::PARTIAL_ONLY`] to save
569    /// just the running recurrent and SWA state, then restore it via
570    /// [`Self::state_seq_set`] to effectively "rewind" the sequence.
571    ///
572    /// The returned [`SeqState`] is opaque — its bytes cannot be
573    /// inspected or forged from safe code, so [`Self::state_seq_set`]
574    /// only ever sees data produced by this method.
575    ///
576    /// Wraps `llama_state_seq_get_data_ext`.
577    ///
578    /// # Errors
579    ///
580    /// Returns [`StateSeqError::SizeMismatch`] if llama.cpp writes a
581    /// different number of bytes than the reported state size.
582    pub fn state_seq_get(
583        &self,
584        seq_id: i32,
585        flags: LlamaStateSeqFlags,
586    ) -> Result<SeqState, crate::StateSeqError> {
587        let size = unsafe {
588            llama_cpp_sys_2::llama_state_seq_get_size_ext(self.context.as_ptr(), seq_id, flags.0)
589        };
590        let mut bytes = vec![0u8; size];
591        let n = unsafe {
592            llama_cpp_sys_2::llama_state_seq_get_data_ext(
593                self.context.as_ptr(),
594                bytes.as_mut_ptr(),
595                size,
596                seq_id,
597                flags.0,
598            )
599        };
600        if n != size {
601            return Err(crate::StateSeqError::SizeMismatch {
602                expected: size,
603                actual: n,
604            });
605        }
606        Ok(SeqState { bytes, flags })
607    }
608
609    /// Restore sequence state previously captured by [`Self::state_seq_get`]
610    /// into `seq_id`.
611    ///
612    /// Cross-sequence restore (`seq_id` different from the sequence the
613    /// state was captured from) is supported — llama.cpp treats the
614    /// destination sequence id independently of the source.
615    ///
616    /// Wraps `llama_state_seq_set_data_ext`.
617    ///
618    /// # Errors
619    ///
620    /// Returns [`StateSeqError::SizeMismatch`] if llama.cpp reads a
621    /// different number of bytes than the state buffer contains — this
622    /// covers shape mismatches (different `n_ctx`, `n_layer`, quantization,
623    /// etc.) that llama.cpp's own deserializer detects and aborts on.
624    pub fn state_seq_set(
625        &mut self,
626        state: &SeqState,
627        seq_id: i32,
628    ) -> Result<(), crate::StateSeqError> {
629        let n = unsafe {
630            llama_cpp_sys_2::llama_state_seq_set_data_ext(
631                self.context.as_ptr(),
632                state.bytes.as_ptr(),
633                state.bytes.len(),
634                seq_id,
635                state.flags.0,
636            )
637        };
638        if n != state.bytes.len() {
639            return Err(crate::StateSeqError::SizeMismatch {
640                expected: state.bytes.len(),
641                actual: n,
642            });
643        }
644        Ok(())
645    }
646}
647
648/// Opaque, immutable snapshot of a single sequence's state.
649///
650/// Produced by [`LlamaContext::state_seq_get`] and consumed by
651/// [`LlamaContext::state_seq_set`]. Bytes cannot be constructed, forged,
652/// or mutated from safe code — the only way to obtain a `SeqState` is
653/// via a get call, which guarantees the payload came from llama.cpp
654/// itself. Combined with llama.cpp's own shape validation on the
655/// deserialize path, this closes the "arbitrary bytes into C parser"
656/// unsoundness that motivates the unsafe raw setters.
657///
658/// The state carries the flag set it was captured with; restore uses
659/// those same flags automatically so the byte layout always matches.
660#[derive(Clone)]
661pub struct SeqState {
662    bytes: Vec<u8>,
663    flags: LlamaStateSeqFlags,
664}
665
666impl SeqState {
667    /// The flag set that was used to capture this state.
668    #[must_use]
669    pub fn flags(&self) -> LlamaStateSeqFlags {
670        self.flags
671    }
672
673    /// Size in bytes of the serialized state.
674    #[must_use]
675    pub fn byte_len(&self) -> usize {
676        self.bytes.len()
677    }
678}
679
680impl Debug for SeqState {
681    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
682        f.debug_struct("SeqState")
683            .field("byte_len", &self.bytes.len())
684            .field("flags", &self.flags)
685            .finish()
686    }
687}
llama_cpp_2/context/session.rs

llama_cpp_2/context/
session.rs
