Grouped-Query Attention (GQA)

Ainslie, Lee-Thorp, de Jong, Zemlyanskiy, Lebrón & Sanghai, "GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints", Google Research, 2023

Learning Objectives

After completing this chapter, you will be able to:

1. Explain why the KV-cache is the primary memory bottleneck during autoregressive inference and why MHA makes it $H\times$ worse than necessary
2. Describe how GQA partitions $H$ query heads into $G$ groups where each group shares a single set of keys and values
3. Compute the group assignment function $g(i) = \lfloor i \,/\, (H/G) \rfloor$ and understand how it generalizes MHA ($G = H$) and MQA ($G = 1$)
4. Perform a full GQA forward pass by hand on the shared “The cat sat on mat” example, comparing $G = 2$ and $G = 1$
5. Calculate KV-cache memory savings for production models like LLaMA 3 70B and Mistral 7B
6. Understand the “uptrain” procedure that converts an existing MHA checkpoint to GQA with only 5% of original training compute
7. Implement GQA from scratch in both NumPy and PyTorch

The Problem: The KV-Cache Dilemma

What Is the KV-Cache?

When a transformer generates text token by token (autoregressive decoding), it must recompute the attention scores for every previously generated token at every step. To avoid this $O(N^2)$ recomputation, modern systems cache the Key and Value tensors from all previous positions. This “KV-cache” trades memory for speed: once a token's $K$ and $V$ are computed, they are stored and reused for all subsequent tokens.

The problem is that the KV-cache grows with three multipliers: the number of attention heads $H$, the sequence length $N$, and the per-head dimension $d_k$. For each transformer layer, the cache stores $2 \times H \times N \times d_k$ floating-point values (the factor of 2 accounts for both K and V).

MHA's Memory Wall

In standard Multi-Head Attention (Chapter 2), every head has its own independent $K$ and $V$ projections. For LLaMA 3 70B with $H = 64$ heads, $d_k = 128$, and 80 layers in FP16:

At a 128K context window, that is 320 GB of KV-cache alone — far exceeding the memory of any single GPU. This memory wall is the reason that long-context inference requires multi-GPU setups or aggressive compression.

MQA's Quality Cliff

Chapter 5 introduced Multi-Query Attention (MQA, Shazeer 2019), which shares a single $K, V$ pair across all $H$ heads. This reduces the KV-cache by a factor of $H$ — from 320 GB to just 5 GB in our LLaMA example. The speed improvement is dramatic.

But the quality cost is real. When all heads are forced to attend through the same key-value lens, the model loses the representational diversity that makes multi-head attention powerful. Benchmarks show measurable degradation on reasoning-heavy tasks, summarization, and multi-hop question answering (Ainslie et al., 2023, Table 1).

The Core Tension: MHA gives maximum quality but is memory-hungry. MQA gives maximum speed but sacrifices quality. Is there a sweet spot between these extremes?

The Intuition Behind Grouped-Query Attention

The Library Analogy

Imagine a research institute with 64 researchers (query heads). Each researcher needs to look up information in reference books (keys and values) to answer questions.

• MHA: Every researcher has their own private library. Maximum research diversity, but the institute needs 64 separate library buildings. Expensive.
• MQA: Everyone shares a single library. Cheap to maintain, but when 64 researchers all need the same book simultaneously, the bottleneck is severe. Popular topics get congested; niche topics are underserved.
• GQA: Researchers are organized into departments of 8. Each department has its own library. Eight libraries instead of 64 — an $8\times$ cost reduction — but each department can specialize its collection. The biology department keeps different books than the physics department.

This is exactly how GQA works. Instead of $H$ separate KV sets (MHA) or 1 shared KV set (MQA), GQA uses $G$ groups, where $1 \le G \le H$. Each group of $H/G$ query heads shares one set of keys and values.

The Paper: Ainslie et al. 2023

Grouped-Query Attention was introduced by Joshua Ainslie and colleagues at Google Research in 2023. Their paper, “GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints,” made two key contributions:

1. The GQA mechanism itself: a clean generalization that places MHA and MQA as special cases on a single continuum controlled by $G$
2. The uptrain procedure: a practical method to convert existing MHA models to GQA without training from scratch, requiring only ~5% of the original training compute

