How do you implement grouped query attention (GQA)?

In this article

1. Before you start
2. Why interviewers ask this
3. Grouped-query attention from first principles
4. Worked example
5. Trade-offs and edge cases
6. How to answer it in an interview
7. Common mistakes
8. Follow-up questions
9. Sources

Before you start

This page builds on the machine learning (ML) coding round guide and on attention you can already write.

Why interviewers ask this

Candidates report being asked to write grouped-query attention, often as the follow-up after multi-head attention in an ML coding round, so the question is common enough to rehearse. It separates people who memorized the attention formula from people who know why modern decoders change it.

A weak answer says "GQA is a faster attention" and cannot say what is shared. A strong answer names the shared tensors (keys and values), states the shapes, writes the repeat step, and explains that the benefit is a smaller cache and less data read per generated token. A Staff-level answer also says what does not change: the number of query heads, the score arithmetic and the causal mask. The interviewer then tests that with one follow-up, such as what happens at the two ends of the head count.

Grouped-query attention from first principles

Grouped-query attention

Grouped-query attention is an attention variant that uses fewer key-value heads than query heads, so each key-value head is shared by a group of query heads, with the head count more than one and fewer than the number of query heads.

Multi-query attention

Multi-query attention is an attention variant that uses a single key-value head shared by every query head, which shrinks the keys and values that decoding must read at each step.

Key-value head

A key-value head is a pair of projections, one for keys and one for values, in an attention layer that produces the keys and values one or more query heads attend to.

In plain words

Picture a classroom where every student (a query head) asks a different question. In multi-head attention each student has a private reference book (a key-value head). In multi-query attention the whole room shares one book. In grouped-query attention the room splits into groups, say 8, and each group shares one book. The analogy breaks in one place: the students still ask different questions, because the query projections are never shared.

Let H be the number of query heads and G the number of key-value heads. G divides H. Query head h reads key-value head h divided by H/G, rounded down.

Shazeer (2019) proposed multi-query attention, where the keys and values are shared across all heads, to cut the memory bandwidth cost of incremental decoding. Memory bandwidth is the rate at which a chip can move data from memory to its compute units, and it limits decoding because every step rereads the cached keys and values.

Step 1: Queries keep H heads, keys and values get G

The layer projects the input of shape [batch, T, model width] three times. The query projection outputs H heads of size d, giving Q of shape [batch, H, T, d]. The key and value projections output only G heads each, giving K and V of shape [batch, G, T, d]. Setting G to H recovers multi-head attention, and setting G to 1 recovers multi-query attention.

Input: the hidden states. Output: one large Q and two smaller K and V. You now know where the saving starts: the key and value projection matrices have G over H of their usual size. Multi-head attention itself comes from Vaswani et al. (2017).

Step 2: Each key-value head serves a group of query heads

The H query heads are split into G groups of H/G consecutive heads. Query heads 0 to H/G minus 1 use key-value head 0, the next H/G use head 1, and so on. Ainslie et al. (2023) describe grouped-query attention as an intermediate number of key-value heads, more than one and fewer than the number of query heads.

Input: the head index. Output: the group index. The mapping is plain integer division, and it must match how you repeat the tensors in the next step, or heads read the wrong group. For example, with 8 query heads and 2 key-value heads, each group holds 4 consecutive query heads: heads 0 to 3 read key-value head 0 and heads 4 to 7 read key-value head 1. Writing the mapping down before you write any tensor code is the quickest way to avoid a silent mismatch.

Step 3: Repeat the keys and values, then attend as usual

To reuse the standard attention code, copy each key-value head H/G times along the head axis, so K and V become [batch, H, T, d]. In PyTorch that is repeat_interleave on dimension 1, which keeps copies of one head next to each other, matching the grouping of Step 2. Scores, the causal mask, softmax and the weighted sum of values are then exactly multi-head attention.

Input: Q, repeated K and V. Output: the attended values, merged back to the model width. In plain PyTorch, repeat_interleave always makes a temporary copy that is thrown away; a fused kernel never materializes the copy and indexes key-value head h divided by H/G, rounded down, directly.

Step 4: The cache stores G heads, not H

The KV cache keeps the keys and values from before the repeat, with G heads, and the repeat happens each step on read. That is the saving: per token and per layer the cache holds 2 times G times d numbers instead of 2 times H times d. The query side and the score arithmetic are unchanged, so the benefit is memory and data movement, not fewer multiplications.

