Lesson 15 of 20Training and Sampling

KV Caching: Making Generation Fast

One practical wrinkle in generation: producing each new token requires computing attention against every token so far, which means computing K and V vectors for the whole sequence. Doing that from scratch at every single step means position 1's K/V get recomputed at step 2, step 3, step 4... every step, forever, even though position 1 never changes. That's wasted work, and it grows fast: generating N tokens naively costs roughly 1+2+3+...+N K/V computations, quadratic in sequence length.

Generating token by token — which K/V vectors get (re)computed?

The

cat

sat

the

mat

Without cache

With cache

The fix is exactly what it sounds like: a KV cache. Compute each position's K and V vectors once, store them, and on every later step only compute K/V for the newest token, reusing the cached ones for everything before it. That turns the quadratic cost into a linear one, and it's why real chat models can hold long conversations without generation grinding to a halt as the conversation grows.

Python

tokens = ["The", "cat", "sat", "on", "the", "mat"]

total_without_cache = 0
total_with_cache = 0
print(f"{'step':>4}  {'without cache':>14}  {'with cache':>11}")
for step in range(len(tokens)):
    recompute_count_without = step + 1  # every position gets recomputed
    recompute_count_with = 1            # only the new token
    total_without_cache += recompute_count_without
    total_with_cache += recompute_count_with
    print(f"{step:>4}  {total_without_cache:>14}  {total_with_cache:>11}")

print()
print(f"total K/V computations without cache: {total_without_cache}")
print(f"total K/V computations with cache:    {total_with_cache}")

tokens = ["The", "cat", "sat", "on", "the", "mat"]

total_without_cache = 0
total_with_cache = 0
print(f"{'step':>4}  {'without cache':>14}  {'with cache':>11}")
for step in range(len(tokens)):
    recompute_count_without = step + 1  # every position gets recomputed
    recompute_count_with = 1            # only the new token
    total_without_cache += recompute_count_without
    total_with_cache += recompute_count_with
    print(f"{step:>4}  {total_without_cache:>14}  {total_with_cache:>11}")

print()
print(f"total K/V computations without cache: {total_without_cache}")
print(f"total K/V computations with cache:    {total_with_cache}")

What problem does a KV cache solve during text generation?

You now have every piece: tokenize, embed, attend, assemble into blocks, train with gradient descent, and sample. The capstone puts them all together into one small program that actually runs.

← Sampling Strategies: Temperature, Top-k & Top-p Fine-Tuning Intuition (LoRA) →