KV cache: making a GPT write faster

Intro

0:00 Hello and welcome back to episode four, where we're going to make our model write faster, with a KV cache, a cache of the keys and values that attention computes. Let's begin with a quick look at where we're going to be by the end of the episode.

0:15 On my laptop's CPU, our fine-tuned model now writes its stories about four times faster, and they come out word for word as they did before. Without the cache, each new token takes longer as the story grows, and with it, the time per token stays flat.

0:31 On the website of the series, you can find more information about the project. The code is open source on GitHub, tagged at each step of each episode, so that you can follow along, and the docs get a page on the KV cache today. And each episode has a page here as well, with its chapters and its code.

Last time

0:49 Now, before we get to today's work, let's quickly go over what we did last time. We fine-tuned our model on TinyStories-Instruct, with each document's instructions parsed into one order, and a loss on the story alone, so the model is scored on the story and not on the instruction. After ten minutes of fine-tuning, its stories used about twenty-seven percent of the words an instruction asked for, up from about three. And our runs now log to TensorBoard, which we read back in Python to plot. Today, we're going to make our model write faster, without changing what it writes. All right, let's get started with how it writes a story, a token at a time.

Writing a token at a time

1:26 We wrote generate in episode one, and it runs a loop with a step for each new token. To write the first token, it runs the prompt through the model, once upon a time, and samples there from the logits at the last position. For the next token, it runs once upon a time there, the prompt and the new token, and then once upon a time there was, and so on, until the story ends. At each step, the model computes all the tokens before the new one again, through all six blocks, even though they come out with the numbers they had on the step before. For two hundred new tokens after our four-token prompt, that's four tokens through the model, then five, then six, and so on up to two hundred and three, twenty thousand seven hundred in all. If we kept what we computed for the earlier tokens, we'd run the prompt once, and then each new token once, two hundred and three tokens in all, about a hundred times fewer. Keeping those numbers is the idea behind a KV cache, and today we're going to see which ones we can keep, add the cache to our model, and measure how much faster it gets.

2:26 Jay Alammar's post, The Illustrated Transformer, draws attention a step at a time, and it's one of the clearest explanations of it, used in courses at Stanford, MIT and others. Here, each word's embedding is multiplied by three matrices of weights, which gives its query, its key and its value. The query of the word being processed is scored against the key of each word, and its output is a mix of their values. Keep the keys and the values in mind, as they're what we're going to cache.

Keys and values that don't change

2:54 So, let's go back to attention, and see what it computes for a new token. In episode one, we saw attention as a formula over all the tokens at once. For the newest token, t, it's this: the token's channels are projected into its query, its key and its value, in each head of each block. Its query is compared with the keys of all the tokens up to it, and its output is the weighted sum of their values. As attention is causal, a token's channels in each block depend on that token and the ones before it, and on none after it. So writing a new token doesn't change what we computed for the tokens before it, in any block. Their keys and values come out as they did on the step before, so we can keep them in a cache, and at each step, compute the query, the key and the value of the new token alone, and add its key and value to the cache. We don't keep the earlier tokens' queries, as a query is used to predict the token after it, and we've predicted those already. So the cache holds the keys and the values, and that's where the name KV cache comes from. The MLP, the norms and the head run on the new token alone too, as they work on each token on its own.

4:05 Each of our six blocks has its own attention, with its own keys and values, so the cache holds a set per block. As the new token goes through a block, attention reads that block's cache for the earlier tokens, and adds the new token's key and value to it. The head then gives the logits for the next token, from the new token's position. And for a batch of stories, the cache holds each story's keys and values side by side, like the rest of a batch.

A cache in attention

