Understanding and Coding the KV Cache in LLMs from Scratch
Source
Sebastian Raschka, PhD
Author
Sebastian Raschka, PhD
Date
Key takeaways · AI-distilled
A KV cacheThe memory a model keeps about text it has already read, so generating each new token doesn't require reprocessing the whole conversation.Full definition → stores the key and value vectors already computed for earlier tokens, so each generation step only projects the newest token instead of re-encoding the whole sequence. The saving grows with sequence length.
The cache is inferenceRunning a trained model to get answers — the phase where AI is actually used, as opposed to trained.Full definition →-only. Training processes a full sequence in one pass with no autoregressive reuse, so the technique buys nothing during a training run and only costs memory.
The core implementation is two registered buffers on the attentionThe mechanism that lets a model weigh which earlier words matter for the word it's currently processing — the core operation of a transformer.Full definition → module: concatenate the new keys and values onto them each step, then read the concatenated tensors instead of the freshly computed ones.
Forgetting to clear the cache between generations is a silent bug. The new prompt's queries attend to leftover keys from the previous sequence, so the model conditions on irrelevant context windowThe maximum amount of text a model can consider at once — its working memory for the current conversation or task.Full definition → and the output goes incoherent.
Terms in this piece · Glossary
inference — Running a trained model to get answers — the phase where AI is actually used, as opposed to trained.
KV cache — The memory a model keeps about text it has already read, so generating each new token doesn't require reprocessing the whole conversation.
attention — The mechanism that lets a model weigh which earlier words matter for the word it's currently processing — the core operation of a transformer.
context window — The maximum amount of text a model can consider at once — its working memory for the current conversation or task.
Why it matters
If you want to actually understand how KV caching accelerates LLM inference rather than just calling an API, this walks through a human-readable from-scratch PyTorch implementation and demonstrates a concrete ~5x speedup on a 124M model. It's from the author of 'Build a Large Language Model From Scratch,' so the code is self-contained and pedagogically clear.
Key quotes
“In short, a KV cache stores intermediate key (K) and value (V) computations for reuse during inference (after training), which results in a substantial speed-up when generating text.”
“So, as we can see, we already get a ~5x speed-up with a small 124 M parameter model and a short 200-token sequence length.”
“Unfortunately, the speed advantages disappear on CUDA devices as this is a tiny model, and the device transfer and communication outweigh the benefits of a KV cache for this small model.”
“Continuously concatenating tensors via torch.cat , as shown earlier, leads to performance bottlenecks due to frequent memory allocation and reallocation.”
articleRecent Developments in LLM Architectures: KV Sharing, mHC, and Compressed AttentionSebastian Raschka, PhD
articleEfficient Memory Management for Large Language Model Serving with PagedAttentionWoosuk Kwon et al.
articleSnapkv Llm Knows What You Are Looking For Before Generation 2024 04 22Cohere editorial sitemap