The Uptrain Insight

The uptrain procedure works by mean-pooling the KV projection weights within each group. If an MHA model has 64 separate $W^K$ matrices and you want $G = 8$ groups, you average heads 0–7 into group 0's $W^K$, heads 8–15 into group 1's $W^K$, and so on. Then you fine-tune the model for a small fraction of the original training budget to recover any quality loss.

This is how LLaMA 2 70B was converted from MHA to GQA. The resulting model achieved near-MHA quality at near-MQA inference speed, and GQA has since become the default architecture for virtually all production large language models.

Mathematical Formulation

The Group Assignment Function

The heart of GQA is a simple function that maps each query head $i$ to its KV group:

where $\lfloor \cdot \rfloor$ is the floor function (integer division). With $H$ query heads and $G$ KV groups, each group serves $H/G$ heads. Heads 0 through $H/G - 1$ map to group 0, heads $H/G$ through $2H/G - 1$ map to group 1, and so on.

The GQA Equation

Using the group assignment, each head computes standard scaled dot-product attention but uses its group's K and V rather than its own:

The subscript $g(i)$ on $W^K$ and $W^V$ is what makes GQA different: multiple heads share the same projection matrices for keys and values.

Symbol-by-Symbol Breakdown

The Spectrum: MQA ↔ GQA ↔ MHA

GQA is a single mechanism that continuously interpolates between MQA and MHA:

Interactive: Group Visualizer

Use the controls below to explore how changing $G$ affects the head-to-group mapping. Try the model presets to see real configurations used in production:

KV-Cache Memory Analysis

The Memory Formula

The KV-cache size per layer in bytes is:

The total cache across all $L$ layers is $L$ times this value. The key insight: the cache scales with $G$, not $H$. By choosing $G < H$, we directly reduce memory consumption.

Interactive: KV-Cache Comparison

Select a model and sequence length to see how GQA reduces KV-cache memory compared to MHA and MQA:

Step-by-Step Calculation

Setup: Shared Example

We use the same $Q, K, V \in \mathbb{R}^{5 \times 4}$ matrices as every chapter, with $d_{\text{model}} = 4$ and $H = 2$ query heads ($d_k = 2$ per head). We compare two configurations:

• G = 2 (one group per head = MHA behavior): Head 0 uses $K[:,0{:}2], V[:,0{:}2]$; Head 1 uses $K[:,2{:}4], V[:,2{:}4]$
• G = 1 (one group for all = MQA behavior): Both heads use $K[:,0{:}2], V[:,0{:}2]$

GQA G = 2 (MHA Mode)

With $G = 2$ and $H = 2$, $\text{heads\_per\_group} = 1$. Each head gets its own KV group. We trace the computation for “The” (row 0).

Head 0 → Group 0

$g(0) = \lfloor 0/1 \rfloor = 0$. Head 0 uses $K_{g0} = K[:,0{:}2]$, $V_{g0} = V[:,0{:}2]$.

Query: $Q_{h0}[0] = Q[0, 0{:}2] = [1, 0]$

Scaled scores ($Q_{h0}[0] \cdot K_{g0}^T / \sqrt{2}$):

Softmax: $[0.1237,\; 0.2509,\; 0.2509,\; 0.1237,\; 0.2509]$

Output: $O_{h0}[0] = \sum w_j \cdot V_{g0}[j] = [0.2491,\; 0.3764]$

Head 1 → Group 1

$g(1) = \lfloor 1/1 \rfloor = 1$. Head 1 uses $K_{g1} = K[:,2{:}4]$, $V_{g1} = V[:,2{:}4]$.

Query: $Q_{h1}[0] = Q[0, 2{:}4] = [1, 0]$

Scaled scores ($Q_{h1}[0] \cdot K_{g1}^T / \sqrt{2}$):

Softmax: $[0.1337,\; 0.2711,\; 0.1337,\; 0.2711,\; 0.1904]$

Output: $O_{h1}[0] = [0.2289,\; 0.3663]$

Concatenation

The final output for “The” concatenates both head outputs:

