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    /// - `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                std::ptr::null_mut(),
410                log_level,
411            )
412        };
413
414        // llama_rs_fit_params returns common_params_fit_status: 0 = success, 1 = failure, 2 = error.
415        match status {
416            0 => {}
417            1 => return Err(FitError::Failure),
418            _ => return Err(FitError::Error),
419        }
420
421        // Wire the owned buffers into the raw params.
422        self.params.tensor_split = self.tensor_split.as_ptr();
423        self.params.tensor_buft_overrides = self.buft_overrides.as_ptr();
424
425        Ok(FitResult {
426            n_ctx: cparams.context_params.n_ctx,
427        })
428    }
429}
430
431impl LlamaModelParams {
432    /// Get the number of layers to offload to the GPU.
433    #[must_use]
434    pub fn n_gpu_layers(&self) -> i32 {
435        self.params.n_gpu_layers
436    }
437
438    /// The GPU that is used for scratch and small tensors
439    #[must_use]
440    pub fn main_gpu(&self) -> i32 {
441        self.params.main_gpu
442    }
443
444    /// only load the vocabulary, no weights
445    #[must_use]
446    pub fn vocab_only(&self) -> bool {
447        self.params.vocab_only
448    }
449
450    /// use mmap if possible
451    ///
452    /// `use_mmap`/`use_mlock` were replaced by a single `load_mode` enum
453    /// (`args: refactor mlock/mmap/directio into load-mode`, #20834); this
454    /// getter decodes the combined mode back into the two independent flags
455    /// this crate's public API still exposes.
456    #[must_use]
457    pub fn use_mmap(&self) -> bool {
458        // `LLAMA_LOAD_MODE_AUTO` generally translates to `use_mmap`.
459        // FIXME(madsmtm): Change this API to expose the load mode directly,
460        // instead of making it two methods.
461        matches!(
462            self.params.load_mode,
463            llama_cpp_sys_2::LLAMA_LOAD_MODE_AUTO
464                | llama_cpp_sys_2::LLAMA_LOAD_MODE_MMAP
465                | llama_cpp_sys_2::LLAMA_LOAD_MODE_MMAP_MLOCK
466        )
467    }
468
469    /// force system to keep model in RAM
470    #[must_use]
471    pub fn use_mlock(&self) -> bool {
472        matches!(
473            self.params.load_mode,
474            llama_cpp_sys_2::LLAMA_LOAD_MODE_MLOCK | llama_cpp_sys_2::LLAMA_LOAD_MODE_MMAP_MLOCK
475        )
476    }
477
478    /// get the split mode
479    ///
480    /// # Errors
481    /// Returns `LlamaSplitModeParseError` if the unknown split mode is encountered.
482    pub fn split_mode(&self) -> Result<LlamaSplitMode, LlamaSplitModeParseError> {
483        LlamaSplitMode::try_from(self.params.split_mode)
484    }
485
486    /// get the devices
487    #[must_use]
488    pub fn devices(&self) -> Vec<usize> {
489        let mut backend_devices = Vec::new();
490        for i in 0..unsafe { llama_cpp_sys_2::ggml_backend_dev_count() } {
491            let dev = unsafe { llama_cpp_sys_2::ggml_backend_dev_get(i) };
492            backend_devices.push(dev);
493        }
494        let mut devices = Vec::new();
495        for &dev in self.devices.iter() {
496            if dev.is_null() {
497                break;
498            }
499            if let Some((index, _)) = backend_devices
500                .iter()
501                .enumerate()
502                .find(|&(_i, &d)| d == dev)
503            {
504                devices.push(index);
505            }
506        }
507        devices
508    }
509
510    /// sets the number of gpu layers to offload to the GPU.
511    /// ```
512    /// # use llama_cpp_2::model::params::LlamaModelParams;
513    /// let params = LlamaModelParams::default();
514    /// let params = params.with_n_gpu_layers(1);
515    /// assert_eq!(params.n_gpu_layers(), 1);
516    /// ```
517    #[must_use]
518    pub fn with_n_gpu_layers(mut self, n_gpu_layers: u32) -> Self {
519        // The only way this conversion can fail is if u32 overflows the i32 - in which case we set
520        // to MAX
521        let n_gpu_layers = i32::try_from(n_gpu_layers).unwrap_or(i32::MAX);
522        self.params.n_gpu_layers = n_gpu_layers;
523        self
524    }
525
526    /// sets the main GPU
527    ///
528    /// To enable this option, you must set `split_mode` to `LlamaSplitMode::None` to enable single GPU mode.
529    #[must_use]
530    pub fn with_main_gpu(mut self, main_gpu: i32) -> Self {
531        self.params.main_gpu = main_gpu;
532        self
533    }
534
535    /// sets `vocab_only`
536    #[must_use]
537    pub fn with_vocab_only(mut self, vocab_only: bool) -> Self {
538        self.params.vocab_only = vocab_only;
539        self
540    }
541
542    /// sets `use_mmap`
543    #[must_use]
544    pub fn with_use_mmap(mut self, use_mmap: bool) -> Self {
545        self.params.load_mode = load_mode_from_flags(use_mmap, self.use_mlock());
546        self
547    }
548
549    /// sets `use_mlock`
550    #[must_use]
551    pub fn with_use_mlock(mut self, use_mlock: bool) -> Self {
552        self.params.load_mode = load_mode_from_flags(self.use_mmap(), use_mlock);
553        self
554    }
555
556    /// sets `split_mode`
557    #[must_use]
558    pub fn with_split_mode(mut self, split_mode: LlamaSplitMode) -> Self {
559        self.params.split_mode = split_mode.into();
560        self
561    }
562
563    /// sets `devices`
564    ///
565    /// The devices are specified as indices that correspond to the ggml backend device indices.
566    ///
567    /// The maximum number of devices is 16.
568    ///
569    /// You don't need to specify CPU or ACCEL devices.
570    ///
571    /// # Errors
572    /// Returns `LlamaCppError::BackendDeviceNotFound` if any device index is invalid.
573    pub fn with_devices(mut self, devices: &[usize]) -> Result<Self, LlamaCppError> {
574        for dev in self.devices.iter_mut() {
575            *dev = std::ptr::null_mut();
576        }
577        // Check device count
578        let max_devices = crate::max_devices().min(LLAMA_CPP_MAX_DEVICES);
579        if devices.len() > max_devices {
580            return Err(LlamaCppError::MaxDevicesExceeded(max_devices));
581        }
582        for (i, &dev) in devices.iter().enumerate() {
583            if dev >= unsafe { llama_cpp_sys_2::ggml_backend_dev_count() } {
584                return Err(LlamaCppError::BackendDeviceNotFound(dev));
585            }
586            let backend_dev = unsafe { llama_cpp_sys_2::ggml_backend_dev_get(dev) };
587            self.devices[i] = backend_dev;
588        }
589        if self.devices.is_empty() {
590            self.params.devices = std::ptr::null_mut();
591        } else {
592            self.params.devices = self.devices.as_mut_ptr();
593        }
594        Ok(self)
595    }
596
597    /// Set `no_alloc`
598    ///
599    /// If this parameter is true, don't allocate memory for the tensor data
600    ///
601    /// You can't use `no_alloc` with `use_mmap`, so this also sets `use_mmap` to false.
602    #[must_use]
603    pub fn with_no_alloc(mut self, no_alloc: bool) -> Self {
604        self.params.no_alloc = no_alloc;
605        if no_alloc {
606            self = self.with_use_mmap(false);
607        }
608        self
609    }
610
611    /// Get `no_alloc`
612    ///
613    /// If this parameter is true, don't allocate memory for the tensor data
614    #[must_use]
615    pub fn no_alloc(&self) -> bool {
616        self.params.no_alloc
617    }
618
619    /// Sets whether to load bundled multi-token prediction (MTP) tensors.
620    #[must_use]
621    pub fn with_load_mtp(mut self, load_mtp: bool) -> Self {
622        self.params.load_mtp = load_mtp;
623        self
624    }
625
626    /// Returns whether bundled multi-token prediction (MTP) tensors are loaded.
627    #[must_use]
628    pub fn load_mtp(&self) -> bool {
629        self.params.load_mtp
630    }
631
632    /// Sets a callback invoked during loading with progress in `0.0..=1.0`.
633    /// Returning `false` aborts the load (it then fails with `NullResult`).
634    #[must_use]
635    pub fn with_progress_callback<F: FnMut(f32) -> bool + 'static>(mut self, callback: F) -> Self {
636        unsafe extern "C" fn trampoline<F: FnMut(f32) -> bool>(
637            progress: f32,
638            user_data: *mut c_void,
639        ) -> bool {
640            let callback = unsafe { &mut *user_data.cast::<F>() };
641            callback(progress)
642        }
643
644        let mut callback = Box::new(callback);
645        self.params.progress_callback_user_data =
646            std::ptr::from_mut(&mut *callback).cast::<c_void>();
647        self.params.progress_callback = Some(trampoline::<F>);
648        self.progress_callback = Some(callback);
649        self
650    }
651}
652
653/// Default parameters for `LlamaModel`. (as defined in llama.cpp by `llama_model_default_params`)
654/// ```
655/// # use llama_cpp_2::model::params::LlamaModelParams;
656/// use llama_cpp_2::model::params::LlamaSplitMode;
657/// let params = LlamaModelParams::default();
658/// assert_eq!(params.n_gpu_layers(), -1, "n_gpu_layers should be -1");
659/// assert_eq!(params.main_gpu(), 0, "main_gpu should be 0");
660/// assert_eq!(params.vocab_only(), false, "vocab_only should be false");
661/// assert_eq!(params.use_mmap(), true, "use_mmap should be true");
662/// assert_eq!(params.use_mlock(), false, "use_mlock should be false");
663/// assert_eq!(params.split_mode(), Ok(LlamaSplitMode::Layer), "split_mode should be LAYER");
664/// assert_eq!(params.devices().len(), 0, "devices should be empty");
665/// assert_eq!(params.no_alloc(), false, "no_alloc should be false");
666/// assert_eq!(params.load_mtp(), false, "load_mtp should be false");
667/// ```
668impl Default for LlamaModelParams {
669    fn default() -> Self {
670        let default_params = unsafe { llama_cpp_sys_2::llama_model_default_params() };
671        LlamaModelParams {
672            params: default_params,
673            // push the next one to ensure we maintain the iterator invariant of ending with a 0
674            kv_overrides: vec![llama_cpp_sys_2::llama_model_kv_override {
675                key: [0; 128],
676                tag: 0,
677                __bindgen_anon_1: llama_cpp_sys_2::llama_model_kv_override__bindgen_ty_1 {
678                    val_i64: 0,
679                },
680            }],
681            buft_overrides: vec![llama_cpp_sys_2::llama_model_tensor_buft_override {
682                pattern: std::ptr::null(),
683                buft: std::ptr::null_mut(),
684            }],
685            devices: Box::pin([std::ptr::null_mut(); 16]),
686            tensor_split: Vec::new(),
687            progress_callback: None,
688        }
689    }
690}
691
692#[cfg(test)]
693mod tests {
694    use super::kv_overrides::ParamOverrideValue;
695    use super::{LlamaModelParams, LlamaSplitMode};
696    use std::pin::pin;
697
698    #[test]
699    fn tensor_buft_override_patterns_empty_by_default() {
700        // Fresh params carry only the null-terminator entry, so no patterns are reported.
701        assert!(LlamaModelParams::default()
702            .tensor_buft_override_patterns()
703            .is_empty());
704    }
705
706    #[test]
707    fn tensor_buft_override_patterns_reads_back_added_override() {
708        // The getter is the read-only counterpart to the setter: the override added is reported and
709        // the trailing null terminator is skipped. This mirrors how `fit_params` populates the same
710        // buffer for the auto-fit's MoE expert offload (`add_cpu_moe_override` uses the same expert
711        // tensor pattern shape the fit emits).
712        let mut params = pin!(LlamaModelParams::default());
713        params.as_mut().add_cpu_moe_override();
714        assert_eq!(
715            params.tensor_buft_override_patterns(),
716            vec!["\\.ffn_(up|down|gate)_(ch|)exps".to_owned()],
717        );
718    }
719
720    #[test]
721    fn tensor_buft_override_appends_on_second_call() {
722        // add_cpu_buft_override is documented as appending. After the first call the
723        // vector is [filled, null-terminator]; writing the second override into the
724        // terminator (then pushing a new one) is what keeps the list null-terminated.
725        let mut params = pin!(LlamaModelParams::default());
726        params
727            .as_mut()
728            .add_cpu_buft_override(c"blk\\.0\\.ffn_(up|down|gate)_(ch|)exps");
729        params
730            .as_mut()
731            .add_cpu_buft_override(c"blk\\.1\\.ffn_(up|down|gate)_(ch|)exps");
732
733        assert_eq!(
734            params.tensor_buft_override_patterns(),
735            vec![
736                "blk\\.0\\.ffn_(up|down|gate)_(ch|)exps".to_owned(),
737                "blk\\.1\\.ffn_(up|down|gate)_(ch|)exps".to_owned(),
738            ],
739        );
740    }
741
742    #[test]
743    fn kv_override_appends_on_second_call() {
744        // Same [filled, null-terminator] layout as the buffer overrides: the second call has to
745        // write into the terminator, not back into slot 0.
746        let mut params = pin!(LlamaModelParams::default());
747        params
748            .as_mut()
749            .append_kv_override(c"first", ParamOverrideValue::Int(1));
750        params
751            .as_mut()
752            .append_kv_override(c"second", ParamOverrideValue::Int(2));
753
754        let kv_overrides = params.kv_overrides().into_iter().collect::<Vec<_>>();
755        assert_eq!(kv_overrides.len(), 2);
756        assert_eq!(kv_overrides[0].0.to_bytes(), b"first");
757        assert_eq!(kv_overrides[1].0.to_bytes(), b"second");
758        assert_eq!(kv_overrides[1].1, ParamOverrideValue::Int(2));
759    }
760
761    #[test]
762    fn tensor_split_mode_round_trips() {
763        assert_eq!(
764            LlamaSplitMode::try_from(llama_cpp_sys_2::LLAMA_SPLIT_MODE_TENSOR),
765            Ok(LlamaSplitMode::Tensor)
766        );
767        assert_eq!(
768            u32::from(LlamaSplitMode::Tensor),
769            llama_cpp_sys_2::LLAMA_SPLIT_MODE_TENSOR as u32
770        );
771        assert_eq!(
772            i32::from(LlamaSplitMode::Tensor),
773            llama_cpp_sys_2::LLAMA_SPLIT_MODE_TENSOR as i32
774        );
775    }
776
777    #[test]
778    fn load_mtp_round_trips() {
779        let params = LlamaModelParams::default();
780        assert!(!params.load_mtp());
781
782        let params = params.with_load_mtp(true);
783        assert!(params.load_mtp());
784
785        let params = params.with_load_mtp(false);
786        assert!(!params.load_mtp());
787    }
788
789    #[test]
790    fn progress_callback_round_trips_and_can_abort() {
791        use super::LlamaModelParams;
792        use std::cell::Cell;
793        use std::rc::Rc;
794
795        let calls = Rc::new(Cell::new(0_u32));
796        let counter = Rc::clone(&calls);
797        let params = LlamaModelParams::default().with_progress_callback(move |_progress| {
798            counter.set(counter.get() + 1);
799            false
800        });
801
802        assert!(params.params.progress_callback.is_some());
803        assert!(!params.params.progress_callback_user_data.is_null());
804
805        let trampoline = params.params.progress_callback.unwrap();
806        let user_data = params.params.progress_callback_user_data;
807        let first = unsafe { trampoline(0.5, user_data) };
808        let second = unsafe { trampoline(1.0, user_data) };
809
810        assert!(!first && !second, "returning false signals an abort");
811        assert_eq!(calls.get(), 2);
812    }
813}
llama_cpp_2/model/params.rs

llama_cpp_2/model/
params.rs
