CS 194-16 · Scalable AI · Final Project
written version

EchoKV: watching a language model decide its own precision

A closed-loop controller that adjusts each layer's attention precision while the model is generating — so memory savings come from the model itself, not a manual tuning pass.

By Nazar Ospanov, Zoir Imomaliev, Jason Trinh, and Sanjay Adhikesaven · Berkeley CS 194-16, Spring 2026

The problem in one sentence

When a large language model serves a chat or a long document, it keeps the entire past conversation in fast GPU memory. That working memory — the KV cache — grows with the conversation, and at 128,000 tokens on a 70-billion-parameter model it can take more space than the model weights themselves.

The standard fix is to compress the cache: store every number in 4 bits instead of 16 bits, cutting memory 4×. That works, but it's blunt — every layer of the model gets the same compression whether or not it needs to. We wanted to know: can a model decide, on the fly, how aggressively to compress each of its own layers?

How EchoKV works

EchoKV has four pieces that talk to each other while the model is generating text:

DECODE STEP  (every step) K (FP16) key projection Kref ring buffer FP16, last 64 tokens (K only) ~10 MB — 0.1% of cache captured BEFORE quant quantb(K) to current bit-width paged KV cache INT2 pool INT4 pool INT8 pool canary fires ~1 in 64 steps CANARY ESTIMATOR for each layer ℓ and each candidate b ∈ {2, 4, 8}:
$\hat{\delta}_{\ell,b} = D_{\mathrm{KL}}\!\bigl(\mathrm{softmax}(Q K_{\text{ref}}^\top) \,\|\, \mathrm{softmax}(Q\,\mathrm{quant}_b(K_{\text{ref}})^\top)\bigr)$
attention distribution 0 p key idx → FP16 softmax quant softmax |FP16 − quant| ≈ KL
A small $\hat{\delta}$ means the quantized cache barely changes attention — that layer can stay at the lower bit-width.
PID CONTROLLER  (per layer) 1. exponential-moving-average smoothing:
$\mathrm{EMA}_{\ell,b} \leftarrow (1-\alpha)\,\mathrm{EMA}_{\ell,b} + \alpha\,\hat{\delta}_{\ell,b}$
2. argmin under budget τ:
$\hat{b}_\ell = \min\{\,b \in \{2,4,8\} : \mathrm{EMA}_{\ell,b} < \tau\,\}$
3. broadcast new tier vector to TP ranks 4. allocator routes new K blocks to the chosen pool
if $\mathrm{EMA}_{\ell,b_\ell} > \tau$  →  upgrade layer $\ell$ next batch
$\hat{\delta}_{\ell,b}$
b ∈ {2,4,8} new tier vector b̂ ∈ {2, 4, 8} routes NEW K blocks ↑
EchoKV's mechanism. Every decode step, the model's key projection $K_\ell$ is written into a per-layer FP16 ring buffer $K_{\text{ref}}$ before being quantized into the paged cache. This buffer is the load-bearing trick (paper §4 / fig 2): it gives the canary a clean FP16 reference, so the measurement doesn't depend on the controller's own past decisions. On about 1 in 64 steps the canary fires, computing the KL divergence between full-precision attention and three candidate quantizations $b \in \{2, 4, 8\}$. The PID controller smooths those values with an EMA and picks each layer's bit-width as the lowest $b$ whose smoothed divergence stays under the budget $\tau$. The new tier vector routes future K blocks to the right pool; already-allocated blocks are never modified.

See it run

The animation below is a simulation of EchoKV running inside a 20-layer language model.* Most of the time, the model just generates text. Periodically a small "canary" measurement fires: for every layer, it computes how much the attention output would drift if that layer's cache were compressed to INT2, INT4, or INT8. Layers whose drift stays below a threshold can use the more aggressive compression; layers whose drift exceeds the threshold get upgraded to keep accuracy intact. Press play, or change the workload, to watch it adjust.

EchoKV controller — interactive simulation

Decoding…
step: 0
workload: GSM8K-Hard
FP16 (16 bits, initial) INT2 (2 bits) INT4 (4 bits) INT8 (8 bits)
16.00avg bits / element
0%memory saved vs FP16
Animation speed Workload