GQA G = 1 (MQA Mode)

Now set $G = 1$. With $\text{heads\_per\_group} = 2$, both heads share the same KV group. Head 0 is unchanged (it already used group 0). But Head 1 now also uses group 0: $K[:,0{:}2]$ and $V[:,0{:}2]$ instead of $K[:,2{:}4]$ and $V[:,2{:}4]$.

Head 1 Changes KV Source

Query: $Q_{h1}[0] = [1, 0]$ (unchanged — queries are always head-specific)

Scaled scores (now against $K_{g0}$ instead of $K_{g1}$):

New softmax: $[0.1237,\; 0.2509,\; 0.2509,\; 0.1237,\; 0.2509]$

New output: $O_{h1}^{\text{MQA}}[0] = [0.2491,\; 0.3764]$

Critical observation: Head 1's output for “The” is now $[0.2491, 0.3764]$ — exactly the same as Head 0's output! When both heads share the same K, V, and their queries happen to produce similar attention patterns, the heads become redundant. This is the quality cost of MQA that GQA exists to mitigate.

Output Comparison: G = 2 vs G = 1

Dims 0–1 (Head 0) are identical in both configurations. Dims 2–3 (Head 1) differ because Head 1 sees different keys and values. The maximum change is 0.186 for “on” in dimension 3. In our toy example with $H = 2$, there is no middle ground — $G$ can only be 1 or 2. The real power of GQA emerges with $H = 64$ where $G$ can be 2, 4, 8, 16, or 32.

Full Attention Weight Matrices

Averaged Weights — GQA G = 2 (MHA)

Full Output Matrices

Output — GQA G = 2 (MHA)

Output — GQA G = 1 (MQA)

Interactive: Attention Heatmap

Toggle between G = 2 (MHA) and G = 1 (MQA) to see how sharing KV groups changes the attention patterns. Cells highlighted in amber indicate weights that changed from the MHA baseline:

Applications Across Domains

GQA has become the de facto standard for all new large models. The few exceptions (very small models under 1B parameters) use MHA because the KV-cache savings are negligible at small scale.

Connection to Modern Systems

• Flash Attention + GQA: Flash Attention (Chapter 13) is an IO-aware implementation trick that does not change the mathematical result. It composes perfectly with GQA — the reduced number of KV heads means fewer memory reads during tiled computation, amplifying Flash Attention's speed benefits. All major inference engines (vLLM, TGI, TensorRT-LLM) implement this combination.
• Paged Attention + GQA: vLLM's PagedAttention manages KV-cache as virtual memory pages. With GQA, the page tables are $H/G \times$ smaller, enabling higher batch sizes and better GPU utilization for serving.
• KV-Cache Quantization: GQA reduces the number of KV heads that need to be stored, and each KV head can be further compressed via INT8 or INT4 quantization. The combination gives $(H/G) \times 2$ to $(H/G) \times 4$ total compression versus MHA FP16.
• Speculative Decoding: Draft models in speculative decoding pipelines often use aggressive GQA ($G = 1$ or $G = 2$) because inference speed matters more than quality for the draft — the verifier catches errors.
• Positional Encodings (Chapters 7–9): RoPE and ALiBi are applied to Q and K vectors. In GQA, the rotation/bias is applied to each query head independently but to each KV group only once. This is a minor efficiency gain but matters at scale.

Complexity Analysis

Note that the compute cost is identical across all three — GQA does not reduce FLOPs. The savings come entirely from memory bandwidth: less data needs to be loaded from GPU memory during each attention operation, which translates to higher throughput during inference.

Python Implementation

The full GQA class with the shared example. The _group_index method is the only addition beyond standard multi-head attention — one line of integer division controls the entire KV sharing behavior.

      d0   d1   d2   d3
The  1.0  0.0  1.0  0.0
cat  0.0  2.0  0.0  1.0
sat  1.0  1.0  1.0  0.0
on   0.0  0.0  1.0  1.0
mat  1.0  0.0  0.0  1.0

      d0   d1   d2   d3
