Grouped Query Attention vs Multi-Query Attention in PyTorch

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 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 for deployment-level analysis.

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

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.

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

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 is a fused CUDA kernel that computes exact attention with tiled IO [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 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 - Ashutosh Sharma (2026)