llama_cpp_2/model/
params.rs

1//! A safe wrapper around `llama_model_params`.
2
3use crate::context::params::LlamaContextParams;
4use crate::model::params::kv_overrides::KvOverrides;
5use crate::LlamaCppError;
6use std::ffi::{c_char, c_void, CStr};
7use std::fmt::{Debug, Formatter};
8use std::pin::Pin;
9use std::ptr::null;
10
11pub mod kv_overrides;
12
13/// Result of [`LlamaModelParams::fit_params`], containing the fitted context size.
14#[cfg(feature = "common")]
15#[derive(Debug, Clone)]
16pub struct FitResult {
17    /// The context size after fitting (may have been reduced from the requested value).
18    pub n_ctx: u32,
19}
20
21/// Error returned by [`LlamaModelParams::fit_params`].
22#[cfg(feature = "common")]
23#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
24pub enum FitError {
25    /// Could not find allocations that are projected to fit available memory.
26    #[error("could not find allocations that fit available memory")]
27    Failure,
28    /// A hard error occurred during fitting (e.g. model not found at the specified path).
29    #[error("hard error during parameter fitting")]
30    Error,
31}
32
33#[allow(clippy::cast_possible_wrap)]
34#[allow(clippy::cast_possible_truncation)]
35const LLAMA_SPLIT_MODE_NONE: i8 = llama_cpp_sys_2::LLAMA_SPLIT_MODE_NONE as i8;
36#[allow(clippy::cast_possible_wrap)]
37#[allow(clippy::cast_possible_truncation)]
38const LLAMA_SPLIT_MODE_LAYER: i8 = llama_cpp_sys_2::LLAMA_SPLIT_MODE_LAYER as i8;
39#[allow(clippy::cast_possible_wrap)]
40#[allow(clippy::cast_possible_truncation)]
41const LLAMA_SPLIT_MODE_ROW: i8 = llama_cpp_sys_2::LLAMA_SPLIT_MODE_ROW as i8;
42#[allow(clippy::cast_possible_wrap)]
43#[allow(clippy::cast_possible_truncation)]
44const LLAMA_SPLIT_MODE_TENSOR: i8 = llama_cpp_sys_2::LLAMA_SPLIT_MODE_TENSOR as i8;
45
46/// A rusty wrapper around `llama_split_mode`.
47#[repr(i8)]
48#[derive(Copy, Clone, Debug, PartialEq, Eq)]
49pub enum LlamaSplitMode {
50    /// Single GPU
51    None = LLAMA_SPLIT_MODE_NONE,
52    /// Split layers and KV across GPUs
53    Layer = LLAMA_SPLIT_MODE_LAYER,
54    /// Split layers and KV across GPUs, use tensor parallelism if supported
55    Row = LLAMA_SPLIT_MODE_ROW,
56    /// Experimental tensor parallelism across GPUs
57    Tensor = LLAMA_SPLIT_MODE_TENSOR,
58}
59
60/// An error that occurs when unknown split mode is encountered.
61#[derive(Debug, Clone, Copy, PartialEq, Eq)]
62pub struct LlamaSplitModeParseError(pub i32);
63
64/// Create a `LlamaSplitMode` from a `i32`.
65///
66/// # Errors
67/// Returns `LlamaSplitModeParseError` if the value does not correspond to a valid `LlamaSplitMode`.
68impl TryFrom<i32> for LlamaSplitMode {
69    type Error = LlamaSplitModeParseError;
70
71    fn try_from(value: i32) -> Result<Self, Self::Error> {
72        let i8_value = value
73            .try_into()
74            .map_err(|_| LlamaSplitModeParseError(value))?;
75        match i8_value {
76            LLAMA_SPLIT_MODE_NONE => Ok(Self::None),
77            LLAMA_SPLIT_MODE_LAYER => Ok(Self::Layer),
78            LLAMA_SPLIT_MODE_ROW => Ok(Self::Row),
79            LLAMA_SPLIT_MODE_TENSOR => Ok(Self::Tensor),
80            _ => Err(LlamaSplitModeParseError(value)),
81        }
82    }
83}
84
85/// Create a `LlamaSplitMode` from a `u32`.
86///
87/// # Errors
88/// Returns `LlamaSplitModeParseError` if the value does not correspond to a valid `LlamaSplitMode`.
89impl TryFrom<u32> for LlamaSplitMode {
90    type Error = LlamaSplitModeParseError;
91
92    fn try_from(value: u32) -> Result<Self, Self::Error> {
93        let i8_value = value
94            .try_into()
95            .map_err(|_| LlamaSplitModeParseError(value.try_into().unwrap_or(i32::MAX)))?;
96        match i8_value {
97            LLAMA_SPLIT_MODE_NONE => Ok(Self::None),
98            LLAMA_SPLIT_MODE_LAYER => Ok(Self::Layer),
99            LLAMA_SPLIT_MODE_ROW => Ok(Self::Row),
100            LLAMA_SPLIT_MODE_TENSOR => Ok(Self::Tensor),
101            _ => Err(LlamaSplitModeParseError(
102                value.try_into().unwrap_or(i32::MAX),
103            )),
104        }
105    }
106}
107
108/// Create a `i32` from a `LlamaSplitMode`.
109impl From<LlamaSplitMode> for i32 {
110    fn from(value: LlamaSplitMode) -> Self {
111        match value {
112            LlamaSplitMode::None => LLAMA_SPLIT_MODE_NONE.into(),
113            LlamaSplitMode::Layer => LLAMA_SPLIT_MODE_LAYER.into(),
114            LlamaSplitMode::Row => LLAMA_SPLIT_MODE_ROW.into(),
115            LlamaSplitMode::Tensor => LLAMA_SPLIT_MODE_TENSOR.into(),
116        }
117    }
118}
119
120/// Create a `u32` from a `LlamaSplitMode`.
121impl From<LlamaSplitMode> for u32 {
122    fn from(value: LlamaSplitMode) -> Self {
123        match value {
124            LlamaSplitMode::None => LLAMA_SPLIT_MODE_NONE as u32,
125            LlamaSplitMode::Layer => LLAMA_SPLIT_MODE_LAYER as u32,
126            LlamaSplitMode::Row => LLAMA_SPLIT_MODE_ROW as u32,
127            LlamaSplitMode::Tensor => LLAMA_SPLIT_MODE_TENSOR as u32,
128        }
129    }
130}
131
132/// The default split mode is `Layer` in llama.cpp.
133impl Default for LlamaSplitMode {
134    fn default() -> Self {
135        LlamaSplitMode::Layer
136    }
137}
138
139/// The maximum number of devices supported.
140///
141/// The real maximum number of devices is the lesser one of this value and the value returned by
142/// `llama_cpp_2::max_devices()`.
143pub const LLAMA_CPP_MAX_DEVICES: usize = 16;
144
145/// Combines the two independent `use_mmap`/`use_mlock` flags this crate's public
146/// API exposes into the single `load_mode` enum llama.cpp now stores them as.
147fn load_mode_from_flags(use_mmap: bool, use_mlock: bool) -> llama_cpp_sys_2::llama_load_mode {
148    match (use_mmap, use_mlock) {
149        (false, false) => llama_cpp_sys_2::LLAMA_LOAD_MODE_NONE,
150        (true, false) => llama_cpp_sys_2::LLAMA_LOAD_MODE_MMAP,
151        (false, true) => llama_cpp_sys_2::LLAMA_LOAD_MODE_MLOCK,
152        (true, true) => llama_cpp_sys_2::LLAMA_LOAD_MODE_MMAP_MLOCK,
153    }
154}
155
156/// A safe wrapper around `llama_model_params`.
157#[allow(clippy::module_name_repetitions)]
158pub struct LlamaModelParams {
159    pub(crate) params: llama_cpp_sys_2::llama_model_params,
160    kv_overrides: Vec<llama_cpp_sys_2::llama_model_kv_override>,
161    buft_overrides: Vec<llama_cpp_sys_2::llama_model_tensor_buft_override>,
162    devices: Pin<Box<[llama_cpp_sys_2::ggml_backend_dev_t; LLAMA_CPP_MAX_DEVICES]>>,
163    tensor_split: Vec<f32>,
164    progress_callback: Option<Box<dyn FnMut(f32) -> bool>>,
165}
166
167impl Debug for LlamaModelParams {
168    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
169        f.debug_struct("LlamaModelParams")
170            .field("n_gpu_layers", &self.params.n_gpu_layers)
171            .field("main_gpu", &self.params.main_gpu)
172            .field("vocab_only", &self.params.vocab_only)
173            .field("use_mmap", &self.use_mmap())
174            .field("use_mlock", &self.use_mlock())
175            .field("split_mode", &self.split_mode())
176            .field("devices", &self.devices)
177            .field("kv_overrides", &"vec of kv_overrides")
178            .finish()
179    }
180}
181
182impl LlamaModelParams {
183    /// See [`KvOverrides`]
184    ///
185    /// # Examples
186    ///
187    /// ```rust
188    /// # use llama_cpp_2::model::params::LlamaModelParams;
189    /// let params = Box::pin(LlamaModelParams::default());
190    /// let kv_overrides = params.kv_overrides();
191    /// let count = kv_overrides.into_iter().count();
192    /// assert_eq!(count, 0);
193    /// ```
194    #[must_use]
195    pub fn kv_overrides<'a>(&'a self) -> KvOverrides<'a> {
196        KvOverrides::new(self)
197    }
198
199    /// Appends a key-value override to the model parameters. It must be pinned as this creates a self-referential struct.
200    ///
201    /// # Examples
202    ///
203    /// ```rust
204    /// # use std::ffi::{CStr, CString};
205    /// use std::pin::pin;
206    /// # use llama_cpp_2::model::params::LlamaModelParams;
207    /// # use llama_cpp_2::model::params::kv_overrides::ParamOverrideValue;
208    /// let mut params = pin!(LlamaModelParams::default());
209    /// let key = CString::new("key").expect("CString::new failed");
210    /// params.as_mut().append_kv_override(&key, ParamOverrideValue::Int(50));
211    ///
212    /// let kv_overrides = params.kv_overrides().into_iter().collect::<Vec<_>>();
213    /// assert_eq!(kv_overrides.len(), 1);
214    ///
215    /// let (k, v) = &kv_overrides[0];
216    /// assert_eq!(v, &ParamOverrideValue::Int(50));
217    ///
218    /// assert_eq!(k.to_bytes(), b"key", "expected key to be 'key', was {:?}", k);
219    /// ```
220    #[allow(clippy::missing_panics_doc)] // panics are just to enforce internal invariants, not user errors
221    pub fn append_kv_override(
222        mut self: Pin<&mut Self>,
223        key: &CStr,
224        value: kv_overrides::ParamOverrideValue,
225    ) {
226        let kv_override = self
227            .kv_overrides
228            .last_mut()
229            .expect("kv_overrides did not have a next allocated");
230
231        assert_eq!(kv_override.key[0], 0, "last kv_override was not empty");
232
233        // There should be some way to do this without iterating over everything.
234        for (i, &c) in key.to_bytes_with_nul().iter().enumerate() {
235            kv_override.key[i] = c_char::try_from(c).expect("invalid character in key");
236        }
237
238        kv_override.tag = value.tag();
239        kv_override.__bindgen_anon_1 = value.value();
240
241        // set to null pointer for panic safety (as push may move the vector, invalidating the pointer)
242        self.params.kv_overrides = null();
243
244        // push the next one to ensure we maintain the iterator invariant of ending with a 0
245        self.kv_overrides
246            .push(llama_cpp_sys_2::llama_model_kv_override {
247                key: [0; 128],
248                tag: 0,
249                __bindgen_anon_1: llama_cpp_sys_2::llama_model_kv_override__bindgen_ty_1 {
250                    val_i64: 0,
251                },
252            });
253
254        // set the pointer to the (potentially) new vector
255        self.params.kv_overrides = self.kv_overrides.as_ptr();
256
257        eprintln!("saved ptr: {:?}", self.params.kv_overrides);
258    }
259}
260
261impl LlamaModelParams {
262    /// Adds buffer type overides to move all mixture-of-experts layers to CPU.
263    pub fn add_cpu_moe_override(self: Pin<&mut Self>) {
264        self.add_cpu_buft_override(c"\\.ffn_(up|down|gate)_(ch|)exps");
265    }
266
267    /// Appends a buffer type override to the model parameters, to move layers matching pattern to CPU.
268    /// It must be pinned as this creates a self-referential struct.
269    pub fn add_cpu_buft_override(mut self: Pin<&mut Self>, key: &CStr) {
270        let buft_override = self
271            .buft_overrides
272            .last_mut()
273            .expect("buft_overrides did not have a next allocated");
274
275        assert!(
276            buft_override.pattern.is_null(),
277            "last buft_override was not empty"
278        );
279
280        // There should be some way to do this without iterating over everything.
281        for &c in key.to_bytes_with_nul().iter() {
282            c_char::try_from(c).expect("invalid character in key");
283        }
284
285        buft_override.pattern = key.as_ptr();
286        buft_override.buft = unsafe { llama_cpp_sys_2::ggml_backend_cpu_buffer_type() };
287
288        // set to null pointer for panic safety (as push may move the vector, invalidating the pointer)
289        self.params.tensor_buft_overrides = null();
290
291        // push the next one to ensure we maintain the iterator invariant of ending with a 0
292        self.buft_overrides
293            .push(llama_cpp_sys_2::llama_model_tensor_buft_override {
294                pattern: std::ptr::null(),
295                buft: std::ptr::null_mut(),
296            });
297
298        // set the pointer to the (potentially) new vector
299        self.params.tensor_buft_overrides = self.buft_overrides.as_ptr();
300    }
301
302    /// Returns the tensor-name patterns of the buffer-type overrides currently set on these
303    /// parameters, in order.
304    ///
305    /// This is the read-only counterpart to [`add_cpu_buft_override`](Self::add_cpu_buft_override)
306    /// and [`add_cpu_moe_override`](Self::add_cpu_moe_override). After
307    /// [`fit_params`](Self::fit_params) it reflects the overrides the auto-fit chose — for example
308    /// the routed-expert tensors (`blk.<N>.ffn_(up|down|gate_up|gate)_(ch|)exps`) a
309    /// mixture-of-experts fit assigns to the CPU buffer type to make the model fit. Returns an empty
310    /// vector when no overrides are set. The trailing null-terminator entry the override list
311    /// carries is skipped; only entries with a non-null pattern are returned.
312    #[must_use]
313    pub fn tensor_buft_override_patterns(&self) -> Vec<String> {
314        self.buft_overrides
315            .iter()
316            .filter(|o| !o.pattern.is_null())
317            .map(|o| {
318                // SAFETY: a non-null `pattern` is a NUL-terminated C string. For fit-produced
319                // overrides it points into process-lifetime function-local `static` storage in
320                // llama.cpp's `common/fit.cpp`, so it is always valid to read here. For overrides set
321                // via `add_cpu_buft_override` the pointer is borrowed from the caller's `&CStr` with
322                // no lifetime tie recorded on the params, so that setter's callers are responsible
323                // for keeping the string alive at least as long as the params; every in-tree caller
324                // passes a `'static` literal. In both cases the string outlives this `&self` borrow.
325                unsafe { CStr::from_ptr(o.pattern) }
326                    .to_string_lossy()
327                    .into_owned()
328            })
329            .collect()
330    }
331}
332
333#[cfg(feature = "common")]
334impl LlamaModelParams {
335    /// Automatically fit model parameters to available device memory.
336    ///
337    /// Wraps llama.cpp's `common_fit_params` (libcommon), which determines optimal `n_gpu_layers`,
338    /// `tensor_split`, and `tensor_buft_overrides` based on available VRAM. On success
339    /// the model and context params are updated in place.
340    ///
341    /// # Requirements
342    ///
343    /// Per the C API docstring, only parameters that still hold their default value
344    /// are modified. In practice this means:
345    /// - `n_gpu_layers` must be at its default (`-1`). Do not call
346    ///   [`with_n_gpu_layers`](Self::with_n_gpu_layers) before this.
347    /// - No `tensor_buft_overrides` may be set. Do not call
348    ///   [`add_cpu_buft_override`](Self::add_cpu_buft_override) or
349    ///   [`add_cpu_moe_override`](Self::add_cpu_moe_override) before this.
350    /// - No tensor split may be set. A split from
351    ///   [`with_tensor_split`](Self::with_tensor_split) is discarded and replaced by the fitted one.
352    /// - `cparams.n_ctx` is only auto-selected if it is `0`; otherwise it is left alone.
353    ///
354    /// # Arguments
355    ///
356    /// - `model_path` — path to the GGUF model file.
357    /// - `cparams` — context parameters; `n_ctx` may be modified (see above).
358    /// - `margins` — memory margin per device in bytes. Must have at least
359    ///   `llama_max_devices()` elements.
360    /// - `n_ctx_min` — minimum context size to preserve when reducing memory usage.
361    /// - `log_level` — minimum log level for fitting output; lower levels are routed
362    ///   to the debug log.
363    ///
364    /// # Thread safety
365    ///
366    /// This function is **not** thread safe: the underlying C call mutates the global
367    /// llama logger state.
368    ///
369    /// # Errors
370    ///
371    /// Returns [`FitError::Failure`] if no fitting allocation could be found, or
372    /// [`FitError::Error`] on a hard error (e.g. the model file could not be read).
373    pub fn fit_params(
374        mut self: Pin<&mut Self>,
375        model_path: &CStr,
376        cparams: &mut LlamaContextParams,
377        margins: &mut [usize],
378        n_ctx_min: u32,
379        log_level: llama_cpp_sys_2::ggml_log_level,
380    ) -> Result<FitResult, FitError> {
381        let max_devices = unsafe { llama_cpp_sys_2::llama_max_devices() };
382        let max_buft = unsafe { llama_cpp_sys_2::llama_max_tensor_buft_overrides() };
383
384        // Allocate tensor_split output buffer.
385        self.tensor_split.clear();
386        self.tensor_split.resize(max_devices, 0.0);
387
388        // Reset and resize buft_overrides for fit output (null-terminated).
389        self.buft_overrides.clear();
390        self.buft_overrides.resize(
391            max_buft + 1,
392            llama_cpp_sys_2::llama_model_tensor_buft_override {
393                pattern: std::ptr::null(),
394                buft: std::ptr::null_mut(),
395            },
396        );
397
398        // Clear pointers before the call — fit writes directly into the buffers above.
399        self.params.tensor_split = null::<f32>();
400        self.params.tensor_buft_overrides = null();
401
402        let status = unsafe {
403            llama_cpp_sys_2::llama_rs_fit_params(
404                model_path.as_ptr(),
405                &raw mut self.params,
406                &raw mut cparams.context_params,
407                self.tensor_split.as_mut_ptr(),
408                self.buft_overrides.as_mut_ptr(),
409                margins.as_mut_ptr(),
410                n_ctx_min,
411                std::ptr::null_mut(),
412                log_level,
413            )
414        };
415
416        // llama_rs_fit_params returns common_params_fit_status: 0 = success, 1 = failure, 2 = error.
417        match status {
418            0 => {}
419            1 => return Err(FitError::Failure),
420            _ => return Err(FitError::Error),
421        }
422
423        // Wire the owned buffers into the raw params.
424        self.params.tensor_split = self.tensor_split.as_ptr();
425        self.params.tensor_buft_overrides = self.buft_overrides.as_ptr();
426
427        Ok(FitResult {
428            n_ctx: cparams.context_params.n_ctx,
429        })
430    }
431}
432
433impl LlamaModelParams {
434    /// Get the number of layers to offload to the GPU.
435    #[must_use]
436    pub fn n_gpu_layers(&self) -> i32 {
437        self.params.n_gpu_layers
438    }
439
440    /// The GPU that is used for scratch and small tensors
441    #[must_use]
442    pub fn main_gpu(&self) -> i32 {
443        self.params.main_gpu
444    }
445
446    /// only load the vocabulary, no weights
447    #[must_use]
448    pub fn vocab_only(&self) -> bool {
449        self.params.vocab_only
450    }
451
452    /// use mmap if possible
453    ///
454    /// `use_mmap`/`use_mlock` were replaced by a single `load_mode` enum
455    /// (`args: refactor mlock/mmap/directio into load-mode`, #20834); this
456    /// getter decodes the combined mode back into the two independent flags
457    /// this crate's public API still exposes.
458    #[must_use]
459    pub fn use_mmap(&self) -> bool {
460        // `LLAMA_LOAD_MODE_AUTO` generally translates to `use_mmap`.
461        // FIXME(madsmtm): Change this API to expose the load mode directly,
462        // instead of making it two methods.
463        matches!(
464            self.params.load_mode,
465            llama_cpp_sys_2::LLAMA_LOAD_MODE_AUTO
466                | llama_cpp_sys_2::LLAMA_LOAD_MODE_MMAP
467                | llama_cpp_sys_2::LLAMA_LOAD_MODE_MMAP_MLOCK
468        )
469    }
470
471    /// force system to keep model in RAM
472    #[must_use]
473    pub fn use_mlock(&self) -> bool {
474        matches!(
475            self.params.load_mode,
476            llama_cpp_sys_2::LLAMA_LOAD_MODE_MLOCK | llama_cpp_sys_2::LLAMA_LOAD_MODE_MMAP_MLOCK
477        )
478    }
479
480    /// get the split mode
481    ///
482    /// # Errors
483    /// Returns `LlamaSplitModeParseError` if the unknown split mode is encountered.
484    pub fn split_mode(&self) -> Result<LlamaSplitMode, LlamaSplitModeParseError> {
485        LlamaSplitMode::try_from(self.params.split_mode)
486    }
487
488    /// get the devices
489    #[must_use]
490    pub fn devices(&self) -> Vec<usize> {
491        let mut backend_devices = Vec::new();
492        for i in 0..unsafe { llama_cpp_sys_2::ggml_backend_dev_count() } {
493            let dev = unsafe { llama_cpp_sys_2::ggml_backend_dev_get(i) };
494            backend_devices.push(dev);
495        }
496        let mut devices = Vec::new();
497        for &dev in self.devices.iter() {
498            if dev.is_null() {
499                break;
500            }
501            if let Some((index, _)) = backend_devices
502                .iter()
503                .enumerate()
504                .find(|&(_i, &d)| d == dev)
505            {
506                devices.push(index);
507            }
508        }
509        devices
510    }
511
512    /// sets the number of gpu layers to offload to the GPU.
513    /// ```
514    /// # use llama_cpp_2::model::params::LlamaModelParams;
515    /// let params = LlamaModelParams::default();
516    /// let params = params.with_n_gpu_layers(1);
517    /// assert_eq!(params.n_gpu_layers(), 1);
518    /// ```
519    #[must_use]
520    pub fn with_n_gpu_layers(mut self, n_gpu_layers: u32) -> Self {
521        // The only way this conversion can fail is if u32 overflows the i32 - in which case we set
522        // to MAX
523        let n_gpu_layers = i32::try_from(n_gpu_layers).unwrap_or(i32::MAX);
524        self.params.n_gpu_layers = n_gpu_layers;
525        self
526    }
527
528    /// sets the main GPU
529    ///
530    /// To enable this option, you must set `split_mode` to `LlamaSplitMode::None` to enable single GPU mode.
531    #[must_use]
532    pub fn with_main_gpu(mut self, main_gpu: i32) -> Self {
533        self.params.main_gpu = main_gpu;
534        self
535    }
536
537    /// sets `vocab_only`
538    #[must_use]
539    pub fn with_vocab_only(mut self, vocab_only: bool) -> Self {
540        self.params.vocab_only = vocab_only;
541        self
542    }
543
544    /// sets `use_mmap`
545    #[must_use]
546    pub fn with_use_mmap(mut self, use_mmap: bool) -> Self {
547        self.params.load_mode = load_mode_from_flags(use_mmap, self.use_mlock());
548        self
549    }
550
551    /// sets `use_mlock`
552    #[must_use]
553    pub fn with_use_mlock(mut self, use_mlock: bool) -> Self {
554        self.params.load_mode = load_mode_from_flags(self.use_mmap(), use_mlock);
555        self
556    }
557
558    /// sets `split_mode`
559    #[must_use]
560    pub fn with_split_mode(mut self, split_mode: LlamaSplitMode) -> Self {
561        self.params.split_mode = split_mode.into();
562        self
563    }
564
565    /// Sets each device's share of the model, parallel to
566    /// [`with_devices`](Self::with_devices).
567    ///
568    /// Shares are relative, so they need not sum to one. An empty slice clears the split and
569    /// restores llama.cpp's default, which distributes in proportion to free memory. The
570    /// slice is padded with zeroes to `llama_max_devices()`, as the C API requires.
571    ///
572    /// The pointer handed to llama.cpp addresses the `Vec`'s heap buffer, not this struct, so
573    /// it survives the moves that builder chaining performs.
574    ///
575    /// # Errors
576    /// Returns `LlamaCppError::MaxDevicesExceeded` if `split` is longer than
577    /// `llama_max_devices()`.
578    pub fn with_tensor_split(mut self, split: &[f32]) -> Result<Self, LlamaCppError> {
579        if split.is_empty() {
580            self.tensor_split.clear();
581            self.params.tensor_split = null::<f32>();
582            return Ok(self);
583        }
584        let max_devices = crate::max_devices();
585        if split.len() > max_devices {
586            return Err(LlamaCppError::MaxDevicesExceeded(max_devices));
587        }
588        self.tensor_split.clear();
589        self.tensor_split.extend_from_slice(split);
590        self.tensor_split.resize(max_devices, 0.0);
591        self.params.tensor_split = self.tensor_split.as_ptr();
592        Ok(self)
593    }
594
595    /// Each device's share of the model, as last set by
596    /// [`with_tensor_split`](Self::with_tensor_split). Empty when the split is llama.cpp's
597    /// free-memory default.
598    #[must_use]
599    pub fn tensor_split(&self) -> &[f32] {
600        &self.tensor_split
601    }
602
603    /// sets `devices`
604    ///
605    /// The devices are specified as indices that correspond to the ggml backend device indices.
606    ///
607    /// The maximum number of devices is 16.
608    ///
609    /// You don't need to specify CPU or ACCEL devices.
610    ///
611    /// # Errors
612    /// Returns `LlamaCppError::BackendDeviceNotFound` if any device index is invalid.
613    pub fn with_devices(mut self, devices: &[usize]) -> Result<Self, LlamaCppError> {
614        for dev in self.devices.iter_mut() {
615            *dev = std::ptr::null_mut();
616        }
617        // Check device count
618        let max_devices = crate::max_devices().min(LLAMA_CPP_MAX_DEVICES);
619        if devices.len() > max_devices {
620            return Err(LlamaCppError::MaxDevicesExceeded(max_devices));
621        }
622        for (i, &dev) in devices.iter().enumerate() {
623            if dev >= unsafe { llama_cpp_sys_2::ggml_backend_dev_count() } {
624                return Err(LlamaCppError::BackendDeviceNotFound(dev));
625            }
626            let backend_dev = unsafe { llama_cpp_sys_2::ggml_backend_dev_get(dev) };
627            self.devices[i] = backend_dev;
628        }
629        if self.devices.is_empty() {
630            self.params.devices = std::ptr::null_mut();
631        } else {
632            self.params.devices = self.devices.as_mut_ptr();
633        }
634        Ok(self)
635    }
636
637    /// Set `no_alloc`
638    ///
639    /// If this parameter is true, don't allocate memory for the tensor data
640    ///
641    /// You can't use `no_alloc` with `use_mmap`, so this also sets `use_mmap` to false.
642    #[must_use]
643    pub fn with_no_alloc(mut self, no_alloc: bool) -> Self {
644        self.params.no_alloc = no_alloc;
645        if no_alloc {
646            self = self.with_use_mmap(false);
647        }
648        self
649    }
650
651    /// Get `no_alloc`
652    ///
653    /// If this parameter is true, don't allocate memory for the tensor data
654    #[must_use]
655    pub fn no_alloc(&self) -> bool {
656        self.params.no_alloc
657    }
658
659    /// Sets whether to load bundled multi-token prediction (MTP) tensors.
660    #[must_use]
661    pub fn with_load_mtp(mut self, load_mtp: bool) -> Self {
662        self.params.load_mtp = load_mtp;
663        self
664    }
665
666    /// Returns whether bundled multi-token prediction (MTP) tensors are loaded.
667    #[must_use]
668    pub fn load_mtp(&self) -> bool {
669        self.params.load_mtp
670    }
671
672    /// Sets a callback invoked during loading with progress in `0.0..=1.0`.
673    /// Returning `false` aborts the load (it then fails with `NullResult`).
674    #[must_use]
675    pub fn with_progress_callback<F: FnMut(f32) -> bool + 'static>(mut self, callback: F) -> Self {
676        unsafe extern "C" fn trampoline<F: FnMut(f32) -> bool>(
677            progress: f32,
678            user_data: *mut c_void,
679        ) -> bool {
680            let callback = unsafe { &mut *user_data.cast::<F>() };
681            callback(progress)
682        }
683
684        let mut callback = Box::new(callback);
685        self.params.progress_callback_user_data =
686            std::ptr::from_mut(&mut *callback).cast::<c_void>();
687        self.params.progress_callback = Some(trampoline::<F>);
688        self.progress_callback = Some(callback);
689        self
690    }
691}
692
693/// Default parameters for `LlamaModel`. (as defined in llama.cpp by `llama_model_default_params`)
694/// ```
695/// # use llama_cpp_2::model::params::LlamaModelParams;
696/// use llama_cpp_2::model::params::LlamaSplitMode;
697/// let params = LlamaModelParams::default();
698/// assert_eq!(params.n_gpu_layers(), -1, "n_gpu_layers should be -1");
699/// assert_eq!(params.main_gpu(), 0, "main_gpu should be 0");
700/// assert_eq!(params.vocab_only(), false, "vocab_only should be false");
701/// assert_eq!(params.use_mmap(), true, "use_mmap should be true");
702/// assert_eq!(params.use_mlock(), false, "use_mlock should be false");
703/// assert_eq!(params.split_mode(), Ok(LlamaSplitMode::Layer), "split_mode should be LAYER");
704/// assert_eq!(params.devices().len(), 0, "devices should be empty");
705/// assert_eq!(params.no_alloc(), false, "no_alloc should be false");
706/// assert_eq!(params.load_mtp(), false, "load_mtp should be false");
707/// ```
708impl Default for LlamaModelParams {
709    fn default() -> Self {
710        let default_params = unsafe { llama_cpp_sys_2::llama_model_default_params() };
711        LlamaModelParams {
712            params: default_params,
713            // push the next one to ensure we maintain the iterator invariant of ending with a 0
714            kv_overrides: vec![llama_cpp_sys_2::llama_model_kv_override {
715                key: [0; 128],
716                tag: 0,
717                __bindgen_anon_1: llama_cpp_sys_2::llama_model_kv_override__bindgen_ty_1 {
718                    val_i64: 0,
719                },
720            }],
721            buft_overrides: vec![llama_cpp_sys_2::llama_model_tensor_buft_override {
722                pattern: std::ptr::null(),
723                buft: std::ptr::null_mut(),
724            }],
725            devices: Box::pin([std::ptr::null_mut(); 16]),
726            tensor_split: Vec::new(),
727            progress_callback: None,
728        }
729    }
730}
731
732#[cfg(test)]
733mod tests {
734    use super::kv_overrides::ParamOverrideValue;
735    use super::{LlamaModelParams, LlamaSplitMode};
736    use std::pin::pin;
737
738    #[test]
739    fn tensor_buft_override_patterns_empty_by_default() {
740        // Fresh params carry only the null-terminator entry, so no patterns are reported.
741        assert!(LlamaModelParams::default()
742            .tensor_buft_override_patterns()
743            .is_empty());
744    }
745
746    #[test]
747    fn tensor_buft_override_patterns_reads_back_added_override() {
748        // The getter is the read-only counterpart to the setter: the override added is reported and
749        // the trailing null terminator is skipped. This mirrors how `fit_params` populates the same
750        // buffer for the auto-fit's MoE expert offload (`add_cpu_moe_override` uses the same expert
751        // tensor pattern shape the fit emits).
752        let mut params = pin!(LlamaModelParams::default());
753        params.as_mut().add_cpu_moe_override();
754        assert_eq!(
755            params.tensor_buft_override_patterns(),
756            vec!["\\.ffn_(up|down|gate)_(ch|)exps".to_owned()],
757        );
758    }
759
760    #[test]
761    fn tensor_buft_override_appends_on_second_call() {
762        // add_cpu_buft_override is documented as appending. After the first call the
763        // vector is [filled, null-terminator]; writing the second override into the
764        // terminator (then pushing a new one) is what keeps the list null-terminated.
765        let mut params = pin!(LlamaModelParams::default());
766        params
767            .as_mut()
768            .add_cpu_buft_override(c"blk\\.0\\.ffn_(up|down|gate)_(ch|)exps");
769        params
770            .as_mut()
771            .add_cpu_buft_override(c"blk\\.1\\.ffn_(up|down|gate)_(ch|)exps");
772
773        assert_eq!(
774            params.tensor_buft_override_patterns(),
775            vec![
776                "blk\\.0\\.ffn_(up|down|gate)_(ch|)exps".to_owned(),
777                "blk\\.1\\.ffn_(up|down|gate)_(ch|)exps".to_owned(),
778            ],
779        );
780    }
781
782    #[test]
783    fn kv_override_appends_on_second_call() {
784        // Same [filled, null-terminator] layout as the buffer overrides: the second call has to
785        // write into the terminator, not back into slot 0.
786        let mut params = pin!(LlamaModelParams::default());
787        params
788            .as_mut()
789            .append_kv_override(c"first", ParamOverrideValue::Int(1));
790        params
791            .as_mut()
792            .append_kv_override(c"second", ParamOverrideValue::Int(2));
793
794        let kv_overrides = params.kv_overrides().into_iter().collect::<Vec<_>>();
795        assert_eq!(kv_overrides.len(), 2);
796        assert_eq!(kv_overrides[0].0.to_bytes(), b"first");
797        assert_eq!(kv_overrides[1].0.to_bytes(), b"second");
798        assert_eq!(kv_overrides[1].1, ParamOverrideValue::Int(2));
799    }
800
801    #[test]
802    fn tensor_split_round_trips_and_pads() {
803        let max = crate::max_devices();
804        let params = LlamaModelParams::default()
805            .with_tensor_split(&[0.25, 0.75])
806            .expect("two devices is under any max");
807        assert_eq!(params.tensor_split().len(), max);
808        assert_eq!(&params.tensor_split()[..2], &[0.25, 0.75]);
809        assert!(params.tensor_split()[2..].iter().all(|share| *share == 0.0));
810
811        // The pointer must survive the moves builder chaining performs.
812        let params = params.with_n_gpu_layers(99);
813        assert_eq!(&params.tensor_split()[..2], &[0.25, 0.75]);
814
815        let cleared = params.with_tensor_split(&[]).expect("empty clears");
816        assert!(cleared.tensor_split().is_empty());
817    }
818
819    #[test]
820    fn tensor_split_rejects_more_devices_than_llama_supports() {
821        let too_many = vec![1.0_f32; crate::max_devices() + 1];
822        assert!(LlamaModelParams::default()
823            .with_tensor_split(&too_many)
824            .is_err());
825    }
826
827    #[test]
828    fn tensor_split_mode_round_trips() {
829        assert_eq!(
830            LlamaSplitMode::try_from(llama_cpp_sys_2::LLAMA_SPLIT_MODE_TENSOR),
831            Ok(LlamaSplitMode::Tensor)
832        );
833        assert_eq!(
834            u32::from(LlamaSplitMode::Tensor),
835            llama_cpp_sys_2::LLAMA_SPLIT_MODE_TENSOR as u32
836        );
837        assert_eq!(
838            i32::from(LlamaSplitMode::Tensor),
839            llama_cpp_sys_2::LLAMA_SPLIT_MODE_TENSOR as i32
840        );
841    }
842
843    #[test]
844    fn load_mtp_round_trips() {
845        let params = LlamaModelParams::default();
846        assert!(!params.load_mtp());
847
848        let params = params.with_load_mtp(true);
849        assert!(params.load_mtp());
850
851        let params = params.with_load_mtp(false);
852        assert!(!params.load_mtp());
853    }
854
855    #[test]
856    fn progress_callback_round_trips_and_can_abort() {
857        use super::LlamaModelParams;
858        use std::cell::Cell;
859        use std::rc::Rc;
860
861        let calls = Rc::new(Cell::new(0_u32));
862        let counter = Rc::clone(&calls);
863        let params = LlamaModelParams::default().with_progress_callback(move |_progress| {
864            counter.set(counter.get() + 1);
865            false
866        });
867
868        assert!(params.params.progress_callback.is_some());
869        assert!(!params.params.progress_callback_user_data.is_null());
870
871        let trampoline = params.params.progress_callback.unwrap();
872        let user_data = params.params.progress_callback_user_data;
873        let first = unsafe { trampoline(0.5, user_data) };
874        let second = unsafe { trampoline(1.0, user_data) };
875
876        assert!(!first && !second, "returning false signals an abort");
877        assert_eq!(calls.get(), 2);
878    }
879}
llama_cpp_2/model/params.rs

llama_cpp_2/model/
params.rs
