# Grouped Query Attention vs Multi-Query Attention in PyTorch

> Compare grouped query attention vs multi query attention in PyTorch. Learn how KV-cache memory, throughput, head settings, and model quality differ.

- Author: Swarnava Dutta (https://swarnava.dev)
- Published: 2026-08-30
- Tags: Multi-head Attention, KV Cache, Flashattention-2
- Reading time: 7 min (1540 words)
- Canonical: https://swarnava.dev/blogs/grouped-query-attention-vs-mqa

---

![Illustration of grouped query attention vs multi query attention: A wide water manifold: many small intake valves on the](/images/blogs/grouped-query-attention-vs-mqa-hero.webp)

Every key-value head you allocate during inference burns GPU memory proportional to sequence length. Consider a 70B model with 80 layers, 8 KV heads, `head_dim=128`, BF16, at 128K context: that single-user cache runs to roughly 80 × 128000 × 8 × 128 × 2 × 2 ≈ 33 GB. Grouped query attention vs multi-query attention is the architectural decision controlling that bill - how many KV heads you keep and what quality you sacrifice.

## Quick answer

Grouped query attention (GQA) shares each key-value head across a group of query heads; multi-query attention (MQA) uses a single KV head for all queries. GQA cuts KV-cache size proportionally to the group ratio while retaining near-MHA quality. MQA maximizes cache reduction but risks degradation on long-context tasks. Choose GQA for balanced throughput and quality; reserve MQA for extreme memory pressure with compatible trained weights.

## How Grouped Query Attention Works

Standard multi-head [attention](/blogs/transformer-attention-mechanism-explained) pairs every query head with its own key and value head. GQA breaks that pairing. You divide `num_attention_heads` query heads into equal groups, each sharing one KV head. Llama 3 70B uses a grouped-query attention configuration with multiple query heads sharing each KV head.

K and V projections output only `num_key_value_heads` heads instead of the full count. During the dot-product step, each KV head broadcasts across its query group. Every query head still computes independent attention scores, so heads within a group attend to different positions.

The KV cache stores one entry per KV head per layer per token. An 8:1 ratio cuts cache size by 8× versus MHA. Compared to MQA's single KV head, GQA preserves multiple distinct key-value representations, encoding more positional and semantic variety.

## Grouped Query Attention vs Multi-Query Attention and MHA

| Criterion | MHA | GQA | MQA |
|---|---|---|---|
| KV heads per layer | = `num_attention_heads` | 2-16 typical | 1 |
| Query-head sharing | None; 1:1 | Group per KV head | All queries share one KV head |
| KV-cache size (relative) | 1× | `n_kv / n_q` | `1 / n_q` |
| Memory-bandwidth pressure | Highest | Moderate | Lowest |
| Decode throughput | Baseline | Higher than MHA | Highest |
| Quality retention | Full | Near-MHA | Degrades on long-context |
| Implementation complexity | Standard | One broadcast op | Same as GQA with `n_kv=1` |

Actual throughput gains depend on GPU bandwidth, batch size, and kernel selection. On bandwidth-bound hardware at large batch sizes, MQA's advantage surfaces clearly; at batch size 1 on compute-bound GPUs, the gap narrows. See the [KV-cache optimization guide](/blogs/kv-cache-optimization-llm-inference) for deployment-level analysis.

![Comparison diagram of grouped query attention vs multi query attention and MHA showing query heads mapped to differing counts of key-value heads](/images/blogs/grouped-query-attention-vs-mqa-diagram-1.jpg "How query heads share key-value heads across MHA, GQA, and M")

## How num_key_value_heads Changes the Architecture

Set `num_key_value_heads` equal to `num_attention_heads` for MHA, to 1 for MQA, anything between for GQA. One constraint: `num_attention_heads` must be evenly divisible by `num_key_value_heads`. 32 query heads with 8 KV heads works; 32 with 5 fails.

For `hidden_size=4096`, 32 query heads, `head_dim=128`: Q projects 4096→4096 always. K and V each project to `num_key_value_heads × 128`. With 8 KV heads, that's 1024 dimensions - one quarter of MHA's K/V parameters.

Llama 3 8B and Mistral 7B v0.1 both set `num_key_value_heads=8` with 32 query heads. Always verify the actual `config.json` in your checkpoint - fine-tune releases sometimes alter these values.

Changing `num_key_value_heads` in a pretrained config without retraining produces garbage. The original GQA paper describes an uptraining procedure to convert MHA checkpoints; skipping it means misaligned weight matrices.

## Implement GQA and MQA with PyTorch Attention

```python
import torch
import torch.nn.functional as F

class FlexAttention(torch.nn.Module):
    def __init__(self, hidden: int, n_q_heads: int, n_kv_heads: int):
        super().__init__()
        assert n_q_heads % n_kv_heads == 0
        self.n_q_heads = n_q_heads
        self.n_kv_heads = n_kv_heads
        self.head_dim = hidden // n_q_heads
        self.q_proj = torch.nn.Linear(hidden, n_q_heads * self.head_dim, bias=False)
        self.k_proj = torch.nn.Linear(hidden, n_kv_heads * self.head_dim, bias=False)
        self.v_proj = torch.nn.Linear(hidden, n_kv_heads * self.head_dim, bias=False)
        self.o_proj = torch.nn.Linear(hidden, hidden, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, S, _ = x.shape
        q = self.q_proj(x).view(B, S, self.n_q_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(B, S, self.n_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, S, self.n_kv_heads, self.head_dim).transpose(1, 2)
        out = F.scaled_dot_product_attention(q, k, v, is_causal=True, enable_gqa=True)
        return self.o_proj(out.transpose(1, 2).contiguous().view(B, S, -1))
```

Pass `n_kv_heads=n_q_heads` for MHA, any divisor for GQA, or `1` for MQA.

### Using enable_gqa and Explicit KV-Head Expansion

The `enable_gqa=True` flag on `scaled_dot_product_attention` landed in PyTorch 2.5. The kernel handles mismatched Q/KV head counts internally without duplicating tensors. PyTorch 2.3 and earlier reject the flag.

For older builds, expand KV heads before the SDPA call with `k.repeat_interleave(group_size, dim=1)`. This physically materializes repeated tensors, erasing the memory savings GQA exists to deliver. The `enable_gqa` path broadcasts inside the fused kernel, preserving the advantage.

## Calculate KV-Cache Memory for MHA, GQA, and MQA

Per-layer formula: `batch × seq_len × num_kv_heads × head_dim × 2 × bytes_per_element`. Multiply by `num_layers` for the full model.

Assume a 32-layer model, 32 query heads, `head_dim=128`, BF16 (2 bytes), batch 8, sequence 4096. MHA (32 KV heads): 8 × 4096 × 32 × 128 × 2 × 2 × 32 layers = ~4.3 GB. GQA with 8 KV heads: ~1.1 GB. MQA with 1 KV head: ~134 MB.

![Diagram comparing grouped query attention vs multi query attention kv-cache memory footprint alongside standard multi-head attention across batch and sequence](/images/blogs/grouped-query-attention-vs-mqa-diagram-2.jpg "Kv-cache size shrinks as kv-head count drops")

Fewer KV heads shrink only cached K/V per token. Query activations stay at `num_attention_heads × head_dim` regardless - they exist transiently, not in the persistent cache. Halving cache size lets you double context length or concurrent batch within the same GPU memory envelope.

## Benchmark Decode Throughput, Memory, and Quality

Vary only `num_key_value_heads` while holding model width, layer count, and dtype constant. Fix your matrix: batch size (1, 8, 32), prompt length (512, 2048), generated tokens (128), BF16 on a specific GPU. Record PyTorch version, CUDA version, and which SDPA backend dispatched.

Separate prefill from decode throughput. Prefill is compute-bound; KV-head count barely moves it. Decode is bandwidth-bound, so fewer KV heads translate directly into higher tokens-per-second. Report peak memory via `torch.cuda.max_memory_allocated()` and p50/p99 latency per token.

Quality comparisons demand checkpoints trained with the target head count. Slicing KV heads from a pretrained MHA model tests nothing except broken weight alignment.

### A Reproducible PyTorch Benchmark Harness

```python
import torch

def bench_decode(n_q: int, n_kv: int, hidden: int = 4096, seq: int = 2048,
                 batch: int = 8, warmup: int = 3, trials: int = 10):
    from __main__ import FlexAttention
    model = FlexAttention(hidden, n_q, n_kv).cuda().bfloat16().eval()
    x = torch.randn(batch, seq, hidden, device="cuda", dtype=torch.bfloat16)
    with torch.inference_mode():
        for _ in range(warmup):
            model(x)
        torch.cuda.synchronize()
        torch.cuda.reset_peak_memory_stats()
        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        start.record()
        for _ in range(trials):
            model(x)
        end.record()
        torch.cuda.synchronize()
    ms = start.elapsed_time(end) / trials
    peak_mb = torch.cuda.max_memory_allocated() / 1e6
    return {"n_kv": n_kv, "ms": round(ms, 2), "peak_mb": round(peak_mb, 1)}
```

Call `bench_decode(32, 32)` for MHA, `(32, 8)` for GQA, `(32, 1)` for MQA. Watch for `repeat_interleave` inflating memory on PyTorch < 2.5 - your "GQA" run measures MHA memory in that case.

## FlashAttention-2 and Sliding-Window Attention Fit Different Problems

[FlashAttention-2](/blogs/install-flash-attention-2) is a fused CUDA kernel that computes exact attention with tiled IO [[1]](#ref-1). It changes *how* attention executes. GQA and MQA change the architecture - how heads share keys and values. GQA support arrived in flash-attn v2.3.0; earlier versions reject mismatched Q/KV head counts.

Sliding-window attention restricts each token to a fixed window of `w` preceding positions. [Mistral 7B](/blogs/mistral-ai-review-2026) v0.1 uses a window of 4096. Cache entries beyond `w` steps are evicted, capping memory regardless of total generated length.

GQA shrinks storage per token (fewer KV heads). Sliding-window shrinks the token count stored (bounded window). Combined - as in Mistral 7B - you get both reductions. The tradeoff: the model cannot attend beyond the window boundary, so full-context retrieval tasks lose information.

## Choosing Between GQA, MQA, and Multi-Head Attention

Start with your checkpoint. No budget for uptraining means MHA weights stay MHA. Training from scratch or selecting a pretrained GQA checkpoint, the decision reduces to GPU memory, target context length, batch size, and quality tolerance.

Run the KV-cache formula with your deployment parameters. If the cache fits at maximum context and batch, MHA gives full head independence. If it exceeds your GPU budget, GQA with 4-8 KV heads recovers enough memory while preserving near-MHA quality. Reserve MQA for extreme memory pressure with a checkpoint trained at one KV head.

- Inspect `config.json` to confirm `num_key_value_heads` matches your variant.
- Calculate cache demand at your serving dtype and maximum sequence length.
- Confirm `enable_gqa` (PyTorch ≥ 2.5) or flash-attn ≥ 2.3.0 for fused GQA kernels.
- Benchmark decode latency and peak memory on the deployment GPU.

## FAQ

### What is grouped query attention?

GQA is a multi-head attention variant that shares each key-value head across a group of query heads. Llama 3 70B uses 8 KV heads for 64 query heads, cutting KV-cache size by 8× versus standard MHA while retaining near-equivalent output quality. It sits between full MHA and single-head MQA on the memory-quality spectrum.

### How does grouped query attention work?

Keys and values project into fewer heads than queries. Each KV head broadcasts to its assigned query group during the dot-product step. Every query head still computes independent attention scores, so heads within a group attend to different positions. The KV cache stores entries only for the reduced head count, shrinking memory linearly with the group ratio.

### What is sliding window attention?

Sliding window attention restricts each token's attention to a fixed window of `w` preceding positions. Mistral 7B v0.1 uses `w=4096`. Cache entries older than `w` steps are evicted, capping memory regardless of total sequence length - at the cost of losing access to context beyond the window boundary.

### How does sliding window attention work?

The attention kernel masks all key-value positions older than `w` steps from the current token. During decoding, out-of-window cache entries are evicted since no future token attends to them. This bounds memory and bandwidth cost per decode step. GQA and sliding-window compose independently - GQA reduces per-token storage while the window reduces token count.


## References

1. [TileMaxSim: IO-Aware GPU MaxSim Scoring with Dimension Tiling and Fused Product Quantization](https://arxiv.org/abs/2606.26439v1) - Ashutosh Sharma (2026)