ik_llama_cpp_2/context/
params.rs

1//! Context parameters ([`LlamaContextParams`]) over ik's `llama_context_params`.
2
3use std::num::NonZeroU32;
4
5use ik_llama_cpp_sys as sys;
6
7/// The kind of context to create.
8///
9/// Mirrors `llama-cpp-2`'s `context::params::LlamaContextType`. ik has no
10/// `llama_context_type` field — instead it toggles MTP via a `bool mtp`, so
11/// this shim maps [`LlamaContextType::Mtp`] onto that flag.
12#[derive(Copy, Clone, Debug, PartialEq, Eq)]
13pub enum LlamaContextType {
14    /// Default decoder context.
15    Default,
16    /// Multi-token-prediction (NextN) draft context.
17    Mtp,
18}
19
20/// Parameters controlling a [`crate::LlamaContext`].
21///
22/// Starts from `llama_context_default_params()`. Note ik keeps a `seed` field in
23/// the context params (stock removed it) and uses a `bool flash_attn` plus the
24/// MTP fields `mtp` / `mtp_op_type`. There is no `ctx_type` in ik.
25#[derive(Debug, Clone)]
26pub struct LlamaContextParams {
27    pub(crate) params: sys::llama_context_params,
28}
29
30impl Default for LlamaContextParams {
31    fn default() -> Self {
32        Self {
33            params: unsafe { sys::llama_context_default_params() },
34        }
35    }
36}
37
38impl LlamaContextParams {
39    /// Context size (tokens); `None` means "take it from the model" (0).
40    ///
41    /// Takes `Option<NonZeroU32>` to match `llama-cpp-2`.
42    #[must_use]
43    pub fn with_n_ctx(mut self, n_ctx: Option<NonZeroU32>) -> Self {
44        self.params.n_ctx = n_ctx.map_or(0, NonZeroU32::get);
45        self
46    }
47
48    /// Logical batch size.
49    #[must_use]
50    pub fn with_n_batch(mut self, n_batch: u32) -> Self {
51        self.params.n_batch = n_batch;
52        self
53    }
54
55    /// Physical (micro) batch size.
56    #[must_use]
57    pub fn with_n_ubatch(mut self, n_ubatch: u32) -> Self {
58        self.params.n_ubatch = n_ubatch;
59        self
60    }
61
62    /// Maximum number of sequences (distinct recurrent states).
63    ///
64    /// `llama-cpp-2`'s NextN fork exposes this as `n_rs_seq`; ik's equivalent is
65    /// `n_seq_max`, which this sets.
66    #[must_use]
67    pub fn with_n_rs_seq(mut self, n_rs_seq: u32) -> Self {
68        self.params.n_seq_max = n_rs_seq;
69        self
70    }
71
72    /// Select the context kind. [`LlamaContextType::Mtp`] enables ik's MTP path
73    /// (equivalent to [`Self::with_mtp(true)`](Self::with_mtp)).
74    #[must_use]
75    pub fn with_context_type(mut self, context_type: LlamaContextType) -> Self {
76        self.params.mtp = matches!(context_type, LlamaContextType::Mtp);
77        self
78    }
79
80    /// RNG seed (ik retains this in the context params).
81    #[must_use]
82    pub fn with_seed(mut self, seed: u32) -> Self {
83        self.params.seed = seed;
84        self
85    }
86
87    /// Threads used for generation.
88    #[must_use]
89    pub fn with_n_threads(mut self, n_threads: u32) -> Self {
90        self.params.n_threads = n_threads;
91        self
92    }
93
94    /// Threads used for batch/prompt processing.
95    #[must_use]
96    pub fn with_n_threads_batch(mut self, n_threads_batch: u32) -> Self {
97        self.params.n_threads_batch = n_threads_batch;
98        self
99    }
100
101    /// Enable flash attention.
102    #[must_use]
103    pub fn with_flash_attn(mut self, flash_attn: bool) -> Self {
104        self.params.flash_attn = flash_attn;
105        self
106    }
107
108    /// Activate the MTP path (requires a model loaded with `.with_mtp(true)`).
109    #[must_use]
110    pub fn with_mtp(mut self, mtp: bool) -> Self {
111        self.params.mtp = mtp;
112        self
113    }
114
115    /// Produce embeddings on decode (sets `params.embeddings`).
116    ///
117    /// Required for embedding and reranker models; usually paired with
118    /// [`Self::with_pooling_type`].
119    #[must_use]
120    pub fn with_embeddings(mut self, embeddings: bool) -> Self {
121        self.params.embeddings = embeddings;
122        self
123    }
124
125    /// Set the pooling strategy for embeddings (sets `params.pooling_type`).
126    ///
127    /// Takes the raw `sys::llama_pooling_type` (e.g. `LLAMA_POOLING_TYPE_MEAN`,
128    /// `LLAMA_POOLING_TYPE_CLS`, or `LLAMA_POOLING_TYPE_LAST`) to avoid
129    /// introducing a new enum.
130    #[must_use]
131    pub fn with_pooling_type(mut self, pooling_type: sys::llama_pooling_type) -> Self {
132        self.params.pooling_type = pooling_type;
133        self
134    }
135
136    /// Access the raw params (advanced/escape hatch).
137    #[must_use]
138    pub fn as_raw(&self) -> &sys::llama_context_params {
139        &self.params
140    }
141}
ik_llama_cpp_2/context/params.rs

ik_llama_cpp_2/context/
params.rs
