Expand description
W5b — KV-vector capture hook for sparse-dictionary calibration.
Sibling of crate::llm_awq (AWQ activation capture): a gated forward-pass hook that records the
engine’s real per-layer K and V vectors during a calibration forward over the eval corpus. The
sparse-dictionary go/no-go (crate::wgsl_forge::calibration) needs OUR engine’s actual KV
geometry — GQA layout, RoPE convention, and layer shapes are engine-specific, so this cannot come
from synthetic data or another runtime.
Capture point: native attention runs through gguf_bridge::…::cpu_attention_pass (the wasm-proven
CPU SDPA the native path routes through), which writes each token’s post-RoPE K and pre-RoPE V into
the KV cache. We tap those exact vectors there — post-RoPE K is what the int8 KV cache would
quantize, so this is an apples-to-apples source for the int8-vs-dictionary comparison.
Gated + bounded: off in production (one relaxed atomic load on the attention path). When on, it
appends into a per-layer buffer under a mutex, capped at max_per_layer vectors per layer per
stream (K, V) so memory stays bounded regardless of corpus length. Calibration-only — never a
production hot path.
Structs§
- KvCapture
- One layer’s captured vectors, split into
head_dim-length rows.
Functions§
- clear
- Drop the capture buffer (free the calibration-only heap).
- disable
- enable
- Begin a KV capture for a model with up to
n_layerlayers, keeping at mostmax_per_layerK andmax_per_layerV vectors per layer.head_dimself-sizes on the first recorded vector. - is_
enabled - record
- Record every head’s K (
k_not_v = true) or V vector from a projection slice.projholdsn_kvcontiguous head vectors ofhead_dimeach ([h0…][h1…]…). No-op (one atomic load) when disabled; stops appending to a layer/stream once itsmax_per_layercap is hit. - snapshot
- Copy out the captured vectors (splitting the flat per-layer blobs into rows). Returns
Noneif nothing was captured (capture never enabled, or the forward never hit the CPU attention path).