Tutorial 07: Multi-Head and Grouped-Query Attention (Full Layer)

The previous tutorial covered the attention math. This tutorial shows how it fits into a complete attention layer — with input/output projections, RoPE, QK-norm, and the KV cache interface. This is the actual code that runs in the transformer.

---

1. The Full Attention Layer Architecture

A complete attention layer does this:

Input x: (batch, seq_len, hidden_size)
    │
    ├──→ Q = W_q @ x    → (batch, seq_len, num_heads × head_dim)
    ├──→ K = W_k @ x    → (batch, seq_len, num_kv_heads × head_dim)
    ├──→ V = W_v @ x    → (batch, seq_len, num_kv_heads × head_dim)
    │
    ├──→ Reshape to multi-head: (batch, num_heads, seq_len, head_dim)
    ├──→ QK-norm (Qwen3 specific)
    ├──→ Apply RoPE
    │
    ├──→ Attention(Q, K, V, mask)
    │
    ├──→ Reshape back: (batch, seq_len, num_heads × head_dim)
    └──→ Output = W_o @ attention_output

Each projection (W_q, W_k, W_v, W_o) is a learned weight matrix.

---

2. Projection Dimensions

For a model with hidden_size=4096, num_heads=32, and head_dim=128:

W_q shape: (num_heads × head_dim, hidden_size) = (4096, 4096)
W_k shape: (num_kv_heads × head_dim, hidden_size) = (1024, 4096)  ← GQA: fewer KV heads
W_v shape: (num_kv_heads × head_dim, hidden_size) = (1024, 4096)
W_o shape: (hidden_size, num_heads × head_dim) = (4096, 4096)

Note that W_k and W_v are smaller when using GQA. If num_kv_heads=8 and head_dim=128, then W_k has shape (1024, 4096) instead of (4096, 4096).

---

3. QK-Norm (Qwen3 Specific)

Qwen3 applies RMSNorm per head to queries and keys before RoPE. This stabilizes training by preventing query/key vectors from growing too large:

# After projection and reshape, but before RoPE:
projection_q = self.q_norm(projection_q)
projection_k = self.k_norm(projection_q)

Not all models do this. Llama 3, Mistral, and Qwen2.5 skip QK-norm. It is a Qwen3 design choice.

Each head's query and key is independently normalized. The norm has dim=head_dim (not hidden_size) — it operates within each head, not across heads.

---

4. The Complete Attention Layer

