ik_llama_cpp_2/
context.rs

1//! Safe wrapper around `llama_context` ([`LlamaContext`]).
2
3pub mod embeddings;
4pub mod kv_cache;
5pub mod params;
6pub mod session;
7
8use std::marker::PhantomData;
9use std::ptr::NonNull;
10
11use ik_llama_cpp_sys as sys;
12
13use crate::llama_batch::LlamaBatch;
14use crate::model::LlamaModel;
15use crate::speculative::MtpOpType;
16use crate::LlamaError;
17
18pub use params::LlamaContextParams;
19
20/// An inference context bound to a [`LlamaModel`].
21///
22/// The lifetime `'a` ties the context to its model (the model must outlive it).
23#[derive(Debug)]
24pub struct LlamaContext<'a> {
25    pub(crate) context: NonNull<sys::llama_context>,
26    n_vocab: i32,
27    /// The raw params the context was built with (needed to derive the MTP
28    /// companion context in the speculative glue; only read under `common`).
29    #[cfg_attr(not(feature = "common"), allow(dead_code))]
30    pub(crate) raw_params: sys::llama_context_params,
31    _model: PhantomData<&'a LlamaModel>,
32}
33
34impl<'a> LlamaContext<'a> {
35    /// Create a context for `model`.
36    pub fn new(model: &'a LlamaModel, params: &LlamaContextParams) -> Result<Self, LlamaError> {
37        // SAFETY: valid model ptr + initialized params.
38        let raw = unsafe { sys::llama_init_from_model(model.model.as_ptr(), params.params) };
39        NonNull::new(raw)
40            .map(|context| Self {
41                context,
42                n_vocab: model.n_vocab(),
43                raw_params: params.params,
44                _model: PhantomData,
45            })
46            .ok_or(LlamaError::ContextCreation)
47    }
48
49    /// Run a decode over `batch`. Errors on a non-zero status from `llama_decode`.
50    pub fn decode(&mut self, batch: &mut LlamaBatch) -> Result<(), LlamaError> {
51        // SAFETY: valid ctx + batch owned by the caller.
52        let ret = unsafe { sys::llama_decode(self.context.as_ptr(), batch.as_raw()) };
53        if ret != 0 {
54            return Err(LlamaError::Decode(ret));
55        }
56        Ok(())
57    }
58
59    /// Logits for batch index `i` of the last decode (length = `n_vocab`).
60    ///
61    /// `i` is the **batch position** (0..n_tokens), not a logits-rank; that entry
62    /// must have been added with `logits = true` (else ik returns null → `NoLogits`).
63    pub fn get_logits_ith(&self, i: i32) -> Result<&[f32], LlamaError> {
64        // SAFETY: ptr valid until the next decode; slice length is n_vocab.
65        let ptr = unsafe { sys::llama_get_logits_ith(self.context.as_ptr(), i) };
66        if ptr.is_null() {
67            return Err(LlamaError::NoLogits(i));
68        }
69        Ok(unsafe { std::slice::from_raw_parts(ptr, self.n_vocab as usize) })
70    }
71
72    /// Candidate token data for batch index `i` (id + logit, `p = 0`).
73    ///
74    /// Returns an empty iterator if logits are unavailable at `i` (e.g. that
75    /// position was not decoded with `logits = true`).
76    pub fn candidates_ith(
77        &self,
78        i: i32,
79    ) -> impl Iterator<Item = crate::token::data::LlamaTokenData> + '_ {
80        let logits = self.get_logits_ith(i).unwrap_or(&[]);
81        (0_i32..).zip(logits).map(|(idx, &logit)| {
82            crate::token::data::LlamaTokenData::new(crate::token::LlamaToken::new(idx), logit, 0.0)
83        })
84    }
85
86    /// Build a [`LlamaTokenDataArray`](crate::token::data_array::LlamaTokenDataArray)
87    /// from the logits at batch index `i` (matches `llama-cpp-2`). Empty if
88    /// logits are unavailable at `i`.
89    #[must_use]
90    pub fn token_data_array_ith(&self, i: i32) -> crate::token::data_array::LlamaTokenDataArray {
91        crate::token::data_array::LlamaTokenDataArray::from_iter(self.candidates_ith(i), false)
92    }
93
94    /// The context size (tokens).
95    #[must_use]
96    pub fn n_ctx(&self) -> u32 {
97        unsafe { sys::llama_n_ctx(self.context.as_ptr()) }
98    }
99
100    /// Set the MTP operation mode for subsequent decodes.
101    ///
102    /// Only meaningful for a context created with `.with_mtp(true)` on a model
103    /// loaded with `.with_mtp(true)` (a NextN model). Drives the low-level MTP
104    /// state machine (`llama_set_mtp_op_type`).
105    pub fn set_mtp_op_type(&mut self, op: MtpOpType) {
106        unsafe { sys::llama_set_mtp_op_type(self.context.as_ptr(), op.to_raw()) };
107    }
108
109    /// Vocabulary size (cached from the model at creation).
110    #[must_use]
111    pub fn n_vocab(&self) -> i32 {
112        self.n_vocab
113    }
114
115    /// Raw context pointer (used by the sampling module; escape hatch).
116    #[must_use]
117    pub(crate) fn as_ptr(&self) -> *mut sys::llama_context {
118        self.context.as_ptr()
119    }
120}
121
122impl Drop for LlamaContext<'_> {
123    fn drop(&mut self) {
124        unsafe { sys::llama_free(self.context.as_ptr()) };
125    }
126}
ik_llama_cpp_2/context.rs

ik_llama_cpp_2/
context.rs
