1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
//! Calibration-time Hessian capture for GPTQ / error-feedback transfer.
//!
//! While `begin()` is active, every `QTensor::matmat` (the batched prefill
//! path) folds its input activations into a per-tensor second-moment
//! `H = Σ_t x_t·x_tᵀ` and per-channel `Σ x²` (for the activation RMS field
//! of the two-field outlier score). One global hook covers every linear in
//! the model — attention projections, FFN, experts — with no per-layer
//! wiring. Zero cost when off (one relaxed atomic load). Single-threaded
//! accumulation from the caller's thread, so it never races the parallel
//! matmat that follows.
use std::collections::HashMap;
use std::sync::Mutex;
use std::sync::atomic::{AtomicBool, Ordering};
static ON: AtomicBool = AtomicBool::new(false);
/// Accumulate the full `H = X·Xᵀ` (needed only for the GPTQ fold). Off ⇒
/// diagonal-only (`Σx²`), which is all the ternary + per-row correction
/// path needs — and the only thing that fits for a 12B (full H would be
/// ~100 GB across the model).
static FULL_H: AtomicBool = AtomicBool::new(true);
static REG: Mutex<Option<HashMap<String, HessianAcc>>> = Mutex::new(None);
/// Accumulated input statistics of one linear layer over the calibration set.
pub struct HessianAcc {
pub cols: usize,
/// `H = Σ_t x_t·x_tᵀ`, dense `[cols·cols]`, f64.
pub h: Vec<f64>,
/// `Σ_t x_t²` per input channel `[cols]` — RMS = sqrt(sumsq / count).
pub sumsq: Vec<f64>,
/// Token-position samples folded in.
pub count: usize,
}
impl HessianAcc {
/// Per-input-channel activation RMS.
pub fn rms(&self) -> Vec<f32> {
let n = self.count.max(1) as f64;
self.sumsq.iter().map(|&s| (s / n).sqrt() as f32).collect()
}
}
/// Start capturing (clears any prior registry). `full_h` = accumulate the
/// dense Hessian (fold path); false = diagonal-only (ternary + correction).
pub fn begin(full_h: bool) {
*REG.lock().unwrap() = Some(HashMap::new());
FULL_H.store(full_h, Ordering::SeqCst);
ON.store(true, Ordering::SeqCst);
}
/// Stop capturing and take the accumulated Hessians.
pub fn end() -> HashMap<String, HessianAcc> {
ON.store(false, Ordering::SeqCst);
REG.lock().unwrap().take().unwrap_or_default()
}
#[inline]
pub fn capturing() -> bool {
ON.load(Ordering::Relaxed)
}
/// Fold a batch of `b` input vectors (`xs = [b·cols]`) into `name`'s Hessian.
pub fn accumulate(name: &str, xs: &[f32], b: usize, cols: usize) {
let mut guard = REG.lock().unwrap();
let Some(map) = guard.as_mut() else {
return;
};
let full = FULL_H.load(Ordering::Relaxed);
let acc = map.entry(name.to_string()).or_insert_with(|| HessianAcc {
cols,
h: if full {
vec![0.0; cols * cols]
} else {
Vec::new()
},
sumsq: vec![0.0; cols],
count: 0,
});
if acc.cols != cols {
return; // shape mismatch — skip defensively
}
for bi in 0..b {
let x = &xs[bi * cols..(bi + 1) * cols];
for i in 0..cols {
let xi = x[i] as f64;
if xi == 0.0 {
continue;
}
acc.sumsq[i] += xi * xi;
if full {
let hrow = &mut acc.h[i * cols..i * cols + cols];
for (j, &xj) in x.iter().enumerate() {
hrow[j] += xi * xj as f64;
}
}
}
acc.count += 1;
}
}