KV Cache Quantization Without Silent Drift

KV cache quantization can turn long-context memory from a concurrency ceiling into usable capacity, while quietly perturbing every later decode step. A credible rollout measures where the error lands—key or value, early or late layer, recent or distant token—not only how many gigabytes disappear.

This guide builds a layer-position heatmap, live allocation ledger, paired service benchmark, and task-level drift gate. It ends with mixed precision and an exact reference fallback instead of one boastful bit count.

KV cache quantization begins with a tensor ledger

KV cache quantization changes the state reused at every decode step, so its first deliverable is a byte-accurate ledger rather than a single quality score. Record layers, key/value heads, head dimension, sequence length, batch shape, element width, scales, zero points, residual buffers, allocator rounding, and any unquantized prefix. The theoretical four-bit total is not the live device footprint if kernels stage dequantized tiles or metadata dominates short requests.

The Hugging Face cache documentation distinguishes cache implementations and describes quantized caches as a memory trade with latency costs. Treat those choices as serving modes with separate receipts. A model using grouped-query attention has fewer cache heads than attention heads; confusing them can double the estimate before one token is generated.

Build the ledger for a concrete Llama-like fixture: 24 layers, eight KV heads, head dimension 128, and 2,048 tokens. FP16 keys and values consume 384 MiB; ideal four-bit storage consumes 96 MiB before scales and packing. KV cache quantization earns attention when the measured allocation approaches that reduction and the extra kernel work does not erase the concurrency gain.

Measure layerwise sensitivity instead of one average

Quantizing every layer identically assumes that cache error has uniform effect. It does not. Early layers shape broad representations, middle layers may carry routing-like distinctions, and later layers directly influence logits.

For each layer, replace only that layer's cache with its quantized form while keeping all others in the reference precision. Compare tensor cosine similarity, key and value error, logit divergence, token agreement, and task outcomes.

KIVI introduced asymmetric two-bit key and value treatment after observing different channel patterns; the KIVI paper is a useful primary reference for why cache quantizers deserve tensor-specific design. Reproduce its relevant premise on your own model rather than copying its bit setting. Kernels, model families, and workload lengths can move the best point.

A per-layer quantization sweep often reveals a small sensitive set. Preserve those layers at eight or sixteen bits while compressing the rest, then rerun the full stack because isolated effects can interact. Publish the layer map beside memory saved. A claim such as 4-bit cache is incomplete when six layers remain high precision and account for a meaningful part of the allocation.

Token position turns small errors into trajectories

Decode depends on every cached position, and quantization error may differ between a recent token, a long-range instruction, and a high-magnitude outlier near the prefix. Slice metrics by absolute position, distance from the current token, prompt versus generated region, and semantic role. Probe 128, 512, 2,048, and the maximum supported length with the same content planted at multiple positions.

The PyTorch quantization documentation provides official background on observers, schemes, and quantized representations, even when a custom cache kernel performs the final packing. Keep the observer's calibration population explicit. A scale fitted on short conversational prompts can saturate on long code blocks, tables, or tool traces.

Track token-position error as a heatmap whose rows are layers and columns are position buckets. The worked fixture flags layer 23 at position 2,048 because cosine similarity falls to 0.982 and KL divergence rises to 0.041. It still passes the illustrative gate, but it identifies the cell to attack with adversarial prompts and longer continuations rather than hiding inside an average.

Runnable artifact: The fixture checks a layer-position error gate and reconciles ideal cache bytes for a concrete model shape.

Save this proof as kv-layer-drift.test.mjs and run node kv-layer-drift.test.mjs. Expected final line: PASS: layer-position gates.

import assert from "node:assert/strict";
const rows=[{layer:0,pos:32,cos:.999,kl:.002},{layer:12,pos:512,cos:.994,kl:.012},{layer:23,pos:2048,cos:.982,kl:.041}];
const gate=r=>r.cos>=.98&&r.kl<=.05; assert.equal(rows.every(gate),true);
assert.equal(rows.sort((a,b)=>b.kl-a.kl)[0].layer,23);
const bytes=(layers,tokens,heads,dim,bits)=>layers*tokens*heads*dim*2*bits/8;
assert.equal(bytes(24,2048,8,128,4),50331648);
console.log("PASS: layer-position gates");

Separate key geometry from value reconstruction

Keys decide attention similarity; values carry the content mixed after weights are formed. Quantization can therefore distort who is attended to or what is retrieved after attention. Measure them independently.

For keys, compare attention-logit and top-attended-position changes. For values, replay reference attention weights against quantized values to isolate reconstruction error. Then run the combined path to see their interaction.

