1use 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#[cfg(feature = "common")]
15#[derive(Debug, Clone)]
16pub struct FitResult {
17 pub n_ctx: u32,
19}
20
21#[cfg(feature = "common")]
23#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
24pub enum FitError {
25 #[error("could not find allocations that fit available memory")]
27 Failure,
28 #[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#[repr(i8)]
48#[derive(Copy, Clone, Debug, PartialEq, Eq)]
49pub enum LlamaSplitMode {
50 None = LLAMA_SPLIT_MODE_NONE,
52 Layer = LLAMA_SPLIT_MODE_LAYER,
54 Row = LLAMA_SPLIT_MODE_ROW,
56 Tensor = LLAMA_SPLIT_MODE_TENSOR,
58}
59
60#[derive(Debug, Clone, Copy, PartialEq, Eq)]
62pub struct LlamaSplitModeParseError(pub i32);
63
64impl 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
85impl 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
108impl 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
120impl 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
132impl Default for LlamaSplitMode {
134 fn default() -> Self {
135 LlamaSplitMode::Layer
136 }
137}
138
139pub const LLAMA_CPP_MAX_DEVICES: usize = 16;
144
145#[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 #[must_use]
184 pub fn kv_overrides<'a>(&'a self) -> KvOverrides<'a> {
185 KvOverrides::new(self)
186 }
187
188 #[allow(clippy::missing_panics_doc)] 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 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 self.params.kv_overrides = null();
232
233 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 self.params.kv_overrides = self.kv_overrides.as_ptr();
245
246 eprintln!("saved ptr: {:?}", self.params.kv_overrides);
247 }
248}
249
250impl LlamaModelParams {
251 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 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 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 self.params.tensor_buft_overrides = null();
279
280 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 self.params.tensor_buft_overrides = self.buft_overrides.as_ptr();
289 }
290}
291
292#[cfg(feature = "common")]
293impl LlamaModelParams {
294 pub fn fit_params(
331 mut self: Pin<&mut Self>,
332 model_path: &CStr,
333 cparams: &mut LlamaContextParams,
334 margins: &mut [usize],
335 n_ctx_min: u32,
336 log_level: llama_cpp_sys_2::ggml_log_level,
337 ) -> Result<FitResult, FitError> {
338 let max_devices = unsafe { llama_cpp_sys_2::llama_max_devices() };
339 let max_buft = unsafe { llama_cpp_sys_2::llama_max_tensor_buft_overrides() };
340
341 self.tensor_split.clear();
343 self.tensor_split.resize(max_devices, 0.0);
344
345 self.buft_overrides.clear();
347 self.buft_overrides.resize(
348 max_buft + 1,
349 llama_cpp_sys_2::llama_model_tensor_buft_override {
350 pattern: std::ptr::null(),
351 buft: std::ptr::null_mut(),
352 },
353 );
354
355 self.params.tensor_split = null::<f32>();
357 self.params.tensor_buft_overrides = null();
358
359 let status = unsafe {
360 llama_cpp_sys_2::llama_rs_fit_params(
361 model_path.as_ptr(),
362 &raw mut self.params,
363 &raw mut cparams.context_params,
364 self.tensor_split.as_mut_ptr(),
365 self.buft_overrides.as_mut_ptr(),
366 margins.as_mut_ptr(),
367 n_ctx_min,
368 log_level,
369 )
370 };
371
372 match status {
374 0 => {}
375 1 => return Err(FitError::Failure),
376 _ => return Err(FitError::Error),
377 }
378
379 self.params.tensor_split = self.tensor_split.as_ptr();
381 self.params.tensor_buft_overrides = self.buft_overrides.as_ptr();
382
383 Ok(FitResult {
384 n_ctx: cparams.context_params.n_ctx,
385 })
386 }
387}
388
389impl LlamaModelParams {
390 #[must_use]
392 pub fn n_gpu_layers(&self) -> i32 {
393 self.params.n_gpu_layers
394 }
395
396 #[must_use]
398 pub fn main_gpu(&self) -> i32 {
399 self.params.main_gpu
400 }
401
402 #[must_use]
404 pub fn vocab_only(&self) -> bool {
405 self.params.vocab_only
406 }
407
408 #[must_use]
410 pub fn use_mmap(&self) -> bool {
411 self.params.use_mmap
412 }
413
414 #[must_use]
416 pub fn use_mlock(&self) -> bool {
417 self.params.use_mlock
418 }
419
420 pub fn split_mode(&self) -> Result<LlamaSplitMode, LlamaSplitModeParseError> {
425 LlamaSplitMode::try_from(self.params.split_mode)
426 }
427
428 #[must_use]
430 pub fn devices(&self) -> Vec<usize> {
431 let mut backend_devices = Vec::new();
432 for i in 0..unsafe { llama_cpp_sys_2::ggml_backend_dev_count() } {
433 let dev = unsafe { llama_cpp_sys_2::ggml_backend_dev_get(i) };
434 backend_devices.push(dev);
435 }
436 let mut devices = Vec::new();
437 for &dev in self.devices.iter() {
438 if dev.is_null() {
439 break;
440 }
441 if let Some((index, _)) = backend_devices
442 .iter()
443 .enumerate()
444 .find(|&(_i, &d)| d == dev)
445 {
446 devices.push(index);
447 }
448 }
449 devices
450 }
451
452 #[must_use]
460 pub fn with_n_gpu_layers(mut self, n_gpu_layers: u32) -> Self {
461 let n_gpu_layers = i32::try_from(n_gpu_layers).unwrap_or(i32::MAX);
464 self.params.n_gpu_layers = n_gpu_layers;
465 self
466 }
467
468 #[must_use]
472 pub fn with_main_gpu(mut self, main_gpu: i32) -> Self {
473 self.params.main_gpu = main_gpu;
474 self
475 }
476
477 #[must_use]
479 pub fn with_vocab_only(mut self, vocab_only: bool) -> Self {
480 self.params.vocab_only = vocab_only;
481 self
482 }
483
484 #[must_use]
486 pub fn with_use_mmap(mut self, use_mmap: bool) -> Self {
487 self.params.use_mmap = use_mmap;
488 self
489 }
490
491 #[must_use]
493 pub fn with_use_mlock(mut self, use_mlock: bool) -> Self {
494 self.params.use_mlock = use_mlock;
495 self
496 }
497
498 #[must_use]
500 pub fn with_split_mode(mut self, split_mode: LlamaSplitMode) -> Self {
501 self.params.split_mode = split_mode.into();
502 self
503 }
504
505 pub fn with_devices(mut self, devices: &[usize]) -> Result<Self, LlamaCppError> {
516 for dev in self.devices.iter_mut() {
517 *dev = std::ptr::null_mut();
518 }
519 let max_devices = crate::max_devices().min(LLAMA_CPP_MAX_DEVICES);
521 if devices.len() > max_devices {
522 return Err(LlamaCppError::MaxDevicesExceeded(max_devices));
523 }
524 for (i, &dev) in devices.iter().enumerate() {
525 if dev >= unsafe { llama_cpp_sys_2::ggml_backend_dev_count() } {
526 return Err(LlamaCppError::BackendDeviceNotFound(dev));
527 }
528 let backend_dev = unsafe { llama_cpp_sys_2::ggml_backend_dev_get(dev) };
529 self.devices[i] = backend_dev;
530 }
531 if self.devices.is_empty() {
532 self.params.devices = std::ptr::null_mut();
533 } else {
534 self.params.devices = self.devices.as_mut_ptr();
535 }
536 Ok(self)
537 }
538
539 #[must_use]
545 pub fn with_no_alloc(mut self, no_alloc: bool) -> Self {
546 self.params.no_alloc = no_alloc;
547 if no_alloc {
548 self = self.with_use_mmap(false);
549 }
550 self
551 }
552
553 #[must_use]
557 pub fn no_alloc(&self) -> bool {
558 self.params.no_alloc
559 }
560
561 #[must_use]
564 pub fn with_progress_callback<F: FnMut(f32) -> bool + 'static>(mut self, callback: F) -> Self {
565 unsafe extern "C" fn trampoline<F: FnMut(f32) -> bool>(
566 progress: f32,
567 user_data: *mut c_void,
568 ) -> bool {
569 let callback = unsafe { &mut *user_data.cast::<F>() };
570 callback(progress)
571 }
572
573 let mut callback = Box::new(callback);
574 self.params.progress_callback_user_data = std::ptr::from_mut(&mut *callback).cast::<c_void>();
575 self.params.progress_callback = Some(trampoline::<F>);
576 self.progress_callback = Some(callback);
577 self
578 }
579}
580
581impl Default for LlamaModelParams {
596 fn default() -> Self {
597 let default_params = unsafe { llama_cpp_sys_2::llama_model_default_params() };
598 LlamaModelParams {
599 params: default_params,
600 kv_overrides: vec![llama_cpp_sys_2::llama_model_kv_override {
602 key: [0; 128],
603 tag: 0,
604 __bindgen_anon_1: llama_cpp_sys_2::llama_model_kv_override__bindgen_ty_1 {
605 val_i64: 0,
606 },
607 }],
608 buft_overrides: vec![llama_cpp_sys_2::llama_model_tensor_buft_override {
609 pattern: std::ptr::null(),
610 buft: std::ptr::null_mut(),
611 }],
612 devices: Box::pin([std::ptr::null_mut(); 16]),
613 tensor_split: Vec::new(),
614 progress_callback: None,
615 }
616 }
617}
618
619#[cfg(test)]
620mod tests {
621 use super::LlamaSplitMode;
622
623 #[test]
624 fn tensor_split_mode_round_trips() {
625 assert_eq!(
626 LlamaSplitMode::try_from(llama_cpp_sys_2::LLAMA_SPLIT_MODE_TENSOR),
627 Ok(LlamaSplitMode::Tensor)
628 );
629 assert_eq!(
630 u32::from(LlamaSplitMode::Tensor),
631 llama_cpp_sys_2::LLAMA_SPLIT_MODE_TENSOR as u32
632 );
633 assert_eq!(
634 i32::from(LlamaSplitMode::Tensor),
635 llama_cpp_sys_2::LLAMA_SPLIT_MODE_TENSOR as i32
636 );
637 }
638
639 #[test]
640 fn progress_callback_round_trips_and_can_abort() {
641 use super::LlamaModelParams;
642 use std::cell::Cell;
643 use std::rc::Rc;
644
645 let calls = Rc::new(Cell::new(0_u32));
646 let counter = Rc::clone(&calls);
647 let params = LlamaModelParams::default().with_progress_callback(move |_progress| {
648 counter.set(counter.get() + 1);
649 false
650 });
651
652 assert!(params.params.progress_callback.is_some());
653 assert!(!params.params.progress_callback_user_data.is_null());
654
655 let trampoline = params.params.progress_callback.unwrap();
656 let user_data = params.params.progress_callback_user_data;
657 let first = unsafe { trampoline(0.5, user_data) };
658 let second = unsafe { trampoline(1.0, user_data) };
659
660 assert!(!first && !second, "returning false signals an abort");
661 assert_eq!(calls.get(), 2);
662 }
663}