llama_cpp_4/context/
memory_breakdown.rs

1//! Per-buffer-type memory usage reported by llama.cpp.
2//!
3//! [`MemoryBreakdownEntry`] values are produced by
4//! [`crate::context::LlamaContext::memory_breakdown`] and classify bytes into
5//! model weights, KV / recurrent cache, and temporary compute buffers for each
6//! ggml backend buffer type (e.g. `CUDA0`, `Metal`, `Host`).
7//!
8//! # Examples
9//!
10//! ```no_run
11//! use llama_cpp_4::prelude::*;
12//!
13//! fn main() {
14//!     let backend = LlamaBackend::init().unwrap();
15//!     let model = LlamaModel::load_from_file(&backend, "model.gguf", &LlamaModelParams::default()).unwrap();
16//!     let ctx = model.new_context(&backend, LlamaContextParams::default()).unwrap();
17//!     for entry in ctx.memory_breakdown() {
18//!         println!("{}: {} bytes total", entry.buft_name, entry.total());
19//!     }
20//! }
21//! ```
22
23use std::ffi::CStr;
24
25/// Memory attributed to a single backend buffer type (e.g. CUDA0, Host).
26#[derive(Debug, Clone, PartialEq, Eq)]
27pub struct MemoryBreakdownEntry {
28    /// Human-readable buffer-type name from ggml.
29    pub buft_name: String,
30    /// Bytes used by model weights on this buffer type.
31    pub model: usize,
32    /// Bytes used by the KV / recurrent context cache.
33    pub context: usize,
34    /// Bytes used by temporary compute buffers.
35    pub compute: usize,
36}
37
38impl MemoryBreakdownEntry {
39    /// Sum of model, context, and compute bytes.
40    #[must_use]
41    pub fn total(&self) -> usize {
42        self.model + self.context + self.compute
43    }
44}
45
46fn raw_entry_to_rust(
47    entry: &llama_cpp_sys_4::llama_memory_breakdown_entry,
48) -> MemoryBreakdownEntry {
49    // SAFETY: `buft_name` is a fixed-size `c_char` array inside the entry, so
50    // its own pointer and length describe exactly that array.
51    let bytes: &[u8] = unsafe {
52        std::slice::from_raw_parts(entry.buft_name.as_ptr().cast(), entry.buft_name.len())
53    };
54    let name = CStr::from_bytes_until_nul(bytes)
55        .map(|c| c.to_string_lossy().into_owned())
56        .unwrap_or_default();
57    MemoryBreakdownEntry {
58        buft_name: name,
59        model: entry.model,
60        context: entry.context,
61        compute: entry.compute,
62    }
63}
64
65/// Collect memory breakdown entries for a live context.
66///
67/// Wraps the `ext_shim` helper around `llama_get_memory_breakdown`. Grows the
68/// output buffer until every entry fits. Returns an empty vector when the
69/// context pointer is invalid or no buffers are registered yet.
70///
71/// Prefer [`crate::context::LlamaContext::memory_breakdown`] in application code.
72#[must_use]
73pub(crate) fn collect_memory_breakdown(
74    ctx: *const llama_cpp_sys_4::llama_context,
75) -> Vec<MemoryBreakdownEntry> {
76    if ctx.is_null() {
77        return Vec::new();
78    }
79
80    let mut capacity = 16usize;
81    loop {
82        let mut raw = vec![
83            llama_cpp_sys_4::llama_memory_breakdown_entry {
84                buft_name: [0; 128],
85                model: 0,
86                context: 0,
87                compute: 0,
88            };
89            capacity
90        ];
91
92        // SAFETY: `ctx` is a live context, and `raw` is a Rust buffer whose
93        // capacity is passed alongside its pointer, so the shim cannot overrun it.
94        let n = unsafe {
95            llama_cpp_sys_4::llama_memory_breakdown_collect(ctx, raw.as_mut_ptr(), capacity)
96        };
97
98        if n < capacity {
99            return raw
100                .into_iter()
101                .take(n)
102                .map(|e| raw_entry_to_rust(&e))
103                .collect();
104        }
105
106        capacity = capacity.saturating_mul(2);
107        if capacity > 4096 {
108            return raw.into_iter().map(|e| raw_entry_to_rust(&e)).collect();
109        }
110    }
111}
llama_cpp_4/context/memory_breakdown.rs

llama_cpp_4/context/
memory_breakdown.rs