* The KL values, threshold position, and tier distributions in this animation are illustrative — they are deterministic synthetic patterns chosen so each workload's typical allocation looks plausible, not the actual per-layer measurements from our experiments. Real Llama-3.1-70B has 80 layers, not 20; the canary fires on ~1 in 64 decode steps in production, not every few steps as shown here. The "Animation speed" control sets only the visualization cadence so you can watch each phase — real decoding runs many orders of magnitude faster (real-world EchoKV adds only ~0.09% to time-per-output-token at production canary rate, so the canary fires and tier updates are essentially invisible on the wall clock). The point is to show the mechanism — how a canary firing drives per-layer bit-width decisions — not the exact numbers a real run produces. Real measured results are further down the page.

What you're seeing

Do we store the whole cache in FP16? — No.

The simulation above hides two mechanism details worth zooming in on. The first: "if the canary measures KL against an FP16 reference, doesn't that mean we're paying full-precision memory somewhere?" Answer: no — the FP16 reference is a tiny sliding ring buffer of just the last 64 K vectors per layer, with no V at all (about 10 MB total for the whole 70 B model at 128 K context — the gold sliver in the memory-footprint figure further down the page). Everything else — the actual KV cache the model reads from during attention — is paged INT2/INT4/INT8, sized and routed per layer by the controller.

How the ring buffer slides

Each cell below is one FP16 K vector for a single layer. The left block (gray) is the prompt context — the last few K's from prefill, already in the ring by the time decoding begins. The controller never sees an empty ring; from tick 0 it has real keys to measure against. The gold cells are decode K's, appended on the right as the model generates. When the canary fires (~1 in 64 decode steps in production, more often here for visibility), it measures KL across whatever's in the ring at that moment — up to 64 K's — in FP16 vs. a temporary INT2/INT4/INT8 copy of those same keys. The canary is testing the quantization function, not reconstructing the model's full output.

$K_{\text{ref}}$ ring buffer — 128 decode steps of one layer's timeline

Decoding…
decode step: 0
K's in ring: 16 / 64
newest:

What happens to old blocks when a layer changes tier?

The second detail: "if a layer's bit-width changes from INT4 to INT8 between firings (the simulation above shows this happening live), where do the new 8-bit numbers come from for the old tokens? Don't we need the FP16 originals to re-quantize?" Right — we don't have FP16 originals for old tokens (the ring buffer only holds the last 64). So we don't re-quantize. Tier changes are forward-only: already-stored K blocks stay at their admit-time bit-width forever; the controller's new decision only affects K blocks allocated after the firing. The result is a temporally-striped cache — each layer's storage is a record of every tier decision the controller made over the lifetime of the request. vLLM's paged attention reads each block independently with its stored bit-width, so mixing precisions inside the same layer is fine.

Three representative layers' KV caches as tokens stream in — each layer's canary decisions move independently

Decoding…
token: 0
current admit tiers: L0:INT4 · L1:INT8 · L79:INT4
Show the math — what the canary actually computes

Attention at layer $\ell$ on the full-precision (FP16) key tensor $K_\ell$ produces a softmax distribution over past tokens. If we replaced $K_\ell$ with its quantized version $\mathrm{quant}_b(K_\ell)$, the softmax shifts. The canary measures how much, per-layer and per-candidate-bit-width $b$, using KL divergence:

$$\hat{\delta}_{\ell,b} \;=\; D_{\mathrm{KL}}\!\left(\,\mathrm{softmax}(Q_\ell K_\ell^\top) \;\big\|\; \mathrm{softmax}(Q_\ell\,\mathrm{quant}_b(K_\ell)^\top)\,\right)$$

where $\mathrm{quant}_b(\cdot)$ is KIVI-style asymmetric per-group $b$-bit quantization. The KL is summed over the attention sequence dim and averaged over the query dim, so each canary firing produces one number per (layer, bit-width) pair — $3L$ numbers per firing for $b \in \{2, 4, 8\}$ and $L$ layers.

Single readings are noisy, so each $(\ell, b)$ value goes into an exponential moving average:

$$\mathrm{EMA}_{t+1} \;=\; (1-\alpha)\,\mathrm{EMA}_t \;+\; \alpha\,\hat{\delta}_{\mathrm{obs}}, \qquad \alpha = 0.1$$

The controller then picks each layer's bit-width as the lowest $b$ whose smoothed divergence stays under the budget $\tau$:

$$\hat{b}_\ell \;=\; \min\!\left\{\,b \in \{2, 4, 8\} \;:\; \mathrm{EMA}_{\ell,b} < \tau_{\mathrm{budget}}\,\right\}$$

