qualia_core_db/inference/
inference_awq.rs

1//! W1/AWQ — activation-statistics capture for Activation-aware Weight Quantization (no external libs).
2//!
3//! AWQ's premise: a weight's salience is set by the magnitude of the *activation* it multiplies, so
4//! scaling each input channel by `s_j = max|X_j|` before quantizing preserves the channels that matter
5//! — letting aggressive (ternary) quantization survive. This module is AWQ **step 1, the forward hook**:
6//! during a calibration forward over the eval corpus it records, per FFN layer, the per-input-channel
7//! max |activation| at the FFN input (post-`ffn_norm`, the input to the gate/up projections).
8//!
9//! Lock-free + gated: off in production (one relaxed atomic load on the FFN path). The accumulator is
10//! `fetch_max` over `|x|.to_bits()` — valid because `|x| >= 0`, so u32 bit-order matches float order.
11//! Heap (the stats buffer) is calibration-only, allocated once, never on a production hot path.
12
13use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
14use std::sync::OnceLock;
15
16/// Upper bounds for the fixed stats buffer (covers up to ~8B-param transformer shapes).
17const MAX_LAYERS: usize = 80;
18const MAX_CHAN: usize = 8192;
19
20static ENABLED: AtomicBool = AtomicBool::new(false);
21static N_LAYER: AtomicU32 = AtomicU32::new(0);
22static N_CHAN: AtomicU32 = AtomicU32::new(0);
23/// Per-forward layer index (reset by [`begin_forward`]); the FFN hook increments it 0..n_layer-1.
24static LAYER_CURSOR: AtomicU32 = AtomicU32::new(0);
25
26/// `max |activation|` bits at `[layer * N_CHAN + chan]`. Allocated once at first use.
27fn stats() -> &'static Vec<AtomicU32> {
28    static STATS: OnceLock<Vec<AtomicU32>> = OnceLock::new();
29    STATS.get_or_init(|| {
30        (0..MAX_LAYERS * MAX_CHAN)
31            .map(|_| AtomicU32::new(0))
32            .collect()
33    })
34}
35
36/// Begin an AWQ calibration capture for a model with `n_layer` FFN layers of `n_chan` input channels.
37pub fn enable(n_layer: u32, n_chan: u32) -> Result<(), String> {
38    if n_layer as usize > MAX_LAYERS || n_chan as usize > MAX_CHAN {
39        return Err(format!(
40            "AWQ capture exceeds bounds: n_layer={n_layer} (max {MAX_LAYERS}), n_chan={n_chan} (max {MAX_CHAN})"
41        ));
42    }
43    N_LAYER.store(n_layer, Ordering::Relaxed);
44    N_CHAN.store(n_chan, Ordering::Relaxed);
45    reset();
46    ENABLED.store(true, Ordering::Relaxed);
47    Ok(())
48}
49
50pub fn disable() {
51    ENABLED.store(false, Ordering::Relaxed);
52}
53
54/// Zero the accumulators + reset the per-forward layer cursor.
55pub fn reset() {
56    LAYER_CURSOR.store(0, Ordering::Relaxed);
57    for a in stats().iter() {
58        a.store(0, Ordering::Relaxed);
59    }
60}
61
62/// Call once at the start of each token's forward so the layer cursor tracks layers 0..n_layer-1.
63/// No-op (one atomic load) when capture is off.
64#[inline]
65pub fn begin_forward() {
66    if ENABLED.load(Ordering::Relaxed) {
67        LAYER_CURSOR.store(0, Ordering::Relaxed);
68    }
69}
70
71/// Record one FFN layer's post-norm input channels. Called from the FFN forward; no-op when off.
72#[inline]
73pub fn record_ffn_input(x: &[f32]) {
74    if !ENABLED.load(Ordering::Relaxed) {
75        return;
76    }
77    let n_chan = N_CHAN.load(Ordering::Relaxed) as usize;
78    let n_layer = N_LAYER.load(Ordering::Relaxed) as usize;
79    if n_chan == 0 {
80        return;
81    }
82    let layer = LAYER_CURSOR.fetch_add(1, Ordering::Relaxed) as usize;
83    if layer >= n_layer {
84        return;
85    }
86    let s = stats();
87    let base = layer * n_chan;
88    let lim = x.len().min(n_chan);
89    for c in 0..lim {
90        // |x| >= 0 → u32 bit pattern is monotone in the float value, so fetch_max is a true max.
91        s[base + c].fetch_max(x[c].abs().to_bits(), Ordering::Relaxed);
92    }
93}
94
95/// Per-layer per-channel max |activation| (`[layer][chan]`). Snapshot after the calibration pass.
96pub fn snapshot() -> Vec<Vec<f32>> {
97    let n_layer = N_LAYER.load(Ordering::Relaxed) as usize;
98    let n_chan = N_CHAN.load(Ordering::Relaxed) as usize;
99    if n_chan == 0 {
100        return Vec::new();
101    }
102    let s = stats();
103    (0..n_layer)
104        .map(|l| {
105            (0..n_chan)
106                .map(|c| f32::from_bits(s[l * n_chan + c].load(Ordering::Relaxed)))
107                .collect()
108        })
109        .collect()
110}
111
112#[inline]
113pub fn is_enabled() -> bool {
114    ENABLED.load(Ordering::Relaxed)
115}
qualia_core_db/inference/inference_awq.rs

qualia_core_db/inference/
inference_awq.rs
