09 — KV Cache Quantization

Weights are a fixed cost; the KV cache is a per-token, per-user cost. At 128k context, the KV cache dwarfs the model. This is where quantization moves from “nice speedup” to “the difference between fitting the workload and not.”

The number that motivates everything

Llama-3-8B, one user at 128k context:

KV bytes = 2 (K and V) × layers × kv_heads × head_dim × bytes × seq_len
         = 2 × 32 × 8 × 128 × 2 (fp16) × 131,072
         = 17.2 GB

One user. At batch 8, you need 138 GB of KV, which is more than any single H100 (80GB) can hold. This is why long-context serving is a memory-capacity problem before it’s anything else. FP8 KV cuts it in half; INT4 KV cuts it 4×.

For Llama-3-70B at 128k the number is ~40 GB per user. This is the number that killed “just serve fp16 everywhere.”

The structural fact that makes KV quant work

Key and Value tensors have very different distributions:

  • Keys have per-CHANNEL outlier structure. Some head-dim channels are 50-100× larger than others, consistently across tokens. This is because the same channels of K get compared to the same channels of Q via dot product — the network learns to concentrate signal in specific channels.

  • Values have per-TOKEN outlier structure. Massive-activation tokens (BOS, punctuation, delimiters) produce Value vectors 100-1000× larger than the median.

This asymmetry is the founding insight of KIVI, and it dictates the quant granularity:

Key: quantize per-channel. Value: quantize per-token.

Memorize this. It’s the KV cache quant equivalent of the W4A16-vs-W8A8 rule.

KIVI — the reference paper

KIVI, arxiv <phone_number_or_numberic_id_or_random_id_138> (Liu, Chen, Hu et al., Feb 2024). Repo: github.com/jy-yuan/KIVI.

Core contributions:

  • Per-channel Key quantization with a small residual buffer (last few tokens kept fp16 for the streaming-quant edge case).

  • Per-token Value quantization — tokens with massive activations get their own scale.

  • Tuning-free: no calibration, just the two granularity choices.

  • 2-bit KIVI achieves ~lossless quality on Llama-2 up to 32k context. Throughput 2.35–3.47×.

Bit-level, per K-cache: for each layer × head, per-channel scale and zero (fp16), int2/int4 storage. The residual buffer at the tail is a small hack that matters more than it sounds — online quantization has to wait for a chunk to fill before you know its true max.

KVQuant — the more surgical variant

KVQuant, arxiv <phone_number_or_numberic_id_or_random_id_139> (Hooper, Gholami et al., Jan 2024, Berkeley). Adds four ideas on top of KIVI:

  1. Per-channel Key + pre-RoPE quantization. Applying RoPE after dequant means K channels stay well-behaved (RoPE rotates channels together; if you quantize post-RoPE, the outlier structure gets scrambled).

  2. Non-uniform datatypes. They fit a nonlinear codebook (like NF4 for weights) per K/V channel.

  3. Dense-and-Sparse quantization for outlier tokens. The 1% biggest V-token norms get stored in a separate sparse fp16 buffer.

  4. Enables 10M context on 8×A100 for Llama-2-7B — the paper’s headline demo.

More accurate than KIVI at 2-bit but with more implementation overhead. In production, KIVI’s simpler recipe is usually what ships.

The lazy default: FP8 KV

Before you reach for any of the above, know this: FP8 KV cache is nearly free.

vLLM/SGLang both support --kv-cache-dtype fp8_e4m3 (or fp8_e5m2). This is:

  • 2× smaller than fp16 KV (16 GB → 8 GB in our Llama-3-8B example).

  • <0.1 PPL degradation on standard benchmarks.

  • Zero calibration, zero tooling, one flag.

  • Native on Hopper+ (H100 tensor cores read fp8 KV directly during attention).

On L40S/L20 (Ada), FP8 KV works via software dequant on the attention kernel path, still gives you the memory win, small compute hit.

Zoho translation: if your on-prem CRM deployment is running Llama-3-70B on 4×L40S at 32k context and running out of memory, --kv-cache-dtype fp8_e5m2 is your first move, not a smaller model. This is the single highest-ROI production toggle in this whole phase.

Per-channel vs per-token in the actual attention kernel