class Attention:
    def __init__(self, layer_prefix, config, weights, features):
        self.hidden_size = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.num_kv_heads = config.num_key_value_heads
        self.head_dim = config.head_dim

        # The scale factor for attention scores: 1/sqrt(head_dim)
        # Precomputed because it is constant for all inputs.
        self.scale = self.head_dim ** -0.5

        # Load the four projection weight matrices
        self.wq = weights.linear(f"{layer_prefix}.self_attn.q_proj")
        self.wk = weights.linear(f"{layer_prefix}.self_attn.k_proj")
        self.wv = weights.linear(f"{layer_prefix}.self_attn.v_proj")
        self.wo = weights.linear(f"{layer_prefix}.self_attn.o_proj")

        # Set up RoPE — same for Q and K
        self.rope = RoPE(
            self.head_dim,
            config.max_position_embeddings,
            config.rope_theta,
        )

        # Per-head QK-norm (Qwen3-specific)
        self.q_norm = RMSNorm(
            self.head_dim,
            weights.tensor(f"{layer_prefix}.self_attn.q_norm.weight"),
            eps=config.rms_norm_eps,
        )
        self.k_norm = RMSNorm(
            self.head_dim,
            weights.tensor(f"{layer_prefix}.self_attn.k_norm.weight"),
            eps=config.rms_norm_eps,
        )

    def __call__(self, x, offsets, cache, mask=None):
        """
        Args:
            x: Input tensor. Shape: (batch, seq_len, hidden_size)
            offsets: Position offsets for RoPE.
            cache: KV cache object (handles storing and retrieving past K/V).
            mask: Attention mask — "causal" string, an mx.array, or None.

        Returns:
            Output tensor. Shape: (batch, seq_len, hidden_size)
        """
        B, L, _ = x.shape

        # --- Step 1: Project input to Q, K, V ---
        # Each projection is a linear layer: x @ W^T
        projection_q = linear(x, self.wq).reshape(B, L, self.num_heads, self.head_dim)
        projection_k = linear(x, self.wk).reshape(B, L, self.num_kv_heads, self.head_dim)
        projection_v = linear(x, self.wv).reshape(B, L, self.num_kv_heads, self.head_dim)

        # --- Step 2: Apply per-head QK-norm ---
        # Normalize Q and K within each head before RoPE
        projection_q = self.q_norm(projection_q)
        projection_k = self.k_norm(projection_k)

        # --- Step 3: Apply RoPE to Q and K (not V!) ---
        projection_q = self.rope(projection_q, offset=offsets)
        projection_k = self.rope(projection_k, offset=offsets)

        # --- Step 4: Transpose to multi-head format ---
        # (batch, seq, heads, head_dim) → (batch, heads, seq, head_dim)
        projection_q = projection_q.transpose(0, 2, 1, 3)
        projection_k = projection_k.transpose(0, 2, 1, 3)
        projection_v = projection_v.transpose(0, 2, 1, 3)

        # --- Step 5: Update KV cache and get full context ---
        # The cache stores past K/V and returns the full history
        projection_k, projection_v, _, mask = cache.update_and_fetch(
            projection_k, projection_v, mask_length=L, mask=mask,
        )

        # --- Step 6: Compute attention ---
        x = scaled_dot_product_attention_grouped(
            projection_q.astype(mx.float32),
            projection_k.astype(mx.float32),
            projection_v.astype(mx.float32),
            scale=self.scale,
            mask=mask,
        ).astype(x.dtype)

        # --- Step 7: Reshape back and project to output ---
        # (batch, heads, seq, head_dim) → (batch, seq, heads × head_dim)
        x = x.transpose(0, 2, 1, 3).reshape(B, L, self.num_heads * self.head_dim)

        # Output projection: combines information from all heads
        return linear(x, self.wo)

---

5. Why the Output Projection Exists

After attention, we have num_heads separate outputs, each of dimension head_dim. The output projection W_o has two jobs:

1. Mix information across heads: Each head learned to attend to different things. W_o combines these signals.
2. Project back to the model dimension: The total attention output has size num_heads × head_dim, which equals hidden_size by construction. But W_o allows the model to learn a different combination.

Without W_o, each head's output would be treated independently. With W_o, the model can learn to route information between heads.

---

6. Reshape Flow Through Attention

Here is a concrete example with Qwen3-8B dimensions:

Input: (batch=1, seq=10, hidden=4096)

After W_q projection:
  → (1, 10, 4096)  # still flat

After reshape to heads:
  → (1, 10, 32, 128)  # 32 heads, each 128-dim

After transpose:
  → (1, 32, 10, 128)  # heads in batch dim for parallel compute

After attention (with K/V of shape (1, 8, S, 128)):
  → (1, 32, 10, 128)  # same shape, but values are weighted sums

After transpose back:
  → (1, 10, 32, 128)

After reshape:
  → (1, 10, 4096)  # flat again

After W_o projection:
  → (1, 10, 4096)  # output

---

7. Summary

• The attention layer consists of four projections (Q, K, V, O), RoPE, and the attention computation.
• GQA makes K and V projections smaller than Q when num_kv_heads < num_heads.
• QK-norm (Qwen3-specific) normalizes Q and K per-head before RoPE.
• The KV cache stores past K and V so we do not recompute them for every token.
• The output projection mixes information from all heads back into the model dimension.

Next: Tutorial 08 — Feed-Forward Networks and SiLU Gating