Grouped Query Attention (GQA) - Transformer Architecture

Grouped Query Attention (GQA)

What you'll build

Attention where many query heads share a smaller set of key/value heads. LLaMA 2 70B, LLaMA 3, Mistral, Gemma and most modern LLMs use it to shrink the KV cache.

Background

During generation the model stores every past token's keys and values (the KV cache, see LLM Generation & Decoding). Its size is 2 · n_layers · n_kv_heads · d_head · seq_len numbers per sequence, and reading it every step is what makes decoding memory-bound. GQA keeps all num_q_heads query heads but only num_kv_heads K/V heads, cutting the cache by num_q_heads / num_kv_heads with little quality loss.

Variant	Query heads	KV heads	Used in
MHA	H	H	Original Transformer, GPT-2
GQA	H	G (1 < G < H)	LLaMA 2 70B, LLaMA 3, Mistral
MQA	H	1	PaLM, Falcon

Algorithm

With d_head = Q.shape[1] / num_q_heads and group_size = num_q_heads / num_kv_heads:

1. Split Q into num_q_heads heads and K, V into num_kv_heads heads: shape (heads, seq_len, d_head), head h taking columns h*d_head : (h+1)*d_head.
2. Query head h uses KV head h // group_size: consecutive query heads share a KV head. Repeating each KV head group_size times along the head axis does this.
3. Per head: softmax(q_h @ k_h.T / sqrt(d_head)) @ v_h.
4. Concatenate the query heads back to (seq_len, num_q_heads * d_head).

Inputs

• Q: (seq_len, num_q_heads * d_head)
• K, V: (seq_len, num_kv_heads * d_head)
• num_q_heads, num_kv_heads: ints, num_kv_heads divides num_q_heads

Expected Return Format

output: array (seq_len, num_q_heads * d_head).

Edge cases and conventions

• num_kv_heads == num_q_heads is ordinary multi-head attention; num_kv_heads == 1 is MQA.
• Use np.repeat (each head repeated consecutively), not np.tile (which cycles through the heads and pairs query head 1 with KV head 1 instead of 0).
• No causal mask; stable softmax.

Constraints

NumPy only, vectorised (reshape/transpose, np.repeat, batched matmul).

Follow-up

• LLaMA 2 70B has 64 query heads and 8 KV heads. By how much does GQA shrink its KV cache?
• Can you avoid materialising the repeated K and V? (Reshape the query heads to (num_kv_heads, group_size, ...) and broadcast.)

Input: Q = np.array([[1.0, 0.0, 0.0, 1.0], [0.0, 1.0, 1.0, 0.0]]), K = np.eye(2), V = np.array([[1.0, 2.0], [3.0, 4.0]]), num_q_heads = 2, num_kv_heads = 1
Output: [[1.6605, 2.6605, 2.3395, 3.3395], [2.3395, 3.3395, 1.6605, 2.6605]]

Input: Q (1, 8), K (1, 4), V = np.array([[1.0, 2.0, 3.0, 4.0]]), num_q_heads = 4, num_kv_heads = 2
Output: [[1.0, 2.0, 1.0, 2.0, 3.0, 4.0, 3.0, 4.0]]

import numpy as np

def grouped_query_attention(Q, K, V, num_q_heads, num_kv_heads):
    """
    Grouped Query Attention (query head h uses KV head h // group_size).

    Args:
        Q: Queries (seq_len, num_q_heads * d_head)
        K: Keys (seq_len, num_kv_heads * d_head)
        V: Values (seq_len, num_kv_heads * d_head)
        num_q_heads: Number of query heads
        num_kv_heads: Number of key/value heads (divides num_q_heads)

    Returns:
        output: (seq_len, num_q_heads * d_head)
    """
    seq_len = Q.shape[0]
    d_head = Q.shape[1] // num_q_heads
    group_size = num_q_heads // num_kv_heads

    # Your code here
    pass