Grouped-query attention: a smaller KV cache

Intro

0:00 Hello and welcome back to episode eighteen, where we're going to give our model grouped-query attention, in which the query heads share their keys and values in groups, so that the KV cache gets smaller. Let's begin with a quick look at where we're going to be by the end of the episode.

0:16 By the end of today, our model keeps two heads of keys and values in each block, shared by its six query heads, and summary dot py prints what that saves. Its KV cache takes about one point six megabytes for a full window of two hundred and fifty-six tokens, a third of the four point seven it took with a key and value head for each query head.

0:37 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 grouped-query attention today. And each episode has a page here as well, with its chapters and its code.

0:56 Now, before we get to today's work, let's quickly go over what we did last time. We brought our transformer block up to date, with options in its config that are off by default, so the checkpoints we've trained so far load with the original block. RMSNorm, which we wrote out, took the place of the layer norms, dividing each token's channels by their root mean square. SwiGLU gave the MLP a gate, with three matrices, for as many weights as GELU's two. QK-norm put an RMSNorm on each head's queries and keys, so that the attention scores don't grow with the size of the weights. Dash dash modern turns those on and drops the biases, four changes in all, and it was all code and tests, without training a model. Qwen3's report also listed grouped-query attention, which we left out, and that's what we're adding today.

Keys and values, shared

1:48 Since episode four, our KV cache has kept the keys and values of the tokens written so far. The cache keeps them for each of the six heads, in each of our six blocks. In that episode, we also saw that Llama 3.1's model of seventy billion parameters has sixty-four heads sharing eight sets of keys and values, to keep its cache smaller. Today, we're going to give our model that option, with its query heads in groups. Each group shares a key head and a value head, which is called grouped-query attention, or GQA. With two groups of three, the cache keeps two heads' keys and values in place of six, a third of the memory. The key and value projections get smaller as well. We're going to write it as an option in our config, with tests that check it against the attention we've had so far. There's no training today, so what we measure is what the config decides, and summary dot py is going to print the size of the cache. All right, let's get started with where sharing keys and values comes from.

Multi-query attention

2:48 Sharing keys and values between heads goes back to Fast Transformer Decoding: One Write-Head is All You Need, by Noam Shazeer at Google, from 2019. Its abstract says that training attention layers is fast, as the positions of a sequence run in parallel, while writing a token at a time is often slow. At each step, the model loads the large tensors of keys and values from memory. As we saw in episode thirteen, moving data in and out of memory can take longer than the arithmetic. So in section three, multi-query attention gives each head its own queries, and the heads share a single set of keys and values, which makes those tensors a lot smaller. In table two, on a translation task, its decoder took about four microseconds per token, against forty-six with multi-head attention. Those times are on Google's TPU, a chip built for machine learning. And in table one, its quality came out slightly lower, while models with fewer heads, or smaller ones, did worse.

Grouped-query attention

3:48 Grouped-query attention comes from GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints, by Joshua Ainslie and others at Google Research, from 2023. Its abstract says that multi-query attention, with a single key and value head, speeds up writing a lot, but can lower the quality. So it introduces grouped-query attention, with a number of key and value heads in between, more than one and fewer than the query heads. Figure two draws the three side by side, with eight query heads along the bottom, and the keys and values above them. On the left, multi-head attention gives each query head a key head and a value head of its own. On the right, multi-query attention has one of each for all eight. In the middle, grouped-query attention splits the query heads into four groups of two, and each group shares a key head and a value head. The paper also turns models trained with multi-head attention into grouped ones, by averaging the key and value heads within each group. Then it pretrains them for another five percent of their original steps. Its experiments use T5, Google's model with an encoder and a decoder. In table one, the largest T5 with eight groups scored close to its multi-head version, and took about a fifth of the time per sample, close to multi-query's.

5:04 Meta's Llama 2, by Hugo Touvron and others, from 2023, used grouped-query attention in its bigger models. Its table one lists the family, and the column GQA has a tick for the two largest, of thirty-four and seventy billion parameters. The caption says it's there so that writing with them scales better. The appendix compares the three kinds of attention at thirty billion parameters, with the MLP widened to make up for the weights that sharing removes. Grouped-query attention did about as well as multi-head attention on most tasks, and better than multi-query on average. Figure twenty-four doubles the batch until the model runs out of memory. With a context of two hundred and fifty-six tokens, like ours, the multi-head model ran out at a batch of a thousand and twenty-four, while the other two still ran. And Llama 3.1, a year later, uses eight key and value heads in all three of its sizes.

