Autoregressive decoding and the KV cache

Autoregressive decoding and the KV cache#

Track 01 - Inference engines · Notebook 01 · Runtime: ≈20 min on T4, ≈6 min on CPU

Prerequisites: you have the GPU mental model from 07_gpu/01_gpu_architecture_tour. You’ve written a recursive Fibonacci function at some point and noticed it’s slow.

What you’ll know by the end: why modern LLMs can’t avoid the “read the whole model per token” treadmill, why caching is the one trick that saves us, and how to reason about the memory footprint of that cache in your head.

Why this notebook exists#

If you’ve ever seen an LLM type out a response, you’ve watched a GPU do a staggering amount of pointless work and then undo it with one cache.

Concretely: generating the 1000th token of a response could cost 1000× the compute of generating the first one, because each new token’s attention depends on all previous tokens. That’s a textbook O(N²) trap - the same one you hit when a naive recursive Fibonacci function recomputes fib(5) a hundred times.

The fix is the same as in Fibonacci: memoize the intermediate results you’d otherwise recompute. In LLMs that memo table is called the KV cache, and it is the reason decoding is tractable.

The Fibonacci analogy (trust me)#

Here’s a pattern every CS undergrad has seen:

def fib(n):
    if n < 2: return n
    return fib(n - 1) + fib(n - 2)

Compute fib(30) and your laptop fans spin up. The problem isn’t that addition is slow - it’s that you recompute fib(25) thousands of times. Memoize it and the whole thing runs in microseconds.

Now think about an LLM. Predicting token \(t\) requires running attention over tokens \(1 \ldots t - 1\). Predict token \(t+1\) and you do it again, from scratch, over tokens \(1 \ldots t\). That’s the same redundancy as unmemoized Fibonacci - except instead of wasted additions, you waste weight reads from HBM, which (as we saw in the GPU notebook) is the scarce resource.

The KV cache stashes the per-layer Key and Value tensors for every token the moment they’re computed, so subsequent steps only do work for the newest token. That’s it. Same trick, different vocabulary.

The number to remember: with a good KV cache, generating token \(t{+}1\) costs roughly the same compute as generating token \(1\). Without one, it costs \(t{+}1\) times more.

from llm_systems_cookbook.nb import bootstrap

import math
import time

import torch

from llm_systems_cookbook.models import get as get_model

s = bootstrap("01_inference_01_autoregressive_decoding_kv_cache")

Step 1 - How big is this cache going to be?#

Before we run anything, let’s predict the cache size with a back-of-envelope formula. You’ll use this formula all the time when you’re sizing a deployment, so commit it to memory.

For each layer we store one K tensor and one V tensor. Each is (batch, num_kv_heads, seq_len, head_dim). So the total bytes are:

\[ \text{KV bytes} = 2 \cdot L \cdot H_\text{kv} \cdot D_\text{head} \cdot T \cdot B \cdot \text{dtype\_bytes} \]

The 2 is K and V. Everything else is “storage per element × number of elements.”

A concrete example you can use as a mental anchor: Llama-3.1-8B with 32 layers, 8 KV heads (grouped-query attention), head-dim 128, sequence length 4096, batch 1, fp16. Plug in:

\[ 2 \cdot 32 \cdot 8 \cdot 128 \cdot 4096 \cdot 1 \cdot 2 = 536\,870\,912 \text{ bytes} \approx 512 \text{ MiB} \]

Half a gigabyte. That’s just for a single-sequence, 4k-token cache. Now you see why batched serving with long contexts is such a memory pressure

each concurrent user needs their own cache.

def kv_cache_bytes(
    num_layers: int,
    num_kv_heads: int,
    head_dim: int,
    seq_len: int,
    batch: int = 1,
    dtype_bytes: int = 2,
) -> int:
    '''Closed-form KV-cache size in bytes.'''
    return 2 * num_layers * num_kv_heads * head_dim * seq_len * batch * dtype_bytes


ref = kv_cache_bytes(num_layers=32, num_kv_heads=8, head_dim=128, seq_len=4096)
print(f"Llama-3.1-8B @ 4k, fp16, B=1  -> {ref / 1024**2:.0f} MiB")

