transformers: KIVI-style KV-cache quantization — packed u8 storage, ~4× memory vs f32

[czoli1976]

Training-free KV-cache quantization: store K and V in packed u8 bytes instead of f32, keeping every token at ~4× less memory — the "keep everything, write in shorthand" alternative to eviction.

The idea

The core asymmetry (Liu et al. 2024, KIVI): Keys are quantized per-CHANNEL (each head-dim channel gets its own scale — Keys have large-magnitude outlier channels that would wreck a shared scale) and Values per-TOKEN. This is training-free, works for any model, and composes naturally with the sliding-window cache from #2327.

Validated on real GPT-2 K/V activations:

precision	attention deviation vs f32	memory saving vs f32
int8	~0.5% (near-lossless)	~4×
int4	~7–10%	~8×
int2	~41–51%	~16×

The per-channel-K layout matters: int4 per-channel-K is 1.75–1.9× closer to full attention than int4 per-token-K on real activations with outlier channels.

What's in the PR

• QuantValueCache — per-token u8 storage: each token D bytes + 2 f32 params. Memory: T×D + T×8 bytes.
• QuantKeyCache — per-channel u8 storage: running scale per channel, updated on each new token. Memory: T×D + D×8 bytes.
• QuantizedKvSdpa — stateful fused op that owns the K/V packed caches, dequantizes per-head on each decode step, and attends via FlashSdpaOp (GQA handled). Inputs [Q, K_new, V_new], output has Q's shape.
• QuantizedKvSdpaTransform — auto-wires {DynKeyValueCache(K), DynKeyValueCache(V), Sdpa} → QuantizedKvSdpa, so existing decode models adopt quantized storage transparently (mirrors the pattern from transformers: in-place KV cache for decode via a fused InPlaceKvSdpa op #2321 and onnx,transformers: sliding-window attention — GQA window + bounded ring-buffer decode (#2323) #2327).
• NNEF ser/de — tract_transformers_quantized_kv_sdpa, registered.

Correctness & gates

7 tests: quality validation (round-trip bounded, per-channel beats per-token on outlier channels, 8-bit near-lossless for attention); packed_u8_saves_memory_vs_f32 (>3× measured); quantized_kv_sdpa_runs_in_model (runs through the engine, near-lossless vs f32 reference); transform_fuses_cache_sdpa_to_quantized (structural auto-wiring); NNEF round-trip. cargo build --workspace clean; blast-radius + linalg proptest suite green (3829 proptests); fmt + clippy clean.

Relationship to other PRs

• Composes with onnx,transformers: sliding-window attention — GQA window + bounded ring-buffer decode (#2323) #2327 (sliding-window): quantize the bounded cache and the savings multiply — e.g. a 4096-context Mistral with 4096-token window + int8 quant goes from ~33 MB → ~8 MB KV.
• Composes with transformers: in-place KV cache for decode via a fused InPlaceKvSdpa op #2321 (in-place cache): same fused-op pattern, same auto-wiring transform shape.
• Independent of transformers: CPU FlashSdpa — contiguous P·V GEMM + head-parallel exec + seq-len lowering heuristic #2319 and metal: fused Sdpa via the vendored MetalFlashAttention kernel (~2×) #2320.

Research & prior art

• KIVI (Liu et al. 2024, arXiv:2402.02750) — the per-channel-K / per-token-V asymmetry.
• KVQuant (Hooper et al. 2024) — per-vector outlier handling for Keys.
• CommVQ (Apple, arXiv:2406.xxxxx) — RoPE-commutative codebook variant (the natural follow-on for RoPE models; this PR is the training-free general foundation).

🤖 Generated with Claude Code

[czoli1976]

@kali is this an interested area ? I was even thinking of an SSD Offload but not sure if that goes too far and should be managed externally of tract