Our six heads, three ways

5:59 So, let's draw it for our model, with its six heads. What we have now is multi-head attention, where each of the six query heads has a key head and a value head of its own, and the cache keeps all six. With grouped-query attention and two key and value heads, the query heads form two groups of three. Query heads zero, one and two share key and value head zero, and three, four and five share head one, so the cache keeps two. With a single key and value head, all six share it, which is multi-query attention, and the cache keeps one. In all three, each query head keeps its own queries, so it scores the shared keys in its own way, and takes its own mix of the shared values. Our config is going to take any number of key and value heads that divides the six query heads evenly: six, three, two or one.

6:47 Attention for head h is the formula from episode one, softmax of q, k transposed, over the square root of d k, times v, on that head's sixty-four channels. Its queries are its own, as before, while its keys and values come from head j, the head its group shares. The index j is h divided by the size of a group, rounded down. A group is H over G heads, where H counts the query heads, and G the key and value heads, one per group. For us, H is six and G is two, so a group is three query heads. Query heads zero, one and two read key and value head zero, and three, four and five read head one, counting from zero, as the code does. With G equal to H, each query head is a group of its own, which is the multi-head attention we have. And with G at one, they all read head zero, which is multi-query attention.

Fewer weights, a smaller cache

7:42 Sharing the heads changes two projections in each block, the key projection and the value projection. With last time's modern block, which has no biases, the key projection takes three hundred and eighty-four channels in, and gives six heads of sixty-four out. So it has a hundred and forty-seven thousand, four hundred and fifty-six weights, and the value projection has as many. With two key and value heads, each gives two heads of sixty-four, a hundred and twenty-eight channels, for forty-nine thousand, one hundred and fifty-two weights. So each block has a hundred and ninety-six thousand, six hundred and eight fewer, and our six blocks have about one point one eight million fewer. With a single key and value head, each projection has twenty-four thousand, five hundred and seventy-six. The query and output projections keep their three hundred and eighty-four by three hundred and eighty-four, as there are still six query heads, each with its own queries and its own output. Shazeer's paper and Llama 2's comparison widened the MLP to give those weights back, so that the models they compared had about as many parameters. We're not training a model today, so ours keeps its MLP as it is, with fewer weights.

8:52 In episode four, we counted the cache per token: two, for the key and the value, times six blocks, times three hundred and eighty-four channels, times four bytes a number. Those channels are six heads of sixty-four, and the six heads are what grouped-query attention changes, so a token takes eighteen thousand, four hundred and thirty-two bytes. With two key and value heads, it's six thousand, one hundred and forty-four, a third, and with one, it's three thousand and seventy-two, a sixth. For a full window of two hundred and fifty-six tokens, a story's cache goes from about four point seven megabytes to about one point six with two heads. With one, it takes about seven hundred and ninety kilobytes. Each story in a batch has a cache of its own. So the sixteen full windows of a batch on episode fourteen's server would take about twenty-five megabytes with two heads, where they took about seventy-five.

Grouped-query attention in code

