KV Cache from scratch in nanoVLM

Autoregressive models have an annoying habit: to spit out one new word, they redo almost all the math they already did for every word before it. Each step means another full pass through every transformer layer, recomputing queries, keys, and values for tokens that haven't changed since the last step. That's wasted compute piling up as the sequence grows, and it's the exact problem KV caching exists to solve.

The Hugging Face team decided to build it themselves inside nanoVLM, their compact PyTorch codebase for training vision-language models, rather than just importing a library that already does it. That choice makes the writeup useful beyond nanoVLM itself — the same redundancy shows up in any decoder-only language model, from GPT-style chatbots to whatever multimodal system you're tinkering with in Colab. Using simple tensor checks, they show that keys and values computed for early tokens stay identical no matter how many new tokens get appended afterward. Only the newest token needs fresh K and V. Everything else is just being recalculated for no reason.

Their fix splits generation into two phases. First comes prefill: the model chews through the entire prompt once and stores keys and values for every layer. Then decode kicks in, where each new token only needs its own query vector, appended to the cached keys and values from before. In code, this meant reworking three pieces: the attention block now appends new K/V pairs to a running cache instead of rebuilding from scratch, the language model tracks a per-layer cache plus a start_pos variable so rotary positional embeddings stay correctly aligned, and the generation loop cleanly separates the one-time prefill pass from the token-by-token decode loop that follows.

The payoff: a 38% speedup in generation, with per-token inference cost dropping from quadratic to roughly linear in sequence length. That's not nothing, particularly for anyone running these models on modest hardware where every wasted FLOP costs real time. The tradeoff is memory — you're storing K and V tensors for every layer and every token — and some added code complexity, plus friction if you want fancier decoding tricks like beam search layered on top. Still, this is essentially why long-context chatbots don't grind to a halt after a few hundred words, and now there's a small, readable codebase showing exactly how that trick works under the hood.