The  0.0  1.0  0.0  1.0
cat  1.0  0.0  1.0  0.0
sat  1.0  1.0  0.0  0.0
on   0.0  0.0  1.0  1.0
mat  1.0  0.0  0.5  0.5

      d0   d1   d2   d3
The  1.0  0.0  0.0  0.0
cat  0.0  1.0  0.0  0.0
sat  0.0  0.0  1.0  0.0
on   0.0  0.0  0.0  1.0
mat  0.5  0.5  0.5  0.5

      d0   d1
The  1.0  0.0
cat  0.0  2.0
sat  1.0  1.0
on   0.0  0.0
mat  1.0  0.0

      d2   d3
The  1.0  0.0
cat  0.0  1.0
sat  1.0  0.0
on   1.0  1.0
mat  0.0  1.0

      d0   d1
The  0.0  1.0
cat  1.0  0.0
sat  1.0  1.0
on   0.0  0.0
mat  1.0  0.0

      d2   d3
The  0.0  1.0
cat  1.0  0.0
sat  0.0  0.0
on   1.0  1.0
mat  0.5  0.5

      The     cat     sat      on     mat
The  0.000   0.707   0.707   0.000   0.707
cat  1.414   0.000   1.414   0.000   0.000
sat  0.707   0.707   1.414   0.000   0.707
on   0.000   0.000   0.000   0.000   0.000
mat  0.000   0.707   0.707   0.000   0.707

       The      cat      sat       on      mat
The  0.1237   0.2509   0.2509   0.1237   0.2509
cat  0.3664   0.0891   0.3664   0.0891   0.0891
sat  0.1812   0.1812   0.3673   0.0893   0.1812
on   0.2000   0.2000   0.2000   0.2000   0.2000
mat  0.1237   0.2509   0.2509   0.1237   0.2509

        d0       d1
The  0.2491   0.3764
cat  0.4110   0.1337
sat  0.2718   0.2718
on   0.3000   0.3000
mat  0.2491   0.3764

       d0       d1       d2       d3
The  0.2491   0.3764   0.2289   0.3663
cat  0.4110   0.1337   0.2289   0.3663
sat  0.2718   0.2718   0.2289   0.3663
on   0.3000   0.3000   0.1799   0.4579
mat  0.2491   0.3764   0.2289   0.3663

      d0   d1   d2   d3
The  1.0  0.0  1.0  0.0
cat  0.0  2.0  0.0  1.0
sat  1.0  1.0  1.0  0.0
on   0.0  0.0  1.0  1.0
mat  1.0  0.0  0.0  1.0

      d0   d1   d2   d3
The  0.0  1.0  0.0  1.0
cat  1.0  0.0  1.0  0.0
sat  1.0  1.0  0.0  0.0
on   0.0  0.0  1.0  1.0
mat  1.0  0.0  0.5  0.5

      d0   d1   d2   d3
The  1.0  0.0  0.0  0.0
cat  0.0  1.0  0.0  0.0
sat  0.0  0.0  1.0  0.0
on   0.0  0.0  0.0  1.0
mat  0.5  0.5  0.5  0.5

       d0       d1       d2       d3
The  0.2491   0.3764   0.2289   0.3663
cat  0.4110   0.1337   0.2289   0.3663
sat  0.2718   0.2718   0.2289   0.3663
on   0.3000   0.3000   0.1799   0.4579
mat  0.2491   0.3764   0.2289   0.3663

       d0       d1       d2       d3
The  0.2491   0.3764   0.2491   0.3764
cat  0.4110   0.1337   0.3583   0.2126
sat  0.2718   0.2718   0.2491   0.3764
on   0.3000   0.3000   0.2718   0.2718
mat  0.2491   0.3764   0.3583   0.2126

       d0     d1     d2     d3
The  0.000  0.000  0.020  0.010
cat  0.000  0.000  0.129  0.154
sat  0.000  0.000  0.020  0.010
on   0.000  0.000  0.092  0.186
mat  0.000  0.000  0.129  0.154

