Autoregressive Generation & The KV Cache: Memory Footprint & Scaling Formulas

1. Two Phases of LLM Inference: Prefill vs. Decode

Transformer-based autoregressive models generate text in two radically distinct operational phases:

  Phase 1: Prefill (Prompt Processing)             Phase 2: Decode (Token-by-Token)
  Compute-Bound (High Parallelism)                  Memory-Bandwidth Bound (Sequential)
  ┌─────────────────────────────────────┐           ┌─────────────────────────────────────┐
  │ Processes all prompt tokens (e.g.   │           │ Generates 1 token at a time:        │
  │ 2,048 tokens) concurrently via GEMM │ ──► K, V  │ Attends to all past tokens via      │
  │ matrix multiplication. Saturates GPU│   Cache   │ cached Keys and Values. GPU memory  │
  │ Tensor Cores.                       │           │ bandwidth is the primary bottleneck.│
  └─────────────────────────────────────┘           └─────────────────────────────────────┘

2. Why the KV Cache is Essential

In standard self-attention, generating token $t$ requires computing attention against all tokens $1 \dots t-1$. Without caching, previous key ($K$) and query/value ($V$) states must be recomputed at every step, creating an $O(N^2)$ generation cost.

• The KV Cache stores computed Key and Value vectors for all past tokens in GPU VRAM.
• Each new decoding step only computes Key and Value for the single newly generated token ($O(1)$ per step).

3. KV Cache Memory Sizing Formula

The memory footprint of the KV cache grows linearly with batch size and context sequence length:

$$\text{Memory (bytes)} = 2 \times 2 \times n_{\text{layers}} \times n_{\text{heads}} \times d_{\text{head}} \times \text{seq_len} \times \text{batch_size} \times \text{bytes_per_elem}$$

Where:

• $2 \times$ accounts for both Keys and Values.
• $\text{bytes_per_elem} = 2$ for FP16/BF16, or $1$ for FP8.

Concrete Example (Llama 3 70B, BF16, Context = 8,192, Batch Size = 8):

• Layers = 80, KV Heads = 8 (Grouped-Query Attention), Head Dim = 128
• $\text{Bytes} = 2 \times 80 \times 8 \times 128 \times 8,192 \times 8 \times 2 = 21,474,836,480\text{ bytes} \approx \mathbf{21.47\text{ GB}}$

The KV cache for just 8 concurrent users consumes over 21 GB of GPU VRAM!

def calculate_kv_cache_vram_gb(
    num_layers: int, 
    num_kv_heads: int, 
    head_dim: int, 
    seq_len: int, 
    batch_size: int, 
    dtype_bytes: int = 2
) -> float:
    # 2 for Key and Value
    total_bytes = 2 * num_layers * num_kv_heads * head_dim * seq_len * batch_size * dtype_bytes
    return total_bytes / (1024 ** 3)

# Example: Llama 3 8B (32 layers, 8 KV heads, 128 head dim, 8k tokens, batch 16)
vram_gb = calculate_kv_cache_vram_gb(32, 8, 128, 8192, 16, dtype_bytes=2)
print(f"Required KV Cache VRAM: {vram_gb:.2f} GB")  # ~17.18 GB