The non-trivial part of KV quant is that the attention kernel itself has to dequantize on the fly. This changes FlashAttention’s memory pattern:

  • Normally FA loads a K-tile in fp16 and does QKᵀ in tensor-core fp16.

  • With int4-per-channel K: FA loads the int4 K-tile + the per-channel scale vector, dequantizes to fp16 in shared memory, then feeds tensor cores.

  • With int4-per-token V: same pattern but scales are per-row of V.

This is why FlashInfer (the paged-attention kernel library vLLM/SGLang use) had to add quantized attention paths as a first-class feature. Read paged_attention docs in the FlashInfer repo when you get to Phase 4 to see the actual API.

Fusion insight (important): you cannot dequantize KV into a full fp16 buffer first — that defeats the purpose. The dequant must be fused into the attention kernel’s tile loop. This is a real kernel-engineering constraint, not a nice-to-have.

Recent developments (2024–2025)

Paper

ArXiv

Idea

KIVI

<phone_number_or_numberic_id_or_random_id_140>

Per-channel K, per-token V. The reference.

KVQuant

<phone_number_or_numberic_id_or_random_id_141>

Pre-RoPE K + dense-and-sparse outliers → 10M context.

SKVQ

<phone_number_or_numberic_id_or_random_id_142>

Sliding-window preservation — keep the recent K tokens at fp16, quantize the tail.

QAQ

<phone_number_or_numberic_id_or_random_id_143>

Quality-adaptive: allocate bits per-layer by sensitivity.

KVSink

<phone_number_or_numberic_id_or_random_id_144>

Preserve attention-sink tokens (BOS-adjacent) at higher precision. Beats naive Preserve-First-N.

KVQuant (v2, extended)

<phone_number_or_numberic_id_or_random_id_145>

Follow-up handling more architectures and MoE.

The direction of travel: from tuning-free simple schemes (KIVI) → to per-layer adaptive schemes (QAQ) → to structural-awareness schemes (KVSink recognizes attention sinks; SKVQ recognizes recency effects).

The prescription table

Scenario

KV format

Why

Default Hopper serving, ≤32k context

FP8 E5M2

Nearly free, one flag, no calibration.

Long context 32k–256k on Hopper

FP8 E4M3 or INT8 KIVI-style

Need better precision at low bits due to context length.

1M+ context research

INT4 KIVI or KVQuant

Only way to fit. Ships with quality caveat.

Ampere serving (A100/A6000)

FP8 E5M2 (software) or INT8

H100 fp8 kernels don’t exist on Ampere.

Edge (Apple/CPU)

Model-native (usually GGUF fp16 KV)

Not the bottleneck.

L40S on-prem (your Zoho case)

FP8 E5M2

Ada-generation fp8 works, capacity is the issue.

Anti-patterns

  • Don’t quantize the sink tokens. The first few tokens absorb massive attention scores; quantizing them wrecks long-context quality. KVSink or preserve-first-N handles this.

  • Don’t use per-tensor scales on KV. The whole point is per-channel-K / per-token-V. Per-tensor throws away the structure.

  • Don’t benchmark KV quant on short prompts. Damage compounds with context. Always eval at ≥8k, ideally 32k+.

  • Don’t confuse KV quant with prefix caching. Different problem, different solution, they combine.

The two-sentence study answer

KV cache dwarfs the model at long context (17GB per user at Llama-3-8B/128k fp16), and quantizing it 2–4× is the difference between fitting a workload and not — the trick is the K/V distributional asymmetry: quantize Key per-channel and Value per-token (KIVI, arxiv <phone_number_or_numberic_id_or_random_id_148>). In production, --kv-cache-dtype fp8_e5m2 in vLLM/SGLang is the free-lunch default; drop to INT4 via KIVI/KVQuant only when memory capacity forces it.

Homework

  1. Compute the KV cache size for Llama-3-70B at 32k, 128k, and 1M context in fp16, fp8, int4. Plot GB vs seq_len for all three.

  2. Read the KIVI paper section 3 (“Per-channel Key, Per-token Value”) and reproduce Figure 2 (the K vs V distribution plot) on Llama-3-8B using ~30 lines of hook code.

  3. Run vLLM with and without --kv-cache-dtype fp8_e5m2 on Llama-3.1-8B at 8k prompt. Measure decode tokens/sec, prefill time, WT2 PPL. Report the trilogy: memory saved, speed change, quality delta.

  4. Bonus: read the DeepSeek-V2 MLA derivation (arxiv <phone_number_or_numberic_id_or_random_id_149>). MLA is the architectural fork in the road — KV quantization is one lever, KV architecture is another.