1import numpy as np
2import math
3
4class GroupedQueryAttention:
5    """
6    Grouped-Query Attention (Ainslie et al., 2023)
7
8    Generalizes MHA (G=H) and MQA (G=1) via G KV groups.
9    Each group of H/G query heads shares one set of K, V.
10    Computes: head_i = Attention(Q_h, K_{g(h)}, V_{g(h)})
11    where g(h) = floor(h / (H/G))
12    """
13
14    def __init__(self, d_model: int, H: int, G: int):
15        """
16        Args:
17            d_model: Total model dimension
18            H: Number of query heads
19            G: Number of KV groups (1 <= G <= H)
20        """
21        assert d_model % H == 0
22        assert H % G == 0
23        self.d_model = d_model
24        self.H = H
25        self.G = G
26        self.d_k = d_model // H
27        self.heads_per_group = H // G
28        self.scale = math.sqrt(self.d_k)
29
30    def _softmax(self, x: np.ndarray) -> np.ndarray:
31        """Numerically stable softmax along last axis."""
32        x_shifted = x - np.max(x, axis=-1, keepdims=True)
33        exp_x = np.exp(x_shifted)
34        return exp_x / np.sum(exp_x, axis=-1, keepdims=True)
35
36    def _group_index(self, head_idx: int) -> int:
37        """Maps query head index to its KV group."""
38        return head_idx // self.heads_per_group
39
40    def forward(self, Q: np.ndarray, K: np.ndarray, V: np.ndarray):
41        """
42        Args:
43            Q: Query matrix  (N, d_model)
44            K: Key matrix    (N, d_model)
45            V: Value matrix  (N, d_model)
46        Returns:
47            all_weights: List of H weight matrices, each (N, N)
48            output: Concatenated output (N, d_model)
49        """
50        N = Q.shape[0]
51        head_outputs = []
52        all_weights = []
53
54        for h in range(self.H):
55            g = self._group_index(h)
56
57            Q_h = Q[:, h * self.d_k : (h + 1) * self.d_k]
58            K_g = K[:, g * self.d_k : (g + 1) * self.d_k]
59            V_g = V[:, g * self.d_k : (g + 1) * self.d_k]
60
61            scores = Q_h @ K_g.T / self.scale
62            weights = self._softmax(scores)
63            output = weights @ V_g
64
65            all_weights.append(weights)
66            head_outputs.append(output)
67
68        return all_weights, np.hstack(head_outputs)
69
70    def explain(self, Q, K, V, tokens, query_idx=0):
71        """Print step-by-step trace for one token."""
72        all_weights, full_output = self.forward(Q, K, V)
73        token = tokens[query_idx]
74        mech = "MHA" if self.G == self.H else "MQA" if self.G == 1 else "GQA"
75        print(f"\n=== GQA trace for '{token}' (H={self.H}, G={self.G}) ===")
76        print(f"Mechanism: {mech}  |  Heads/group: {self.heads_per_group}")
77        print(f"KV saving: {(1 - self.G / self.H) * 100:.0f}%\n")
78
79        for h in range(self.H):
80            g = self._group_index(h)
81            print(f"--- Head {h} -> KV Group {g} ---")
82            w = all_weights[h]
83            for j, t in enumerate(tokens):
84                bar = "#" * int(w[query_idx, j] * 40)
85                print(f"  A[{token},{t}] = {w[query_idx, j]:.4f} |{bar}|")
86
87        print(f"\nOutput[{token}] = {np.round(full_output[query_idx], 4)}")
88
89
90# ── Shared Example (used in every chapter) ──
91tokens = ["The", "cat", "sat", "on", "mat"]
92
93Q = np.array([
94    [1.0, 0.0, 1.0, 0.0],   # The
95    [0.0, 2.0, 0.0, 1.0],   # cat
96    [1.0, 1.0, 1.0, 0.0],   # sat
97    [0.0, 0.0, 1.0, 1.0],   # on
98    [1.0, 0.0, 0.0, 1.0],   # mat
99])
100
101K = np.array([
102    [0.0, 1.0, 0.0, 1.0],   # The
103    [1.0, 0.0, 1.0, 0.0],   # cat
104    [1.0, 1.0, 0.0, 0.0],   # sat
105    [0.0, 0.0, 1.0, 1.0],   # on
106    [1.0, 0.0, 0.5, 0.5],   # mat
107])
108
109V = np.array([
110    [1.0, 0.0, 0.0, 0.0],   # The
111    [0.0, 1.0, 0.0, 0.0],   # cat
112    [0.0, 0.0, 1.0, 0.0],   # sat
113    [0.0, 0.0, 0.0, 1.0],   # on
114    [0.5, 0.5, 0.5, 0.5],   # mat
115])
116
117# ── GQA with G=2 (same as MHA for H=2) ──
118gqa_mha = GroupedQueryAttention(d_model=4, H=2, G=2)
119weights_g2, output_g2 = gqa_mha.forward(Q, K, V)
120print("=== GQA G=2 (MHA mode) ===")
121print("Output:\n", np.round(output_g2, 4))
122gqa_mha.explain(Q, K, V, tokens, query_idx=0)
123
124# ── GQA with G=1 (same as MQA for H=2) ──
125gqa_mqa = GroupedQueryAttention(d_model=4, H=2, G=1)
126weights_g1, output_g1 = gqa_mqa.forward(Q, K, V)
127print("\n=== GQA G=1 (MQA mode) ===")
128print("Output:\n", np.round(output_g1, 4))
129gqa_mqa.explain(Q, K, V, tokens, query_idx=0)
130
131# ── Compare ──
132diff = np.abs(output_g2 - output_g1)
133print("\nDifference:\n", np.round(diff, 4))
134print("Max diff:", round(diff.max(), 4))