Input: new K and V of the current token. Output: a cache that grows by G heads per token. Shazeer (2019) frames the target as the memory-bandwidth cost of repeatedly loading large keys and values, which is exactly what a smaller G cuts.

1 · Project2 · Group3 · Attend4 · CacheReplay

Step-by-step description

1. Queries keep H heads; keys and values get G.
2. Each key-value head serves a group of query heads.
3. Repeat keys and values, then attend as usual.
4. The cache stores G heads, not H.

• Queries keep H heads while keys and values get G heads.
• Each key-value head serves H over G query heads.
• The cache stores G heads, not H.

Worked example

The numbers below use one illustrative configuration, not a specific model: 32 query heads, head size 128, model width 4,096, 32 layers, 16-bit numbers (2 bytes each) and a context of 8,192 tokens.¹ The cache bytes per token and layer are 2 (keys and values) times the key-value head count times 128 times 2 bytes.

Per token and layer that is 16,384 bytes with 32 key-value heads (multi-head attention (MHA)), 4,096 with 8 (grouped-query attention) and 512 with 1 (multi-query attention (MQA)).¹ Across 32 layers and 8,192 tokens, one sequence needs 4 GiB, 1 GiB and 0.125 GiB.¹ The same ratios apply to the data read per decoding step: 128 MiB, 32 MiB and 4 MiB per layer at a context of 8,192 tokens, while the score computation stays at 33,554,432 multiply-adds per layer and step in all three cases.¹

Text description of the diagram

• One sequence needs 4 GiB with 32 key-value heads.
• One sequence needs 1 GiB with 8 key-value heads and 0.125 GiB with 1.

Here is the implementation. It handles grouped-query attention, multi-query attention (n_kv_heads=1) and multi-head attention (n_kv_heads=n_heads) with one class, plus a cache and a causal mask.

import math
import torch
import torch.nn as nn

class GroupedQueryAttention(nn.Module):
    def __init__(self, d_model, n_heads, n_kv_heads):
        super().__init__()
        assert d_model % n_heads == 0 and n_heads % n_kv_heads == 0
        self.h, self.g = n_heads, n_kv_heads
        self.d = d_model // n_heads
        self.wq = nn.Linear(d_model, self.h * self.d, bias=False)
        self.wk = nn.Linear(d_model, self.g * self.d, bias=False)
        self.wv = nn.Linear(d_model, self.g * self.d, bias=False)
        self.wo = nn.Linear(self.h * self.d, d_model, bias=False)

    def forward(self, x, past=None):
        b, t, _ = x.shape
        q = self.wq(x).view(b, t, self.h, self.d).transpose(1, 2)  # [b, h, t, d]
        k = self.wk(x).view(b, t, self.g, self.d).transpose(1, 2)  # [b, g, t, d]
        v = self.wv(x).view(b, t, self.g, self.d).transpose(1, 2)  # [b, g, t, d]
        if past is not None:                       # cache keeps g heads, not h
            k = torch.cat([past[0], k], dim=2)
            v = torch.cat([past[1], v], dim=2)
        cache = (k, v)
        rep = self.h // self.g
        k = k.repeat_interleave(rep, dim=1)        # [b, h, s, d]
        v = v.repeat_interleave(rep, dim=1)
        scores = q @ k.transpose(-2, -1) / math.sqrt(self.d)  # [b, h, t, s]
        s = k.shape[2]
        # query i sits at absolute position s - t + i and sees keys up to it
        keep = torch.ones(t, s, dtype=torch.bool, device=x.device).tril(s - t)
        scores = scores.masked_fill(~keep, float("-inf"))
        out = torch.softmax(scores, dim=-1) @ v    # [b, h, t, d]
        out = out.transpose(1, 2).reshape(b, t, self.h * self.d)
        return self.wo(out), cache

Read it against the steps. The three projections are Step 1, repeat_interleave is Steps 2 and 3, and returning cache before the repeat is Step 4. The mask offset s - t makes the same code work for a full prompt (t equals s) and for one new token (t equals 1). For a stable softmax and the mask itself, see the pages on numerically stable softmax and the causal mask.

Trade-offs and edge cases

Grouped-query attention trades a little quality for a lot of cache. Ainslie et al. (2023) report that multi-query attention can lead to quality degradation and that uptrained grouped-query attention achieves quality close to multi-head attention with speed comparable to multi-query attention.

Text description of the diagram

