ferrox_core/
activation_tap.rs

1//! A process-wide observer of the f32 activations fed to every
2//! [`WeightMatrix`] projection: the seam `ferrox imatrix` collects an
3//! importance matrix through.
4//!
5//! llama.cpp collects its imatrix with a scheduler callback
6//! (`cb_eval`) that sees every `MUL_MAT` node and its `src1` input
7//! (`tools/imatrix/imatrix.cpp:219-237`). ferrox has no graph, so the
8//! equivalent seam is the function every projection goes through:
9//! [`WeightMatrix::apply`] for one activation and
10//! [`WeightMatrix::apply_batch_with_acts`] for a batch. Both call
11//! [`observe`] before they touch the weights.
12//!
13//! **A `WeightMatrix` carries no name**, and the loader that knows the
14//! names lives in `ferrox-models`, so the observer is handed the matrix
15//! by reference and it is the installer's job to know which one it is
16//! -- `ferrox imatrix` walks the decoder's public weight fields and
17//! keys on the address. That is why the callback takes `&WeightMatrix`
18//! and not a string: adding a name field would have to be threaded
19//! through thirty construction sites across seven loaders to change
20//! nothing for inference.
21//!
22//! Cost when nothing is installed: one relaxed atomic load per
23//! projection call, not per element. A decode step is a few dozen
24//! calls, so this is not a hot-path concern.
25//!
26//! The tap sees the f32 input EXACTLY as the projection receives it,
27//! before any activation quantization the CPU INT_DOT path does. That
28//! is what llama.cpp's callback sees too (`src1->type == F32` is a
29//! precondition for collection).
30//!
31//! What it does NOT see: a GPU batch that fails to launch and degrades
32//! to per-row [`WeightMatrix::apply`] would fire the observer twice for
33//! the same rows. `ferrox imatrix` pins the CPU backend, where no such
34//! degradation path exists, and checks every dense entry's row count
35//! against the token count so a double observation is a refusal rather
36//! than a silently doubled matrix.
37
38use std::sync::atomic::{AtomicBool, Ordering};
39use std::sync::{Arc, RwLock};
40
41use crate::weight_matrix::WeightMatrix;
42
43/// The observer: the matrix being applied, the activation rows laid
44/// out `[n_rows][cols]`, and `n_rows`.
45pub type Observer = dyn Fn(&WeightMatrix, &[f32], usize) + Send + Sync;
46
47static INSTALLED: AtomicBool = AtomicBool::new(false);
48static OBSERVER: RwLock<Option<Arc<Observer>>> = RwLock::new(None);
49
50/// Installs `observer` for the life of the returned guard. Only one
51/// may be installed at a time: a second install while one is live is
52/// refused, because two observers would each see every row and neither
53/// would know the other was counting.
54pub fn install(observer: Arc<Observer>) -> Result<TapGuard, AlreadyInstalled> {
55    let mut slot = OBSERVER.write().unwrap_or_else(|e| e.into_inner());
56    if slot.is_some() {
57        return Err(AlreadyInstalled);
58    }
59    *slot = Some(observer);
60    INSTALLED.store(true, Ordering::Release);
61    Ok(TapGuard(()))
62}
63
64/// Uninstalls the observer on drop.
65#[must_use = "dropping the guard uninstalls the tap immediately"]
66pub struct TapGuard(());
67
68impl Drop for TapGuard {
69    fn drop(&mut self) {
70        INSTALLED.store(false, Ordering::Release);
71        *OBSERVER.write().unwrap_or_else(|e| e.into_inner()) = None;
72    }
73}
74
75#[derive(Debug, Clone, Copy, PartialEq, Eq)]
76pub struct AlreadyInstalled;
77
78impl std::fmt::Display for AlreadyInstalled {
79    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
80        f.write_str("an activation tap is already installed in this process")
81    }
82}
83
84impl std::error::Error for AlreadyInstalled {}
85
86/// Called by the projection entry points. The fast path is the
87/// `INSTALLED` load; the observer is cloned out of the lock so a slow
88/// observer never holds it across the call.
89#[inline]
90pub(crate) fn observe(matrix: &WeightMatrix, rows: &[f32], n_rows: usize) {
91    if !INSTALLED.load(Ordering::Acquire) {
92        return;
93    }
94    let observer = OBSERVER
95        .read()
96        .unwrap_or_else(|e| e.into_inner())
97        .as_ref()
98        .cloned();
99    if let Some(observer) = observer {
100        observer(matrix, rows, n_rows);
101    }
102}
103
104#[cfg(test)]
105mod tests {
106    use super::*;
107    use crate::tensor::Tensor;
108    use std::sync::Mutex;
109
110    // Tests that install a tap share the one process-wide slot, so
111    // they serialise on this rather than racing each other.
112    static SERIAL: Mutex<()> = Mutex::new(());
113
114    fn small_matrix() -> WeightMatrix {
115        WeightMatrix::F32(Tensor::new(vec![1.0; 8], vec![2, 4]))
116    }
117
118    /// The tap sees each `apply_batch` call once, with every row, and
119    /// can tell which matrix was applied by address. Both halves are
120    /// what `ferrox imatrix` relies on.
121    #[test]
122    fn the_observer_sees_every_batch_row_once_and_the_matrix_identity() {
123        let _serial = SERIAL.lock().unwrap_or_else(|e| e.into_inner());
124        let m = small_matrix();
125        let addr = &m as *const WeightMatrix as usize;
126        type Seen = Mutex<Vec<(usize, Vec<f32>, usize)>>;
127        let seen: Arc<Seen> = Arc::new(Mutex::new(Vec::new()));
128        let sink = seen.clone();
129        let guard = install(Arc::new(move |w: &WeightMatrix, rows: &[f32], n: usize| {
130            sink.lock()
131                .unwrap()
132                .push((w as *const WeightMatrix as usize, rows.to_vec(), n));
133        }))
134        .unwrap();
135        let x: Vec<f32> = (0..12).map(|i| i as f32).collect();
136        let _ = m.apply_batch(&x, 3);
137        drop(guard);
138        let seen = seen.lock().unwrap();
139        assert_eq!(seen.len(), 1, "one batch call, one observation");
140        assert_eq!(seen[0].0, addr);
141        assert_eq!(seen[0].1, x);
142        assert_eq!(seen[0].2, 3);
143    }
144
145    /// After the guard drops nothing is observed, and a second install
146    /// while one is live is refused rather than replacing it.
147    #[test]
148    fn the_tap_is_gone_after_the_guard_drops_and_cannot_be_installed_twice() {
149        let _serial = SERIAL.lock().unwrap_or_else(|e| e.into_inner());
150        let m = small_matrix();
151        let calls = Arc::new(Mutex::new(0usize));
152        let c = calls.clone();
153        let guard = install(Arc::new(move |_: &WeightMatrix, _: &[f32], _| {
154            *c.lock().unwrap() += 1;
155        }))
156        .unwrap();
157        assert_eq!(
158            install(Arc::new(|_: &WeightMatrix, _: &[f32], _| {}))
159                .err()
160                .map(|_| ()),
161            Some(()),
162            "a second install must be refused while the first is live"
163        );
164        let _ = m.apply(&[1.0; 4]);
165        drop(guard);
166        let _ = m.apply(&[1.0; 4]);
167        assert_eq!(*calls.lock().unwrap(), 1);
168    }
169}
ferrox_core/activation_tap.rs

ferrox_core/
activation_tap.rs