Three details that matter for stability:

  • The "absolute" canary. $K_\ell$ in the formula above is read from a per-layer ring buffer of the most recent 64 FP16 K projections (≈ 10 MB total for Llama-3.1-70B at 128 K context — <0.1% of the working quantized cache), captured before any tier-application. Earlier designs measured KL against the running quantized cache, which creates a feedback loop where controller decisions perturb the measurement substrate. Reading against a fixed FP16 reference breaks the loop.
  • Bidirectional adjustment — promote AND demote. The argmin is recomputed from scratch on every canary firing, so a layer isn't locked into its first decision. A layer currently at INT2 whose smoothed KL@INT2 starts exceeding the INT2 budget gets promoted to INT4 or INT8; an INT8 layer whose KL@INT4 has dropped back below $\tau$ gets demoted to INT4. The reason it shifts: every firing sees a different $K_{\text{ref}}$ window (the 64-token ring rotates as decoding proceeds), a different batch of queries, and content-dependent quantization error (KIVI's per-group scale is a function of the group's max/min, so incoming keys with outliers blow up the per-group quant error), all amplified by softmax sharpness — workloads that demand precise attention to specific tokens see larger KL swings under quantization than workloads with diffuse attention. The EMA smooths token-level noise out, but real workload drift propagates through, and the controller follows it bidirectionally.
  • The $\tau(L)$ scaling. Attention sparsity falls as context length grows, so the same per-step divergence translates into different model-output drift at different context lengths. We rescale the budget by $\tau(L) = \tau_0 \sqrt{L/L_0}$ (Pinsker-derived; at $L_0 = 4\text{K}$, $\tau_0 = 0.022$). At 128K context this gives $\tau = 0.124$. Without the rescaling, the controller cascades to nearly-uniform INT8 at long context even when INT4 would have been fine.

At production canary rate p_c = 1/64, this whole pipeline adds ~0.09% to time-per-output-token on Llama-3.1-8B + RULER.

Multi-seed results on Llama-3.1-70B

We tested EchoKV on five workloads at Llama-3.1-70B, running each one with three random seeds and using paired McNemar's test for significance. The plot below puts each substrate on a single bits-vs-accuracy axis — bits per element on x (left = less memory), 3-seed-mean accuracy on y — with vertical error bars showing the accuracy spread across seeds and horizontal error bars showing the bit-budget spread. FP16, INT4, and INT8 have fixed bit-widths so their horizontal bars collapse to a point; EchoKV's bar widens whenever the controller drifts across seeds (most visibly on GSM8K-Hard).

Hardware & serving configuration

All measurements use the same physical host and the same pinned software stack. The only thing that changes between cells is the KV-cache substrate (FP16 / static INT4 / static INT8 / EchoKV); prompts, decoding parameters, and random seeds are paired across substrates.

  • Model: Llama-3.1-70B-Instruct (FP16 weights), via vLLM's LLM.generate path.
  • GPUs: 8 × NVIDIA H100 80 GB HBM3 on a single host.
  • Tensor parallel: TP=4 (4 H100s) for GSM8K-Hard, HotpotQA, and RULER@128K — these fit comfortably and let us run two workloads in parallel on the 8-GPU box. TP=8 (all 8 H100s) for τ-Bench airline and τ-Bench retail because the tool-use agent generates longer trajectories and benefits from the extra throughput.
  • Serving stack: vLLM 0.16.0 + PyTorch 2.9.1 (CUDA 12.8) in a project-pinned conda environment. The EchoKV controller runs as a forward-pre-hook shadow-quant emulation layer on top of the unmodified vLLM PagedAttention path; the static-INT4/INT8 cells use the same emulation layer with a fixed tier vector.
  • Sample sizes (per seed): n=500 for GSM8K-Hard and HotpotQA, n=100 for RULER@128K (its standard eval slice), full test set for τ-Bench airline (50 tasks) and τ-Bench retail (115 tasks).
  • Statistical protocol: three seeds per workload; paired McNemar's test between substrates on per-prompt correctness, plus the per-seed accuracy itself for the small-n τ-Bench cells where McNemar is underpowered.

Per-seed wall-time is roughly 55–65 min for GSM8K-Hard and HotpotQA, 90–120 min for RULER@128K, and ~2.2 h for each τ-Bench split. The same numbers reproduce on a clean clone of the repo using the launcher scripts in scripts/22 through scripts/29.

Measurement protocol — warmup, freeze, measure