# Two sanity checks that pin down the formula's behaviour.
s.check("kv_formula_matches_hand_calc", lambda: ref == 2 * 32 * 8 * 128 * 4096 * 2)
s.check("kv_formula_linear_in_seq",     lambda: kv_cache_bytes(32, 8, 128, 8192) == 2 * ref)

Step 2 - Meet the model#

We’ll use SmolLM2-135M - the smallest model in our registry. Tiny enough to run on a Colab CPU, big enough that the KV cache actually shows up in memory measurements. The architecture is the same family as Llama; what we learn here transfers directly to the 8 B and 70 B models you’ll meet later.

After loading we print the architectural constants because every later calculation will reference them.

from transformers import AutoModelForCausalLM, AutoTokenizer

spec = get_model("smollm2-135m")
print(f"loading {spec.hf_id}")
tok = AutoTokenizer.from_pretrained(spec.hf_id)
# transformers 4.46 renamed torch_dtype -> dtype; fall back for older builds.
try:
    model = AutoModelForCausalLM.from_pretrained(spec.hf_id, dtype=torch.float32).to(DEVICE)
except TypeError:
    model = AutoModelForCausalLM.from_pretrained(spec.hf_id, torch_dtype=torch.float32).to(DEVICE)
model.eval()

cfg = model.config
NUM_LAYERS   = cfg.num_hidden_layers
NUM_KV_HEADS = getattr(cfg, "num_key_value_heads", cfg.num_attention_heads)
HEAD_DIM     = cfg.hidden_size // cfg.num_attention_heads
DTYPE_BYTES  = 2 if model.dtype in (torch.float16, torch.bfloat16) else 4
print(f"L={NUM_LAYERS} H_kv={NUM_KV_HEADS} D={HEAD_DIM} dtype_bytes={DTYPE_BYTES}")

Step 3 - Naive vs cached decoding, side by side#

We’re going to implement two versions of greedy generation and run them on the same prompt:

generate_no_cache - every new token re-runs the model on the entire sequence so far. This is the Fibonacci-without-memoization version. It burns compute in \(O(T^2)\).
generate_with_cache - uses Hugging Face’s DynamicCache object. Prefill runs once on the whole prompt, then each new token sends in just one token ID and pulls cached K/V for everything before it. Compute per token is \(O(1)\) in the context.

Two things must be true if we’ve done this right:

They produce the same tokens. Greedy decoding is deterministic, so the cache is a pure optimisation. If outputs differ, we have a bug.
The cached version is much faster. On any real GPU we expect ≥ 5× throughput; on CPU the speedup is smaller but still visible.

from transformers import DynamicCache


@torch.inference_mode()
def generate_no_cache(prompt_ids: torch.Tensor, n_new: int) -> tuple[torch.Tensor, float]:
    '''Greedy decode with no cache: re-feed the full sequence on every step.'''
    ids = prompt_ids.clone()
    t0 = time.perf_counter()
    for _ in range(n_new):
        out = model(ids)  # full forward over len(ids) tokens.
        next_id = out.logits[:, -1].argmax(-1, keepdim=True)
        ids = torch.cat([ids, next_id], dim=-1)
    return ids, time.perf_counter() - t0


@torch.inference_mode()
def generate_with_cache(
    prompt_ids: torch.Tensor, n_new: int
) -> tuple[torch.Tensor, float, DynamicCache]:
    '''Greedy decode with a DynamicCache: one prefill, then one-token forwards.'''
    cache = DynamicCache()
    ids = prompt_ids.clone()
    t0 = time.perf_counter()
    # Prefill: process the entire prompt once and fill the cache.
    out = model(ids, past_key_values=cache, use_cache=True)
    cache = out.past_key_values
    next_id = out.logits[:, -1].argmax(-1, keepdim=True)
    ids = torch.cat([ids, next_id], dim=-1)
    # Decode: send only the newest token; attention reads K/V from cache.
    for _ in range(n_new - 1):
        out = model(next_id, past_key_values=cache, use_cache=True)
        cache = out.past_key_values
        next_id = out.logits[:, -1].argmax(-1, keepdim=True)
        ids = torch.cat([ids, next_id], dim=-1)
    return ids, time.perf_counter() - t0, cache


prompt = "The key idea behind the Transformer architecture is"
prompt_ids = tok(prompt, return_tensors="pt").input_ids.to(DEVICE)
N_NEW = 32

