KV Cache from scratch in nanoVLM

We have implemented KV Caching from scratch in our nanoVLM repository (a small codebase to train your own Vision Language Model with pure PyTorch). This gave us a 38% speedup in generation. In this blog post we cover KV Caching and all our experiences while implementing it. The lessons learnt are general and can be applied to all autoregressive language model generations. Implementing from scratch on a small codebase is a great learning experience, come along for the ride!

Autoregressive language models generate text by sampling one token at a time . During inference, the model processes a given input sequence, predicts the next token, appends it to the sequence, and repeats this process until some stopping criterion:

This repetition also leads to computational redundancy . In this post, we explore KV Caching , an optimisation technique that mitigates this inefficiency.

Before diving into caching, let’s revisit how attention operates in transformer models. A Transformer language model consists of stacked layers, each composed of:

To understand where KV Caching helps , we focus on the self-attention mechanism, specifically within a single attention head.

Let’s walk through a simple PyTorch implementation to visualise the key computations.

For a sequence of input embeddings represented as , self-attention is computed as: