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            .get_mut(0)
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            .get_mut(0)
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    /// - `cparams.n_ctx` is only auto-selected if it is `0`; otherwise it is left alone.
351    ///
352    /// # Arguments
353    ///
354    /// - `model_path` — path to the GGUF model file.
355    /// - `cparams` — context parameters; `n_ctx` may be modified (see above).
356    /// - `margins` — memory margin per device in bytes. Must have at least
357    ///   `llama_max_devices()` elements.
358    /// - `n_ctx_min` — minimum context size to preserve when reducing memory usage.
359    /// - `log_level` — minimum log level for fitting output; lower levels are routed
360    ///   to the debug log.
361    ///
362    /// # Thread safety
363    ///
364    /// This function is **not** thread safe: the underlying C call mutates the global
365    /// llama logger state.
366    ///
367    /// # Errors
368    ///
369    /// Returns [`FitError::Failure`] if no fitting allocation could be found, or
370    /// [`FitError::Error`] on a hard error (e.g. the model file could not be read).
371    pub fn fit_params(
372        mut self: Pin<&mut Self>,
373        model_path: &CStr,
374        cparams: &mut LlamaContextParams,
375        margins: &mut [usize],
376        n_ctx_min: u32,
377        log_level: llama_cpp_sys_2::ggml_log_level,
378    ) -> Result<FitResult, FitError> {
379        let max_devices = unsafe { llama_cpp_sys_2::llama_max_devices() };
380        let max_buft = unsafe { llama_cpp_sys_2::llama_max_tensor_buft_overrides() };
381
382        // Allocate tensor_split output buffer.
383        self.tensor_split.clear();
384        self.tensor_split.resize(max_devices, 0.0);
385
386        // Reset and resize buft_overrides for fit output (null-terminated).
387        self.buft_overrides.clear();
388        self.buft_overrides.resize(
389            max_buft + 1,
390            llama_cpp_sys_2::llama_model_tensor_buft_override {
391                pattern: std::ptr::null(),
392                buft: std::ptr::null_mut(),
393            },
394        );
395
396        // Clear pointers before the call — fit writes directly into the buffers above.
397        self.params.tensor_split = null::<f32>();
398        self.params.tensor_buft_overrides = null();
399
400        let status = unsafe {
401            llama_cpp_sys_2::llama_rs_fit_params(
402                model_path.as_ptr(),
403                &raw mut self.params,
404                &raw mut cparams.context_params,
405                self.tensor_split.as_mut_ptr(),
406                self.buft_overrides.as_mut_ptr(),
407                margins.as_mut_ptr(),
408                n_ctx_min,
409                log_level,
410            )
411        };
412
413        // llama_rs_fit_params returns common_params_fit_status: 0 = success, 1 = failure, 2 = error.
414        match status {
415            0 => {}
416            1 => return Err(FitError::Failure),
417            _ => return Err(FitError::Error),
418        }
419
420        // Wire the owned buffers into the raw params.
421        self.params.tensor_split = self.tensor_split.as_ptr();
422        self.params.tensor_buft_overrides = self.buft_overrides.as_ptr();
423
424        Ok(FitResult {
425            n_ctx: cparams.context_params.n_ctx,
426        })
427    }
428}
429
430impl LlamaModelParams {
431    /// Get the number of layers to offload to the GPU.
432    #[must_use]
433    pub fn n_gpu_layers(&self) -> i32 {
434        self.params.n_gpu_layers
435    }
436
437    /// The GPU that is used for scratch and small tensors
438    #[must_use]
439    pub fn main_gpu(&self) -> i32 {
440        self.params.main_gpu
441    }
442
443    /// only load the vocabulary, no weights
444    #[must_use]
445    pub fn vocab_only(&self) -> bool {
446        self.params.vocab_only
447    }
448
449    /// use mmap if possible
450    ///
451    /// `use_mmap`/`use_mlock` were replaced by a single `load_mode` enum
452    /// (`args: refactor mlock/mmap/directio into load-mode`, #20834); this
453    /// getter decodes the combined mode back into the two independent flags
454    /// this crate's public API still exposes.
455    #[must_use]
456    pub fn use_mmap(&self) -> bool {
457        matches!(
458            self.params.load_mode,
459            llama_cpp_sys_2::LLAMA_LOAD_MODE_MMAP | llama_cpp_sys_2::LLAMA_LOAD_MODE_MMAP_MLOCK
460        )
461    }
462
463    /// force system to keep model in RAM
464    #[must_use]
465    pub fn use_mlock(&self) -> bool {
466        matches!(
467            self.params.load_mode,
468            llama_cpp_sys_2::LLAMA_LOAD_MODE_MLOCK | llama_cpp_sys_2::LLAMA_LOAD_MODE_MMAP_MLOCK
469        )
470    }
471
472    /// get the split mode
473    ///
474    /// # Errors
475    /// Returns `LlamaSplitModeParseError` if the unknown split mode is encountered.
476    pub fn split_mode(&self) -> Result<LlamaSplitMode, LlamaSplitModeParseError> {
477        LlamaSplitMode::try_from(self.params.split_mode)
478    }
479
480    /// get the devices
481    #[must_use]
482    pub fn devices(&self) -> Vec<usize> {
483        let mut backend_devices = Vec::new();
484        for i in 0..unsafe { llama_cpp_sys_2::ggml_backend_dev_count() } {
485            let dev = unsafe { llama_cpp_sys_2::ggml_backend_dev_get(i) };
486            backend_devices.push(dev);
487        }
488        let mut devices = Vec::new();
489        for &dev in self.devices.iter() {
490            if dev.is_null() {
491                break;
492            }
493            if let Some((index, _)) = backend_devices
494                .iter()
495                .enumerate()
496                .find(|&(_i, &d)| d == dev)
497            {
498                devices.push(index);
499            }
500        }
501        devices
502    }
503
504    /// sets the number of gpu layers to offload to the GPU.
505    /// ```
506    /// # use llama_cpp_2::model::params::LlamaModelParams;
507    /// let params = LlamaModelParams::default();
508    /// let params = params.with_n_gpu_layers(1);
509    /// assert_eq!(params.n_gpu_layers(), 1);
510    /// ```
511    #[must_use]
512    pub fn with_n_gpu_layers(mut self, n_gpu_layers: u32) -> Self {
513        // The only way this conversion can fail is if u32 overflows the i32 - in which case we set
514        // to MAX
515        let n_gpu_layers = i32::try_from(n_gpu_layers).unwrap_or(i32::MAX);
516        self.params.n_gpu_layers = n_gpu_layers;
517        self
518    }
519
520    /// sets the main GPU
521    ///
522    /// To enable this option, you must set `split_mode` to `LlamaSplitMode::None` to enable single GPU mode.
523    #[must_use]
524    pub fn with_main_gpu(mut self, main_gpu: i32) -> Self {
525        self.params.main_gpu = main_gpu;
526        self
527    }
528
529    /// sets `vocab_only`
530    #[must_use]
531    pub fn with_vocab_only(mut self, vocab_only: bool) -> Self {
532        self.params.vocab_only = vocab_only;
533        self
534    }
535
536    /// sets `use_mmap`
537    #[must_use]
538    pub fn with_use_mmap(mut self, use_mmap: bool) -> Self {
539        self.params.load_mode = load_mode_from_flags(use_mmap, self.use_mlock());
540        self
541    }
542
543    /// sets `use_mlock`
544    #[must_use]
545    pub fn with_use_mlock(mut self, use_mlock: bool) -> Self {
546        self.params.load_mode = load_mode_from_flags(self.use_mmap(), use_mlock);
547        self
548    }
549
550    /// sets `split_mode`
551    #[must_use]
552    pub fn with_split_mode(mut self, split_mode: LlamaSplitMode) -> Self {
553        self.params.split_mode = split_mode.into();
554        self
555    }
556
557    /// sets `devices`
558    ///
559    /// The devices are specified as indices that correspond to the ggml backend device indices.
560    ///
561    /// The maximum number of devices is 16.
562    ///
563    /// You don't need to specify CPU or ACCEL devices.
564    ///
565    /// # Errors
566    /// Returns `LlamaCppError::BackendDeviceNotFound` if any device index is invalid.
567    pub fn with_devices(mut self, devices: &[usize]) -> Result<Self, LlamaCppError> {
568        for dev in self.devices.iter_mut() {
569            *dev = std::ptr::null_mut();
570        }
571        // Check device count
572        let max_devices = crate::max_devices().min(LLAMA_CPP_MAX_DEVICES);
573        if devices.len() > max_devices {
574            return Err(LlamaCppError::MaxDevicesExceeded(max_devices));
575        }
576        for (i, &dev) in devices.iter().enumerate() {
577            if dev >= unsafe { llama_cpp_sys_2::ggml_backend_dev_count() } {
578                return Err(LlamaCppError::BackendDeviceNotFound(dev));
579            }
580            let backend_dev = unsafe { llama_cpp_sys_2::ggml_backend_dev_get(dev) };
581            self.devices[i] = backend_dev;
582        }
583        if self.devices.is_empty() {
584            self.params.devices = std::ptr::null_mut();
585        } else {
586            self.params.devices = self.devices.as_mut_ptr();
587        }
588        Ok(self)
589    }
590
591    /// Set `no_alloc`
592    ///
593    /// If this parameter is true, don't allocate memory for the tensor data
594    ///
595    /// You can't use `no_alloc` with `use_mmap`, so this also sets `use_mmap` to false.
596    #[must_use]
597    pub fn with_no_alloc(mut self, no_alloc: bool) -> Self {
598        self.params.no_alloc = no_alloc;
599        if no_alloc {
600            self = self.with_use_mmap(false);
601        }
602        self
603    }
604
605    /// Get `no_alloc`
606    ///
607    /// If this parameter is true, don't allocate memory for the tensor data
608    #[must_use]
609    pub fn no_alloc(&self) -> bool {
610        self.params.no_alloc
611    }
612
613    /// Sets a callback invoked during loading with progress in `0.0..=1.0`.
614    /// Returning `false` aborts the load (it then fails with `NullResult`).
615    #[must_use]
616    pub fn with_progress_callback<F: FnMut(f32) -> bool + 'static>(mut self, callback: F) -> Self {
617        unsafe extern "C" fn trampoline<F: FnMut(f32) -> bool>(
618            progress: f32,
619            user_data: *mut c_void,
620        ) -> bool {
621            let callback = unsafe { &mut *user_data.cast::<F>() };
622            callback(progress)
623        }
624
625        let mut callback = Box::new(callback);
626        self.params.progress_callback_user_data =
627            std::ptr::from_mut(&mut *callback).cast::<c_void>();
628        self.params.progress_callback = Some(trampoline::<F>);
629        self.progress_callback = Some(callback);
630        self
631    }
632}
633
634/// Default parameters for `LlamaModel`. (as defined in llama.cpp by `llama_model_default_params`)
635/// ```
636/// # use llama_cpp_2::model::params::LlamaModelParams;
637/// use llama_cpp_2::model::params::LlamaSplitMode;
638/// let params = LlamaModelParams::default();
639/// assert_eq!(params.n_gpu_layers(), -1, "n_gpu_layers should be -1");
640/// assert_eq!(params.main_gpu(), 0, "main_gpu should be 0");
641/// assert_eq!(params.vocab_only(), false, "vocab_only should be false");
642/// assert_eq!(params.use_mmap(), true, "use_mmap should be true");
643/// assert_eq!(params.use_mlock(), false, "use_mlock should be false");
644/// assert_eq!(params.split_mode(), Ok(LlamaSplitMode::Layer), "split_mode should be LAYER");
645/// assert_eq!(params.devices().len(), 0, "devices should be empty");
646/// assert_eq!(params.no_alloc(), false, "no_alloc should be false");
647/// ```
648impl Default for LlamaModelParams {
649    fn default() -> Self {
650        let default_params = unsafe { llama_cpp_sys_2::llama_model_default_params() };
651        LlamaModelParams {
652            params: default_params,
653            // push the next one to ensure we maintain the iterator invariant of ending with a 0
654            kv_overrides: vec![llama_cpp_sys_2::llama_model_kv_override {
655                key: [0; 128],
656                tag: 0,
657                __bindgen_anon_1: llama_cpp_sys_2::llama_model_kv_override__bindgen_ty_1 {
658                    val_i64: 0,
659                },
660            }],
661            buft_overrides: vec![llama_cpp_sys_2::llama_model_tensor_buft_override {
662                pattern: std::ptr::null(),
663                buft: std::ptr::null_mut(),
664            }],
665            devices: Box::pin([std::ptr::null_mut(); 16]),
666            tensor_split: Vec::new(),
667            progress_callback: None,
668        }
669    }
670}
671
672#[cfg(test)]
673mod tests {
674    use super::{LlamaModelParams, LlamaSplitMode};
675    use std::pin::pin;
676
677    #[test]
678    fn tensor_buft_override_patterns_empty_by_default() {
679        // Fresh params carry only the null-terminator entry, so no patterns are reported.
680        assert!(LlamaModelParams::default()
681            .tensor_buft_override_patterns()
682            .is_empty());
683    }
684
685    #[test]
686    fn tensor_buft_override_patterns_reads_back_added_override() {
687        // The getter is the read-only counterpart to the setter: the override added is reported and
688        // the trailing null terminator is skipped. This mirrors how `fit_params` populates the same
689        // buffer for the auto-fit's MoE expert offload (`add_cpu_moe_override` uses the same expert
690        // tensor pattern shape the fit emits).
691        let mut params = pin!(LlamaModelParams::default());
692        params.as_mut().add_cpu_moe_override();
693        assert_eq!(
694            params.tensor_buft_override_patterns(),
695            vec!["\\.ffn_(up|down|gate)_(ch|)exps".to_owned()],
696        );
697    }
698
699    #[test]
700    fn tensor_split_mode_round_trips() {
701        assert_eq!(
702            LlamaSplitMode::try_from(llama_cpp_sys_2::LLAMA_SPLIT_MODE_TENSOR),
703            Ok(LlamaSplitMode::Tensor)
704        );
705        assert_eq!(
706            u32::from(LlamaSplitMode::Tensor),
707            llama_cpp_sys_2::LLAMA_SPLIT_MODE_TENSOR as u32
708        );
709        assert_eq!(
710            i32::from(LlamaSplitMode::Tensor),
711            llama_cpp_sys_2::LLAMA_SPLIT_MODE_TENSOR as i32
712        );
713    }
714
715    #[test]
716    fn progress_callback_round_trips_and_can_abort() {
717        use super::LlamaModelParams;
718        use std::cell::Cell;
719        use std::rc::Rc;
720
721        let calls = Rc::new(Cell::new(0_u32));
722        let counter = Rc::clone(&calls);
723        let params = LlamaModelParams::default().with_progress_callback(move |_progress| {
724            counter.set(counter.get() + 1);
725            false
726        });
727
728        assert!(params.params.progress_callback.is_some());
729        assert!(!params.params.progress_callback_user_data.is_null());
730
731        let trampoline = params.params.progress_callback.unwrap();
732        let user_data = params.params.progress_callback_user_data;
733        let first = unsafe { trampoline(0.5, user_data) };
734        let second = unsafe { trampoline(1.0, user_data) };
735
736        assert!(!first && !second, "returning false signals an abort");
737        assert_eq!(calls.get(), 2);
738    }
739}
llama_cpp_2/model/params.rs

llama_cpp_2/model/
params.rs
