KV-cache optimization

Why We Need It: Derivation from First Principles

The Naive Attention Cost

In multi-head self-attention, for sequence length $L$, we compute:

Step-by-step cost analysis:

1. During training (teacher forcing): We have all $L$ tokens at once.

   • $Q, K, V \in \mathbb{R}^{L \times d}$ are computed once
   • Attention matrix: $O(L^2 d)$ operations
   • Total per layer: $O(L^2 d + Ld^2)$

2. During autoregressive generation (the problem):

   • Generate token at position $t=1$: compute $K_1, V_1$ ✓
   • Generate token at position $t=2$: compute $Q_2$ (new), but recompute $K_1, V_1, K_2, V_2$ ✗
   • Generate token at position $t=L$: recompute $K_1, \ldots, K_L$ and $V_1, \ldots, V_L$ ✗

Why this step? Each generation step only adds one new token, but without caching, the model treats it as a fresh sequence and recomputes projections for all previous tokens.

Total naive cost for generating $L$ tokens:

This is quadratic in sequence length just for the projection overhead, before even doing attention!

The KV-Cache Solution

Derivation of savings:

Why this works: The causality mask in attention ensures token $t$ only attends to positions $\leq t$. Since $K_i, V_i$ for $i < t$ are fixed (the model is frozen during generation), their values never change—perfect for caching.

Full-model cache (across all layers) = $2 \times n_{\text{layers}} \times L \times h \times d_k \times \text{precision}$.

Example: GPT-3 style model ($n_{\text{layers}}=96$, $h=96$ heads, $d_k=128$, fp16):

• Per token, per layer: $2 \times 96 \times 128 \times 2 = 49{,}152 \text{ bytes} \approx 48\text{ KB}$
• Per token, all 96 layers: $96 \times 48\text{ KB} = 2 \times 96 \times 96 \times 128 \times 2 \text{ bytes} \approx 4.7\text{ MB}$
• For 2048 tokens (all layers): $2048 \times 4.7\text{ MB} \approx 9.6\text{ GB}$

Why this step? We need 2 matrices (K and V), each of size (num_tokens × num_heads × head_dim), stored per layer. The full-model cost multiplies by the number of layers, which is why the per-token figure jumps from ~48 KB (one layer) to ~4.7 MB (all 96 layers).

Worked Examples

Step 1: Previous cache state (after generating tokens 1-2):

Why this step? Each $K_i$ has shape (num_heads, head_dim) = (4, 64). We've generated 2 tokens so far.

Step 2: Compute projections for new token:

Step 3: Update cache (concatenation):

Step 4: Compute attention using only $Q_3$ with all cached keys:

Why this step? $Q_3$ (shape 4×64) multiplies with cached K (shape 3×4×64 → transpose to 4×64×3). Each head gets a 1×3 score vector (attending to 3 positions).

Cost: $O(d^2)$ for one token's projections + $O(t \cdot d)$ for attention scores. No recomputation of $K_1, K_2$.

Generating a 100-token sequence:

Without cache:

• Per layer, per token $t$: Compute $t$ key-value pairs → $O(t \times768^2)$
• Total: $12 \times \sum_{t=1}^{100} t \times 768^2 \approx 3.5 \times 10^{10}$ FLOPs

With cache:

• Per layer, per token: Compute 1 new K, V → $O(768^2)$
• Total: $12 \times 100 \times 768^2 \approx 7.1 \times 10^8$ FLOPs
• Speedup: 50× (approximately matches the $L$ factor for $L=100$)

Memory overhead:

• Cache size: $12 \text{ layers} \times 2 \text{ (K+V)} \times 100 \text{ tokens} \times 12 \text{ heads} \times 64 \text{ dim} \times 2 \text{ bytes} \approx 3.7\text{ MB}$

Why this matters: 3.7 MB is negligible compared to the 500 MB model weights, but saves 98% of projection compute.

Advanced Optimizations

1. Multi-Query Attention (MQA)

Cache reduction:

Why this step? With $h=32$ heads, MQA uses 32× less KV-cache memory. Quality drop is typically 1-2% perplexity.

2. Grouped-Query Attention (GQA)

Middle ground: Group $h$ query heads into $g$ groups ($g < h$), each group shares one K, V head.

Example: Llama-2 70B uses $h=64$, $g=8$ (8× cache reduction vs. standard).

3. Sliding Window / Local Attention

For very long sequences, only cache the last $W$ tokens:

Tradeoff: Constant $O(W)$ memory, but loses long-range dependencies beyond $W$ tokens.

Use case: Summarization with 100K+ token contexts where recent context matters most.

4. FlashAttention Integration

FlashAttention computes attention without materializing the full $O(L^2)$ score matrix by fusing operations and using SRAM tiling. KV-cache + FlashAttention synergize:

