The KV Cache

The Autoregressive Bottleneck Without KV Cache

In Causal Self-Attention:

$$\text{Attention}(Q, K, V) = \text{Softmax}\left( \frac{Q K^T}{\sqrt{d_k}} \right) V$$

At step $t = 100$, we generate token 100. To compute attention for token 100, we need Key vectors $K_{1..100}$ and Value vectors $V_{1..100}$.

• Without KV Cache: Re-compute $K, V$ for tokens $1 \dots 99$ from scratch. Total operations over $N$ tokens = $O(N^3)$.
• With KV Cache: Fetch pre-computed $K_{1..99}, V_{1..99}$ from GPU VRAM; compute $Q_{100}, K_{100}, V_{100}$ for token 100 only. Total operations over $N$ tokens = $O(N^2)$.

  WITHOUT KV CACHE (Redundant Matrix Re-computation)
  Step 1: Compute [T1]
  Step 2: Re-compute [T1], Compute [T2]
  Step 3: Re-compute [T1], Re-compute [T2], Compute [T3]  <-- O(N³) Waste!

  WITH KV CACHE (Append to VRAM Memory)
  Step 1: Compute [T1] -> Store K1, V1 in VRAM
  Step 2: Fetch K1, V1 -> Compute [T2] -> Append K2, V2 to VRAM
  Step 3: Fetch K1..2, V1..2 -> Compute [T3] -> Append K3, V3 to VRAM  <-- O(N²) Efficient!

KV Cache VRAM Formula

For a batch of $B$ sequences, sequence length $L$, $L_{\text{layers}}$ layers, $H_{\text{kv}}$ Key-Value heads, and head dimension $D_{\text{head}}$:

$$\text{KV Cache Memory (Bytes)} = 2 \times B \times L \times L_{\text{layers}} \times H_{\text{kv}} \times D_{\text{head}} \times \text{BytesPerElement}$$

Concrete Example (LLaMA-3 70B FP16)

• $L_{\text{layers}} = 80$, $H_{\text{kv}} = 8$ (GQA), $D_{\text{head}} = 128$, Precision = 16-bit (2 bytes).
• Per token per sequence: $2 \times 80 \times 8 \times 128 \times 2 = 327,680 \text{ Bytes} \approx 328 \text{ KB/token}$.
• For batch size $B = 32$ at sequence length $L = 8,192$:

$$\text{KV Cache VRAM} = 32 \times 8192 \times 328\text{ KB} \approx \mathbf{86.0\text{ GB of VRAM!}}$$

The KV Cache size quickly exceeds the base model weights size at high batch sizes and long contexts!

Architectural Mitigations

  Multi-Head Attention (MHA)            Grouped-Query Attention (GQA)           Multi-Query Attention (MQA)
  N_q Query Heads, N_kv = N_q KV Heads   N_q Query Heads, N_kv = 8 KV Heads      N_q Query Heads, N_kv = 1 KV Head
  [Q1 Q2 Q3 Q4] [K1 K2 K3 K4]            [Q1 Q2 Q3 Q4] ──► [K1]                 [Q1 Q2 Q3 Q4] ──► [K1]
  (Large KV Cache)                       (8x Memory Reduction! - LLaMA-3)       (32x Memory Reduction!)

Say this out loud

"The KV Cache stores Key and Value projection tensors of past tokens in GPU VRAM during autoregressive generation. This eliminates redundant self-attention re-computations, reducing generation complexity per token from O(N²) to O(N). However, KV cache VRAM footprint scales linearly with context length and batch size, requiring Grouped-Query Attention (GQA) and PagedAttention to fit long context serving in GPU memory."

Follow-ups to expect

• What is KV Cache Quantization (INT8 / INT4 KV Cache)? Quantizes cached K and V tensors from FP16 to INT8 or FP8 before storing in VRAM, reducing KV Cache VRAM footprint by 50-75% with minimal accuracy degradation.
• How does Chunked Prefill work with KV Cache? Splits long input prompt prefill processing into smaller chunks interleaved with decode steps, preventing long prompt prefills from starving ongoing decode token latency.