Use per-channel, per-token, or grouped scales according to measured outlier structure and kernel support. Smaller groups reduce quantization error but increase scale bytes and memory traffic. Symmetric coding is simple; asymmetric coding can use range more effectively when distributions shift away from zero. The choice belongs in a cache calibration table with clipping percentile, group size, rounding rule, and scale dtype.

Do not use perplexity alone. A model can preserve average next-token loss while losing exact copying, distant constraints, tool arguments, or rare names. Add retrieval needles, multi-turn instruction retention, long code completion, structured output, and free generation. RoPE scaling without recall loss supplies a position-by-depth pattern that adapts well to quantized cache evaluation.

Benchmark latency on the path users actually run

A smaller cache reduces memory bandwidth and may raise batch capacity, yet dequantization, packing, and unsupported shapes can add work. Measure prefill separately from decode, because some systems quantize only after a residual window fills. Report time to first token, inter-token latency, tokens per second, p50 and p95, peak allocation, power if available, and maximum stable concurrency.

Warm up every kernel, pin software and clocks where possible, and use paired prompts in randomized order. Sweep batch, context length, generation length, and cache reuse. KV cache optimization for predictable speed explains why a hot-path result cannot substitute for cold and eviction behavior. Quantization may improve long sequences but regress a short interactive request whose cache never becomes the bottleneck.

The decode latency receipt needs a break-even point. In one hypothetical A10G run, four-bit cache loses four percent at 128 tokens, breaks even near 700, and gains eighteen percent at 4,096 while doubling feasible batch. The product can enable it only beyond 1,024 cached tokens, or reserve it for memory-pressure mode, rather than declaring one global winner.

Test quality with matched generations and tasks

Use teacher-forced logits for precise local comparison and seeded generation for trajectory-level effects. Log top-one agreement, rank correlation, KL divergence, output edit distance, task accuracy, refusal behavior, and format validity. Seeded sampling does not make GPU execution universally deterministic, so archive prompts, decoded tokens, library versions, kernels, and the tolerance used for paired conclusions.

Stratify by input structure: natural dialogue, code, multilingual text, repetitive tokens, long numbers, retrieved passages, and system-plus-tool messages. Add outlier fixtures with repeated high-magnitude tokens and alternating topics. LLM quantization with a quality budget provides the release discipline, while this cache-specific study adds position and layer coordinates to every failure.

A silent-drift gate fails when an important cohort crosses its task limit even if aggregate logit error remains small. Preserve examples that change a tool argument, citation, safety decision, or exact string. Human review is useful for open responses but should follow deterministic diffs that identify where to look. The goal is not proving numerical identity; it is bounding the behavioral changes the product cannot accept.

Design a mixed-precision fallback before rollout

Implement policies for protected prefix length, sensitive layers, group size, bit width, and escalation to a wider format. The policy should be chosen from the layer-position heatmap, not intuition. If a request exceeds calibrated length, uses an unsupported head shape, or trips an outlier detector, route to eight-bit or FP16. If allocation still fails, reject or shorten honestly rather than silently changing quality.

Stage rollout with mirrored cache construction and reference decoding on sampled requests where cost permits. Compare decisions offline before serving quantized output. A canary gate can bind model hash, quantizer version, and kernel build to progressive exposure. Keep a kill switch that clears incompatible cache blocks; stale packed state must never survive a policy change.

Multi-tenant serving adds identity concerns. Bit width does not alter the requirement that cache blocks belong to an authorized request or explicitly shareable prefix. Prefix cache tenant isolation covers those boundaries. The quantized block header should include tenant or scope, model, adapter, quantizer, group size, sequence range, and checksum before a worker accepts it.

Publish the layer-position-quality frontier

The final KV cache quantization receipt includes tensor shapes, reference dtype, quantizer scheme, bit allocation by layer, residual window, calibration corpus, clipping, group dimensions, scale dtype, kernel and device, memory at several lengths, latency by batch and length, tensor heatmaps, task cohorts, changed examples, unsupported cases, fallback, and rollout alarms. Show the Pareto frontier rather than only the chosen point.

Reject the candidate if savings exist only on paper, p95 decode regresses on its target workload, long-position cells exceed the error gate, critical tool or retrieval tasks drift, or fallback cannot restore the reference path. Re-run after model, kernel, driver, or scheduler changes. A cached representation is coupled to all four, and yesterday's fastest packing may become tomorrow's slow compatibility path.

The useful conclusion may be mixed: four bits for middle layers after a 128-token residual window, eight bits at the beginning and end, and FP16 for short prompts. That answer is less marketable than one bit count, but it is operationally honest. KV cache quantization succeeds when each saved byte is tied to a measured layer, position, latency effect, and quality boundary.

KV cache quantization should enter the release ledger as a tensor policy, not a model-wide adjective. Re-run KV cache quantization evidence whenever kernels, schedulers, context limits, or cache shapes change.