9:42 In model dot py, the config gets a new field, the number of heads of keys and values, each shared by a group of query heads. It's an int or None, and None, the default, means as many as the query heads, the multi-head attention we've had so far. So the configs saved with our checkpoints, which have no such field, load as before, and so do the bigger models of episode twelve, with their twelve heads. The comment at the end of the line says it: heads of keys and values, each shared by a group of query heads, where None is n head. In the KV cache, the comment on what each slot holds changes, as the tensors now have the key and value heads, and not one per query head. In attention's constructor, we keep the number of key and value heads, where the or falls back to n head when the field is None. Above it, the comment reads: grouped-query attention, n kv head heads of keys and values, each shared by a group of query heads. The query heads have to split evenly into the groups, so if n head isn't a multiple of n kv head, we raise a value error, Python's exception for an argument with the wrong value, with both numbers in its message. With six query heads, that allows six, three, two or one, and asking for four stops the model as it's built, with that message, where it would otherwise fail inside attention, on tensors whose shapes don't match. The head size moves up, as the projections need it now: the embedding size over the heads, sixty-four channels. The query projection keeps its three hundred and eighty-four outputs, while the key and value projections go to n kv head times the head size, a hundred and twenty-eight channels with two heads. The output projection stays as it was, as attention still gives an output for each query head. In forward, the keys and the values are viewed as n kv head heads of sixty-four, where the queries keep all six. QK-norm and the rotary embeddings work on each head on its own, so they apply to the fewer heads as they are, and the cache's update stores them, which is where the memory goes down. After the cache, when there are fewer key and value heads than query heads, we repeat them. Repeat interleave repeats each slice of a tensor a number of times, one copy after another, along a dimension. Here, the dimension is one, the heads, and the number is n head over n kv head, three. So key and value head zero becomes heads zero, one and two, and head one becomes three, four and five, which lines each group of query heads up with its shared head. Above the if, the comment says why: the cache keeps the fewer heads; attention needs one per query head, so each is repeated for its group. The values are repeated like the keys, and the repeat comes after the cache, as repeating before it would store the copies, and the cache would be as large as before. PyTorch's scaled dot product attention can also repeat them for us, with enable G Q A set to true, from PyTorch two point five. We repeat by hand, as attention can run through the functions that attend with takes, such as the flash attention kernel from episode thirteen, and those expect as many key and value heads as query heads. Then the tests, in a new file, tests slash test G Q A dot py. It imports pytest and torch, and from model dot py, the GPT, the modern block's settings, the config and the KV cache. Tiny builds a small model to test with, and it seeds PyTorch first, so each run gets identical weights. It has a vocabulary of fifty, two blocks of six heads, forty-eight channels, so heads of eight, no dropout and rotary embeddings. It also takes the key and value heads it's given, and the modern block when asked, and returns the model in evaluation mode. The first test runs twice, with the original block and the modern one, as pytest's parametrize runs a test once for each value in its list. Its name says what it checks: grouped-query attention is multi-head attention, with each key and value head repeated for its group. We build a model with two key and value heads, and a multi-head one, and take the first one's state dict, a dictionary of its weights by name. For each key and value projection, a weight's rows are two heads of eight channels, and repeating each head three times gives six heads, the multi-head model's shape. The comment says so: two heads of eight channels, each repeated for the three query heads of its group, six heads, as in the multi-head model. Load state dict puts those weights into the multi-head model, and on a batch of random tokens, the two models' logits agree to within ten to the minus five. So the repeat in attention gives each group of query heads its own key and value head, in the order the projection's heads come. The second test runs with two key and value heads, and with one, which is multi-query attention. Its name says what it checks: the cache holds fewer heads, and matches running the whole sequence at once. It builds the small model with the modern block, and reads twelve tokens in one call, and then again through a cache, the first eight in one go and the last four one at a time, and their logits agree. The cache now holds the fewer heads, and the repeat runs on what it returns, so this checks the two together. And the last check is the shape of the cache's keys: a batch of two, n kv head heads, twelve tokens and eight channels a head, as the comment above it lists.

The cache's size, and the flags

14:57 To see what we save, model dot py gets a function, KV cache bytes, which works out the cache's size from the config. Per token, the cache keeps a key and a value for each layer, each key and value head and each channel of a head, where n kv head or n head stands in for a missing field again. It returns that times the tokens, times the bytes per number, four by default, for float32. Its comment reads: per token, a key and a value for each layer, key and value head and channel; float32 is four bytes a number. Summary dot py imports the function, and it already takes dash dash rope and dash dash modern. Now it takes dash dash K V heads too, with a number after it, as its comment says: grouped-query attention, with N heads of keys and values. Sys argv is the list of the command's words, so the number is the word after the flag, as an int, or None when the flag isn't there. The config gets it as n kv head, next to rope and the modern block's settings. After the total, a new line prints the KV cache: the bytes for one token, and for a full window of block size tokens, in megabytes of a million bytes, as in episodes four and fourteen. The comment over it says: what the KV cache keeps as the model writes, from the config, per token and for a full window. Pretrain dot py takes the flag too, and main gets kv heads as an argument. The run folder gets dash kv and the number at its end, so a run with two key and value heads saves to a folder of its own. That way, it doesn't overwrite the checkpoints of a model with multi-head attention. The config gets n kv head, and at the bottom, the flag is read like summary dot py's, with a comment on what it is. Pretrain DDP dot py, from episode twelve, gets dash dash kv heads as an option of its argument parser, an int, with a help text, and dash dash name's help now mentions both suffixes. The config gets the number as well, and the run folder's default name, pretrain dash DDP, gets dash modern and dash kv with the number after it. The comment over it now reads: a run with the modern block or fewer heads of keys and values gets a folder of its own, so it doesn't overwrite pretrain dash DDP, the bigger model we trained on four GPUs. Last, a test for the size, and the tests' import gets KV cache bytes. Test KV cache bytes counts what the cache holds: it fills a cache with ten tokens through the small model with two key and value heads, and adds up the bytes of its tensors. Numel is a tensor's number of elements, and element size the bytes of one, four for float32. That sum has to equal KV cache bytes for ten tokens, so the formula counts what's there. Then two checks on our own config, with its default sizes, and above them, the comment reads: our model, eighteen thousand, four hundred and thirty-two bytes a token with a key and value head per query head, and a third of that with two. So the second check asks for six thousand, one hundred and forty-four bytes with two key and value heads.