PyTorch Implementation

The PyTorch version is structurally identical to the NumPy version. Key differences: torch.matmul replaces @ for clarity, F.softmax replaces our manual implementation, and an optional mask parameter supports causal attention (Chapter 3).

1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4import math
5
6class GroupedQueryAttention(nn.Module):
7    """
8    Grouped-Query Attention (Ainslie et al., 2023) - PyTorch
9
10    Raw version: takes pre-computed Q, K, V (no learned projections).
11    Supports GPU, autograd, optional causal mask.
12    """
13
14    def __init__(self, d_model: int, H: int, G: int):
15        super().__init__()
16        assert d_model % H == 0 and H % G == 0
17        self.d_model = d_model
18        self.H = H
19        self.G = G
20        self.d_k = d_model // H
21        self.heads_per_group = H // G
22        self.scale = math.sqrt(self.d_k)
23
24    def forward(
25        self,
26        Q: torch.Tensor,
27        K: torch.Tensor,
28        V: torch.Tensor,
29        mask: torch.Tensor | None = None,
30    ) -> tuple[list[torch.Tensor], torch.Tensor]:
31        """
32        Args:
33            Q: (N, d_model), K: (N, d_model), V: (N, d_model)
34            mask: Optional (N, N) boolean, True = masked
35        Returns:
36            all_weights: list of H tensors, each (N, N)
37            output: (N, d_model) concatenated
38        """
39        head_outputs = []
40        all_weights = []
41
42        for h in range(self.H):
43            g = h // self.heads_per_group
44
45            Q_h = Q[:, h * self.d_k : (h + 1) * self.d_k]
46            K_g = K[:, g * self.d_k : (g + 1) * self.d_k]
47            V_g = V[:, g * self.d_k : (g + 1) * self.d_k]
48
49            scores = torch.matmul(Q_h, K_g.T) / self.scale
50
51            if mask is not None:
52                scores = scores.masked_fill(mask, float("-inf"))
53
54            weights = F.softmax(scores, dim=-1)
55            output = torch.matmul(weights, V_g)
56
57            all_weights.append(weights)
58            head_outputs.append(output)
59
60        return all_weights, torch.cat(head_outputs, dim=-1)
61
62
63# ── Shared Example ──
64tokens = ["The", "cat", "sat", "on", "mat"]
65
66Q = torch.tensor([
67    [1.0, 0.0, 1.0, 0.0],
68    [0.0, 2.0, 0.0, 1.0],
69    [1.0, 1.0, 1.0, 0.0],
70    [0.0, 0.0, 1.0, 1.0],
71    [1.0, 0.0, 0.0, 1.0],
72])
73K = torch.tensor([
74    [0.0, 1.0, 0.0, 1.0],
75    [1.0, 0.0, 1.0, 0.0],
76    [1.0, 1.0, 0.0, 0.0],
77    [0.0, 0.0, 1.0, 1.0],
78    [1.0, 0.0, 0.5, 0.5],
79])
80V = torch.tensor([
81    [1.0, 0.0, 0.0, 0.0],
82    [0.0, 1.0, 0.0, 0.0],
83    [0.0, 0.0, 1.0, 0.0],
84    [0.0, 0.0, 0.0, 1.0],
85    [0.5, 0.5, 0.5, 0.5],
86])
87
88# GQA G=2
89gqa = GroupedQueryAttention(d_model=4, H=2, G=2)
90with torch.no_grad():
91    weights, output = gqa(Q, K, V)
92print("GQA G=2 output:")
93print(output.numpy().round(4))
94
95# GQA G=1
96gqa1 = GroupedQueryAttention(d_model=4, H=2, G=1)
97with torch.no_grad():
98    weights1, output1 = gqa1(Q, K, V)
99print("\nGQA G=1 output:")
100print(output1.numpy().round(4))