Every EchoKV cell on this page is the result of a two-phase run that cleanly separates "the controller deciding its tier vector" from "scoring that tier vector against the static baselines." The two phases share a model load but use disjoint prompts so the controller never sees the prompts it's being graded on.

  1. Phase 1 — warmup (n=16 prompts, controller active). We draw 16 prompts from the workload's dataset at a different seed offset than the measurement set (e.g. for GSM8K-Hard: warmup at seed+100, measurement at seed). The model answers each prompt fully (prefill → generate up to max_new_tokens of reasoning + answer). During the decode loop, the canary fires on a fraction of decode steps, computes $\hat{\delta}_{\ell,b}$ for every (layer, bit-width) pair, and feeds the per-(ℓ, b) EMA. Each layer's bit-width is updated live as the EMA evolves. Because each answer is hundreds of decode tokens long, even 16 prompts give the controller thousands of canary measurements per layer to converge on — far more than the prompt count suggests.
  2. Phase 2 — consolidate & freeze. At the end of warmup, we snapshot the per-layer tier vector from every TP rank, consolidate across ranks (majority vote per layer), and lock that 80-element vector. The canary rate is dropped to ~$10^{-6}$ (effectively off) and freeze_controller is set to True — no further updates.
  3. Phase 3 — measurement (frozen tiers). The measurement prompts (different from warmup) are run with the frozen tier vector. The per-workload sample size is whatever's standard for that benchmark — n=500 for GSM8K-Hard and HotpotQA, n=100 for RULER@128K, the full test sets for τ-Bench airline (n=50 tasks) and τ-Bench retail (n=115 tasks). We score each prompt for correctness, and we run FP16, static INT4, and static INT8 cells on the exact same prompts in the same order with the same seed, then pair them prompt-by-prompt for McNemar's test. The 80-element vectors visualized in the "real per-layer tier vectors" figure below are these frozen post-warmup assignments — one decision per (workload, seed), held constant across all measurement prompts of that seed.

The freeze is an evaluation artifact, not a property of the deployed system. In production EchoKV stays live throughout — the controller keeps updating EMAs and re-deciding bit-widths across every decode step and every back-to-back prompt. We freeze only during measurement because paired McNemar needs a fixed schedule per prompt: if the controller kept adjusting, "EchoKV on prompt i" would depend on the prompt history (EMAs carry state from prompts 1..i−1), so the same prompt at position 50 vs position 300 isn't the same trial, the pairing loses its meaning, and the test's effective n shrinks. Live measurement would also fold controller-induced variance into the EchoKV-vs-uniform contrast, raising the effect size needed to clear p<0.05 on n=500.

The practical consequence is that the numbers below are a lower bound on what a live EchoKV deployment would deliver: production gets more degrees of freedom (within-prompt and across-prompt adaptation) that the frozen protocol throws away. A separate live-controller benchmark — bootstrapped over prompt orderings to characterize the controller's own variance — is the most defensible follow-up, and is listed in Future Work alongside "online τ adaptation."

The picture across all five workloads

workloadINT2INT4INT8FP16EchoKVEchoKV avg bitsEchoKV's choice

What the controller actually chose — real per-layer tier vectors

Each strip below is the actual 80-layer bit-width assignment EchoKV produced on Llama-3.1-70B at one seed of one workload — the frozen post-warmup tier vector that was then held constant across all measurement prompts for that workload (n=500 for GSM8K-Hard / HotpotQA, n=100 for RULER@128K, n=50 for τ-Bench airline, n=115 for τ-Bench retail; see "Measurement protocol" above for what "post-warmup" means). Cells are colored by the chosen precision (red = INT2, blue = INT4, orange = INT8). The pattern is markedly different per workload: RULER stays almost entirely at INT4; τ-bench cascades to almost entirely INT8; GSM8K-Hard is a mix and shifts between seeds.

Real tier vectors from our multi-seed runs (80 layers each)

INT2 INT4 INT8 hover a cell for the layer's tier & workload

Memory cost across workloads — one row per workload

Each colored trace below is one workload, shown as its own horizontal row. Markers along each row sit at the bit-widths of the four substrates (FP16, INT4, INT8, EchoKV). Left = less memory. The bit-width EchoKV actually chose for that workload is shown on the right.

Click workload names in the legend below to toggle visibility. Hover any marker for accuracy and config.

What the controller costs you in memory

The accuracy story is the headline, but EchoKV's whole pitch depends on the FP16 reference being cheap — otherwise we'd just be trading a static KV cache for a live one of the same size. The chart below puts the three storage layers to scale so you can see the cost directly: the FP16 reference EchoKV needs to make a per-layer decision is roughly 1000× smaller than the paged INT-cache it's helping to size, and the paged cache itself is ~75% smaller than the full-FP16 baseline it replaces.

