Grouped-Query Attention

Share the KV cache across attention heads .

Think of attention heads as readers taking notes on the same text. Grouped-Query Attention (GQA) lets a group share one notebook. Each reader has its own query. The group shares its keys and values.

query heads stay separate · key/value heads are shared

Read the paper

Readers can focus on different words

Imagine one reader marking the subject and another marking the action. Attention heads can also give different weights to the same tokens. Select a question to compare the examples.

Readers with different questions

Imagine four readers marking different parts of a sentence. Select a question to see which words they mark.

Thesmallredfoxquicklycrossedthefield

Marked words: fox

This reader marks the subject.

The readers represent attention heads. Real heads learn their own attention patterns. The roles shown here are simplified examples.

Each reader keeps a notebook

In standard multi-head attention, every reader keeps separate notes. Each new token adds a page to every notebook. Longer text means more stored pages.

Longer text means more notes

Each notebook adds a page for every new token. Compare one notebook per reader with one notebook per group.

Tokens saved6

one notebook per reader

8 notebooks

6 tokens × 8 notebooks = 48 saved pages

four readers share each notebook

2 notebooks

6 tokens × 2 notebooks = 12 saved pages

The notebooks represent KV heads. Their saved keys and values form the KV cache. GQA keeps fewer notebooks while retaining entries for every token.

A group can share one notebook

GQA gives each group a shared notebook. Readers still ask different questions, using the same notes to answer them. Choose a layout and follow a reader to its notebook.

Readers sharing notebooks

Choose how many notebooks to keep. Select a reader to see which notebook it shares.

Select a reader

Notebook 1

Q1–Q4

Notebook 2

Q5–Q8

Reader Q4 → Notebook 1

Each reader has its own query. Its group shares the same keys and values, represented by one notebook.

Choose how many notebooks to keep

One notebook per reader is multi-head attention (MHA). One notebook for everyone is multi-query attention (MQA). GQA uses a notebook for each group.

Choose how many notebooks to share

Compare one notebook for everyone, one per group, and one per reader. More notebooks keep more separate views of the text and use more memory.

sharing level2 notebooks

current choice

GQA · group sharing

8 readers · 2 notebooks

memory used25%

separate stored views · illustrativebalanced

Estimate the cache size

The number of readers stays the same. Sharing notebooks reduces how many copies of the notes are stored. Change the context length and notebook count to estimate the memory use.

Estimate the memory for shared notebooks

This model has 32 readers, representing query heads. Choose the context length and number of shared notebooks, representing KV heads.

Calculation: 32 layers, head dimension 128, 16-bit keys and values.

conversation length

shared notebooks

one notebook per reader

32 GB

32 readers · 32 stored notebooks

your shared setup

8.0 GB

32 readers · 8 stored notebooks

without sharing32 GB

with sharing8.0 GB

The model still has 32 readers, and the notebooks use 4× less memory. In GQA: 32 query heads share 8 key/value heads.

Original paper

GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints Ainslie et al. · 2023

Loading diagram…

Reference

← All guides

Grouped-Query Attention

Grouped-Query Attention (GQA) lets several query heads share one key/value head. This reduces the KV cache while keeping separate query projections.

Overview

How it works

A model saves keys and values from earlier tokens during generation. Each query head uses those cached vectors to compute attention weights. GQA reduces how many separate key/value heads must be stored and read.

Details

The KV cache

Without a cache, the model would recompute earlier keys and values at each generation step. Caching avoids that work, but uses memory in proportion to the sequence length, layer count, KV head count and head dimensions.

Details

MHA, GQA and MQA

Multi-head attention (MHA) has one key/value head per query head. Multi-query attention (MQA) shares one key/value head across all query heads. GQA uses several groups, each with its own shared key/value head.

Details

What is shared

Query heads in one group use the same key and value projections. Their query projections remain separate, so they can produce different attention weights over the same cached vectors. The semantic roles shown in the demo are illustrative, not fixed jobs assigned to real heads.

Details

Memory savings and quality

With equal head dimensions and precision, cache size scales with the number of KV heads. For example, reducing 32 KV heads to eight cuts that storage by four. The effect on quality depends on how the model is trained. Fewer heads are not a free change to an arbitrary checkpoint.

Sources

GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints ↗Ainslie et al. · 2023