Skip to main content

Module kv_capture

Module kv_capture 

Source
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_layer layers, keeping at most max_per_layer K and max_per_layer V vectors per layer. head_dim self-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. proj holds n_kv contiguous head vectors of head_dim each ([h0…][h1…]…). No-op (one atomic load) when disabled; stops appending to a layer/stream once its max_per_layer cap is hit.
snapshot
Copy out the captured vectors (splitting the flat per-layer blobs into rows). Returns None if nothing was captured (capture never enabled, or the forward never hit the CPU attention path).