out_no,  dt_no  = generate_no_cache(prompt_ids, N_NEW)
out_yes, dt_yes, final_cache = generate_with_cache(prompt_ids, N_NEW)

print(f"no-cache:   {dt_no * 1000:7.1f} ms  ({N_NEW / dt_no:6.1f} tok/s)")
print(f"with-cache: {dt_yes * 1000:7.1f} ms  ({N_NEW / dt_yes:6.1f} tok/s)")
print(f"generated:  {tok.decode(out_yes[0], skip_special_tokens=True)!r}")

Correctness first, speed second#

Before we celebrate the speedup, we check that the cache didn’t change the semantics. Caching is an implementation detail - the model’s probability distribution over the next token must be bit-for-bit identical. Greedy decoding is a pure function of that distribution, so the two runs must emit the same tokens. Any deviation here is a bug in the attention masking, not a feature.

s.check(
    "greedy_outputs_agree",
    lambda: torch.equal(out_no, out_yes),
    msg="no-cache and cached greedy must produce identical token ids",
)
s.benchmark(
    "cache_speedup_at_least_3x",
    measured=N_NEW / dt_yes,
    baseline=N_NEW / dt_no,
    must_beat=3.0,
    unit="tok/s",
    higher_is_better=True,
)

Step 4 - Memory scales linearly with context#

The formula says cache bytes grow linearly with sequence length. Let’s watch it happen on a real forward pass by measuring max_memory_allocated before and after decoding at several context lengths.

Two caveats:

On CPU, PyTorch doesn’t track peak allocation in the same way, so the memory check gets skipped (the scorer records it as “skipped”, not “failed”). The timing measurements still run.
The measured peak includes everything - weights, activations, intermediate tensors - not just the KV cache. What we really verify is that the difference between two contexts matches the formula’s prediction for the difference in KV bytes. Other tensors cancel out.

IS_CUDA = torch.cuda.is_available()
CONTEXTS = [128, 256, 512, 1024, 2048] if IS_CUDA else [128, 256, 512]

filler = tok("Transformers " * 4096, return_tensors="pt").input_ids[0].to(DEVICE)

mem: dict[int, dict[str, float]] = {}
throughput: dict[int, float] = {}
for ctx in CONTEXTS:
    ids = filler[:ctx].unsqueeze(0)
    predicted_mb = kv_cache_bytes(
        NUM_LAYERS, NUM_KV_HEADS, HEAD_DIM, seq_len=ctx + 16, batch=1,
        dtype_bytes=DTYPE_BYTES,
    ) / 1024**2
    if IS_CUDA:
        torch.cuda.empty_cache()
        torch.cuda.reset_peak_memory_stats()
    _, dt, _ = generate_with_cache(ids, n_new=16)
    peak_mb = torch.cuda.max_memory_allocated() / 1024**2 if IS_CUDA else float("nan")
    mem[ctx] = {"predicted_kv_mb": predicted_mb, "measured_peak_mb": peak_mb}
    throughput[ctx] = 16 / dt
    print(f"ctx={ctx:>4}  predicted KV={predicted_mb:6.1f} MiB  "
          f"measured peak={peak_mb:6.1f} MiB  {throughput[ctx]:5.1f} tok/s")

if IS_CUDA:
    peaks = [mem[c]["measured_peak_mb"] for c in CONTEXTS]
    delta_measured = peaks[-1] - peaks[0]
    delta_predicted = (
        kv_cache_bytes(NUM_LAYERS, NUM_KV_HEADS, HEAD_DIM,
                        CONTEXTS[-1] + 16, dtype_bytes=DTYPE_BYTES)
        - kv_cache_bytes(NUM_LAYERS, NUM_KV_HEADS, HEAD_DIM,
                          CONTEXTS[0] + 16, dtype_bytes=DTYPE_BYTES)
    ) / 1024**2
    print(f"\nΔ measured peak = {delta_measured:.1f} MiB  vs Δ predicted KV = {delta_predicted:.1f} MiB")
    s.assert_close(
        "memory_growth_tracks_kv_formula",
        actual=delta_measured,
        expected=delta_predicted,
        rtol=0.35,
    )
else:
    s.skip("memory_growth_tracks_kv_formula", "peak-memory metric requires CUDA")

s.check(
    "tok_per_s_nondegenerate_all_contexts",
    lambda: all(throughput[c] > 0.5 for c in CONTEXTS),
    msg=f"measured {throughput}",
)