Key Takeaways

1. GQA is a generalization. It places MHA ($G = H$) and MQA ($G = 1$) as endpoints of a continuous spectrum controlled by a single parameter $G$.
2. The math is one line. The entire mechanism is controlled by $g(i) = \lfloor i / (H/G) \rfloor$ — a floor division that maps each query head to its KV group.
3. KV-cache scales with G, not H. Choosing $G = H/8$ gives 87.5% memory savings with negligible quality loss in practice.
4. Uptrain, don't retrain. Existing MHA models can be converted to GQA by mean-pooling KV heads within groups and fine-tuning with ~5% of original compute.
5. It's the industry standard. LLaMA 3, Mistral, Gemma 2, StarCoder 2, and virtually all new production models use GQA.
6. Composes with everything. GQA works seamlessly with Flash Attention, RoPE, ALiBi, PagedAttention, and KV-cache quantization.

Exercises

Exercise 1: Compute for “cat”

Using the shared example with $H = 2, G = 2$, compute the full GQA output for “cat” (row 1). Then repeat with $G = 1$ and compare the two outputs. Which dimensions change and why?

Exercise 2: KV-Cache Budget

You have a 128K context model with $H = 48$ heads, $d_k = 128$, 40 layers, FP16. Your GPU has 24 GB of memory and 8 GB is available for KV-cache. What is the maximum $G$ you can use? What is the minimum?

Exercise 3: Prove the Spectrum

Prove formally that GQA with $G = H$ is mathematically identical to MHA, and that GQA with $G = 1$ is identical to MQA. Start from the group assignment function and show that the KV slicing reduces to the expected behavior in each case.

Exercise 4: Uptrain Design

You have a trained MHA model with $H = 32$ and want to convert it to GQA with $G = 4$. (a) How many heads are in each group? (b) Write pseudocode for mean-pooling the $W^K$ matrices. (c) If the original training took 1000 GPU-hours, how many GPU-hours would the uptrain take (approximately)?

Exercise 5: Flash Attention Synergy

Explain why GQA improves Flash Attention's performance beyond just reducing KV-cache size. Hint: think about memory bandwidth and the ratio of computation to memory access in the tiled attention kernel.

References

1. Ainslie, J., Lee-Thorp, J., de Jong, M., Zemlyanskiy, Y., Lebrón, F., & Sanghai, S. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. EMNLP 2023. arXiv:2305.13245
2. Shazeer, N. (2019). Fast Transformer Decoding: One Write-Head is All You Need. arXiv:1911.02150
3. Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., & Polosukhin, I. (2017). Attention Is All You Need. NeurIPS 2017. arXiv:1706.03762
4. Touvron, H., et al. (2023). LLaMA 2: Open Foundation and Fine-Tuned Chat Models. arXiv:2307.09288
5. Meta AI (2024). The LLaMA 3 Herd of Models. arXiv:2407.21783
6. Jiang, A. Q., et al. (2023). Mistral 7B. arXiv:2310.06825
7. Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arXiv:2307.08691
8. Kwon, W., et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. SOSP 2023. arXiv:2309.06180