A decoder can keep many query heads while sharing fewer key and value heads. Multi-head attention (MHA), grouped-query attention (GQA), and multi-query attention (MQA) choose how much K/V to share. This changes the learned projections and the size of the KV cache.
FlashAttention answers a separate question: how should the selected attention operation run on the GPU? It tiles and fuses computation to reduce transfers to high-bandwidth memory (HBM), without approximating softmax attention. A supported FlashAttention implementation can run MHA, GQA or MQA. Head sharing and an IO-aware kernel can therefore be combined. FlashAttention paper.
Compare the Variants
Move the slider from eight K/V heads to one. The eight query heads stay in place while the wiring and compact cache payload change. This is the same teaching configuration as the GQA explorer: model width 1,024, head width 128, 32 layers, 4,096 cached tokens, one sequence and two bytes per element.
4 query heads share each K/V head.
Every query head still scores T = 4,096 cached keys: Hq × T = 32,768 pairs. Sharing K/V does not share attention weights or remove query heads.
D = 1024 model width · d = 128 per head · Hq = 8 · Hkv = 2
B = 1 sequence · L = 32 layers · T = 4,096 tokens · b = 2 bytes (FP16/BF16)
KV payload = 2 × B × L × T × Hkv × d × b. Excludes weights, workspace, allocator overhead and replicated or expanded K/V.
Read the outcome in two parts. Fewer K/V heads mean fewer K/V projection weights and cached rows. Every query head still computes its own attention scores and output, so sharing does not reduce the number of query-key score pairs by the cache ratio. The figure calculates dimensions and bytes; it does not predict model quality or measured speed.
Detailed Comparison
Architecture Differences
Keep model width and head width distinct:
| Symbol | Meaning | Figure value |
|---|---|---|
| D | Model width, the input and output hidden dimension | 1,024 |
| d | Width of each query, key and value head | 128 |
| Hq | Number of query heads | 8 |
| Hkv | Number of key heads and number of value heads | 8, 4, 2 or 1 |
| B | Number of equal-length sequences in the batch | 1 |
| L | Number of attention layers | 32 |
| T | Cached tokens per sequence | 4,096 |
| b | Bytes per cached element | 2, for FP16 or BF16 |
The figure uses H_q d = D, but the general formulas below do not require that equality. Each input row is projected to H_q d query features and H_{kv} d key and value features:
| Feature | MHA | GQA | MQA |
|---|---|---|---|
| Query heads | Hq | Hq | Hq |
| K/V heads | Hkv = Hq | 1 < Hkv < Hq | Hkv = 1 |
| Queries per K/V head | 1 | Hq / Hkv | Hq |
| Q projection shape | D × (Hq d) | D × (Hq d) | D × (Hq d) |
| K and V shapes, each | D × (Hq d) | D × (Hkv d) | D × d |
For this evenly grouped layout, Hkv divides Hq. Counting Q, K and V weights per layer, excluding biases gives:
The output projection maps the concatenated query-head outputs back to model width. Count it separately:
MHA has P_{QKV} = 3D H_q d, which becomes 3D^2 when H_q d = D. GQA and MQA reduce the K/V portion; the Q and output projections remain the same in this comparison. These counts do not include feed-forward layers, embeddings or other model weights.
Memory Footprint
A compact dense cache stores both K and V for every cached token at every layer:
The factor 2 is for K and V, not the element size. Using the figure's B = 1, L = 32, T = 4,096, d = 128 and b = 2:
| Variant | Hkv | Payload bytes | Payload MiB | Fraction of MHA |
|---|---|---|---|---|
| MHA | 8 | 536,870,912 | 512 | 100% |
| GQA with 4 K/V heads | 4 | 268,435,456 | 256 | 50% |
| GQA with 2 K/V heads | 2 | 134,217,728 | 128 | 25% |
| MQA | 1 | 67,108,864 | 64 | 12.5% |
One MiB is 2^{20} bytes. At fixed B, L, T, d and b, the payload fraction is exactly H_{kv}/H_q. A label such as “GQA-8” specifies eight K/V heads; its ratio also depends on Hq. With 32 query heads it is 25% of MHA; with 64 it is 12.5%.
These are tensor payloads, not total GPU memory. Weights, attention workspace, allocator padding and metadata add memory. Tensor-parallel replication or materializing K/V copies for every query head can also change the stored amount. For variable sequence lengths, count their cached tokens individually; windowed caches use the tokens actually retained.
Mathematical Formulations
For a query head i, head sharing changes which K and V tensors it reads:
MHA: Separate K/V per Query
Each query head has a separate learned K/V projection. There is no sharing between query heads.
GQA: Share within Groups
With eight queries and two K/V heads, queries 0–3 read K/V head 0 and queries 4–7 read K/V head 1. The index is zero-based. Queries in a group share keys and values, but each query produces its own softmax weights and mixed value vector.
MQA: One Shared K/V Head
All query heads read the same K/V history. This is the maximum-sharing endpoint of the same axis.
For one new query token attending to T cached keys, each of Hq query heads still scores T keys:
In the figure that is 8 × 4,096 = 32,768 score pairs at every slider position. Sharing can save K/V projection work and memory traffic while leaving this count unchanged. Actual latency depends on how the kernel reuses the shared data.
Production Model Configurations
These examples describe named model releases, not recommended settings to apply to any checkpoint:
| Release | Head-sharing architecture | Query / K/V heads | Other detail |
|---|---|---|---|
| Llama 2 7B | MHA | 32 / 32 | The original model card marks GQA absent |
| Llama 2 70B | GQA | 64 / 8 | The paper describes eight K/V projections |
| Mistral 7B v0.1 | GQA | 32 / 8 | Configured sliding window of 4,096 tokens |
| Falcon 40B | GQA by this definition | 128 / 8 | Published config uses new_decoder_architecture: true |
Sources: Llama 2 model card, Llama 2 architecture discussion, Mistral 7B v0.1 config, and Falcon 40B config. The Llama query counts follow the original Llama architecture table and the architecture carried into Llama 2; the Llama 2 sources identify its head-sharing change.
Read configuration fields alongside the model implementation. A family name or an informal “multi-query” label is insufficient to determine how many K/V heads a particular release stores.
Optimization Combinations
The mechanisms operate on different parts of attention:
| Choice | What it changes | What it preserves |
|---|---|---|
| MHA, GQA or MQA | Learned Q/K/V projection layout and K/V sharing | The selected number of query heads |
| FlashAttention | Tiling, fusion and transfers between HBM and on-chip memory | The mathematical attention operation and selected mask |
| Sliding-window attention | Which past positions each query may attend to | The chosen head-sharing layout |
| RoPE | Position-dependent rotations of queries and keys | The chosen head-sharing layout |
FlashAttention's “exact” means it computes softmax attention without a sparsity or kernel approximation. Floating-point roundoff can differ with the order of operations; it does not imply bitwise-identical outputs. It also does not shrink the persistent KV cache by the Hkv/Hq ratio. That ratio comes from the architecture, and the implementation must store the shared K/V compactly to realize it. FlashAttention paper.
Quick Decision Matrix
Start from the available checkpoints and workload, then measure the quantity that matters:
| Constraint | What to check |
|---|---|
| KV memory budget | Calculate payload for the checkpoint's heads, dtype, retained tokens and concurrency; measure total allocated memory |
| Task quality | Evaluate the candidate checkpoints on the same task and data; head count alone does not predict quality |
| Decode latency | Benchmark the serving kernel at the intended batch and context lengths |
| Prefill throughput | Benchmark the attention implementation and mask on representative prompt lengths |
| Multiple GPUs | Check K/V sharding and replication across the tensor-parallel layout |
There is no universal “GQA-8 default” or “MQA is lowest latency” rule. In the Llama 2 paper's Appendix A.2.1, the authors compare trained attention variants and discuss serving across eight tensor-parallel GPUs. Replicating a single MQA K/V head across those GPUs can make its aggregate cache as large as eight-head GQA. This illustrates why logical payload ratios alone do not determine deployment speed or memory.
Best Practices
- Read the checkpoint's architecture first. Changing Hkv changes learned projection shapes. Converting a pretrained MHA checkpoint to GQA or MQA requires an architecture conversion and training or uptraining, followed by evaluation; it is not a drop-in inference setting. The GQA paper studies such conversion through uptraining.
- Separate architecture from kernel choice. Choose a supported attention kernel for that checkpoint's head layout, dtype and mask. A FlashAttention option changes execution, not the learned number of K/V heads.
- Keep comparisons controlled. Report the model release, task metric, batch size, prompt and decode lengths, dtype, GPU and serving implementation with any quality, latency or throughput claim.
Related concepts
Trace how grouped-query attention keeps independent query heads while sharing fewer key/value heads, projections, and compact KV-cache rows during LLM decoding.
Understand normalized kernel attention, low-rank sequence projections and their trade-offs, and compare them with sparse attention and exact tiled softmax.
Trace how multi-query attention keeps independent query heads while sharing one key/value head, with exact cache bytes and a shape-correct decoding example.
Explore causal and symmetric local attention with exact window counts, then distinguish direct keys, multi-layer input reach, and inference KV storage.
Compose sparse attention patterns, follow one query through masking and softmax, and separate fewer allowed pairs from actual sparse-kernel savings.
Follow a query-key distance through ALiBi head slopes, attention logits and normalized weights, then understand the evidence and limits of length extrapolation.