18:01 Let's run the tests with pytest, where dash v lists each case. All five pass: grouped-query attention against multi-head attention, with the original block and with the modern one, the cache against the whole sequence with two key and value heads and with one, and the cache's bytes against the formula.

18:19 Now let's print the model, with rotary embeddings and the modern block, first with a key and value head for each query head. Over the six blocks, the key projections hold eight hundred and eighty-four thousand seven hundred and thirty-six weights, and the value projections as many. The total is about fourteen million five hundred thousand parameters, and the KV cache takes eighteen thousand four hundred and thirty-two bytes a token, about four point seven megabytes for a full window of two hundred and fifty-six tokens. With dash dash K V heads two, the key and value projections give one hundred and twenty-eight channels, and over the six blocks, the key projections drop to two hundred and ninety-four thousand nine hundred and twelve weights, a third. The total has about one million one hundred and eighty thousand fewer parameters, and the cache takes six thousand one hundred and forty-four bytes a token, about one point six megabytes for the window. And with a single key and value head, multi-query attention, the key projections have one hundred and forty-seven thousand four hundred and fifty-six weights, and the cache takes three thousand and seventy-two bytes a token, about seven hundred and ninety kilobytes for the window. Those are the numbers we worked out by hand, as both come from the config, and the test we just ran checks the formula against a cache that's been filled.

What's real, and what's scaled down

19:34 The mechanism is the one in large open models: query heads in groups, each group sharing a key head and a value head, and the cache keeping the shared ones. Llama 3.1's model of seventy billion parameters has sixty-four query heads sharing eight key and value heads, in groups of eight, where ours has six sharing two, in groups of three. Smaller models use it too: Qwen3's smallest, of six hundred million parameters, has sixteen query heads sharing eight key and value heads. Llama 3.1's eighty blocks have heads of a hundred and twenty-eight channels, where ours has six blocks with heads of sixty-four. For a full context of about a hundred and thirty-one thousand tokens, at two bytes a number, Llama 3.1's cache takes the forty-three gigabytes we worked out in episode four. With a key and value head per query head, it would take about three hundred and forty. Ours goes from about four point seven megabytes to one point six, which matters little on a laptop, and a lot more as the model, the context and the batch grow. And Llama 3.1 was trained with grouped-query attention, while today we built the option and checked it against the attention we had, without training a model with it.

20:44 Last, the docs get a page on grouped-query attention, with the papers, how the heads are shared and then repeated for attention, the tests, and what it saves in weights and in the cache, worked out from the config. It goes in the list on the docs' home page, and at the end of the navigation. And the README has the commands to print the model with two key and value heads, along with the size of its cache, and to run the tests.

Wrap-up

21:10 So, today we gave our model grouped-query attention, where the query heads work in groups, and each group shares a key head and a value head. The cache keeps the shared heads, and they're repeated for each query head after the cache, right before attention, so what the cache holds stays small. Our tests check it against multi-head attention with the key and value weights repeated for each group. They check writing with the cache against writing without it, for two key and value heads and for one, and the cache's size from the config against what a filled cache holds. With two, the key and value projections and the cache take a third of what they did. A story's full window takes about one point six megabytes, in place of four point seven. And dash dash K V heads sets it on pretrain dot py, pretrain DDP dot py and summary dot py. The default stays a key and value head per query head, so the models we've trained load as before.

22:06 Let me know in the comments how many key and value heads you'd give a model like ours, and whether you'd go as far as a single one. And if you'd like to follow along, subscribe to the channel. I hope you find this useful, and thanks for watching.