KV Cache Optimization

KV Cache Optimization refers to the family of techniques that reduce the memory footprint, computational cost, and latency of storing and retrieving the key-value tensors generated during autoregressive inference in transformer-based language models. It matters because, without caching, every new token generation requires recomputing attention over the entire context, making inference quadratically expensive — a bottleneck that dominates the cost of production LLM serving. The counterintuitive fact: in many large-scale deployments, KV cache memory consumption — not model weights — is the primary constraint on batch size and throughput.

The core insight

During autoregressive decoding, the attention keys and values for all previously generated tokens never change — they are deterministic functions of tokens that have already been committed. Instead of recomputing these tensors on every forward pass, you compute them once and store them; each new token only needs to compute its own K and V vectors and attend over the cached history. The entire science of KV cache optimization is then about answering one question: how do you make that cache smaller, faster, or smarter without hurting model quality?

How It Works

Standard Attention and the Recomputation Problem

In a decoder-only transformer, the scaled dot-product attention at layer $\ell$ for query position $t$ over context positions $1 \dots t$ is:

\quad \text{Attention}(Q_t, K_{1:t}, V_{1:t}) = \text{softmax}\!\left(\frac{Q_t K_{1:t}^\top}{\sqrt{d_k}}\right) V_{1:t}

Explanation: The query vector for the current token attends over all cached key vectors, producing a weighted sum of value vectors — and $d_k$ is the per-head key dimension used to stabilize gradients. Here $Q_t \in \mathbb{R}^{d_k}$ is the query for token $t$ , $K_{1:t} \in \mathbb{R}^{t \times d_k}$ is the stacked key matrix for all tokens seen so far, and $V_{1:t} \in \mathbb{R}^{t \times d_v}$ is the stacked value matrix. Without caching, computing $K_{1:t}$ and $V_{1:t}$ from scratch at every decoding step costs $O(t)$ linear projections per layer, giving $O(L \cdot t)$ total work per token where $L$ is the number of layers.

KV Cache Memory Footprint

The total memory required to store the KV cache for a single sequence of length $T$ is:

\quad M_{\text{KV}} = 2 \times L \times H \times d_k \times T \times \text{sizeof}(\text{dtype})

Explanation: Factor of 2 accounts for both keys and values; $L$ layers each have $H$ attention heads each with $d_k$ -dimensional projections, accumulated over $T$ tokens. For a concrete example: LLaMA-2 70B has $L=80$ layers, $H=64$ heads, $d_k=128$ , using float16 (2 bytes). A single sequence of $T=4096$ tokens requires $2 \times 80 \times 64 \times 128 \times 4096 \times 2 \approx 10.7\text{ GB}$ .

Multi-Query and Grouped-Query Attention

Multi-Query Attention (MQA) and Grouped-Query Attention (GQA) reduce cache size by sharing key and value heads across query heads. With $G$ key-value groups and $H$ total query heads:

\quad M_{\text{GQA}} = 2 \times L \times G \times d_k \times T \times \text{sizeof}(\text{dtype}), \quad G \leq H

Explanation: By sharing K and V across $H/G$ query heads each, you reduce KV cache by a factor of $H/G$ with minimal quality degradation when $G \geq 8$ . MQA sets $G=1$ (maximum compression); GQA uses $G \in \{4, 8\}$ as a quality-efficiency sweet spot — LLaMA-3 and Mistral both use GQA.

Quantized KV Cache

Quantizing cached K and V tensors from float16 to int8 or int4 reduces memory by 2× or 4× respectively. For a quantization scale $s$ and zero-point $z$ , the dequantization is:

\quad \hat{x} = s \cdot (x_q - z)

Explanation: Each stored integer $x_q$ is mapped back to a floating-point approximation $\hat{x}$ using per-channel or per-token scales, introducing bounded quantization error. The attention computation then proceeds on dequantized tensors, adding a small but measurable perplexity penalty that becomes acceptable at int8 precision for most production use cases.

Attention Score Sparsity and Eviction

Empirical studies (Zhang et al., StreamingLLM; Ge et al., FastGen) show that attention weight distributions are highly sparse — only a small fraction of cached tokens receive meaningful attention weight. This motivates eviction policies that drop low-importance tokens. A common importance score for token $i$ at decoding step $t$ is its cumulative attention mass:

\quad \text{score}(i) = \sum_{j=i}^{t} \alpha_{j \to i}, \quad \alpha_{j \to i} = \text{softmax}\!\left(\frac{q_j k_i^\top}{\sqrt{d_k}}\right)

Explanation: Tokens that have accumulated high attention weight across many decoding steps are considered important and retained; tokens with near-zero cumulative weight are candidates for eviction.

researcher · Ashish Vaswani | b. 1983 | Invented the Transformer Architecture

Vaswani led the team at Google Brain that published "Attention Is All You Need" (2017), introducing the scaled dot-product attention mechanism and multi-head attention that underlies every modern LLM. His work directly created the architectural context in which KV caching became both necessary and possible — without the autoregressive transformer, there would be no KV cache to optimize. He later co-founded Inceptive, a company applying transformers to RNA biology.

researcher · Noam Shazeer | b. 1975 | Invented Multi-Query Attention

Shazeer, one of the original "Attention Is All You Need" co-authors, published "Fast Transformer Decoding: One Write-Head is All You Need" in 2019, introducing Multi-Query Attention as a direct solution to KV cache memory pressure during inference. His insight — that query heads need not share their key and value projections — reduced KV cache by up to 64× while preserving most model quality, and became the foundation for GQA. He later co-founded Character.AI.

Figure

KV Cache Data Flow During Autoregressive Decoding

Intuition

The Mental Model

Imagine you are translating a long novel and, on every page, you need to remember the names of all characters introduced so far. Without a notepad (no cache), you would re-read the entire novel from page 1 every time you need to recall a name — obviously insane. With a notepad (the KV cache), you jot down each character name once and only consult your notes when needed. KV cache optimization is then about making that notepad more compact: MQA/GQA is like sharing one notepad among multiple translators working in parallel (they all need the same character list, so why keep separate copies?); quantization is like using shorthand instead of full spellings (saves space with slightly less precision); token eviction is like crossing out minor characters who were mentioned once in chapter 2 and never again, because experience shows you will never need them; and PagedAttention is like using a binder with removable pages instead of a fixed notebook, so you can give each translator exactly as many pages as they currently need, no more, no less, and recycle pages when a translation job is done.

Historical Development

Origins · 1 minhow this was discovered

History

2017

Vaswani et al., "Attention Is All You Need" — introduces scaled dot-product attention and the transformer architecture, implicitly defining the K and V tensors whose caching becomes critical.

2019

Shazeer, "Fast Transformer Decoding: One Write-Head is All You Need" — introduces Multi-Query Attention, reducing KV cache by

H\times

by sharing a single K/V head across all query heads.

2020

Radford et al. (OpenAI) deploy GPT-3 with standard multi-head attention; KV caching is used in inference but cache memory becomes a visible bottleneck at scale for the first time.

2022

Pope et al., "Efficiently Scaling Transformer Inference" (Google) — systematic analysis of inference memory bounds, introducing the arithmetic intensity framework that formalizes when inference is memory-bound