qualia_core_db/inference/
kv_capture.rs

1//! W5b — KV-vector capture hook for sparse-dictionary calibration.
2//!
3//! Sibling of [`crate::llm_awq`] (AWQ activation capture): a gated forward-pass hook that records the
4//! engine's real per-layer K and V vectors during a calibration forward over the eval corpus. The
5//! sparse-dictionary go/no-go ([`crate::wgsl_forge::calibration`]) needs OUR engine's actual KV
6//! geometry — GQA layout, RoPE convention, and layer shapes are engine-specific, so this cannot come
7//! from synthetic data or another runtime.
8//!
9//! Capture point: native attention runs through `gguf_bridge::…::cpu_attention_pass` (the wasm-proven
10//! CPU SDPA the native path routes through), which writes each token's post-RoPE K and pre-RoPE V into
11//! the KV cache. We tap those exact vectors there — post-RoPE K is what the int8 KV cache would
12//! quantize, so this is an apples-to-apples source for the int8-vs-dictionary comparison.
13//!
14//! Gated + bounded: off in production (one relaxed atomic load on the attention path). When on, it
15//! appends into a per-layer buffer under a mutex, capped at `max_per_layer` vectors per layer per
16//! stream (K, V) so memory stays bounded regardless of corpus length. Calibration-only — never a
17//! production hot path.
18
19#![cfg(not(target_arch = "wasm32"))]
20
21use std::sync::atomic::{AtomicBool, Ordering};
22use std::sync::{Mutex, OnceLock};
23
24static ENABLED: AtomicBool = AtomicBool::new(false);
25
26struct KvBuf {
27    /// Head-dim of the captured vectors; 0 until the first record self-sizes it.
28    head_dim: usize,
29    /// Cap on vectors per layer per stream.
30    max_per_layer: usize,
31    n_layer: usize,
32    /// Per-layer flat blobs; `k[layer]` = captured K vectors concatenated (len = count * head_dim).
33    k: Vec<Vec<f32>>,
34    v: Vec<Vec<f32>>,
35}
36
37fn buf() -> &'static Mutex<Option<KvBuf>> {
38    static B: OnceLock<Mutex<Option<KvBuf>>> = OnceLock::new();
39    B.get_or_init(|| Mutex::new(None))
40}
41
42/// Begin a KV capture for a model with up to `n_layer` layers, keeping at most `max_per_layer` K and
43/// `max_per_layer` V vectors per layer. `head_dim` self-sizes on the first recorded vector.
44pub fn enable(n_layer: usize, max_per_layer: usize) {
45    if let Ok(mut g) = buf().lock() {
46        *g = Some(KvBuf {
47            head_dim: 0,
48            max_per_layer,
49            n_layer,
50            k: (0..n_layer).map(|_| Vec::new()).collect(),
51            v: (0..n_layer).map(|_| Vec::new()).collect(),
52        });
53    }
54    ENABLED.store(true, Ordering::Relaxed);
55}
56
57pub fn disable() {
58    ENABLED.store(false, Ordering::Relaxed);
59}
60
61#[inline]
62pub fn is_enabled() -> bool {
63    ENABLED.load(Ordering::Relaxed)
64}
65
66/// Record every head's K (`k_not_v = true`) or V vector from a projection slice. `proj` holds `n_kv`
67/// contiguous head vectors of `head_dim` each (`[h0…][h1…]…`). No-op (one atomic load) when disabled;
68/// stops appending to a layer/stream once its `max_per_layer` cap is hit.
69#[inline]
70pub fn record(layer: usize, k_not_v: bool, proj: &[f32], n_kv: usize, head_dim: usize) {
71    if !ENABLED.load(Ordering::Relaxed) || head_dim == 0 {
72        return;
73    }
74    let Ok(mut g) = buf().lock() else {
75        return;
76    };
77    let Some(b) = g.as_mut() else {
78        return;
79    };
80    if b.head_dim == 0 {
81        b.head_dim = head_dim;
82    }
83    if b.head_dim != head_dim || layer >= b.n_layer {
84        return;
85    }
86    let cap_floats = b.max_per_layer.saturating_mul(head_dim);
87    let dst = if k_not_v {
88        &mut b.k[layer]
89    } else {
90        &mut b.v[layer]
91    };
92    for h in 0..n_kv {
93        if dst.len() >= cap_floats {
94            break;
95        }
96        let s = h * head_dim;
97        if s + head_dim > proj.len() {
98            break;
99        }
100        dst.extend_from_slice(&proj[s..s + head_dim]);
101    }
102}
103
104/// One layer's captured vectors, split into `head_dim`-length rows.
105pub struct KvCapture {
106    pub head_dim: usize,
107    /// `k[layer]` = list of captured K vectors (each `head_dim` long).
108    pub k: Vec<Vec<Vec<f32>>>,
109    pub v: Vec<Vec<Vec<f32>>>,
110}
111
112impl KvCapture {
113    /// Total K vectors captured across all layers.
114    pub fn total_k(&self) -> usize {
115        self.k.iter().map(|l| l.len()).sum()
116    }
117    /// Total V vectors captured across all layers.
118    pub fn total_v(&self) -> usize {
119        self.v.iter().map(|l| l.len()).sum()
120    }
121}
122
123/// Copy out the captured vectors (splitting the flat per-layer blobs into rows). Returns `None` if
124/// nothing was captured (capture never enabled, or the forward never hit the CPU attention path).
125pub fn snapshot() -> Option<KvCapture> {
126    let g = buf().lock().ok()?;
127    let b = g.as_ref()?;
128    if b.head_dim == 0 {
129        return None;
130    }
131    let hd = b.head_dim;
132    let rows =
133        |flat: &Vec<f32>| -> Vec<Vec<f32>> { flat.chunks_exact(hd).map(|c| c.to_vec()).collect() };
134    Some(KvCapture {
135        head_dim: hd,
136        k: b.k.iter().map(rows).collect(),
137        v: b.v.iter().map(rows).collect(),
138    })
139}
140
141/// Drop the capture buffer (free the calibration-only heap).
142pub fn clear() {
143    if let Ok(mut g) = buf().lock() {
144        *g = None;
145    }
146}
qualia_core_db/inference/kv_capture.rs

qualia_core_db/inference/
kv_capture.rs