4:31 In model dot py, the cache gets a class of its own, above attention. KV cache takes the number of blocks, as each block's attention has its own keys and values. It keeps two lists, keys and values, with a slot per block, and each slot starts as None, as the cache starts empty. Above them, we note what each slot holds: per block, the keys and the values of the tokens so far, each a tensor of the batch, by the heads, by the tokens, by sixty-four. Len is the method Python calls when we ask for the length of an object, and for the cache, it's how many tokens it holds: zero while the first block's slot is empty, and otherwise the size of dimension two of its keys, the tokens. Update is how a block adds to its slot. It takes the block's number, and the new tokens' keys and values. If the slot has keys already, torch cat joins the new keys after them, along dimension two, and the values too. Then the slot keeps the result, and update returns it, so attention gets the keys and values of all the tokens so far. Cat copies the slot each time it grows, which is fine at our size, while serving code for large models sets aside room for the whole context window up front, and writes into it. Attention has to know which slot is its own, so its constructor takes the layer, the block's number, and keeps it. Forward takes an optional cache, and the queries, keys and values are computed as before, for the tokens we pass in, which is the newest one alone once the cache holds the prompt. If there's a cache, we hand it the new keys and values, and get back the keys and values of all the tokens so far. S counts the tokens attended to, the cached ones and these T. Then comes the mask, which has three cases, depending on the cache. Without a cache, S equals T, and is causal applies the mask from episode one, so training runs as it did. With a single new token, T is one, and the new token may look at all S tokens, the cached ones and its own, so it doesn't need a mask. And when several new tokens come after cached ones, we build the mask: torch ones makes a T by S grid of trues, and tril keeps its lower triangle, moved to the right by S minus T, so the last new token sees all S. Its comment notes that query i is token S minus T plus i, which sees the keys up to its own, and a single new token sees all S. We don't use that third case today, but it keeps attention correct for any number of new tokens. Scaled dot product attention gets the mask as attn mask, and is causal when S equals T. The block's constructor takes the layer, and passes it on to attention, and its forward passes the cache on too. In the GPT, the blocks are built with their numbers, from range of the number of blocks, so the first block is layer zero, and the sixth is layer five, and each block's attention reads and writes its own slot in the cache. Forward takes an optional cache as well, which sets where the positions start. With a cache, the tokens we pass in come after the ones it holds. So their positions start at the cache's length, and not at zero, and each gets the embedding our model learned for that position. Each block then gets the cache, and the rest of forward stays as it was.

Generating with the cache

7:39 Now, generate puts the cache to use, through a new argument, use cache. Use cache is true by default, so sample dot py, the samples during training and evaluate dot py all get the cache without a change, while false keeps the old way, so we can compare the two. The cache starts as None, before the first step. The first branch is the old way, the whole window through the model, and we take it when we're not using the cache, or once the story is longer than the context window. Past block size tokens, the window slides, which moves each token to a new position and changes its keys and values, so from there on we run the whole window again. With learned positions, a token's position is part of its channels from the first block on, so a key cached at position ten is wrong once that token sits at position nine. On the first step with the cache, we make an empty one, and run the whole prompt through the model at once, which fills the cache with the prompt's keys and values, block by block. That first call does the most work, as it computes the whole prompt, while each step after it computes a single token. From the second step on, the last branch passes the model the newest token alone, along with the cache, which has the keys and values of the ones before it. The logits come back for that one position, and the rest of the loop stays as it was, with the temperature, top k and sampling, as it reads the last position either way.

Checking it with evaluate.py

9:03 The first check is evaluate dot py from last episode, which writes fifty stories with each model and measures them. And time, in front of it, measures how long the whole command takes. On the left, it runs at the end of episode three, without the cache, and on the right, with it. The two tables match, down to the requested words used, how often the stories stop and how long they are, and so do the likeliest next tokens below them. The cache changes how the logits are computed, not what they are, and the seed makes the sampling match. The run took about fifty-one seconds without the cache, and about thirty-six with it, while loading the models and measuring the losses take as long either way.

Checking and timing the cache

