shimmy_llama_cpp_2/model/

params.rs

//! A safe wrapper around `llama_model_params`.

use crate::model::params::kv_overrides::KvOverrides;
use std::ffi::{c_char, CStr, CString};
use std::fmt::{Debug, Formatter};
use std::pin::Pin;
use std::ptr::null;

pub mod kv_overrides;

/// A safe wrapper around `llama_model_params`.
#[allow(clippy::module_name_repetitions)]
pub struct LlamaModelParams {
    pub(crate) params: shimmy_llama_cpp_sys_2::llama_model_params,
    kv_overrides: Vec<shimmy_llama_cpp_sys_2::llama_model_kv_override>,
    /// Storage for tensor buffer override patterns (keeps CStrings alive)
    tensor_override_patterns: Vec<CString>,
    /// Tensor buffer overrides (NULL-terminated array)
    tensor_overrides: Vec<shimmy_llama_cpp_sys_2::llama_model_tensor_buft_override>,
}

impl Debug for LlamaModelParams {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("LlamaModelParams")
            .field("n_gpu_layers", &self.params.n_gpu_layers)
            .field("main_gpu", &self.params.main_gpu)
            .field("vocab_only", &self.params.vocab_only)
            .field("use_mmap", &self.params.use_mmap)
            .field("use_mlock", &self.params.use_mlock)
            .field("kv_overrides", &"vec of kv_overrides")
            .finish()
    }
}

impl LlamaModelParams {
    /// See [`KvOverrides`]
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use llama_cpp_2::model::params::LlamaModelParams;
    /// let params = Box::pin(LlamaModelParams::default());
    /// let kv_overrides = params.kv_overrides();
    /// let count = kv_overrides.into_iter().count();
    /// assert_eq!(count, 0);
    /// ```
    #[must_use]
    pub fn kv_overrides(&self) -> KvOverrides {
        KvOverrides::new(self)
    }

    /// Appends a key-value override to the model parameters. It must be pinned as this creates a self-referential struct.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use std::ffi::{CStr, CString};
    /// use std::pin::pin;
    /// # use llama_cpp_2::model::params::LlamaModelParams;
    /// # use llama_cpp_2::model::params::kv_overrides::ParamOverrideValue;
    /// let mut params = pin!(LlamaModelParams::default());
    /// let key = CString::new("key").expect("CString::new failed");
    /// params.as_mut().append_kv_override(&key, ParamOverrideValue::Int(50));
    ///
    /// let kv_overrides = params.kv_overrides().into_iter().collect::<Vec<_>>();
    /// assert_eq!(kv_overrides.len(), 1);
    ///
    /// let (k, v) = &kv_overrides[0];
    /// assert_eq!(v, &ParamOverrideValue::Int(50));
    ///
    /// assert_eq!(k.to_bytes(), b"key", "expected key to be 'key', was {:?}", k);
    /// ```
    #[allow(clippy::missing_panics_doc)] // panics are just to enforce internal invariants, not user errors
    pub fn append_kv_override(
        mut self: Pin<&mut Self>,
        key: &CStr,
        value: kv_overrides::ParamOverrideValue,
    ) {
        let kv_override = self
            .kv_overrides
            .get_mut(0)
            .expect("kv_overrides did not have a next allocated");

        assert_eq!(kv_override.key[0], 0, "last kv_override was not empty");

        // There should be some way to do this without iterating over everything.
        for (i, &c) in key.to_bytes_with_nul().iter().enumerate() {
            kv_override.key[i] = c_char::try_from(c).expect("invalid character in key");
        }

        kv_override.tag = value.tag();
        kv_override.__bindgen_anon_1 = value.value();

        // set to null pointer for panic safety (as push may move the vector, invalidating the pointer)
        self.params.kv_overrides = null();

        // push the next one to ensure we maintain the iterator invariant of ending with a 0
        self.kv_overrides
            .push(shimmy_llama_cpp_sys_2::llama_model_kv_override {
                key: [0; 128],
                tag: 0,
                __bindgen_anon_1: shimmy_llama_cpp_sys_2::llama_model_kv_override__bindgen_ty_1 {
                    val_i64: 0,
                },
            });

        // set the pointer to the (potentially) new vector
        self.params.kv_overrides = self.kv_overrides.as_ptr();

        eprintln!("saved ptr: {:?}", self.params.kv_overrides);
    }
}

impl LlamaModelParams {
    /// Get the number of layers to offload to the GPU.
    #[must_use]
    pub fn n_gpu_layers(&self) -> i32 {
        self.params.n_gpu_layers
    }

    /// The GPU that is used for scratch and small tensors
    #[must_use]
    pub fn main_gpu(&self) -> i32 {
        self.params.main_gpu
    }

    /// only load the vocabulary, no weights
    #[must_use]
    pub fn vocab_only(&self) -> bool {
        self.params.vocab_only
    }

    /// use mmap if possible
    #[must_use]
    pub fn use_mmap(&self) -> bool {
        self.params.use_mmap
    }

    /// force system to keep model in RAM
    #[must_use]
    pub fn use_mlock(&self) -> bool {
        self.params.use_mlock
    }

    /// sets the number of gpu layers to offload to the GPU.
    /// ```
    /// # use llama_cpp_2::model::params::LlamaModelParams;
    /// let params = LlamaModelParams::default();
    /// let params = params.with_n_gpu_layers(1);
    /// assert_eq!(params.n_gpu_layers(), 1);
    /// ```
    #[must_use]
    pub fn with_n_gpu_layers(mut self, n_gpu_layers: u32) -> Self {
        // The only way this conversion can fail is if u32 overflows the i32 - in which case we set
        // to MAX
        let n_gpu_layers = i32::try_from(n_gpu_layers).unwrap_or(i32::MAX);
        self.params.n_gpu_layers = n_gpu_layers;
        self
    }

