Grouped Query Attention Explained: How GQA Shrinks the KV Cache

The model saves keys and values, short lists of numbers, for every token (a word or piece of a word) in every open chat. With long documents, those saved numbers can outgrow the model itself. Grouped query attention is a design choice inside the model that shrinks them.

Picture a common setup. An insurance team runs a chat assistant for surveyors, the people who inspect damage and write claim reports. A surveyor uploads a long report, say on a flooded textile warehouse, and asks questions about it. The setup is made up, but the problem is real.

The team runs an open model, one whose weights you can download and run yourself. Weights are the numbers the model learned in training, and here they take 16 GB. The GPU, the chip that does the maths, has 80 GB of memory.

A report plus the chat comes to about 32,000 tokens. The pilot went well with three surveyors. Then the monsoon brought a wave of flood claims, and 40 surveyors logged in on the same Monday.

• Only three long chats could run at the same time. The rest waited in a queue.
• When a fourth long report arrived, there was no room left for its saved numbers. A server that never planned for this fails with an out-of-memory error.
• vLLM, a popular open-source serving engine, pauses a chat instead and recomputes it later. That avoids the crash, but paused chats stutter while they rebuild.

The obvious fix is to rent more GPUs, at about Rs 2.5 lakh (250,000) a month each (made up). The cheaper fix is often a model whose attention stores less. This guide builds up to that choice step by step.

Attention is the step where a model decides which earlier tokens matter for the token it is about to write. It scores every earlier token, then blends their information using those scores.

A model writes one token at a time. Each new token goes through attention again, in every layer. A layer is one round of processing, and the team's model stacks 32 of them.

Attention keeps three lists of numbers, called vectors, for each token. Think of them as a question, a label and a content card.

• Query: what the current position is looking for. After the words "the claim amount is", the next position looks for a money figure.
• Key: a label each earlier token carries, saying what it offers. The tokens of an amount like 18.4 lakh rupees carry keys that match money questions well.
• Value: the information an earlier token passes on once it is picked. It is what gets blended into the answer.

The maths is short. Compare the query with every key to get a score. Turn the scores into weights that add up to 1, using a step called softmax. Then take the weighted blend of the values.

The 2017 paper that introduced the transformer, the design behind today's language models, also divides each score by the square root of the key length. Its authors suspected that large scores push softmax into a region where learning barely moves.

Multi-head attention runs several small attention steps side by side, each with its own query, key and value. Each one is called a head, and each can learn to look for something different.

In a survey report, one head might track dates and another might link "the insured" back to a company name. Nobody assigns these jobs. They emerge in training, and many heads do mixed jobs nobody can name.

The 2017 paper used 8 heads in its base model. It split the model's width, the count of numbers describing each token, across them: 512 numbers became 8 heads of 64. The paper notes the total cost stays similar to one full-width head.

Here is the one picture this guide uses. A wedding photography crew has eight photographers. In multi-head style, each carries a full kit bag of lights and lenses. Each can shoot their own way, but the van is packed to the roof.

The van is the GPU's memory. Every head's keys and values must be stored, for every token, in every layer, for every chat. The next section shows how fast that adds up.

The KV cache is GPU memory where the model keeps the keys and values of every token it has already processed, so it does not recompute them. It grows with every token, layer, key-value head and open chat.

Without it, writing token 32,001 would mean recomputing keys and values for all 32,000 tokens before it. With it, the model computes them for the new token only and appends them.

Here is one writing step in one layer. Follow the arrows and find the box where the model reads the whole cache.

The cache size follows one formula. Multiply 2 (a key and a value), the number of layers, the number of key-value heads, the head size and the bytes per number. That gives bytes per token.

The team's first model has 32 layers and 32 heads of size 128, so it has 32 key-value heads. It stores each number in 16-bit form, which is 2 bytes. So each token needs 2 times 32 times 32 times 128 times 2 bytes.

That is 524,288 bytes, about 0.52 MB, for one token. Times 32,000 tokens, one chat needs about 16.8 GB. That is more than the 16 GB of weights.

After the weights and some working space, about 60 GB is left for caches. Three chats fit in 50.3 GB. A fourth needs another 16.8 GB, and the memory runs out.

Look at the formula again. Layers and head size are fixed by the model's design, and bytes per number has its own quality trade-offs. The number of key-value heads is the dial that attention variants turn.

