#2318 KV cache: FP16 sliding window + attention sinks, and an optional 2-bit K / 4-bit V cache scheme (--kv q2k) [EXPERIMENTAL]

Description

# KV cache: FP16 sliding window + attention sinks, and an optional 2-bit K / 4-bit V cache scheme (`--kv q2k`) [EXPERIMENTAL]

> **Status: experimental.** This ships as an opt-in scheme (`--kv q2k`, `--kv-window N`) and needs more testing before it can be considered production-ready: the behavioral A/B and long-context runs below are single-model, single-document measurements, and the window-1024 anomaly is unexplained. Default behavior is unchanged when the flags are off.

## Summary

Two independent, composable KV-cache features, shipped behind opt-in flags (default behavior unchanged):

1. **`--kv-window N`** — an FP16 overlay over the sink cells and the most recent `N` cells of the quantized KV bulk. Recent context and attention sinks read exact; older context stays quantized.
2. **`--kv q2k`** — an optional cache format storing **K at 2 bits per-channel and V at 4 bits per token (q4_0)**: 448 B/cell vs q4_0's 576 B (−22%), 3.6× smaller than FP16.

The scheme choices are measured, not assumed. Every directional claim below carries the number that pins it.

## Why this shape

### The bit-allocation rule came from the real cache, not the textbook

The published 2-bit KV schemes (KIVI, SKVQ) allocate bits symmetrically or K-heavy. We measured every allocation at **equal bytes per cell** with the gate's attention-output probe on real cached activations (12 QSA layers dumped from a live serve; probe = attention output relative RMS vs FP16, layer 3 unless stated):

| arm (equal bytes) | attention error |
|---|---:|
| **K 2-bit per-channel + V 4-bit (`q2k`)** | **1.911%** |
| K 4-bit + V 2-bit per-channel | 3.171% |
| K 2-bit + V 2-bit per-channel (symmetric q2) | 3.508% |
| KIVI split (K per-channel, V per-token, both 2-bit) | 5.188% |

At equal bytes, **the bits belong on V**: K's quantization error is largely common-mode under the softmax (the softmax only responds to relative score differences, so a common-mode score shift cancels in normalization), while V's error enters the output directly through `A·V`. Measured confirmation: raising K to 6 levels improved K's per-element RMS 1.7× but moved attention error < 2%.

Both published intuitions survive, scoped:

- **K's outliers are channel-aligned** (KIVI's premise) — true, and honored by the per-channel K scaling, not by more bits.
- **V is the tensor the model feels** — true, and it is why V gets 4 bits while K takes 2.

### Why per-channel K and per-token V

Per-channel K scales (`block_q2c_0`) absorb K's channel outliers — K's channel RMS spread on real layers runs 5.8× (L3) to 74.9× (L47). Per-token V groups match V's variance structure (channel RMS spread 7.9–14.5× but flat cell spread). This mirrors KIVI's placement, which we reproduce; KIVI's *bit budget* did not transfer (their K2/V1.5 config measures 10.8% on our cache — see papers section).

### Why the FP16 window

2-bit K is a real quality floor: with no window, `q2k` logit KL to FP16 is 2.5–5× q4_0's. The FP16 overlay is the cheap lever — 2048 B/cell over only the recent slice. With the default 128-cell window plus sinks, `q2k` matches the q4_0 reference on every behavioral measure (below), at 22% less memory.

## Measurements

### Gate (device kernel vs host reference, synthetic + real cache)

- Synthetic: 34 checks pass.
- Real-cache ingestion (4 layers, L3/L11/L23/L47): 31 checks each; the device arm lands on the host row to 3 decimals (L3: device 1.910% vs host 1.911%).
- Ring/append restore: FP16 overlay reads exact over the quantized bulk (dedicated parity suite).

### Multi-layer attention error (K2c + V4 vs references)

| layer | q4_0 | q2k |
|---|---:|---:|
| L3 | 0.888% | 1.911% |
| L11 | 0.740% | 2.002% |
| L23 | 0.941% | 2.211% |
| L47 | 1.040% | 1.645% |

Note L47 — the outlier layer (K channel RMS spread 74.9×) — is where q2k is *closest* to q4_0.

### Logits (512-token prompts, teacher-forced, vs FP16; top-1 agreement / mean KL)