The two things to eyeball#

Left: the closed-form KV prediction against the GPU’s measured peak allocation. They shouldn’t be equal - peak includes weights and activations too - but the slope in context length should match. Right: decode throughput vs context. Tokens/sec holds roughly flat because decode is memory-bound on weights, not on the cache.

import matplotlib.pyplot as plt

if IS_CUDA:
    ctxs = list(CONTEXTS)
    predicted = [mem[c]["predicted_kv_mb"] for c in ctxs]
    measured  = [mem[c]["measured_peak_mb"] for c in ctxs]
    tps       = [throughput[c] for c in ctxs]

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(9, 3.2))
    ax1.plot(ctxs, predicted, "o--", label="predicted KV (formula)")
    ax1.plot(ctxs, measured,  "o-",  label="measured peak")
    ax1.set_xlabel("context length")
    ax1.set_ylabel("MiB")
    ax1.set_title("KV memory: formula vs measured peak")
    ax1.legend()
    ax1.grid(True, alpha=0.3)

    ax2.plot(ctxs, tps, "o-", color="tab:red")
    ax2.set_xlabel("context length")
    ax2.set_ylabel("tok/s")
    ax2.set_title("cached decode throughput")
    ax2.grid(True, alpha=0.3)

    fig.tight_layout()
    plt.show()
else:
    print("skipped - peak-memory metric requires CUDA.")

Step 5 - Why decode is memory-bound (and prefill isn’t)#

You now have enough to answer a question that confuses a lot of engineers the first time they think about it:

“We’re doing attention, which is O(N²) in context length. How can it possibly be memory-bound rather than compute-bound?”

Here’s the intuition. During decode (generating one new token), attention only has one query row but needs to read all \(T\) cached keys and values. That’s a lot of bytes moved for very little math - roughly 1-5 FLOPs per byte. Way below the ridge of any GPU. So the chip finishes the compute in nanoseconds and then waits for the next byte to arrive.

During prefill (processing a long prompt), every token attends to every other token in one giant matmul. Arithmetic intensity shoots up to hundreds of FLOPs per byte - compute-bound, and this is where your GPU’s TFLOP rating finally gets spent.

The practical consequence is that decode latency is dominated by weight reads, not by attention math. Making attention faster by a factor of two barely helps; making weight reads faster (via quantisation, speculative decoding, batching) is where the big wins are.

NUM_PARAMS = sum(p.numel() for p in model.parameters())

# Decode AI: read every weight once + a KV tile for the single new token.
DECODE_AI  = 2 / DTYPE_BYTES                   # ≈1 for fp16; far below any GPU ridge.
# Prefill AI: same weight read amortised across T tokens of parallel work.
PREFILL_AI = 2 * 512 / DTYPE_BYTES             # ≈512 for T=512 prefill.

print(f"params = {NUM_PARAMS/1e6:.1f}M")
print(f"decode  arithmetic intensity ≈ {DECODE_AI:.1f}  FLOPs/byte  (memory-bound)")
print(f"prefill arithmetic intensity ≈ {PREFILL_AI:.0f}  FLOPs/byte  (compute-bound)")

s.check("decode_is_memory_bound",   lambda: DECODE_AI < 50,   msg=f"AI={DECODE_AI:.1f}")
s.check("prefill_is_compute_bound", lambda: PREFILL_AI > 100, msg=f"AI={PREFILL_AI:.0f}")

Exercises#

Without running it, estimate: how much KV memory does a batch of 32 concurrent 2k-token conversations need on Llama-3.1-8B fp16? What about fp8? Check your answer against the formula.
Quantise the cache. Store K and V as int8 instead of fp16. Memory drops 2×. Will the quality drop measurably on a short prompt? (The answer is surprisingly “barely” - which is why int8 KV is a standard technique.)
Break the cache. Modify generate_with_cache to reset the cache between steps and rerun the correctness check. It must fail: you’ve turned cached decoding back into no-cache decoding. Watch the tok/s collapse.

What’s next#

Notebook 02 in this track (attention from scratch and roofline) takes the attention operation apart and measures its arithmetic intensity empirically. Notebook 03 (PagedAttention block allocator) fixes the memory fragmentation you implicitly accepted when we stored KVs in contiguous tensors.

s.summary()
s.save()