Grouped query attention splits the query heads into groups, and each group shares one key head and one value head. Every head still asks its own question, but far fewer keys and values are stored.

Back to the crew. The eight photographers form two teams of four, and each team shares one kit bag. Every photographer still frames their own shots.

Mapped to the model, each photographer is a query head and their framing is its query. The shared kit bag is one key-value head. The van space saved is KV cache memory.

Where the picture breaks: photographers pick their kit, but heads do not. In training, the model learns what the shared keys and values should hold so that every head in the group can use them.

Here is a small version with four query heads and two key-value heads. Notice that two arrows reach each shared box, and that only those two boxes are stored.

The 2023 paper that named GQA treats it as a dial between two ends. With one group, all heads share one set, which is multi-query attention. With one group per head, nothing is shared, which is multi-head attention.

Its authors used 8 key-value heads in their main tests and called that a favourable middle ground. Whether a model you can download does the same is written in its config file, as a later section shows.

Grouped query attention in Python

This sketch computes attention for one new token in one layer. q holds one query per head. K and V hold the cached keys and values, one set per key-value head. It uses NumPy, the standard Python library for arrays.

import numpy as np

def gqa_one_step(q, K, V):
    # q: (query_heads, d)   K and V: (kv_heads, tokens, d)
    query_heads, d = q.shape
    kv_heads = K.shape[0]
    assert query_heads % kv_heads == 0, "query heads must split evenly into groups"
    group = query_heads // kv_heads          # query heads per key-value head
    out = np.empty_like(q)
    for h in range(query_heads):
        kv = h // group                       # the shared key-value head
        scores = K[kv] @ q[h] / np.sqrt(d)    # one score per earlier token
        weights = np.exp(scores - scores.max())
        weights /= weights.sum()              # softmax: weights add up to 1
        out[h] = weights @ V[kv]              # blend of the stored values
    return out

# Check: copy each key-value head 4 times and run plain multi-head attention
rng = np.random.default_rng(0)
q = rng.standard_normal((32, 128))
K = rng.standard_normal((8, 1000, 128))
V = rng.standard_normal((8, 1000, 128))
gqa = gqa_one_step(q, K, V)
mha = gqa_one_step(q, np.repeat(K, 4, axis=0), np.repeat(V, 4, axis=0))
print(gqa.shape, np.abs(gqa - mha).max())
print("cache numbers stored:", K.size + V.size, "vs", 4 * (K.size + V.size))

What just happened? group is how many query heads share each key-value head: 32 query heads over 8 key-value heads gives 4. Then kv = h // group sends heads 0 to 3 to set 0, heads 4 to 7 to set 1, and so on.

The check at the bottom copies each key-value head four times with np.repeat, which turns the input into plain multi-head attention. On NumPy 2.5.4 it prints (32, 128) 0.0, so both give the same answer.

The second line prints 2048000 vs 8192000 stored numbers: a quarter of the cache for the same result. Give K and V 32 sets and the function is multi-head attention. Give them 1 set and it is multi-query attention.

Multi-query attention keeps a separate query for every head but stores only one key head and one value head, shared by all of them. It gives the smallest KV cache, at some cost in quality.

In crew terms, all eight photographers share one kit bag. The van is nearly empty, but every shot must be made with the same few lenses.

The idea comes from a 2019 paper. It argued that writing one token at a time is slow mainly because the large key and value tensors (blocks of numbers) are loaded from memory again and again.

Its translation tests ran the decoder, the part that writes the output, on a batch of 1,024 sentences at once. It used greedy decoding, always picking the likeliest next token. Chip time per output token, averaged over that batch, fell from 46 to 3.8 microseconds, about 12 times less. Quality dipped slightly: BLEU, a translation score, went from 26.7 to 26.5 on the development set, sentences held back from training.

The 2023 GQA paper reports that multi-query attention can lead to quality loss and unstable training. It also notes that larger models usually have more heads, so cutting to one key-value head is a harsher cut for them.

GQA keeps the same proportional saving as models grow. A model with 64 query heads and 8 key-value heads still keeps 8 different views of the earlier text.

Grouped query attention cuts the KV cache by the ratio of query heads to key-value heads. With 32 query heads and 8 key-value heads, the cache is four times smaller.

The team compares three models of the same depth and width that differ only in attention. The helper below is plain arithmetic from the formula. Nothing in it touches a GPU.