• Multi-query attention can lead to quality degradation.
• Grouped-query attention uses more than one and fewer than the number of query heads.

Edge cases to name in an interview:

• Divisibility. H must be a multiple of G; the assert in the code enforces it.
• Mismatched repeat. Using repeat (which tiles the whole head list) instead of repeat_interleave pairs query heads with the wrong key-value heads and silently changes the model.
• Position encodings. Rotary or other position encodings apply to the G key heads before caching; see positional encodings from scratch.
• Converting a model. Ainslie et al. (2023), §2.1 convert a multi-head checkpoint by mean-pooling the key and value projection matrices of the heads in each group, then uptrain with a small share of the original pre-training compute (5% in the abstract's multi-query recipe), rather than training a separate model.
• No savings in floating-point operations (FLOPs). Per the worked example, the score arithmetic is the same, so training speed gains are small compared with decoding gains.
• Tensor parallelism. When tensor parallelism splits heads across devices and the device count exceeds G, each device still needs one key-value head, so heads are replicated and the per-device cache stops shrinking once G equals the device count. Pick G with the serving layout in mind.
• Where the field moved. Multi-head latent attention goes further than sharing heads: DeepSeek-AI (2024) report that it "reduces the KV cache by 93.3%" against DeepSeek 67B, so grouped-query attention is a common baseline, not the last word on cache size.

How to answer it in an interview

The spoken structure: say what is shared, give the shapes, say where the repeat happens, then state the cache saving and the two endpoints.

Common mistakes

• "GQA reduces the number of query heads." (wrong) Only keys and values are shared. (right) "The queries keep all H heads; the key-value heads drop to G."
• "GQA makes attention need fewer multiplications." (wrong) The score work is the same after the repeat. (right) "It shrinks the cache and the data read per decoding step."
• "I cache the repeated keys and values." (wrong) That gives back the whole saving. (right) "I cache G heads and repeat on read."
• "GQA is the same as multi-query attention." (wrong) Multi-query attention is the case G equals one. (right) "GQA generalizes it; one head is the extreme, H heads is multi-head."

Follow-up questions

How is grouped query attention different from multi head attention?

Multi-head attention gives every query head its own key-value head, so G equals H. Grouped-query attention lets groups of query heads share one, so the cache and the key and value projections shrink by H/G while the queries stay the same.

How does grouped query attention differ from multi query attention?

Multi-query attention uses one key-value head for all query heads. Grouped-query attention uses an intermediate number, which recovers quality close to multi-head attention at a speed comparable to multi-query attention, according to Ainslie et al. (2023).

Does grouped query attention reduce the number of query heads?

No. The query projection still outputs H heads, and only the key and value projections shrink to G heads.

How do you convert a multi-head checkpoint to grouped-query attention?

Mean-pool the key and value projection matrices of the original heads in each group into one head, then uptrain: continue training for a small proportion of the original pre-training steps, as described in Ainslie et al. (2023), §2.1. Their abstract gives 5% of original compute for the multi-query recipe.

How do you handle the KV cache and a causal mask in one implementation?

Concatenate new keys and values onto the cached G-head tensors, repeat after concatenation, and build the mask with an offset of the cache length, as in the code above. The generative pre-trained transformer (GPT) from scratch page shows the surrounding decoder block.

Which models use grouped query attention, such as Llama 3?

Llama 2's 34B and 70B models use grouped-query attention, according to Touvron et al. (2023), §2.2. The configuration file on the Llama 3 8B model card lists 32 attention heads and 8 key-value heads (as of October 2026), so four query heads share each key-value head. A model's configuration lists both head counts; if they differ, it uses grouped-query or multi-query attention.

Sources

1. GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints (arXiv 2305.13245v3) — arXiv (Ainslie et al., Google Research); accessed 2026-10-26, 23 Dec 2023
2. Fast Transformer Decoding: One Write-Head is All You Need (arXiv 1911.02150) — arXiv (Shazeer); accessed 2026-10-26, 6 Nov 2019
3. Attention Is All You Need (arXiv 1706.03762) — arXiv (Vaswani et al.), 12 Jun 2017
4. Llama 2: Open Foundation and Chat Models from Meta (arXiv 2307.09288) — arXiv (Touvron et al.), 18 Jul 2023
5. Meta Llama 3 8B model card and configuration — Meta, Hugging Face; accessed 2026-10-26, 18 Apr 2024
6. DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model (arXiv 2405.04434) — arXiv (DeepSeek-AI), 7 May 2024