GQA and MQA: Shrinking the KV Cache

The KV cache is the dragon of long-context inference. Standard MHA stores one K and V per query head, scaling linearly with both context length and head count. Grouped-Query Attention (GQA) and its extreme cousin Multi-Query Attention (MQA) attack this directly by sharing K and V across multiple Q heads.

The shapes

Variant	Q heads	KV heads	KV cache reduction	Used by
MHA	H	H	1× (baseline)	BERT, GPT-3, GPT-2
GQA	H	G < H	H/G ×	Llama 3, Mistral, Gemma 3, Mixtral
MQA	H	1	H ×	Early PaLM, Falcon

Llama 3.3 70B has 64 query heads and 8 KV heads — a GQA group size of 8. Each group of 8 query heads shares one K and one V projection. The KV cache shrinks by 8x. Quality stays close to full MHA because GQA has more KV diversity than MQA but most of the cache savings.

Why this is now the default

Empirically, going from MHA to GQA with 8 KV heads loses very little quality but saves enormous memory at long context. MQA is even more aggressive but shows quality degradation in some tasks, so the field settled on GQA as the default. You can also uptrain an MHA checkpoint into GQA: average the KV heads within each group and continue training for ~5% of the original pretraining compute. Llama 2 70B-chat → GQA was done this way.

Code

GQA in PyTorch (Llama-style)·python

class GroupedQueryAttention(nn.Module):
    def __init__(self, d_model, n_q_heads, n_kv_heads):
        super().__init__()
        assert n_q_heads % n_kv_heads == 0
        self.n_q_heads, self.n_kv_heads = n_q_heads, n_kv_heads
        self.d_head = d_model // n_q_heads
        self.W_q = nn.Linear(d_model, n_q_heads * self.d_head, bias=False)
        self.W_k = nn.Linear(d_model, n_kv_heads * self.d_head, bias=False)
        self.W_v = nn.Linear(d_model, n_kv_heads * self.d_head, bias=False)
        self.W_o = nn.Linear(d_model, d_model, bias=False)
        self.repeat = n_q_heads // n_kv_heads
    def forward(self, x, mask=None):
        B, L, _ = x.shape
        Q = self.W_q(x).view(B, L, self.n_q_heads, self.d_head)
        K = self.W_k(x).view(B, L, self.n_kv_heads, self.d_head)
        V = self.W_v(x).view(B, L, self.n_kv_heads, self.d_head)
        # Repeat KV heads to match query heads
        K = K.repeat_interleave(self.repeat, dim=2)
        V = V.repeat_interleave(self.repeat, dim=2)
        Q, K, V = Q.transpose(1,2), K.transpose(1,2), V.transpose(1,2)
        out = F.scaled_dot_product_attention(Q, K, V, attn_mask=mask)
        return self.W_o(out.transpose(1,2).contiguous().view(B, L, -1))

GQA and MQA: Shrinking the KV Cache

The shapes

Why this is now the default

Code

External links

Exercise

Progress

Comments 0