DL S042 · KV Cache — Why Inference Is 100× Faster With It

Series: Deep Learning & LLMs From Scratch — 80 sessions · Session 42 / 80 · Module M07 · ~1.7 hours

The story

You wrote generate() in S040. Look at it again. Every generation step, you feed the entire growing context back through the model — all N positions get recomputed from scratch, all their Q/K/V projections, all their attention, all their FFNs. To generate 500 tokens starting from a 1-token prompt, you do work proportional to 1² + 2² + 3² + ... + 500² ≈ 42 million position-computes. To generate the same 500 tokens with a KV cache, you do 500 position-computes. That's ~85,000× fewer.

The insight: at generation step t, you're only producing a new key, value, and query for position t. All previous positions' K and V vectors were computed on earlier steps and haven't changed. So cache them. On step t, compute only q_t, k_t, v_t, append k_t and v_t to a running cache, and compute attention as softmax(q_t · [K_cache; k_t]ᵀ) · [V_cache; v_t]. Attention scores are now shape (1, t+1) instead of (t+1, t+1). FFN runs on one position instead of t+1.

This is the trick that makes LLM chatbots economically viable. vLLM, TensorRT-LLM, TGI, and every commercial serving framework live and die by KV cache management. Today we build the basic version; in S066 and S071 we cover PagedAttention (vLLM's contribution) and continuous batching.

Prerequisites

• S040, S041 — the model and generation loop we'll optimise.

1 · Why naive generation is wasteful

Recall the S040 generate loop:

Every iteration feeds idx_cond (length t) through the model, computes Q, K, V for every position, computes the full (t, t) attention matrix, runs the FFN on all t positions. Then we grab the last position's logits and throw away all the intermediate work for positions 0..t-2 — even though they were computed identically on the previous step.

Total work for generating N tokens: Σ_{t=1}^{N} O(t²) ≈ O(N³). That's cubic. For N = 2048, that's ~8 billion operations of overhead you didn't need to do.

The KV cache brings it down to O(N²) total (each step is O(N) in the attention). And in practice, most of the "N" here is memory bandwidth, not compute — you're reading a growing KV tensor, not multiplying huge things.

2 · The core observation

For a decoder-only causal transformer, during generation:

• K and V for past positions are FIXED once computed. Position t's key vector k_t depends only on the input token at position t (and layer inputs at position t, which also depend only on earlier positions). Feeding token t+1 doesn't change k_t or v_t.
• Q for past positions is IRRELEVANT. We only need attention from the new position — i.e., q_t · K[:t+1]. We don't need q_{t-1} any more; that computation was for producing the previous token.
• FFN for past positions is IRRELEVANT. The residual stream at position < t doesn't affect the output at position t in a decoder-only model, once its K/V have been cached.

So on step t, the compute is:

1. Take the new token embedding (single position).
2. Run it through the model: for each layer, compute q_t, k_t, v_t (single position each), append k_t, v_t to that layer's cache.
3. Attention: q_t · [K_cache; k_t]ᵀ, softmax, times [V_cache; v_t]. Result is a single new position's attention output.
4. FFN on single position. Residual add. Onto the next layer.
5. Final LN + head, single position. Sample the next token.

Per step: O(t) for the attention (dot product against growing cache), O(1) for everything else. Total for N tokens: O(N²).

3 · The implementation

Modify CausalSelfAttention.forward to accept an optional cache:

class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        # ... same as S040 ...
 
    def forward(self, x, kv_cache=None):
        B, T, C = x.size()
        d_k = C // self.n_head
        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
        q = q.view(B, T, self.n_head, d_k).transpose(1, 2)   # (B, h, T, d_k)
        k = k.view(B, T, self.n_head, d_k).transpose(1, 2)
        v = v.view(B, T, self.n_head, d_k).transpose(1, 2)
 
        # Append to cache
        if kv_cache is not None:
            past_k, past_v = kv_cache
            k = torch.cat([past_k, k], dim=2)               # (B, h, T_past + T, d_k)
            v = torch.cat([past_v, v], dim=2)
        new_cache = (k, v)
 
        # Attention
        # When using cache, T=1 (new token only), T_total = T_past + 1
        # No causal mask needed — the new query can attend to all cached positions
        att = (q @ k.transpose(-2, -1)) / math.sqrt(d_k)     # (B, h, T, T_total)
        # Note: if T > 1 (e.g. processing a full prompt), we still need a causal mask
        # on the "new" positions relative to each other. Standard trick: apply mask only
        # to the T × T sub-block on the right.
        if T > 1:
            T_total = k.size(2)
            mask = torch.tril(torch.ones(T, T, device=x.device))
            # Position i in the new block can see cached positions (all) + new positions 0..i
            full_mask = torch.ones(T, T_total, device=x.device)
            full_mask[:, T_total-T:] = mask
            att = att.masked_fill(full_mask == 0, float('-inf'))
        att = F.softmax(att, dim=-1)
        y = att @ v                                          # (B, h, T, d_k)
 
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.c_proj(y), new_cache

The GPT forward propagates the cache through all layers:

def forward(self, idx, kv_caches=None):
    B, T = idx.size()
    past_len = 0 if kv_caches is None else kv_caches[0][0].size(2)
    pos = torch.arange(past_len, past_len + T, device=idx.device)
 
    tok = self.transformer.wte(idx)
    pos_emb = self.transformer.wpe(pos)
    x = self.transformer.drop(tok + pos_emb)
 
    new_caches = []
    for i, block in enumerate(self.transformer.h):
        past = None if kv_caches is None else kv_caches[i]
        x, new = block(x, kv_cache=past)
        new_caches.append(new)
 
    x = self.transformer.ln_f(x)
    logits = self.lm_head(x[:, [-1], :])
    return logits, new_caches

Note the pos computation — the new tokens' positions are past_len, past_len+1, …, NOT 0, 1, …. Get this wrong and generation gets scrambled after the first cached step.

And the generate loop:

@torch.no_grad()
def generate_kv(self, idx, max_new_tokens, temperature=1.0, top_k=None):
    kv_caches = None
    # Process the prompt in one shot (fills the cache with T positions)
    logits, kv_caches = self(idx, kv_caches=kv_caches)
    for _ in range(max_new_tokens):
        logits = logits[:, -1, :] / temperature
        if top_k is not None:
            v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
            logits[logits < v[:, [-1]]] = -float('Inf')
        probs = F.softmax(logits, dim=-1)
        idx_next = torch.multinomial(probs, num_samples=1)
        idx = torch.cat((idx, idx_next), dim=1)
        # Only pass the NEW token through the model
        logits, kv_caches = self(idx_next, kv_caches=kv_caches)
    return idx

The critical line: self(idx_next, kv_caches=kv_caches) — passing only the ONE new token, not the whole context. That's what makes each step O(N) instead of O(N²).

4 · Benchmark — measure the speedup

import time
torch.manual_seed(0)
prompt = torch.zeros((1, 1), dtype=torch.long, device='cuda')
 
t0 = time.time()
_ = model.generate(prompt, max_new_tokens=1000, temperature=1.0)      # naive
print(f"Naive: {time.time() - t0:.2f}s")
 
t0 = time.time()
_ = model.generate_kv(prompt, max_new_tokens=1000, temperature=1.0)   # cached
print(f"KV cache: {time.time() - t0:.2f}s")

Typical results (small model on A100):

• Naive: 12.4 s
• KV cache: 1.6 s
• Speedup: ~8×

At max_new_tokens = 100, speedup is only ~2×. At 4096, it's ~30×. The ratio grows with sequence length because naive is O(N²) in the length and cached is O(N).

import time
lengths = [64, 128, 256, 512, 1024, 2048]
for N in lengths:
    for tag, use_cache in [("naive", False), ("cache", True)]:
        t0 = time.perf_counter()
        _ = model.generate(prompt, max_new_tokens=N, use_cache=use_cache)
        dt = time.perf_counter() - t0
        print(f"N={N:5d}  {tag:>5s}  {dt:6.3f}s")

5 · The memory cost

KV cache is not free. Memory per token, per layer, per batch element:

(the 2 is for K and V; the head/dim factors collapse to n_embd).

For Llama-2-7B (n_embd = 4096, n_layer = 32) in fp16 (2 bytes):

For a 4096-token context: 0.5 MB × 4096 = 2 GB PER SEQUENCE.

A single 40-GB A100 running Llama-2-7B has ~13 GB left after model weights, so it can hold ~6 concurrent 4k-context sequences before running out of KV cache memory. This is why KV cache is the memory bottleneck for LLM serving, not model weights.

vLLM's PagedAttention (S066) attacks this by allocating KV cache in fixed-size "pages" like an OS virtual memory system, dramatically reducing fragmentation and enabling ~4× more concurrent requests. Same math, different memory management.

6 · Pitfalls

7 · Beyond basic KV cache — a preview

Real-world LLM serving has many KV-cache optimizations we'll cover later:

Every one of these builds on the basic KV cache you just implemented.

8 · Modern-2025 twist — what actually runs in production

The basic KV cache you just wrote is table stakes. Every 2025 open-model serving stack layers a stack of tricks on top, each solving a different pathology of the naive approach:

1. PagedAttention — [vLLM (Kwon et al., SOSP 2023, arXiv:2309.06180)]. Fragmentation kills naive KV cache: if you allocate contiguous buffers of max_len per request, ~60% of KV memory is wasted on unused padding. vLLM allocates in 16-token "pages" like OS virtual memory, hits ~96% utilisation, serves ~2–4× more concurrent requests. Every serious open-source serving stack has now adopted the pattern (vLLM v1, SGLang, TensorRT-LLM, TGI).

2. Prefix caching — same system prompt for 10,000 users? Compute the KV cache once, share it. RadixAttention in SGLang (Zheng et al., 2024, arXiv:2312.07104) stores prefix caches in a radix tree keyed on token prefix. For chatbot workloads with long system prompts this is a 5–10× throughput win.

3. GQA / MQA / MLA — covered in S037's modern-twist section. Shrinks the KV cache 4–64× by sharing K/V across query heads (GQA) or compressing them (MLA). This is why Llama 3 70B (GQA g=8) can serve 4k contexts on a single H100 while an equivalent MHA model can only serve a few.

4. Quantized KV cache — store K and V in int8 or int4 instead of fp16. KIVI (Liu et al., 2024, arXiv:2402.02750) shows per-channel int2 KV for K and per-token int2 for V hits ~4× memory reduction with <1% quality loss. TensorRT-LLM ships int8 KV out of the box.

5. Speculative decoding (S063 preview) — use a small draft model to propose K tokens per step, verify with the big model in one batched forward. Realistic 2–4× wall-clock speedup on top of KV cache. Chinese labs (DeepSeek-V3 with multi-token prediction, Zhipu GLM-4) increasingly train models with native MTP heads to skip the draft-model dance.

6. Sliding window + attention sinks — StreamingLLM (Xiao et al., 2023, arXiv:2309.17453). Keep only the first few "sink" tokens plus a sliding window of recent tokens. Enables infinite-length inference at bounded memory, at the cost of losing mid-history detail. Used in Mistral 7B's original release and several 2024 forks.

Cheat-sheet KV cache size, per token per layer, K+V together, fp16:

This table is worth memorising: it's the difference between "we can serve 20 users on this GPU" and "we can serve 2."

Further reading:

• vLLM PagedAttention paper
• Aleksa Gordić, "KV Cache Explained" (2024) — clean intuitive walk-through.
• Character.AI's engineering blog, "Optimizing Inference" (2024) — real production numbers for a 20B chatbot.

# 1. Batch scaling. If decoding were compute-bound, per-step time
#    should scale with batch. Measure it.
for bs in (1, 2, 4, 8, 16, 32):
    t = time_one_decode_step(model, batch=bs, ctx=128)
    print(bs, t, t / bs)     # watch time-per-sequence fall sharply

# 2. Context scaling at fixed batch. Per-step FLOPs grow only
#    with the attention term, which is small -- but the cache read
#    grows linearly with ctx.
for ctx in (128, 512, 2048, 8192):
    print(ctx, time_one_decode_step(model, batch=1, ctx=ctx))

# Near-flat time across batch + rising time with context
# = memory bandwidth, not arithmetic.

Retention scaffold

One-line summary (write it in your own words): _______________________________________________

Spaced review: re-read §5 (memory math) and §8 (the production tricks) in 24 hours. On day 7, memorise the KV-cache cheat-sheet table for Llama 3 sizes.

Next session (S043): swap sinusoidal PE for RoPE, and explain FlashAttention-3's tiling trick — how a single kernel can compute softmax(QKᵀ)V without ever materialising the full (T, T) score matrix, saving huge memory on long contexts. We also close the loop on MLA.

Sticky note (keep on your desk): KV cache: past K,V frozen → cache them, only compute new q_t, k_t, v_t. O(N³) → O(N²). Memory dominates weights past 2k tokens. GQA/MLA + PagedAttention are how you make production math work.

Previous: ← DL S041 · Next: DL S043 →