Still Mini KV Compactor

A laptop-scale measurement harness for KV cache compaction in LLM inference: it implements a mini Perceiver compactor in the spirit of Baseten's Still paper, trains it, and benchmarks it head-to-head against selection heuristics — tracking retained memory, latency, and answer quality on a MacBook Air M4.

What this is: a clean, reproducible benchmark + a faithful mini-implementation of Still's architecture, with an honest result. What this is not: a reproduction of Still's full experiments, a claim that learned compaction beats selection at 0.5B, or production code.

Inspired by Still: Amortized KV Cache Compaction in a Single Forward Pass (O'Neill et al., arXiv 2606.07878; Baseten writeup).

The honest headline

At 0.5B on synthetic long-context MCQ, simple selection (sink+recent) retains most of the quality; the learned Perceiver compactor collapses to chance. The interesting part is why, and it points straight at the hard problem Still actually solves.

Quality vs memory (aggregate, 8× compression, 24 examples/cell)

ctx	method	accuracy	Δ vs full	retained KV (MiB)	slots
1024	full	0.67	+0.00	11.81	1011
1024	no_context	0.58	−0.08	0.00	0
1024	sink_recent	0.54	−0.12	1.50	126
1024	learned	0.12	−0.54	1.50	128
1024	random	0.12	−0.54	1.48	126
2048	full	0.71	+0.00	23.51	2009
2048	no_context	0.58	−0.12	0.00	0
2048	sink_recent	0.50	−0.21	2.93	251
2048	learned	0.25	−0.46	1.50	128
2048	random	0.21	−0.50	2.93	251

Reading it:

• The benchmark has real signal. Full context beats no-context (0.67 vs 0.58 at 1024; 0.71 vs 0.58 at 2048), so there's a quality gap for compaction to compete in.
• Sink+recent is the strongest compaction — at 8× it keeps the model near its no-context baseline while holding only ~1.5 MiB of cache.
• The learned compactor lands at chance (0.12). Its training loss converged to ~1.45 CE, which is ln(1/0.25) — i.e. it converged to a degenerate ~25% solution and never learned to preserve retrievable context.

Why the learned compactor fails (the actual finding)

This is the part worth a conversation. The mini-compactor operates in rotated (RoPE) key space: it cross-attends learned latents into the already-RoPE-rotated keys and emits 128 compact slots. But those compact slots have no coherent positional structure for the model's queries to retrieve against — the model applies RoPE to the 128 slots as if they were positions 0–127, while the real content spanned 1024–2048 positions. The geometry is broken, so attention over the compact cache is worse than having no cache at all.

That is exactly the hard problem Still solves (a trained, position-aware synthesis that stays in the model's retrieval geometry). This benchmark reproduces the easy half (the architecture and the measurement) and isolates the half that actually matters — handling position correctly under compaction. The next experiment is the pre- vs post-rotation variant (below).

What's solid here

• A clean measurement harness. Every row records the same schema; timing is split into prefill / compaction / forward; memory is measured from real tensor bytes, not just process RSS. Memory math is exact: full 1024-ctx cache = 11.81 MiB → 128 slots = 1.50 MiB (8×).
• A faithful mini-Still Perceiver (src/perceiver_compactor.py, 1.3M params): learned latent queries cross-attend into the full KV cache, self-attend, and emit compact K/V in one forward pass.
• A real training pipeline (src/train_main.py): distillation from the frozen base model (teacher = full context, student = compact cache), CE + KL loss; only the compactor trains.
• Honest baselines and an honest negative result, with the failure diagnosed rather than hidden.

Setup

Machine	MacBook Air M4, 16 GB RAM
Model	Qwen/Qwen2.5-0.5B-Instruct (24 layers, 2 KV heads, GQA)
Backend	PyTorch MPS; frozen model in float16, compactor trains in float32
Methods	no_context, full, random, sink_recent, learned (128-slot Perceiver)

Run it

python -m venv .venv && source .venv/bin/activate && pip install -r requirements.txt

# 1) train the 128-slot compactor (frozen fp16 model + fp32 compactor; ~15 min on M4)
BENCH_DTYPE=float16 HF_HUB_OFFLINE=1 python -m src.train_main --examples 100 --epochs 5 --context 1024

# 2) benchmark all methods where there's quality signal (ctx >= 1024)
BENCH_DTYPE=float16 HF_HUB_OFFLINE=1 python -m src.benchmark \
  --methods full no_context random sink_recent learned --contexts 1024 2048 --compressions 8 --examples 24

# 3) charts
python src/plot_results.py results/benchmark_results.csv

A trained checkpoint (checkpoints/compactor_slots_128.pt) is included so you can reproduce the learned rows without retraining.

Limitations (read before drawing conclusions)

1. 0.5B scale. KV-cache memory isn't the real bottleneck at 494M params; results don't transfer to 7B+ where compaction actually matters.
2. Rotated-cache MVP. The compactor sees post-RoPE keys — the likely cause of the chance-level learned result. The pre-rotation variant is deferred (see Next).
3. Single-forward MCQ eval. No autoregressive decode throughput (decode_tok_s is always 0.0).
4. Synthetic data. 5 deterministic task families, not real workloads.
5. GQA. Qwen2.5-0.5B has 2 KV heads (cache already 7× smaller than non-GQA), so compaction savings are less dramatic than on a full-KV model.

Next

• Pre- vs post-rotation compaction (the load-bearing experiment): compact pre-RoPE keys and re-apply rotation to the compact slots, so the geometry survives. This directly targets the diagnosed failure and is the cleanest path from "chance" to a real result.
• Scale to 1.5B, longer contexts, and a real long-context benchmark subset (e.g. RULER).
• Real decode-throughput measurement with an autoregressive generate() loop.

License

MIT.