9:42 To check the cache and time it, we write bench dot py. It imports argparse, json and os, statistics, for the median, and time, for the clock, along with torch, load checkpoint, and from model dot py, the KV cache and pick device. The prompt is once upon a time, and the story is two sentences in TinyStories' style, to compare logits on. Sync is there for timing on a GPU, and on the CPU, it returns straight away. If the device is CUDA, NVIDIA's GPUs, it calls torch cuda synchronize, and if it's m p s, the GPU in the Apple chip of my laptop, it calls torch m p s synchronize. When PyTorch gives a GPU work, it queues the work and carries on, while the GPU runs it in the background. So reading the clock straight after would time the queueing, not the work, and we wait for the GPU to finish first. On the CPU, the work is done by the time the call returns, so there's no need to wait. Largest difference runs the model on a sequence of tokens twice, without gradients. Full is a single call to the model, on all the ids, the way training runs it, and above it, we write what we compare: the logits for all the tokens at once, and for one token at a time through the cache. They should come out alike, and a position off by one, or a wrong mask, would make them differ by far more than rounding. For those, we make an empty cache, and go through the sequence a position at a time, with t from zero to its length, passing the slice from t to t plus one, a single token, along with the cache, and keeping the logits each call returns. Cat joins them along the tokens, so the two have one shape, and we return the largest absolute difference between them. They won't agree to the last digit, as the arithmetic happens in a different order, but the difference should be tiny. Step times measures how long each new token takes as the story grows. It makes a cache if we use one, and None if we don't, and starts with the prompt as the ids, the prompt as x, the tokens to run, and an empty list of times. While the story is shorter than the context window, it waits for the device, reads the clock with perf counter, Python's clock for short durations, runs the model on x, with the cache, waits again, and keeps the time in milliseconds. We pick the likeliest next token with argmax, as we care about the time here, not the story, and append it to the ids. X, the tokens for the next step, is where the two ways differ: with the cache, the next step runs the new token alone, and without it, all of them. Once the story fills the window, step times returns the list, with a time for each new token. Tokens per second times generate, and it repeats the prompt once per story, as a batch, waits for the device, reads the clock, writes two hundred tokens with generate, with the sampling settings we use elsewhere, and waits again. It returns the tokens written, the batch times two hundred, over the seconds it took. Generate doesn't wait for the GPU between its steps, so this is how long writing a story takes in practice. In main, an argument parser describes the script, which checks the KV cache and measures how fast the model writes with and without it. The parser takes the checkpoint, the device, which defaults to the one pick device chooses, and a folder to write the timings to, as JSON, which is optional, so a quick run can skip it. We load the model and its tokenizer from the checkpoint, on that device, put it in evaluation mode, so dropout is off, and encode the prompt. Difference is the largest difference in the logits for the story, which we print. Then we write two stories from one seed, without the cache and with it, and torch equal checks that they match, token for token. Then come the timings, where lengths are how many tokens there are at each step, from the prompt's four up to the window. We do five runs each, and take the median at each length, as a single run is noisy. Without is five runs of step times without the cache, where zip lines the five lists up, token by token, and median takes the middle value at each length, and with cache does that with the cache. We print a header, then a row each time the story grows by fifty tokens, with the two times. Then the speed, for a batch of one story, and of eight stories at once, the way a server would write for several people. For each, a short run of twenty tokens without the cache and with it comes first, to warm up, as PyTorch sets things up the first time it runs an operation on a device, which would slow the first measurement down. Slow is the speed without the cache, and fast with it, and we keep both in speeds, along with how many times faster the cache is, and print them as a row. And with dash dash out, we make the folder, where exist ok means a folder that's already there isn't an error, and list the timings at each length as rows, with the tokens so far and the two times, rounded to a thousandth of a millisecond. Then we write a JSON file named after the device, for our charts, with the device, the difference, the match, the timings and the speeds. And as in our other scripts, main runs when we call bench dot py from the command line.

Timing it on the CPU

14:42 Let's run it on the CPU first, with our fine-tuned model. The largest difference in the logits is about one point eight times ten to the minus five, while the logits are in the tens, so the two ways agree up to rounding, and the stories match. Without the cache, each new token takes longer as the story grows, from about two point eight milliseconds to about eight point seven at the end of the window. With it, a token takes about one point five milliseconds, wherever it is in the story. Writing a story, that's about four point two times faster, and with eight stories at once, about four point nine.

