feat(candle-nn): QuantizedKvCache — INT8 KV cache with attention sinks (TurboQuant)

[aryanputta]

Summary

Adds QuantizedKvCache to candle-nn::kv_cache, implementing the KV-cache quantisation component of TurboQuant (Google Research, ICLR 2024).

Closes #3425 (partially — see "What's not included" below).

Motivation

KV cache is the dominant memory bottleneck for long-context LLM inference. TurboQuant shows that symmetric INT8 quantisation of keys and values with per-head scales achieves ~4× memory reduction with negligible accuracy loss. The crucial insight is that the first few tokens ("attention sinks") must be kept at full precision — they receive disproportionately large attention scores and quantising them causes measurable degradation.

Implementation

Quantisation scheme

For a tensor x of shape (B, H, S, D), per-token per-head symmetric INT8 mapped to U8:

scale = clamp(max(|x|, dim=D), min=1e-6)   // shape (B, H, S, 1)
q     = clip(round(x/scale) + 128, 0, 255) // U8, same shape as x
x'    = (q.f32() - 128) * scale            // dequant to original dtype

Memory layout

Field	Type	Shape
k_sink, v_sink	original dtype	(B, H, n_sink, D)
k_q, v_q	U8	(B, H, bulk, D)
k_scale, v_scale	F32	(B, H, bulk, 1)

API (drop-in for existing caches)

let mut cache = QuantizedKvCache::new(/*dim=*/2, /*n_sink_tokens=*/4);

// Append during prefill or decode — returns dequantised tensors ready for attention.
let (k, v) = cache.append(&k_new, &v_new)?;

// Standard lifecycle methods.
cache.current_seq_len()
cache.is_empty()
cache.reset()

The interface is identical to KvCache / ConcatKvCache / RotatingKvCache, so it can be swapped in at the model level without touching attention code.

Tests

Five CPU-only unit tests:

• test_quantized_kv_cache_roundtrip — max error < 0.02 (INT8 tolerance)
• test_quantized_kv_cache_sinks_full_precision — sink tokens are lossless
• test_quantized_kv_cache_mixed_sink_and_bulk — correct shapes with sink+bulk split
• test_quantized_kv_cache_incremental_append — autoregressive append grows seq correctly
• test_quantized_kv_cache_reset — reset clears all state

What's not included

This PR implements the KV quantisation + attention sink component of TurboQuant. The third component — cross-layer KV sharing (sharing K/V pairs across adjacent layers, which provides the bulk of the 6× aggregate compression) — requires per-model architectural changes and is a separate PR. This primitive is the foundation that makes cross-layer sharing composable.