Why Cache Anything

Notes

KV caching is the inference-time memoization of the key and value projections for all previously generated tokens in an autoregressive Transformer, so step $t$ computes only $q_t,k_t,v_t$ for the new token and attends $q_t$ to cached $K_{1:t},V_{1:t}$ instead of recomputing $K,V$ for the entire prefix. For one attention layer, uncached decoding repeatedly evaluates $K=X{W}_K,V=X{W}_V$ for all tokens seen so far and forms attention for the growing prefix; cached decoding updates $K_{\mathrm{cache}}\leftarrow [K_{\mathrm{cache}};x_t{W}_K]$ and $V_{\mathrm{cache}}\leftarrow [V_{\mathrm{cache}};x_t{W}_V]$, then computes $\mathrm{Attn}(q_t,K_{\mathrm{cache}},V_{\mathrm{cache}})=\mathrm{softmax}(q_t{K}_{\mathrm{cache}}^{\top }/\sqrt{d_h})V_{\mathrm{cache}}$. The useful measurement is therefore not “tokens per second” in isolation, but $S(L)=T_{\mathrm{uncached}}(L)/T_{\mathrm{cached}}(L)$ for the same prompt and output length $L$, because the avoided work grows with the decoded suffix. GPT-2 is a good lab model precisely because the difference is visible on a laptop GPU while still small enough that framework overhead can obscure the first few tokens.

For a walkthrough benchmark, run GPT-2 once with `use_cache=False` and once with `use_cache=True`, keeping greedy decoding, batch size $1$, fixed prompt, identical max new tokens, warmed kernels, CUDA synchronization around the timed region, and no sampling overhead. If the prompt length is $P$ and the model emits $L$ tokens, a naive uncached loop feeds lengths $P,P+1,\dots ,P+L-1$ through the full Transformer, so the repeated-sequence part scales roughly like ${\sum }_{i=0}^{L-1}(P+i)=LP+L(L-1)/2$. The cached loop pays one prefill over $P$ tokens and then $L$ single-token decode steps, with attention still reading a growing cache but MLP and projection work no longer repeated for old tokens. Thus a short generation such as $L=4$ often shows almost nothing: the one-time Python dispatch, tokenizer, CUDA graph warmup failure, and prefill dominate, while the avoided term $L(L-1)/2=6$ token-equivalents is tiny. At $L=128$, the avoided repeated suffix is $8128$ token-equivalents, so the speedup becomes obvious even on GPT-2.

A concrete expectation check helps students avoid mystical explanations. GPT-2 small has $12$ layers, hidden size $768$, $12$ heads, and FP16 KV entries of $2$ bytes; its KV cache footprint per token is $2\text{ tensors}\times 12\text{ layers}\times 12\text{ heads}\times 64\text{ dim}\times 2\text{ B}=36,864\text{ B}$, about $36$ KiB/token. For a $P=64,L=128$ run, the cache at the end is only $192\times 36,864=7,077,888$ B, about $6.75$ MiB, so caching is not memory-capacity limited. The uncached loop, however, reprocesses ${\sum }_{i=0}^{127}(64+i)=16,320$ token positions, while the cached path processes a prefill of $64$ plus $128$ decode positions, or $192$ positions through the non-attention projections and MLPs. The ratio $16,320/192=85$ is an upper bound for the repeatable block work; measured end-to-end speedup will be far lower because single-token decode has poor matmul occupancy and still streams weights every token, but it explains why $L=128$ looks qualitatively different from $L=4$.

On production-sized models, the same mechanism trades compute recomputation for KV memory bandwidth and capacity. For a $70$B decoder in FP16 with grouped-query attention using $8$ KV heads, head dimension $128$, and say $80$ layers, the KV footprint per token is $2\times 80\times 8\times 128\times 2=327,680$ B, about $320$ KiB/token. An 80GB H100 has about $3.35$ TB/s HBM bandwidth, so a single request at context length $8192$ must read roughly $8192\times 327,680=2.684$ GB of KV per generated token, giving a bandwidth lower bound of $2.684/3.35\approx 0.80$ ms/token just to stream cached KV, before weights and logits. Without a cache, the model would instead redo all layers for thousands of old tokens every step, which is catastrophically worse; with the cache, decode becomes memory-bandwidth dominated rather than FLOP dominated. This is why FlashAttention from Dao et al. attacks attention IO, PagedAttention from the vLLM paper attacks KV allocation and fragmentation, and TensorRT-LLM’s in-flight batching plus paged KV cache try to keep the H100 busy despite single-token decode.

The technique stops helping when the saved recomputation is smaller than the cost of maintaining and reading the cache, or when the cache forces a worse execution regime. Concretely, for GPT-2 with $L\le 4$ and a short prompt, `use_cache=True` can benchmark equal or slower because `past_key_values` object handling, extra memory writes, and less favorable kernel shapes exceed the avoided work. For large models, very long contexts or large batch sizes can make KV cache bandwidth and capacity the bottleneck: the $70$B example at $8192$ tokens already reads gigabytes of KV per token per sequence, and multiplying by dozens of active sequences can hit HBM bandwidth or evict blocks, causing scheduler stalls. Caching also interacts badly with training-style full-sequence kernels: FlashAttention is excellent for prefill because it tiles the $Q{K}^{\top }$ computation without materializing the attention matrix, but decode with one query per sequence is a different low-arithmetic-intensity kernel. Prefixes that are never reused, speculative branches that are rejected, or beam searches with aggressive branching can store KV that is soon discarded.

The real systems students will see are engineering around this exact cache, not replacing it. Hugging Face exposes the basic `past_key_values` switch used in a GPT-2 lab; vLLM’s 2023 PagedAttention paper virtualizes KV blocks so continuous batching does not require contiguous per-request allocations; SGLang’s RadixAttention paper/release reuses KV across shared prompt prefixes in structured serving programs; TensorRT-LLM’s 2023-2024 releases implement paged KV cache, inflight batching, and quantized KV paths; NVIDIA Dynamo is the 2025 disaggregated serving stack that separates prefill and decode workers because their cache and compute profiles differ; llm-d is the Kubernetes-native distributed inference project that routes and schedules around KV locality; FlashAttention’s papers define the IO-aware attention kernels used heavily in prefill; and EAGLE is speculative decoding, where draft tokens only become useful if their KV survives verification. In the GPT-2 exercise, the student should report a table of $L$ versus $T_{\mathrm{uncached}}$, $T_{\mathrm{cached}}$, and $S(L)$, then explain the curve by the arithmetic above: near-flat at tiny $L$, rising once repeated old-token work dominates, and eventually limited by memory traffic and decode overhead rather than recomputation.