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/// A safe wrapper around `llama_model_params`.
146#[allow(clippy::module_name_repetitions)]
147pub struct LlamaModelParams {
148    pub(crate) params: llama_cpp_sys_2::llama_model_params,
149    kv_overrides: Vec<llama_cpp_sys_2::llama_model_kv_override>,
150    buft_overrides: Vec<llama_cpp_sys_2::llama_model_tensor_buft_override>,
151    devices: Pin<Box<[llama_cpp_sys_2::ggml_backend_dev_t; LLAMA_CPP_MAX_DEVICES]>>,
152    tensor_split: Vec<f32>,
153    progress_callback: Option<Box<dyn FnMut(f32) -> bool>>,
154}
155
156impl Debug for LlamaModelParams {
157    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
158        f.debug_struct("LlamaModelParams")
159            .field("n_gpu_layers", &self.params.n_gpu_layers)
160            .field("main_gpu", &self.params.main_gpu)
161            .field("vocab_only", &self.params.vocab_only)
162            .field("use_mmap", &self.params.use_mmap)
163            .field("use_mlock", &self.params.use_mlock)
164            .field("split_mode", &self.split_mode())
165            .field("devices", &self.devices)
166            .field("kv_overrides", &"vec of kv_overrides")
167            .finish()
168    }
169}
170
171impl LlamaModelParams {
172    /// See [`KvOverrides`]
173    ///
174    /// # Examples
175    ///
176    /// ```rust
177    /// # use llama_cpp_2::model::params::LlamaModelParams;
178    /// let params = Box::pin(LlamaModelParams::default());
179    /// let kv_overrides = params.kv_overrides();
180    /// let count = kv_overrides.into_iter().count();
181    /// assert_eq!(count, 0);
182    /// ```
183    #[must_use]
184    pub fn kv_overrides<'a>(&'a self) -> KvOverrides<'a> {
185        KvOverrides::new(self)
186    }
187
188    /// Appends a key-value override to the model parameters. It must be pinned as this creates a self-referential struct.
189    ///
190    /// # Examples
191    ///
192    /// ```rust
193    /// # use std::ffi::{CStr, CString};
194    /// use std::pin::pin;
195    /// # use llama_cpp_2::model::params::LlamaModelParams;
196    /// # use llama_cpp_2::model::params::kv_overrides::ParamOverrideValue;
197    /// let mut params = pin!(LlamaModelParams::default());
198    /// let key = CString::new("key").expect("CString::new failed");
199    /// params.as_mut().append_kv_override(&key, ParamOverrideValue::Int(50));
200    ///
201    /// let kv_overrides = params.kv_overrides().into_iter().collect::<Vec<_>>();
202    /// assert_eq!(kv_overrides.len(), 1);
203    ///
204    /// let (k, v) = &kv_overrides[0];
205    /// assert_eq!(v, &ParamOverrideValue::Int(50));
206    ///
207    /// assert_eq!(k.to_bytes(), b"key", "expected key to be 'key', was {:?}", k);
208    /// ```
209    #[allow(clippy::missing_panics_doc)] // panics are just to enforce internal invariants, not user errors
210    pub fn append_kv_override(
211        mut self: Pin<&mut Self>,
212        key: &CStr,
213        value: kv_overrides::ParamOverrideValue,
214    ) {
215        let kv_override = self
216            .kv_overrides
217            .get_mut(0)
218            .expect("kv_overrides did not have a next allocated");
219
220        assert_eq!(kv_override.key[0], 0, "last kv_override was not empty");
221
222        // There should be some way to do this without iterating over everything.
223        for (i, &c) in key.to_bytes_with_nul().iter().enumerate() {
224            kv_override.key[i] = c_char::try_from(c).expect("invalid character in key");
225        }
226
227        kv_override.tag = value.tag();
228        kv_override.__bindgen_anon_1 = value.value();
229
230        // set to null pointer for panic safety (as push may move the vector, invalidating the pointer)
231        self.params.kv_overrides = null();
232
233        // push the next one to ensure we maintain the iterator invariant of ending with a 0
234        self.kv_overrides
235            .push(llama_cpp_sys_2::llama_model_kv_override {
236                key: [0; 128],
237                tag: 0,
238                __bindgen_anon_1: llama_cpp_sys_2::llama_model_kv_override__bindgen_ty_1 {
239                    val_i64: 0,
240                },
241            });
242
243        // set the pointer to the (potentially) new vector
244        self.params.kv_overrides = self.kv_overrides.as_ptr();
245
246        eprintln!("saved ptr: {:?}", self.params.kv_overrides);
247    }
248}
249
250impl LlamaModelParams {
251    /// Adds buffer type overides to move all mixture-of-experts layers to CPU.
252    pub fn add_cpu_moe_override(self: Pin<&mut Self>) {
253        self.add_cpu_buft_override(c"\\.ffn_(up|down|gate)_(ch|)exps");
254    }
255
256    /// Appends a buffer type override to the model parameters, to move layers matching pattern to CPU.
257    /// It must be pinned as this creates a self-referential struct.
258    pub fn add_cpu_buft_override(mut self: Pin<&mut Self>, key: &CStr) {
259        let buft_override = self
260            .buft_overrides
261            .get_mut(0)
262            .expect("buft_overrides did not have a next allocated");
263
264        assert!(
265            buft_override.pattern.is_null(),
266            "last buft_override was not empty"
267        );
268
269        // There should be some way to do this without iterating over everything.
270        for &c in key.to_bytes_with_nul().iter() {
271            c_char::try_from(c).expect("invalid character in key");
272        }
273
274        buft_override.pattern = key.as_ptr();
275        buft_override.buft = unsafe { llama_cpp_sys_2::ggml_backend_cpu_buffer_type() };
276
277        // set to null pointer for panic safety (as push may move the vector, invalidating the pointer)
278        self.params.tensor_buft_overrides = null();
279
280        // push the next one to ensure we maintain the iterator invariant of ending with a 0
281        self.buft_overrides
282            .push(llama_cpp_sys_2::llama_model_tensor_buft_override {
283                pattern: std::ptr::null(),
284                buft: std::ptr::null_mut(),
285            });
286
287        // set the pointer to the (potentially) new vector
288        self.params.tensor_buft_overrides = self.buft_overrides.as_ptr();
289    }
290
291    /// Returns the tensor-name patterns of the buffer-type overrides currently set on these
292    /// parameters, in order.
293    ///
294    /// This is the read-only counterpart to [`add_cpu_buft_override`](Self::add_cpu_buft_override)
295    /// and [`add_cpu_moe_override`](Self::add_cpu_moe_override). After
296    /// [`fit_params`](Self::fit_params) it reflects the overrides the auto-fit chose — for example
297    /// the routed-expert tensors (`blk.<N>.ffn_(up|down|gate_up|gate)_(ch|)exps`) a
298    /// mixture-of-experts fit assigns to the CPU buffer type to make the model fit. Returns an empty
299    /// vector when no overrides are set. The trailing null-terminator entry the override list
300    /// carries is skipped; only entries with a non-null pattern are returned.
301    #[must_use]
302    pub fn tensor_buft_override_patterns(&self) -> Vec<String> {
303        self.buft_overrides
304            .iter()
305            .filter(|o| !o.pattern.is_null())
306            .map(|o| {
307                // SAFETY: a non-null `pattern` is a NUL-terminated C string. For fit-produced
308                // overrides it points into process-lifetime function-local `static` storage in
309                // llama.cpp's `common/fit.cpp`, so it is always valid to read here. For overrides set
310                // via `add_cpu_buft_override` the pointer is borrowed from the caller's `&CStr` with
311                // no lifetime tie recorded on the params, so that setter's callers are responsible
312                // for keeping the string alive at least as long as the params; every in-tree caller
313                // passes a `'static` literal. In both cases the string outlives this `&self` borrow.
314                unsafe { CStr::from_ptr(o.pattern) }
315                    .to_string_lossy()
316                    .into_owned()
317            })
318            .collect()
319    }
320}
321
322#[cfg(feature = "common")]
323impl LlamaModelParams {
324    /// Automatically fit model parameters to available device memory.
325    ///
326    /// Wraps llama.cpp's `common_fit_params` (libcommon), which determines optimal `n_gpu_layers`,
327    /// `tensor_split`, and `tensor_buft_overrides` based on available VRAM. On success
328    /// the model and context params are updated in place.
329    ///
330    /// # Requirements
331    ///
332    /// Per the C API docstring, only parameters that still hold their default value
333    /// are modified. In practice this means:
334    /// - `n_gpu_layers` must be at its default (`-1`). Do not call
335    ///   [`with_n_gpu_layers`](Self::with_n_gpu_layers) before this.
336    /// - No `tensor_buft_overrides` may be set. Do not call
337    ///   [`add_cpu_buft_override`](Self::add_cpu_buft_override) or
338    ///   [`add_cpu_moe_override`](Self::add_cpu_moe_override) before this.
339    /// - `cparams.n_ctx` is only auto-selected if it is `0`; otherwise it is left alone.
340    ///
341    /// # Arguments
342    ///
343    /// - `model_path` — path to the GGUF model file.
344    /// - `cparams` — context parameters; `n_ctx` may be modified (see above).
345    /// - `margins` — memory margin per device in bytes. Must have at least
346    ///   `llama_max_devices()` elements.
347    /// - `n_ctx_min` — minimum context size to preserve when reducing memory usage.
348    /// - `log_level` — minimum log level for fitting output; lower levels are routed
349    ///   to the debug log.
350    ///
351    /// # Thread safety
352    ///
353    /// This function is **not** thread safe: the underlying C call mutates the global
354    /// llama logger state.
355    ///
356    /// # Errors
357    ///
358    /// Returns [`FitError::Failure`] if no fitting allocation could be found, or
359    /// [`FitError::Error`] on a hard error (e.g. the model file could not be read).
360    pub fn fit_params(
361        mut self: Pin<&mut Self>,
362        model_path: &CStr,
363        cparams: &mut LlamaContextParams,
364        margins: &mut [usize],
365        n_ctx_min: u32,
366        log_level: llama_cpp_sys_2::ggml_log_level,
367    ) -> Result<FitResult, FitError> {
368        let max_devices = unsafe { llama_cpp_sys_2::llama_max_devices() };
369        let max_buft = unsafe { llama_cpp_sys_2::llama_max_tensor_buft_overrides() };
370
371        // Allocate tensor_split output buffer.
372        self.tensor_split.clear();
373        self.tensor_split.resize(max_devices, 0.0);
374
375        // Reset and resize buft_overrides for fit output (null-terminated).
376        self.buft_overrides.clear();
377        self.buft_overrides.resize(
378            max_buft + 1,
379            llama_cpp_sys_2::llama_model_tensor_buft_override {
380                pattern: std::ptr::null(),
381                buft: std::ptr::null_mut(),
382            },
383        );
384
385        // Clear pointers before the call — fit writes directly into the buffers above.
386        self.params.tensor_split = null::<f32>();
387        self.params.tensor_buft_overrides = null();
388
389        let status = unsafe {
390            llama_cpp_sys_2::llama_rs_fit_params(
391                model_path.as_ptr(),
392                &raw mut self.params,
393                &raw mut cparams.context_params,
394                self.tensor_split.as_mut_ptr(),
395                self.buft_overrides.as_mut_ptr(),
396                margins.as_mut_ptr(),
397                n_ctx_min,
398                log_level,
399            )
400        };
401
402        // llama_rs_fit_params returns common_params_fit_status: 0 = success, 1 = failure, 2 = error.
403        match status {
404            0 => {}
405            1 => return Err(FitError::Failure),
406            _ => return Err(FitError::Error),
407        }
408
409        // Wire the owned buffers into the raw params.
410        self.params.tensor_split = self.tensor_split.as_ptr();
411        self.params.tensor_buft_overrides = self.buft_overrides.as_ptr();
412
413        Ok(FitResult {
414            n_ctx: cparams.context_params.n_ctx,
415        })
416    }
417}
418
419impl LlamaModelParams {
420    /// Get the number of layers to offload to the GPU.
421    #[must_use]
422    pub fn n_gpu_layers(&self) -> i32 {
423        self.params.n_gpu_layers
424    }
425
426    /// The GPU that is used for scratch and small tensors
427    #[must_use]
428    pub fn main_gpu(&self) -> i32 {
429        self.params.main_gpu
430    }
431
432    /// only load the vocabulary, no weights
433    #[must_use]
434    pub fn vocab_only(&self) -> bool {
435        self.params.vocab_only
436    }
437
438    /// use mmap if possible
439    #[must_use]
440    pub fn use_mmap(&self) -> bool {
441        self.params.use_mmap
442    }
443
444    /// force system to keep model in RAM
445    #[must_use]
446    pub fn use_mlock(&self) -> bool {
447        self.params.use_mlock
448    }
449
450    /// get the split mode
451    ///
452    /// # Errors
453    /// Returns `LlamaSplitModeParseError` if the unknown split mode is encountered.
454    pub fn split_mode(&self) -> Result<LlamaSplitMode, LlamaSplitModeParseError> {
455        LlamaSplitMode::try_from(self.params.split_mode)
456    }
457
458    /// get the devices
459    #[must_use]
460    pub fn devices(&self) -> Vec<usize> {
461        let mut backend_devices = Vec::new();
462        for i in 0..unsafe { llama_cpp_sys_2::ggml_backend_dev_count() } {
463            let dev = unsafe { llama_cpp_sys_2::ggml_backend_dev_get(i) };
464            backend_devices.push(dev);
465        }
466        let mut devices = Vec::new();
467        for &dev in self.devices.iter() {
468            if dev.is_null() {
469                break;
470            }
471            if let Some((index, _)) = backend_devices
472                .iter()
473                .enumerate()
474                .find(|&(_i, &d)| d == dev)
475            {
476                devices.push(index);
477            }
478        }
479        devices
480    }
481
482    /// sets the number of gpu layers to offload to the GPU.
483    /// ```
484    /// # use llama_cpp_2::model::params::LlamaModelParams;
485    /// let params = LlamaModelParams::default();
486    /// let params = params.with_n_gpu_layers(1);
487    /// assert_eq!(params.n_gpu_layers(), 1);
488    /// ```
489    #[must_use]
490    pub fn with_n_gpu_layers(mut self, n_gpu_layers: u32) -> Self {
491        // The only way this conversion can fail is if u32 overflows the i32 - in which case we set
492        // to MAX
493        let n_gpu_layers = i32::try_from(n_gpu_layers).unwrap_or(i32::MAX);
494        self.params.n_gpu_layers = n_gpu_layers;
495        self
496    }
497
498    /// sets the main GPU
499    ///
500    /// To enable this option, you must set `split_mode` to `LlamaSplitMode::None` to enable single GPU mode.
501    #[must_use]
502    pub fn with_main_gpu(mut self, main_gpu: i32) -> Self {
503        self.params.main_gpu = main_gpu;
504        self
505    }
506
507    /// sets `vocab_only`
508    #[must_use]
509    pub fn with_vocab_only(mut self, vocab_only: bool) -> Self {
510        self.params.vocab_only = vocab_only;
511        self
512    }
513
514    /// sets `use_mmap`
515    #[must_use]
516    pub fn with_use_mmap(mut self, use_mmap: bool) -> Self {
517        self.params.use_mmap = use_mmap;
518        self
519    }
520
521    /// sets `use_mlock`
522    #[must_use]
523    pub fn with_use_mlock(mut self, use_mlock: bool) -> Self {
524        self.params.use_mlock = use_mlock;
525        self
526    }
527
528    /// sets `split_mode`
529    #[must_use]
530    pub fn with_split_mode(mut self, split_mode: LlamaSplitMode) -> Self {
531        self.params.split_mode = split_mode.into();
532        self
533    }
534
535    /// sets `devices`
536    ///
537    /// The devices are specified as indices that correspond to the ggml backend device indices.
538    ///
539    /// The maximum number of devices is 16.
540    ///
541    /// You don't need to specify CPU or ACCEL devices.
542    ///
543    /// # Errors
544    /// Returns `LlamaCppError::BackendDeviceNotFound` if any device index is invalid.
545    pub fn with_devices(mut self, devices: &[usize]) -> Result<Self, LlamaCppError> {
546        for dev in self.devices.iter_mut() {
547            *dev = std::ptr::null_mut();
548        }
549        // Check device count
550        let max_devices = crate::max_devices().min(LLAMA_CPP_MAX_DEVICES);
551        if devices.len() > max_devices {
552            return Err(LlamaCppError::MaxDevicesExceeded(max_devices));
553        }
554        for (i, &dev) in devices.iter().enumerate() {
555            if dev >= unsafe { llama_cpp_sys_2::ggml_backend_dev_count() } {
556                return Err(LlamaCppError::BackendDeviceNotFound(dev));
557            }
558            let backend_dev = unsafe { llama_cpp_sys_2::ggml_backend_dev_get(dev) };
559            self.devices[i] = backend_dev;
560        }
561        if self.devices.is_empty() {
562            self.params.devices = std::ptr::null_mut();
563        } else {
564            self.params.devices = self.devices.as_mut_ptr();
565        }
566        Ok(self)
567    }
568
569    /// Set `no_alloc`
570    ///
571    /// If this parameter is true, don't allocate memory for the tensor data
572    ///
573    /// You can't use `no_alloc` with `use_mmap`, so this also sets `use_mmap` to false.
574    #[must_use]
575    pub fn with_no_alloc(mut self, no_alloc: bool) -> Self {
576        self.params.no_alloc = no_alloc;
577        if no_alloc {
578            self = self.with_use_mmap(false);
579        }
580        self
581    }
582
583    /// Get `no_alloc`
584    ///
585    /// If this parameter is true, don't allocate memory for the tensor data
586    #[must_use]
587    pub fn no_alloc(&self) -> bool {
588        self.params.no_alloc
589    }
590
591    /// Sets a callback invoked during loading with progress in `0.0..=1.0`.
592    /// Returning `false` aborts the load (it then fails with `NullResult`).
593    #[must_use]
594    pub fn with_progress_callback<F: FnMut(f32) -> bool + 'static>(mut self, callback: F) -> Self {
595        unsafe extern "C" fn trampoline<F: FnMut(f32) -> bool>(
596            progress: f32,
597            user_data: *mut c_void,
598        ) -> bool {
599            let callback = unsafe { &mut *user_data.cast::<F>() };
600            callback(progress)
601        }
602
603        let mut callback = Box::new(callback);
604        self.params.progress_callback_user_data =
605            std::ptr::from_mut(&mut *callback).cast::<c_void>();
606        self.params.progress_callback = Some(trampoline::<F>);
607        self.progress_callback = Some(callback);
608        self
609    }
610}
611
612/// Default parameters for `LlamaModel`. (as defined in llama.cpp by `llama_model_default_params`)
613/// ```
614/// # use llama_cpp_2::model::params::LlamaModelParams;
615/// use llama_cpp_2::model::params::LlamaSplitMode;
616/// let params = LlamaModelParams::default();
617/// assert_eq!(params.n_gpu_layers(), -1, "n_gpu_layers should be -1");
618/// assert_eq!(params.main_gpu(), 0, "main_gpu should be 0");
619/// assert_eq!(params.vocab_only(), false, "vocab_only should be false");
620/// assert_eq!(params.use_mmap(), true, "use_mmap should be true");
621/// assert_eq!(params.use_mlock(), false, "use_mlock should be false");
622/// assert_eq!(params.split_mode(), Ok(LlamaSplitMode::Layer), "split_mode should be LAYER");
623/// assert_eq!(params.devices().len(), 0, "devices should be empty");
624/// assert_eq!(params.no_alloc(), false, "no_alloc should be false");
625/// ```
626impl Default for LlamaModelParams {
627    fn default() -> Self {
628        let default_params = unsafe { llama_cpp_sys_2::llama_model_default_params() };
629        LlamaModelParams {
630            params: default_params,
631            // push the next one to ensure we maintain the iterator invariant of ending with a 0
632            kv_overrides: vec![llama_cpp_sys_2::llama_model_kv_override {
633                key: [0; 128],
634                tag: 0,
635                __bindgen_anon_1: llama_cpp_sys_2::llama_model_kv_override__bindgen_ty_1 {
636                    val_i64: 0,
637                },
638            }],
639            buft_overrides: vec![llama_cpp_sys_2::llama_model_tensor_buft_override {
640                pattern: std::ptr::null(),
641                buft: std::ptr::null_mut(),
642            }],
643            devices: Box::pin([std::ptr::null_mut(); 16]),
644            tensor_split: Vec::new(),
645            progress_callback: None,
646        }
647    }
648}
649
650#[cfg(test)]
651mod tests {
652    use super::{LlamaModelParams, LlamaSplitMode};
653    use std::pin::pin;
654
655    #[test]
656    fn tensor_buft_override_patterns_empty_by_default() {
657        // Fresh params carry only the null-terminator entry, so no patterns are reported.
658        assert!(LlamaModelParams::default()
659            .tensor_buft_override_patterns()
660            .is_empty());
661    }
662
663    #[test]
664    fn tensor_buft_override_patterns_reads_back_added_override() {
665        // The getter is the read-only counterpart to the setter: the override added is reported and
666        // the trailing null terminator is skipped. This mirrors how `fit_params` populates the same
667        // buffer for the auto-fit's MoE expert offload (`add_cpu_moe_override` uses the same expert
668        // tensor pattern shape the fit emits).
669        let mut params = pin!(LlamaModelParams::default());
670        params.as_mut().add_cpu_moe_override();
671        assert_eq!(
672            params.tensor_buft_override_patterns(),
673            vec!["\\.ffn_(up|down|gate)_(ch|)exps".to_owned()],
674        );
675    }
676
677    #[test]
678    fn tensor_split_mode_round_trips() {
679        assert_eq!(
680            LlamaSplitMode::try_from(llama_cpp_sys_2::LLAMA_SPLIT_MODE_TENSOR),
681            Ok(LlamaSplitMode::Tensor)
682        );
683        assert_eq!(
684            u32::from(LlamaSplitMode::Tensor),
685            llama_cpp_sys_2::LLAMA_SPLIT_MODE_TENSOR as u32
686        );
687        assert_eq!(
688            i32::from(LlamaSplitMode::Tensor),
689            llama_cpp_sys_2::LLAMA_SPLIT_MODE_TENSOR as i32
690        );
691    }
692
693    #[test]
694    fn progress_callback_round_trips_and_can_abort() {
695        use super::LlamaModelParams;
696        use std::cell::Cell;
697        use std::rc::Rc;
698
699        let calls = Rc::new(Cell::new(0_u32));
700        let counter = Rc::clone(&calls);
701        let params = LlamaModelParams::default().with_progress_callback(move |_progress| {
702            counter.set(counter.get() + 1);
703            false
704        });
705
706        assert!(params.params.progress_callback.is_some());
707        assert!(!params.params.progress_callback_user_data.is_null());
708
709        let trampoline = params.params.progress_callback.unwrap();
710        let user_data = params.params.progress_callback_user_data;
711        let first = unsafe { trampoline(0.5, user_data) };
712        let second = unsafe { trampoline(1.0, user_data) };
713
714        assert!(!first && !second, "returning false signals an abort");
715        assert_eq!(calls.get(), 2);
716    }
717}
llama_cpp_2/model/params.rs

llama_cpp_2/model/
params.rs