• Cache eliminates redundant K, V projections
• FlashAttention eliminates memory overhead of attention scores

Combined speedup: 10-20× for long sequences (both reduce different bottlenecks).

Common Mistakes

Why it feels right: $Q$ also requires a matrix multiply, so caching seems logical.

Steel-man: In a typical forward pass, caching $Q$ would help. But in autoregressive generation, each new token generates a fresh query that attends to all previous positions. We need a new $Q_t$ for token $t$ because it encodes "what the new token is looking for."

The fix: Only K and V are reusable because they represent "what previous tokens offer." The query represents "what the current token needs"—it must be recomputed.

Why it feels right: The cache has the data, why not use it?

Steel-man: During training with full sequences, we use a causal mask. It's easy to forget this is still needed during generation.

The fix: Even with caching, apply causal mask:

This ensures token $t$ only sees positions $1, \ldots, t$ (even though physically $K_{t+1}, \ldots$ don't exist yet in autoregressive mode, this matters for batched generation).

Why it feels right: The notation says $O(L)$ memory.

Steel-man: The notation hides the constants. For a 32-layer, 32-head model with $d_k=128$ in fp16:

For $L=100{,}000$ tokens: ~105 GB just for the cache (exceeds most GPU memory).

The fix: Use hybrid strategies like sliding window cache, offloading old cache to CPU, or sparse attention for ultra-long contexts.

Implementation Details

Typical PyTorch pattern:

class CachedAttention:
    def __init__(self):
        self.kv_cache = None  # Will hold (batch, num_heads, seq_len, head_dim)
    def forward(self, x, use_cache=True):
        B, T, C = x.shape
        Q = self.W_q(x)  # Always compute fresh
        K_new = self.W_k(x)
        V_new = self.W_v(x)
        if use_cache and self.kv_cache is not None:
            K_cached, V_cached = self.kv_cache
            K = torch.cat([K_cached, K_new], dim=2)  # Concat along seq_len
            V = torch.cat([V_cached, V_new], dim=2)
        else:
            K, V = K_new, V_new
        # Standard attention computation
        scores = (Q @ K.transpose(-2, -1)) / math.sqrt(self.head_dim)
        atn = F.softmax(scores, dim=-1)
        out = atn @ V
        
        if use_cache:
            self.kv_cache = (K.detach(), V.detach())  # Detach to avoid gradient retention
        
        return out

Why detach? Prevents backprop through the entire cached history (only needed during training; during inference, no gradients anyway).

When to Use What

The slow way: Every time you write a new word, you re-read the entire story from the beginning and take notes about each word. Writing 100 words means reading the story 100 times!

The KV-cache way: The first time you read each word, you write down notes on sticky notes (that's the "cache"). When you write word #50, you just glance at your 49 sticky notes instead of re-reading the whole story. You only add one new sticky note for word #50.

The tradeoff: You need space on your desk for all the sticky notes (memory), but you save tons of time not re-reading. If your desk gets too crowded (very long story), you might only keep the last 20 sticky notes and throw away older ones (sliding window).

Connections

• 6.1.1-transformer-architecture - The self-attention mechanism that KV-cache optimizes
• 6.1.4-multi-head-attention - How multiple heads multiply cache requirements
• 6.2.3-inference-optimization - Broader category of deployment speedups
• 6.3.1-model-quantization - Reducing cache memory via lower precision (int8, fp16)
• 5.4.2-beam-search - Generation algorithm that becomes cache-intensive with multiple beams
• 6.1.11-sparse-attention - Alternative approach to reduce $O(L^2)$ attention cost
• 7.2.1-llm-deployment - Real-world systems where cache management is critical

#flashcards/ai-ml

What is the primary problem KV-cache solves? :: During autoregressive generation, preventing redundant recomputation of key and value projections for previously generated tokens. Without caching, we recompute K and V for all prior tokens at every step, leading to $O(L^2 d^2)$ cost instead of $O(Ld^2)$.

Derive the KV-cache memory formula per layer :: Cache size = $2 \times L \times h \times d_k \times \text{bytes}$. The 2 accounts for separate K and V matrices, $L$ is sequence length, $h$ is number of attention heads, $d_k$ is head dimension, and bytes is precision (2 for fp16, 4 for fp32). Multiply by $n_{\text{layers}}$ for the full model.

Calculate KV-cache size for GPT-2 Small generating 512 tokens :: GPT-2 Small: 12 layers, 12 heads, $d_k=64$, fp16. Cache = $12 \times 2 \times 512 \times 12 \times 64 \times 2$ bytes = 18.9 MB. This is per-sample; batch size multiplies this.

Concept Map

Hinglish (regional understanding)