09 — Architecture

KV-Cache#

Memory Optimization✓ Mathematical
◆ The PatternCache once, reuse forever — the key to fast autoregressive generation

During generation, each new token only needs to compute its own Q, K, V — but it attends to all previous K and V vectors. Without caching, we'd recompute K,V for all prior tokens at every step. The KV-cache stores these, turning generation from O(n²) to O(n) per step.

Step t: K_cache = [K₁, K₂, ..., K_t], V_cache = [V₁, V₂, ..., V_t]
Append new K_t, V_t each step. Only compute attention for the new query against cached K,V.
Memory: 2 · n_layers · seq_len · n_kv_heads · d_k · bytes_per_param
For LLaMA-70B with 4K context in FP16: ~2.5 GB per request just for KV cache.

KV-cache is why batch size during inference is heavily memory-constrained. GQA (fewer KV heads) directly reduces this cost. This is the main motivation behind MQA and GQA research.

KV-cache memory scales linearly with sequence length × batch size. For long contexts (128K+), a single request can consume 40+ GB. This dominates GPU memory during serving.
Interactive — KV-cache growth during generation

Python — KV-cache in generation loop#

def generate_with_cache(model, prompt_ids, max_new=50):
    past_kv = None  # will hold cached K,V tensors
    input_ids = prompt_ids

    for _ in range(max_new):
        # Only feed new token(s) — cache handles the rest
        logits, past_kv = model(input_ids, past_key_values=past_kv)
        next_id = logits[:, -1, :].argmax(dim=-1, keepdim=True)
        input_ids = next_id  # only the new token
        if next_id.item() == eos_token_id:
            break
    return all_generated_ids
Pattern bridge: Caching previously computed keys and values to avoid recomputation. The same memory accumulation that RNN hidden states perform. In markets, support/resistance levels are cached price memories the market doesn’t recompute.
← Previous
Decoder-Only Models
Open in the full reader, with the topic sidebar →