Inference & KV Cache

Generation is autoregressive: emit a token, append it, feed everything back in, repeat. Done naively this is quadratic. kv_cache.py is the same GPT — identical weights — with one change to forward() that makes generation linear per step.

Summary

Script: kv_cache.py  ·  Loads: ckpt_sft.pt (falls back to ckpt_base.pt)  ·  Run: python kv_cache.py

The problem: O(T²) generation

base_model.GPT.generate() re-runs the entire context through all six blocks for every new token. To emit token T+1 it recomputes the keys and values for tokens 0..T — which it already computed the previous step. Step t costs O(t), so a full generation is O(T²) of wasted work.

The fix: cache keys and values

Keys and values for past tokens never change, so compute them once and keep them. Each attention layer returns the (k, v) it computed; generate() stores them and passes them back on the next step, which then feeds in only the new token:

if past_kv is not None:
    pk, pv = past_kv                # (B, nh, past_len, hd)
    k = torch.cat([pk, k], dim=2)   # THE cache step: reuse, don't recompute
    v = torch.cat([pv, v], dim=2)

Per-step cost drops from "the whole prefix" to "one token." The first step still processes the full prompt (filling the cache in one pass); every step after feeds a single position.

Two things must respect the cache

1. Position embeddings — the new token sits at absolute position past_len, not 0, so positions are offset by the cache length:

   pos = torch.arange(past_len, past_len + T, device=idx.device)
   

2. The causal mask — a single query at global position past_len may attend to all cached keys. SDPA's is_causal=True assumes the query at row i sits at position i, which is wrong once a cache offsets the queries, so the mask is built explicitly and sliced at the global row index:

   attn_mask = self.mask[S - T:S, :S]     # S = past_len + T total keys
   

Weight-compatible by design

The parameter names and shapes are identical to base_model.GPT, so any checkpoint from pretrain / sft / dpo loads here unchanged — only forward() differs. The cache is an inference optimization, never a retraining requirement.

What main() demonstrates

Running the script shows both correctness and speed:

• Correctness — feeding a prompt token-by-token through the cache reproduces a single full forward pass to within float noise (max |full − incremental| logit difference ≈ 1e-5, not exact zero, because the math is identical but the batching differs).
• Speed — the same 200-token generation is timed with and without the cache, printing tokens/s for each so the quadratic-vs-linear gap is visible.

python kv_cache.py

---

📄 Source: kv_cache.py