llama_cpp_4/context/
tensor_capture.rs

1//! Capture intermediate tensor outputs during [`crate::LlamaContext::decode`].
2//!
3//! llama.cpp builds a computation graph for each forward pass. Every node has a
4//! string name — for transformer blocks the layer output is typically
5//! `"l_out-{N}"` (e.g. `"l_out-13"`), attention norms are `"attn_norm-{N}"`, and
6//! the final norm is `"result_norm"`.
7//!
8//! The graph evaluation callback (`cb_eval`) runs in two phases for each node:
9//!
10//! | Phase | `ask` | Behaviour |
11//! |---|---|---|
12//! | Ask | `true` | Return `true` to request a copy of this tensor's data. |
13//! | Data | `false` | Tensor is computed; data is copied via `ggml_backend_tensor_get`. |
14//!
15//! [`TensorCapture`] implements that callback and stores matching tensors in a
16//! [`HashMap`] you can read after `decode()` finishes.
17//!
18//! # Typical use cases
19//!
20//! - **Layer probing** — inspect hidden states at specific depths.
21//! - **EAGLE / distillation** — read draft-model anchor layers (see `examples/eagle`).
22//! - **Debugging** — dump norms or attention outputs with [`TensorCapture::for_prefix`].
23//!
24//! # Setup
25//!
26//! 1. Build a [`TensorCapture`] with the filter you need ([`TensorCapture::for_layers`]
27//!    is the common case).
28//! 2. Pass it to the legacy unsafe
29//!    [`LlamaContextParams::with_tensor_capture`](crate::LlamaContextParams::with_tensor_capture).
30//!    The capture must remain at a stable address and outlive the
31//!    [`LlamaContext`](crate::LlamaContext). Prefer the owned, checked
32//!    [`TensorTransactions`](crate::TensorTransactions) API in new code.
33//! 3. Run [`LlamaContext::decode`](crate::LlamaContext::decode) as usual.
34//! 4. Read [`CapturedTensor`] values via [`TensorCapture::get_layer`],
35//!    [`TensorCapture::get`], or [`TensorCapture::iter`].
36//!
37//! Call [`TensorCapture::clear`](crate::TensorCapture::clear) before reusing the same capture on another batch.
38//!
39//! # Example
40//!
41//! ```no_run
42//! use llama_cpp_4::prelude::*;
43//! use std::num::NonZeroU32;
44//!
45//! fn main() {
46//!     let backend = LlamaBackend::init().unwrap();
47//!     let model = LlamaModel::load_from_file(
48//!         &backend,
49//!         "model.gguf",
50//!         &LlamaModelParams::default(),
51//!     )
52//!     .unwrap();
53//!
54//!     let mut capture = TensorCapture::for_layers(&[13, 20, 27]);
55//!     let ctx_params = unsafe {
56//!         LlamaContextParams::default()
57//!             .with_n_ctx(NonZeroU32::new(512))
58//!             .with_tensor_capture(&mut capture)
59//!     };
60//!     let mut ctx = model.new_context(&backend, ctx_params).unwrap();
61//!
62//!     let tokens = model.str_to_token("Hello", AddBos::Always).unwrap();
63//!     let mut batch = LlamaBatch::new(512, 1);
64//!     for (i, &tok) in tokens.iter().enumerate() {
65//!         batch
66//!             .add(tok, i as i32, &[0], i == tokens.len() - 1)
67//!             .unwrap();
68//!     }
69//!     ctx.decode(&mut batch).unwrap();
70//!
71//!     for &layer in &[13, 20, 27] {
72//!         if let Some(t) = capture.get_layer(layer) {
73//!             println!(
74//!                 "l_out-{layer}: {} tokens × {} dims",
75//!                 t.n_tokens(),
76//!                 t.n_embd()
77//!             );
78//!             if let Some(vec) = t.token_embedding(0) {
79//!                 println!("  first token, first 3 dims: {:?}", &vec[..3.min(vec.len())]);
80//!             }
81//!         }
82//!     }
83//! }
84//! ```
85//!
86//! # Tensor layout
87//!
88//! Each [`CapturedTensor`] stores a flat `f32` buffer with
89//! `data[token_idx * n_embd + dim_idx]` (ggml row-major: `ne0` = embedding dim,
90//! `ne1` = token count). Use [`CapturedTensor::token_embedding`] to slice one row.
91
92use std::collections::HashMap;
93
94/// A single tensor copied out of the decode graph.
95///
96/// Produced by [`TensorCapture`] after a successful [`crate::LlamaContext::decode`].
97/// For layer outputs (`"l_out-N"`), [`Self::layer`] is set to `N`.
98#[derive(Debug, Clone)]
99pub struct CapturedTensor {
100    /// Graph node name (e.g. `"l_out-13"`, `"result_norm"`).
101    pub name: String,
102    /// Layer index when `name` is `"l_out-{N}"`, otherwise `None`.
103    pub layer: Option<usize>,
104    /// First dimension (typically `n_embd` / hidden size).
105    pub ne0: usize,
106    /// Second dimension (typically number of tokens in the batch position).
107    pub ne1: usize,
108    /// Flattened `ne0 * ne1` values in ggml row-major order.
109    ///
110    /// Index as `data[token_idx * ne0 + dim_idx]`.
111    pub data: Vec<f32>,
112}
113
114impl CapturedTensor {
115    /// Number of embedding dimensions (alias for [`Self::ne0`]).
116    #[inline]
117    #[must_use]
118    pub fn n_embd(&self) -> usize {
119        self.ne0
120    }
121
122    /// Number of token positions (alias for [`Self::ne1`]).
123    #[inline]
124    #[must_use]
125    pub fn n_tokens(&self) -> usize {
126        self.ne1
127    }
128
129    /// Hidden-state vector for one token index.
130    ///
131    /// Returns `None` when `token_idx >= n_tokens()`.
132    #[must_use]
133    pub fn token_embedding(&self, token_idx: usize) -> Option<&[f32]> {
134        if token_idx >= self.ne1 {
135            return None;
136        }
137        let start = token_idx * self.ne0;
138        Some(&self.data[start..start + self.ne0])
139    }
140}
141
142/// Strategy for selecting which tensors to capture.
143#[derive(Debug, Clone)]
144enum CaptureFilter {
145    /// `"l_out-{N}"` for each listed layer index `N`.
146    Layers(Vec<usize>),
147    /// Exact graph node names.
148    Names(Vec<String>),
149    /// Names starting with a prefix (e.g. `"attn_out"`).
150    Prefix(String),
151    /// Every node (can be very large — debug only).
152    All,
153}
154
155/// Captures intermediate tensors during [`crate::LlamaContext::decode`].
156///
157/// Attach with the legacy unsafe
158/// [`LlamaContextParams::with_tensor_capture`](crate::LlamaContextParams::with_tensor_capture)
159/// before creating the context. The same instance can be reused across decodes
160/// if you call [`Self::clear`] between passes.
161///
162/// # Lifetime
163///
164/// The capture must remain at a stable address and outlive the
165/// [`crate::LlamaContext`] it is wired into. This requirement is not enforced
166/// by the returned params; callers must uphold the unsafe constructor's
167/// contract. Prefer [`crate::TensorTransactions`], whose state is pinned and
168/// owned by the context.
169pub struct TensorCapture {
170    filter: CaptureFilter,
171    captured: HashMap<String, CapturedTensor>,
172}
173
174impl std::fmt::Debug for TensorCapture {
175    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
176        f.debug_struct("TensorCapture")
177            .field("filter", &self.filter)
178            .field("captured_count", &self.captured.len())
179            .field("captured_keys", &self.captured.keys().collect::<Vec<_>>())
180            .finish()
181    }
182}
183
184impl TensorCapture {
185    /// Capture transformer layer outputs `"l_out-{N}"` for the given indices.
186    ///
187    /// This is the usual choice for hidden-state extraction. EAGLE-3 draft models
188    /// often use three layers at ~50%, 75%, and 100% depth — e.g. `[13, 20, 27]`
189    /// on a 28-layer model.
190    #[must_use]
191    pub fn for_layers(layer_indices: &[usize]) -> Self {
192        Self {
193            filter: CaptureFilter::Layers(layer_indices.to_vec()),
194            captured: HashMap::new(),
195        }
196    }
197
198    /// Capture tensors whose graph names match exactly.
199    ///
200    /// Example names: `"result_norm"`, `"l_out-27"`.
201    #[must_use]
202    pub fn for_names(names: &[&str]) -> Self {
203        Self {
204            filter: CaptureFilter::Names(
205                names.iter().map(std::string::ToString::to_string).collect(),
206            ),
207            captured: HashMap::new(),
208        }
209    }
210
211    /// Capture every tensor whose name starts with `prefix`.
212    ///
213    /// Useful for families like `"attn_out-*"` or `"attn_norm-*"`.
214    #[must_use]
215    pub fn for_prefix(prefix: &str) -> Self {
216        Self {
217            filter: CaptureFilter::Prefix(prefix.to_string()),
218            captured: HashMap::new(),
219        }
220    }
221
222    /// Capture **all** graph nodes.
223    ///
224    /// Warning: memory use scales with model size and sequence length. Prefer
225    /// [`Self::for_layers`] or [`Self::for_names`] in production code.
226    #[must_use]
227    pub fn all() -> Self {
228        Self {
229            filter: CaptureFilter::All,
230            captured: HashMap::new(),
231        }
232    }
233
234    /// Drop captured tensors but keep the filter (safe to call before another decode).
235    pub fn clear(&mut self) {
236        self.captured.clear();
237    }
238
239    /// Lookup by full graph name (e.g. `"l_out-13"`).
240    #[must_use]
241    pub fn get(&self, name: &str) -> Option<&CapturedTensor> {
242        self.captured.get(name)
243    }
244
245    /// Lookup a layer output (`"l_out-{layer_idx}"`).
246    #[must_use]
247    pub fn get_layer(&self, layer_idx: usize) -> Option<&CapturedTensor> {
248        self.captured.get(&format!("l_out-{layer_idx}"))
249    }
250
251    /// Whether `"l_out-{layer_idx}"` was captured in the last decode.
252    #[must_use]
253    pub fn has_layer(&self, layer_idx: usize) -> bool {
254        self.captured.contains_key(&format!("l_out-{layer_idx}"))
255    }
256
257    /// Number of tensors stored from the most recent decode.
258    #[must_use]
259    pub fn len(&self) -> usize {
260        self.captured.len()
261    }
262
263    /// `true` when [`Self::len`] is zero.
264    #[must_use]
265    pub fn is_empty(&self) -> bool {
266        self.captured.is_empty()
267    }
268
269    /// Iterate `(name, tensor)` pairs from the last decode.
270    pub fn iter(&self) -> impl Iterator<Item = (&str, &CapturedTensor)> {
271        self.captured.iter().map(|(k, v)| (k.as_str(), v))
272    }
273
274    /// Sorted layer indices present among captured `"l_out-*"` tensors.
275    #[must_use]
276    pub fn captured_layers(&self) -> Vec<usize> {
277        let mut layers: Vec<usize> = self.captured.values().filter_map(|ct| ct.layer).collect();
278        layers.sort_unstable();
279        layers.dedup();
280        layers
281    }
282
283    fn matches(&self, name: &str) -> bool {
284        match &self.filter {
285            CaptureFilter::Layers(indices) => {
286                if let Some(suffix) = name.strip_prefix("l_out-") {
287                    if let Ok(idx) = suffix.parse::<usize>() {
288                        return indices.contains(&idx);
289                    }
290                }
291                false
292            }
293            CaptureFilter::Names(names) => names.iter().any(|n| n == name),
294            CaptureFilter::Prefix(prefix) => name.starts_with(prefix.as_str()),
295            CaptureFilter::All => true,
296        }
297    }
298
299    fn store(&mut self, name: String, ne0: usize, ne1: usize, data: Vec<f32>) {
300        let layer = name
301            .strip_prefix("l_out-")
302            .and_then(|s| s.parse::<usize>().ok());
303
304        self.captured.insert(
305            name.clone(),
306            CapturedTensor {
307                name,
308                layer,
309                ne0,
310                ne1,
311                data,
312            },
313        );
314    }
315}
316
317/// `cb_eval` callback installed by [`LlamaContextParams::with_tensor_capture`](crate::LlamaContextParams::with_tensor_capture).
318///
319/// # Safety
320///
321/// `user_data` must point to a live [`TensorCapture`] for the context lifetime.
322pub(crate) unsafe extern "C" fn tensor_capture_callback(
323    t: *mut llama_cpp_sys_4::ggml_tensor,
324    ask: bool,
325    user_data: *mut std::ffi::c_void,
326) -> bool {
327    if t.is_null() || user_data.is_null() {
328        return false;
329    }
330
331    let name_bytes = &(*t).name;
332    let len = name_bytes
333        .iter()
334        .position(|&b| b == 0)
335        .unwrap_or(name_bytes.len());
336    let name = std::str::from_utf8_unchecked(std::slice::from_raw_parts(
337        name_bytes.as_ptr().cast::<u8>(),
338        len,
339    ));
340
341    let state = &mut *user_data.cast::<TensorCapture>();
342
343    if !state.matches(name) {
344        return false;
345    }
346
347    if ask {
348        return true;
349    }
350
351    let ne0 = usize::try_from((*t).ne[0]).expect("tensor ne[0] must be non-negative");
352    let ne1 = usize::try_from((*t).ne[1]).expect("tensor ne[1] must be non-negative");
353    let n_elements = ne0 * ne1;
354
355    let mut buf = vec![0f32; n_elements];
356    llama_cpp_sys_4::ggml_backend_tensor_get(
357        t,
358        buf.as_mut_ptr().cast::<std::ffi::c_void>(),
359        0,
360        n_elements * std::mem::size_of::<f32>(),
361    );
362
363    state.store(name.to_string(), ne0, ne1, buf);
364
365    true
366}
367
368#[cfg(test)]
369mod tests {
370    use super::*;
371
372    #[test]
373    fn test_for_layers_matching() {
374        let capture = TensorCapture::for_layers(&[13, 20, 27]);
375        assert!(capture.matches("l_out-13"));
376        assert!(capture.matches("l_out-20"));
377        assert!(capture.matches("l_out-27"));
378        assert!(!capture.matches("l_out-0"));
379        assert!(!capture.matches("l_out-14"));
380        assert!(!capture.matches("attn_norm-13"));
381        assert!(!capture.matches("result_norm"));
382    }
383
384    #[test]
385    fn test_for_names_matching() {
386        let capture = TensorCapture::for_names(&["result_norm", "l_out-27"]);
387        assert!(capture.matches("result_norm"));
388        assert!(capture.matches("l_out-27"));
389        assert!(!capture.matches("l_out-13"));
390        assert!(!capture.matches("result_output"));
391    }
392
393    #[test]
394    fn test_for_prefix_matching() {
395        let capture = TensorCapture::for_prefix("attn_out");
396        assert!(capture.matches("attn_out-0"));
397        assert!(capture.matches("attn_out-27"));
398        assert!(!capture.matches("attn_norm-0"));
399        assert!(!capture.matches("l_out-0"));
400    }
401
402    #[test]
403    fn test_all_matching() {
404        let capture = TensorCapture::all();
405        assert!(capture.matches("l_out-13"));
406        assert!(capture.matches("result_norm"));
407        assert!(capture.matches("anything"));
408    }
409
410    #[test]
411    fn test_store_and_get() {
412        let mut capture = TensorCapture::for_layers(&[13]);
413        let data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
414        capture.store("l_out-13".to_string(), 3, 2, data.clone());
415
416        assert_eq!(capture.len(), 1);
417        assert!(!capture.is_empty());
418
419        let ct = capture.get("l_out-13").unwrap();
420        assert_eq!(ct.name, "l_out-13");
421        assert_eq!(ct.layer, Some(13));
422        assert_eq!(ct.n_embd(), 3);
423        assert_eq!(ct.n_tokens(), 2);
424        assert_eq!(ct.data, data);
425
426        let ct2 = capture.get_layer(13).unwrap();
427        assert_eq!(ct2.name, ct.name);
428        assert!(capture.has_layer(13));
429        assert!(!capture.has_layer(14));
430    }
431
432    #[test]
433    fn test_token_embedding() {
434        let mut capture = TensorCapture::for_layers(&[5]);
435        let data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
436        capture.store("l_out-5".to_string(), 3, 2, data);
437
438        let ct = capture.get_layer(5).unwrap();
439        assert_eq!(ct.token_embedding(0), Some(&[1.0, 2.0, 3.0][..]));
440        assert_eq!(ct.token_embedding(1), Some(&[4.0, 5.0, 6.0][..]));
441        assert_eq!(ct.token_embedding(2), None);
442    }
443
444    #[test]
445    fn test_captured_layers() {
446        let mut capture = TensorCapture::for_layers(&[5, 10, 20]);
447        capture.store("l_out-10".to_string(), 2, 1, vec![0.0, 0.0]);
448        capture.store("l_out-5".to_string(), 2, 1, vec![0.0, 0.0]);
449        assert_eq!(capture.captured_layers(), vec![5, 10]);
450    }
451
452    #[test]
453    fn test_clear() {
454        let mut capture = TensorCapture::for_layers(&[5]);
455        capture.store("l_out-5".to_string(), 2, 1, vec![0.0, 0.0]);
456        assert_eq!(capture.len(), 1);
457        capture.clear();
458        assert_eq!(capture.len(), 0);
459        assert!(capture.is_empty());
460    }
461
462    #[test]
463    fn test_non_layer_tensor() {
464        let mut capture = TensorCapture::for_names(&["result_norm"]);
465        capture.store("result_norm".to_string(), 4, 3, vec![0.0; 12]);
466        let ct = capture.get("result_norm").unwrap();
467        assert_eq!(ct.layer, None);
468        assert_eq!(ct.n_embd(), 4);
469        assert_eq!(ct.n_tokens(), 3);
470    }
471}
llama_cpp_4/context/tensor_capture.rs

llama_cpp_4/context/
tensor_capture.rs
