Context Windows & KV Cache

Post four is the hands-on lab. The previous posts built the theory; now we run real, copy-pasteable code to make context windows and the KV cache concrete. Everything here uses Hugging Face transformers with GPT-2, a small model that runs on a laptop CPU, so nothing requires special hardware. The point isn't GPT-2 specifically — it's that the mechanisms you've learned are directly observable.

We'll count tokens, check the window size, generate with the cache, benchmark the speedup, inspect the cache tensors, drive the loop manually, and guard the window limit. By the end, the abstractions from posts one through three will be things you've measured with your own hands.

This first block establishes the environment and grounds the two core concepts in real numbers. We load GPT-2's tokenizer and model, then tokenize a short prompt and print the token count — four tokens for 'The context window is.' Crucially, we also print model.config.n_positions, which is GPT-2's context window: 1024 tokens. That's the hard ceiling for this model.

The lesson is that the window isn't an abstract setting buried in documentation — it's a concrete attribute on the model config you can read in one line. Every model exposes its maximum context this way (the attribute name varies by architecture), and checking it should be reflexive before you build anything that streams long inputs.

Here we generate text with the KV cache explicitly enabled via use_cache=True. In practice this is the default, but stating it makes the contrast in the next slide clean. We set do_sample=False for greedy, deterministic decoding so the output is reproducible, and seed the RNG for good measure. The model generates 60 new tokens continuing the prompt.

Under the hood, this single call does a prefill pass over the four prompt tokens to seed the cache, then 60 decode steps, each appending one K/V pair and reusing everything before it. You don't see the cache in this snippet, but it's working on every one of those 60 steps — which is exactly what makes the next benchmark interesting.

This is the money snippet: a direct A/B benchmark of the cache. The timed function generates 120 tokens and returns the elapsed seconds, and we call it once with use_cache=True and once with use_cache=False. With the cache off, the model recomputes Keys and Values for the entire growing sequence at every single step — the quadratic cost from post three, made real.

The gap is dramatic and grows with output length: at 120 tokens cache-off is several times slower than cache-on, and the ratio widens further at longer generations. Running this yourself is the most convincing demonstration of why the KV cache exists. It's not a micro-optimization; it's the difference between practical and impractical generation.

This block opens the black box and shows the cache's actual structure. Calling the model with use_cache=True returns past_key_values, a tuple with one entry per transformer layer — len(past) is 12 for GPT-2's twelve layers. Each entry is a pair of tensors: the Keys and the Values for that layer. Indexing past[0] gives layer zero's K and V.

The printed key shape, (1, 12, 4, 64), maps exactly to the formula from post three: batch size 1, 12 attention heads, sequence length 4 (our four prompt tokens), and head dimension 64. Seeing the formula's factors appear as literal tensor dimensions is what makes the memory math click. Generate more tokens and watch the third dimension grow.

The pipeline diagram frames what the code in this post actually does, end to end. Raw text is tokenized into integer IDs. Those IDs go through prefill, which fills the KV cache. Then decode runs repeatedly, each step reusing the cache to produce one token. Finally the generated IDs are detokenized back into readable text.

This four-stage view is the skeleton of essentially every text-generation system, from a laptop GPT-2 demo to a production serving stack handling thousands of requests. The implementations differ wildly in optimization, but the stages are the same. Holding this pipeline in mind helps you locate any performance or correctness issue in the right phase.

This snippet drives the generation loop by hand to demystify what generate() does internally. We start with no cache and the prompt IDs, then loop: call the model passing the previous past_key_values, capture the updated cache, and take only the most recent token's logits to pick the next token. Critically, after the first iteration we feed just the new token forward, not the whole sequence — the cache holds the rest.

The final print confirms the mechanism: past[0][0].shape[2], the sequence-length dimension of layer zero's Keys, grows by one on every iteration. Watching that number climb is direct evidence that decode appends to the cache one position at a time. This manual loop is exactly the conceptual code from post three, now actually executing.

This guard function operationalizes the window limit so generation fails loudly instead of silently. It reads the model's maximum context once, then before generating it checks whether the prompt length plus the requested new tokens would exceed the ceiling, raising a clear error if so. The key subtlety it captures: output tokens count against the same budget as input, so you must reserve room for what you're about to generate.

Skipping this check is a classic source of confusing runtime crashes — generation proceeds fine until the sequence hits n_positions mid-stream and the model errors or produces garbage. A simple upfront guard turns a baffling failure into an actionable message, and it's the kind of defensive code every production generation path should have.

These four takeaways distill what the code demonstrated. use_cache=True being the default reflects that the cache is essential, not optional, for practical speed. The speedup grows with sequence length because the cost the cache avoids is quadratic. The cache shape [batch, heads, seq, head_dim] is the formula made tangible. And the cache length growing by exactly one per generated token is the literal definition of autoregressive decode.

If you ran the snippets, you saw all four directly: timed the speedup, printed the shape, and watched the sequence dimension increment. That hands-on confirmation is worth far more than taking the theory on faith.

These are the most common ways people break generation in code. Re-feeding the full sequence while also passing a cache double-counts tokens and corrupts the output — you feed either the new token with the cache, or the full sequence with no cache, never both. Forgetting the cache is per-request leads to expecting cross-call memory that doesn't exist. Ignoring n_positions until generation crashes is the avoidable failure the guard function fixes.

And assuming token count equals word count, the recurring theme of this whole day, quietly breaks budgeting and window checks. Each of these is easy to fall into and easy to prevent once you've seen the cache work explicitly, which is the real value of writing the code by hand at least once.

That's the lab complete. You've counted tokens, read the window size off the config, generated with the cache, benchmarked the cache-on versus cache-off speedup, inspected the real cache tensors, driven the decode loop manually, and guarded the window limit.

The final post steps back to the production view: the recurring context and cache mistakes that bite real teams — treating the window as infinite, miscounting tokens, expecting memory across calls, burying facts in the middle, and letting the cache exhaust GPU memory — along with the concrete fixes for each.