def kv_cache_bytes(tokens, layers=32, kv_heads=32, head_dim=128, bytes_per_number=2):
    per_token = 2 * layers * kv_heads * head_dim * bytes_per_number  # 2 = key + value
    return tokens * per_token

GB = 1_000_000_000
free_memory = 60 * GB   # GPU memory left after the model weights load
for name, kv_heads in [("multi-head", 32), ("grouped, 8", 8), ("multi-query", 1)]:
    per_chat = kv_cache_bytes(32_000, kv_heads=kv_heads)
    print(name, round(per_chat / GB, 1), "GB per chat,",
          free_memory // per_chat, "chats at once")

It prints 16.8, 4.2 and 0.5 GB per chat, and 3, 14 and 114 chats at once. Only kv_heads changes between runs. The // is whole-number division, because half a chat cannot run.

Units matter when you compare with a GPU tool. This guide uses decimal gigabytes of 1,000,000,000 bytes. A tool that reports GiB (1,073,741,824 bytes) shows the same 16.8 GB as 15.6 GiB.

Fourteen chats on one GPU changes the Monday rush. Forty surveyors do not all type at once, so the queue mostly disappears and no second GPU is needed yet.

Why not jump to multi-query and fit 114? The team ran 50 saved surveyor questions with known answers. The multi-query model missed more cross-references, like linking a page 3 photo to a page 97 repair estimate.

Yes, mostly while the model writes its answer, because each new token reads less from GPU memory. The number of attention scores it computes stays the same.

Generation has two phases. Prefill reads the whole prompt in one pass, which keeps the GPU's maths units busy for long prompts. Decode then writes the answer one token at a time.

During decode, the GPU does little maths per token but must fetch the weights and the entire cache from memory. Memory bandwidth is how many bytes per second the GPU can move. When fetching is the slow part, a step is called memory-bound.

Both papers make this point. The 2019 paper blames slow decoding on repeatedly loading keys and values. The 2023 paper says decoding is held back by loading the weights and all keys and values at every step.

The bytes per step, in numbers

Take the team's three chats at 32,000 tokens each. With multi-head attention, one decode step reads about 50.3 GB of cache plus 16 GB of weights, so 66.3 GB in total.

With GQA-8, the same step reads 12.6 GB of cache plus about 14.4 GB of weights, so 27 GB. At a round 3 TB per second of bandwidth, that is at least 22 ms versus 9 ms per step.

Now fill the GPU with 14 GQA chats. The step reads about 73 GB, close to the multi-head step for 3 chats. It takes about 10% longer, at least 24 ms, but serves almost five times as many chats.

These are lower bounds from arithmetic, not measurements. Real steps also spend time on maths, scheduling and the work between layers.

Why does decode stay memory-bound even with GQA? For each cached token, one key-value head is 512 bytes to read. Each query head sharing it does about 512 maths operations on it.

So multi-head attention does about 1 operation per byte, GQA with groups of 4 does about 4, and multi-query with 32 heads does 32. A GPU doing 1,000 trillion operations a second on 3 TB per second needs about 330 per byte to stay busy.

Those GPU figures are round numbers; check your card's data sheet. Even multi-query attention sits far below that line. GQA does not change what limits decode. It just moves fewer bytes.

• Faster per token: less cache to read on each decode step.
• More chats at once: each chat's cache is smaller, so batches, the groups of chats served in one step, can be bigger.
• Longer documents: the same memory holds a longer report per chat.
• Not much faster prefill: for long prompts, prefill is limited by maths, and GQA keeps the same number of scores.

Serving engines pile other tricks on top. vLLM stores the cache in small blocks, an idea borrowed from operating-system paging, and its paper reports near-zero wasted cache memory. Our guide to LLM inference system design covers batching and paging.

Read the model's config.json, the settings file most open models ship with their weights. Compare the number of attention heads with the number of key-value heads: equal means multi-head, one means multi-query, anything between means GQA.

The script below does that check and the cache sum together. It uses only Python's built-in json module. The config inside it is a made-up example shaped like a typical file.

import json

# Paste the model's config.json here, or use: cfg = json.load(open("config.json"))
cfg = json.loads("""{"hidden_size": 4096, "num_hidden_layers": 32,
  "num_attention_heads": 32, "num_key_value_heads": 8}""")

q_heads = cfg["num_attention_heads"]
kv_heads = cfg.get("num_key_value_heads", q_heads)    # missing means multi-head
head_dim = cfg.get("head_dim") or cfg["hidden_size"] // q_heads
kind = "MHA" if kv_heads == q_heads else "MQA" if kv_heads == 1 else "GQA"
per_token = 2 * cfg["num_hidden_layers"] * kv_heads * head_dim * 2   # 2 bytes each
print(kind, q_heads, "query heads,", kv_heads, "key-value heads")
print(per_token, "bytes per token =", per_token * 32_000 / 1e9, "GB per 32,000-token chat")

It prints GQA 32 query heads, 8 key-value heads, then 131072 bytes per token = 4.194304 GB per 32,000-token chat. Paste in a real config to size the model you plan to serve. This rule fits classic multi-head, grouped and multi-query attention only.

Some newer designs compress keys and values into one smaller shared vector per token, or limit some layers to a recent window. Their config can look multi-head while the real cache is far smaller. The script then overestimates, so trust the startup log below.

What vLLM 0.32.0 does with these numbers

We read the vLLM 0.32.0 source, released on 18 October 2026, to see how it finds the key-value heads. It looks for num_key_value_heads and a few other names that some model families use.

In its default rules, if none is present, it assumes multi-head and uses the attention-head count. Some model families get their own overrides. In its default rules, a missing head_dim means hidden_size divided by the head count. Our script follows those two rules.

Splitting a model across GPUs, called tensor parallelism, adds a twist. vLLM gives each GPU the key-value heads divided by the GPU count, but never fewer than one.

So a model with 8 key-value heads on 16 GPUs stores each head twice. A multi-query model on 8 GPUs stores its single head eight times. The source comment calls this replication.

To check on a real server, read vLLM's startup log. Version 0.32.0 prints the GPU KV cache size in tokens and the maximum concurrency for your maximum request length.

Start it with --max-model-len 32000, and a concurrency of about 14.00x would mean fourteen full-length chats fit. The cache comes from the memory share set by gpu_memory_utilization, 0.92 by default in 0.32.0.

PyTorch has a shortcut too. Its scaled_dot_product_attention function takes an enable_gqa flag. The 2.14 documentation marks GQA support as experimental, as of 19 October 2026.

Yes, but not with a setting. You merge the key and value heads in each group, then train the model a little more so it adjusts. The GQA paper calls this uptraining.

The paper mean-pooled, or averaged, the key and value weights of each group's heads into one head. That worked better than keeping a single head, which beat starting a fresh one. Uptraining then used about 5% of the original pre-training compute.

The uptrained GQA models reached quality close to multi-head attention at speed close to multi-query attention. Those tests used encoder-decoder models, which have a separate reader (encoder) and writer (decoder). GQA replaced only the decoder's attention, not the encoder's self-attention.

Five percent of a big training run is still a lot of compute. The paper reports about 600 accelerator chip-days, one chip running for one day each, for its 5% uptraining run. For a small team, picking a model trained with GQA from the start is far cheaper.

Here is the team's choice drawn as a decision. Start at the top and follow the answers that match your own traffic.

The team swapped models and reran the 50 saved questions with the same prompts. They switched once the GQA model matched the old one on 48 of them.

Explain it as a memory problem first: the KV cache limits how many long chats a GPU can hold. Then present GQA as the dial that shrinks that cache.

1. 1

   State the bottleneck

   Say that decoding reads the whole KV cache for each new token. Long chats and many users make that cache, not the weights, the limit.

2. 2

   Do the sum out loud

   Give the formula: 2 times layers times key-value heads times head size times bytes. Multiply by tokens and by chats at once.

3. 3

   Introduce the dial

   Multi-head stores one key-value head per query head. Multi-query stores one in total. Grouped query stores a few, shared by groups.

4. 4

   Name the trade

   Fewer key-value heads means a smaller cache, faster decoding and bigger batches. Too few can hurt quality, so GQA sits in the middle.

5. 5

   Say what it does not fix

   It barely shrinks the weights or speeds up prefill. If you split the model across more GPUs than it has key-value heads, vLLM copies heads and the saving shrinks.

6. 6

   Close with measurement

   Choose by testing your own questions and measuring chats per GPU and tokens per second. Paper results are a starting point, not proof for your task.