KV Cache from scratch in nanoVLM

TL;DR

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!

Introduction

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 step-by-step generation is inherently sequential:

• To generate token ti+1 t_{i+1} , the model must consider the entire sequence from t0 t_0 to ti t_i . From the first instance in the above example ti+1 t_{i+1} would be the , while all the previous tokens t0 t_0 to ti t_i would be [What, is, in].
• Although transformers are internally parallel, each new prediction requires a full forward pass through all transformer layers, which incurs a quadratic memory/compute in terms of the sequence length.

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

Table of contents:

• Revisiting the Transformer Architecture
• Where Redundancy Creeps In
• How KV Caching Fixes It
• KV Caching in nanoVLM: From Theory to Practice
• Summary: Why KV Caching Matters

Revisiting the Transformer Architecture

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

• Multi-head self-attention
• Feed-forward network (MLP)
• Residual connections and layer normalisation

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