Memory footprint, Llama-3.1-70B at 128 K context (80 layers × 8 KV heads × 128 head dim) Bars drawn to scale — note how invisible the gold FP16 reference is next to the working cache. Full FP16 KV cache, K+V, all 128 K tokens — (what we'd pay without compression) 41.9 GB Paged KV cache @ INT4, K+V, all 128 K tokens — (what we actually store for attention) 10.5 GB ~75% cut vs. FP16 FP16 Kref ring buffer, K only, last 64 tokens — (the canary's reference; the ONLY FP16 storage) 10.5 MB ~1000× smaller than the working cache
Three storage layers shown to scale. Only the gold sliver ($K_{\text{ref}}$) is FP16 — and it is a per-layer sliding ring of just 64 recent K vectors, ≈ 10 MB total for the whole 70 B model at 128 K context. The blue bar (≈ 10 GB) is what attention actually reads from on every decode step; it is INT4 / INT8 / INT2 per layer, chosen by the canary using the gold reference as ground truth. The gray bar (≈ 42 GB) is the hypothetical "store everything in FP16" baseline that EchoKV never pays.

Key takeaways

Conventional KV-cache quantization fixes the bit-width before the workload is known. You ship a model at static INT4 (or INT8) and live with that choice across every task it sees: hard math, multi-hop QA, long-context retrieval, tool-use agents, short chat. But the per-layer sensitivity to quantization is not a property of the model alone — it depends on what the model is actually doing. A layer that is safe at INT2 on a 100-token summarization request may be the layer that breaks a 50-turn τ-bench retail agent.

EchoKV inverts that contract. The bit-width assignment is decided at serving time, per layer, from the live KL signal the canary measures. The same deployment that runs at near-uniform INT4 on RULER@128K cascades to near-uniform INT8 on τ-bench retail — automatically, without re-quantizing the model, and without an operator having to know in advance which task is "hard." Adaptation isn't just across workloads either: within a single workload the controller re-decides bits per prompt distribution (on GSM8K-Hard the per-seed tier vector swings from 22 INT8 layers up to 67 INT8 layers across the three seeds, depending on which prompts each seed happened to sample). The controller's behavior is content-driven — adaptive to the specific domain or task mix the model is being asked to handle — so as the serving traffic shifts (more agentic tool-use this quarter, more retrieval next quarter, hard math vs. easy summarization within the same endpoint), the per-layer assignment moves with it.

This matters because LLM-serving traffic is heterogeneous by default. A single static bit-width is a bet that one number works for every request the model will ever see. Adaptive per-layer quantization replaces that bet with a measurement.

Future work

Several threads we left open, in order of how interesting we think they are:

Code & reproducing

Everything in this post — the EchoKV controller, the canary estimator, the per-workload multi-seed scripts, and the raw aggregate JSONs — lives in our public repository:

github.com/nazar-ospanov/echokv

Show the reproduction recipe

Each workload has one tmux-launchable script under scripts/ that runs the full three-seed paired protocol with all four substrates (FP16, static INT2, static INT4, static INT8, EchoKV) on identical prompts and writes an aggregate.json with paired McNemar comparisons. Each script is self-contained.

# 1. clone and enter the repo
git clone https://github.com/nazar-ospanov/echokv.git
cd echokv

# 2. activate the pinned environment (vLLM 0.16.0 + torch 2.9.1)
conda activate echokv

# 3. reproduce each workload — every script runs the full
#    FP16 / INT2 / INT4 / INT8 / EchoKV sweep across 3 seeds
bash scripts/22_section63_emulation_multiseed.sh    # GSM8K-Hard
bash scripts/23_section64_hotpotqa_multiseed.sh     # HotpotQA
bash scripts/24_section65_ruler128k_multiseed.sh    # RULER@128K, τ=0.124
bash scripts/25_taubench_airline_multiseed.sh       # τ-Bench airline (TP=8)
bash scripts/26_taubench_retail_multiseed.sh        # τ-Bench retail (TP=8)

# 4. results land under results/serving/<workload>_repro/aggregate.json
#    (raw seed JSONs are alongside; the blog charts read from them)

The aggregate JSONs contain per-seed FP16 / INT2 / INT4 / INT8 / EchoKV accuracies, the per-seed bit-width budget, the per-seed 80-element tier vector, and the paired McNemar p-values used in the results section above. The exact files driving every chart on this page are checked into the repo.