15:14 Without the cache, the time per new token on the CPU climbs with the length of the story, as each step runs all the tokens so far through the six blocks, and computes logits for each of them, while we use the last. With the cache, each step runs a single token, and attends over the cached keys and values, which grow too, but cost little next to running the blocks, so the line stays almost flat. The time to write a whole story is the area under each line.

Timing it on the GPU

15:41 Now on the GPU, with m p s, the logits and the stories check out here too. Without the cache, a token takes from about one point six milliseconds to about three point six, and with it, about one point seven wherever it is in the story. Writing one story, the cache makes it about one point six times faster, and with eight stories at once, about three point five.

16:02 On the GPU, both lines start at about one point seven milliseconds, and the one with the cache stays there. A step on the four tokens of the prompt takes about as long as a step on a single token, so the arithmetic is a small part of that time. Most of it goes to launching the model's operations one after another, each a small job for the GPU, and waiting for the result. A model this small, with one story, doesn't give the GPU enough work to hide that. The cache saves the time a step takes beyond those one point seven milliseconds, which grows with the story, so for one story it saves less than on the CPU.

16:40 Writing one story, the cache puts the CPU ahead of the GPU, at about seven hundred tokens a second against about four hundred and ninety. With eight stories at once, each step has eight times the work for about as many operations, and the GPU's overhead is shared between them. So the GPU pulls ahead, at about thirty-two hundred tokens a second with the cache, about three point five times what it managed without it. Which is faster depends on the model, the batch and the device, so we measure on the machine we're going to run on, such as a server without a GPU.

What changes at scale

17:14 Our numbers are from a small model, so let's see what changes at scale. Our cache holds a key and a value per token, in each of six blocks, of three hundred and eighty-four numbers each, at four bytes a number, which is about eighteen kilobytes per token. A full context window of two hundred and fifty-six tokens comes to about four point seven megabytes per story, next to the model's fifty-eight megabytes. Meta's Llama 3.1, with seventy billion parameters, has eighty blocks, and keeps its numbers in two bytes each. Its sixty-four heads share eight sets of keys and values, of a hundred and twenty-eight numbers each, which is called grouped-query attention, and it's there to make the cache smaller. Its cache still takes about three hundred and twenty kilobytes per token, and a full context window of a hundred and thirty-one thousand tokens takes about forty-three gigabytes, for one conversation. At that scale, the cache limits how many conversations a GPU can serve at once, and serving systems such as vLLM manage it in pages, the way an operating system manages memory.

18:18 Last, the docs get a page on the KV cache, with what it keeps and why, the three cases of the mask, why the cache stops at the context window, and the timings from today's runs, along with the size of the cache at scale. The architecture page mentions the mask attention uses with a cache, where it describes attention, and its section on generating now says that generate runs the prompt once, and then each new token on its own, with a link to the new page. The home page of the docs lists the new page, and so does the navigation, in mkdocs dot yml. And the README has the command to run the bench.

Wrap-up

18:53 So, today we made our model write faster, with a KV cache. Generating ran the whole story through the model for each new token, while the keys and values of the earlier tokens don't change, so we keep them, a set per block, and run the new token alone. The logits agree up to float rounding, and the stories and evaluate dot py's numbers match. On my laptop's CPU, the model writes about four times faster. On the GPU, one story gets about one point six times faster, as with a model this small, most of each step goes to launching operations. And eight stories at once get about three point five times faster. Large models keep a cache like ours too, and at their scale, it can take tens of gigabytes for a single conversation.

19:38 If you run the bench on your own machine, let me know in the comments how much faster the cache makes it, on the CPU and on the GPU, as that depends a lot on the device. And if you'd like to follow along, subscribe to the channel. I hope you find this useful, and thanks for watching.