| text | q4_0 | q2 | **q2k** |
|---|---|---|---|
| code | 88.7% / 0.052 | 72.7% / 0.311 | **79.3% / 0.238** |
| document | 94.1% / 0.028 | 79.3% / 0.196 | **84.0% / 0.135** |
| chat | 94.7% / 0.027 | 82.0% / 0.176 | **85.5% / 0.129** |

### Behavioral A/B (5K-token buried-fact probe)

| arm | mid | tail | short answers | loops | decode |
|---|---|---|---|---|---|
| q4_0 (reference) | 4/4 | 4/4 | 3/3 | 0 | 157 t/s |
| q2 (symmetric) | 4/4 | 4/4 | **2/3** | 0 | 136 t/s |
| **q2k + window 128** | 4/4 | 4/4 | **3/3** | 0 | 137 t/s |

q2k matches the q4_0 reference on every measure at 22% less KV memory; symmetric q2 is the arm that loses a short-answer recall.

### Long-context recall at depth (experimental run)

~200K-token operations log (199,337 tokens after the chat template) served with `--kv q2k --max-context 262144`; three fact needles planted in filler blocks at **25% / 50% / 90%** context depth:

| needle | depth | result |
|---|---|---|
| WALRUS-CANARY-2247 | ~25% | **recalled exactly** |
| QUARTZ-DELTA-8813 | ~50% | **recalled exactly** |
| TUNDRA-FALCON-6631 | ~90% | **recalled exactly** |

Cold prefill of the full 200K prompt ran at **2574 tok/s** through the q2k cache; the answers were read back from the q2k-quantized ~200K KV state (the verification pass rode the serve's conversation cache — the quantized cache is the same either way). Caveat: one model, one document, one run; treat as a smoke test of depth recall, not a benchmark.

A small multi-turn agent task (find and fix two planted bugs across two files, iterate pytest to green) also completed correctly against the served q2k/262144 endpoint.

### Window sizing

| window | short answers | decode |
|---|---|---|
| 128 (default) | 3/3 | 128 t/s |
| 512 | 3/3 | 136 t/s |
| 1024 | 3/3 (but mid/tail recall regressed in the sweep) | 129 t/s |

512 is behaviorally identical to 128 and fastest; 1024 showed a mid/tail regression we could not attribute (logged as an open question, not a recommendation).

### Memory math (262144 ctx, 12 QSA layers, 2 KV heads × 256 dims)

| format | B/cell | KV total |
|---|---:|---:|
| FP16 | 2048 | 6.00 GiB |
| q4_0 | 576 | 1.69 GiB |
| **q2k** | **448** | **1.31 GiB** |

## Reference papers

- **KIVI** (arXiv 2402.02750) — per-channel K / per-token V placement. Their *placement* reproduces on our cache; their *bit budget* does not (K2/V1.5 measures 10.8% attention error on L3 vs our 1.9% for K2/V4 at the same total bytes as their K2/V2).
- **SKVQ** (arXiv 2405.06219) — channel reordering + clipping for low-bit V. Reordering is a no-op for per-channel-scaled K (a relabeling — proven in the gate, not assumed); it can only act through V's per-token groups, where the measured benefit did not materialize.
- **KVQuant** (arXiv 2401.18079) — non-uniform codebooks; our 2-bit K uses a Lloyd-Max codebook over block RMS, and the gate's 1M-sample fit pins the shipped constants to the true Gaussian grid (0.100% gap).

## What is deliberately not in this PR

- `q2t` (top-q mixed-level K), fractional-bit V, joint channel×token scales, asymmetric zero-point arms — measured, did not win, and are excluded to keep the shipping surface to the one format the data picked.
- The `kv_paper_domain` un-rotation harness and the 2.5-bit/1.5-bit width sweeps — measurement-only tooling.

## Known limits

- 2-bit K is the quality floor: logit KL stays 2.5–5× q4_0's without the window; the window recovers behavior, not logits.
- Greedy decode is not run-to-run reproducible under QSA tie-breaks + atomic partial-group KV writes; decode-path measurements are reported as 5-run means.
- The window-1024 mid/tail regression is unexplained (documented above).

## Commits

1. `kv_window` — FP16 sliding window + attention sinks over the quantized KV cache
2. `kv_q2` — 2-bit per-channel KV format, parity gate, measurement tooling
3. `kv_q2k` — K 2-bit per-channel + V 4-bit mixed-width format