    /// sets the main GPU
    #[must_use]
    pub fn with_main_gpu(mut self, main_gpu: i32) -> Self {
        self.params.main_gpu = main_gpu;
        self
    }

    /// sets `vocab_only`
    #[must_use]
    pub fn with_vocab_only(mut self, vocab_only: bool) -> Self {
        self.params.vocab_only = vocab_only;
        self
    }

    /// sets `use_mlock`
    #[must_use]
    pub fn with_use_mlock(mut self, use_mlock: bool) -> Self {
        self.params.use_mlock = use_mlock;
        self
    }

    /// Offload ALL MoE (Mixture of Experts) expert tensors to CPU.
    ///
    /// This reduces VRAM usage for large MoE models (e.g., GPT-OSS, Qwen) by keeping
    /// expert tensors in system RAM while attention layers remain on GPU.
    ///
    /// Matches llama.cpp `--cpu-moe` flag behavior.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use llama_cpp_2::model::params::LlamaModelParams;
    /// let params = LlamaModelParams::default()
    ///     .with_n_gpu_layers(999)  // Offload all non-expert layers to GPU
    ///     .with_cpu_moe_all();     // But keep expert tensors in CPU
    /// ```
    #[must_use]
    pub fn with_cpu_moe_all(mut self) -> Self {
        self.push_tensor_override(r"\.ffn_(up|down|gate)_exps");
        self
    }

    /// Offload the first N MoE layers' expert tensors to CPU.
    ///
    /// This allows fine-grained control over VRAM usage by offloading only some expert layers.
    /// Typically used when you have limited VRAM and want to balance GPU/CPU usage.
    ///
    /// Matches llama.cpp `--n-cpu-moe N` flag behavior.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use llama_cpp_2::model::params::LlamaModelParams;
    /// let params = LlamaModelParams::default()
    ///     .with_n_gpu_layers(999)  // Offload layers to GPU
    ///     .with_n_cpu_moe(10);     // Except first 10 MoE expert layers -> CPU
    /// ```
    #[must_use]
    pub fn with_n_cpu_moe(mut self, n: usize) -> Self {
        for i in 0..n {
            let pattern = format!(r"blk\.{}\.ffn_(up|down|gate)_exps", i);
            self.push_tensor_override(&pattern);
        }
        self
    }

    /// Internal: Push a tensor buffer override pattern.
    ///
    /// Patterns are regex expressions matched against tensor names in the model.
    /// Matched tensors are allocated using CPU buffer type instead of GPU.
    fn push_tensor_override(&mut self, pattern: &str) {
        // Create and store CString to keep it alive
        let c_pattern = CString::new(pattern)
            .expect("pattern must not contain NUL bytes");
        
        self.tensor_override_patterns.push(c_pattern);
        
        // Get CPU buffer type
        let cpu_buft = unsafe { shimmy_llama_cpp_sys_2::ggml_backend_cpu_buffer_type() };
        
        // Create override entry pointing to the stored CString
        let override_entry = shimmy_llama_cpp_sys_2::llama_model_tensor_buft_override {
            pattern: self.tensor_override_patterns.last().unwrap().as_ptr(),
            buft: cpu_buft,
        };
        
        // Remove old NULL terminator if it exists
        if let Some(last) = self.tensor_overrides.last() {
            if last.pattern.is_null() {
                self.tensor_overrides.pop();
            }
        }
        
        // Add new entry
        self.tensor_overrides.push(override_entry);
        
        // Re-add NULL terminator (pattern=NULL signals end of array to C)
        self.tensor_overrides.push(shimmy_llama_cpp_sys_2::llama_model_tensor_buft_override {
            pattern: std::ptr::null(),
            buft: std::ptr::null_mut(),
        });
        
        // Update C params pointer
        self.params.tensor_buft_overrides = self.tensor_overrides.as_ptr();
    }
}

/// Default parameters for `LlamaModel`. (as defined in llama.cpp by `llama_model_default_params`)
/// ```
/// # use llama_cpp_2::model::params::LlamaModelParams;
/// let params = LlamaModelParams::default();
/// assert_eq!(params.n_gpu_layers(), 999, "n_gpu_layers should be 999");
/// assert_eq!(params.main_gpu(), 0, "main_gpu should be 0");
/// assert_eq!(params.vocab_only(), false, "vocab_only should be false");
/// assert_eq!(params.use_mmap(), true, "use_mmap should be true");
/// assert_eq!(params.use_mlock(), false, "use_mlock should be false");
/// ```
impl Default for LlamaModelParams {
    fn default() -> Self {
        let default_params = unsafe { shimmy_llama_cpp_sys_2::llama_model_default_params() };
        LlamaModelParams {
            params: default_params,
            // push the next one to ensure we maintain the iterator invariant of ending with a 0      
            kv_overrides: vec![shimmy_llama_cpp_sys_2::llama_model_kv_override {
                key: [0; 128],
                tag: 0,
                __bindgen_anon_1: shimmy_llama_cpp_sys_2::llama_model_kv_override__bindgen_ty_1 {
                    val_i64: 0,
                },
            }],
            tensor_override_patterns: Vec::new(),
            tensor_overrides: Vec::new(),
        }
    }
}