Causal Masks and Attention Heads

In Attention from Numbers: Queries, Keys, and Values, you derived scaled dot-product attention, traced every tensor shape, and computed a complete numerical example. You learned that attention computes softmax((Q @ K^T) / sqrt(d_k)) @ V and that each token attends to every other token in the sequence. Now we add two critical pieces: causal masking to prevent future-token leakage during autoregressive generation, and attention head variants that change how Query and Key/Value heads are paired — directly impacting the KV cache memory required for serving.

---

Why Autoregressive Models Need Causal Masking

In autoregressive text generation, the model predicts one token at a time from left to right. When predicting token t, the model must only attend to tokens 1 ... t. It cannot see future tokens t+1, t+2, ....

During training, all T tokens of a training sequence are processed in parallel (teacher forcing). Without a causal mask, token 2 could simply look ahead at token 3, trivially predicting it without learning the underlying language distribution. This is information leakage — the model cheats by peeking at the answer.

Causal masking forces the attention mechanism to respect the left-to-right ordering: each token may attend only to itself and tokens to its left (earlier positions). Future positions are masked out.

---

How Causal Masking Works Mathematically

The raw score matrix S_raw = (Q @ K^T) / sqrt(d_k) has shape [T, T]. Each row i represents the attention scores from query token i to all key tokens j.

An additive upper-triangular mask M is defined as:

M_ij = 0      if j <= i   (past or current token: allowed)
M_ij = -inf   if j > i    (future token: masked)

For a 3-token prompt:

M = [[  0,  -inf,  -inf ],
     [  0,     0,  -inf ],
     [  0,     0,     0 ]]

When added to raw scores before softmax:

Masked_S = S_raw + M

Why -inf? Because exp(-inf) = 0.0 exactly!

In softmax, any element with -inf evaluates to an exponential of 0. Its numerator is 0, so its attention weight is exactly 0.0000 (0%). The future token receives precisely zero probability mass.

Let us trace this concretely.

---

Hands-On Exercise 1: Causal Masked Softmax Hand Calculation

Problem: For token at position 1 (second token) in a 3-token sequence with raw scaled scores S = [2.0, 3.0, 5.0].

Causal mask: M = [0.0, 0.0, -inf].

Masked scores: [2.0, 3.0, -inf].

Step-by-step Solution:

Step 1: Apply the mask

Masked_S = S + M = [2.0 + 0.0, 3.0 + 0.0, 5.0 + (-inf)]
            = [2.0, 3.0, -inf]

Step 2: Compute exponentials

exp(2.0)  = 7.389056...
exp(3.0)  = 20.085537...
exp(-inf) = 0.0          (exactly zero!)

Step 3: Sum the exponentials (denominator)

sum = 7.389056 + 20.085537 + 0.0 = 27.474593

Step 4: Compute softmax probabilities

p_0 = 7.389056 / 27.474593 = 0.2689...
p_1 = 20.085537 / 27.474593 = 0.7311...
p_2 = 0.0 / 27.474593 = 0.0000    (exactly zero!)

Result: p = [0.2689, 0.7311, 0.0000]

Verification: Sum = 0.2689 + 0.7311 + 0.0000 = 1.0000 ✓

The future token (position 2) receives exactly 0.0% attention. It contributes nothing to the weighted Value blend for token 1.

---

The Three Attention Head Topologies

In the previous lesson, we worked with a single head. Real transformer models use multiple heads in parallel. The topology of how Query, Key, and Value heads are arranged determines both the model’s expressivity and the KV cache memory footprint during inference.

1. Multi-Head Attention (MHA — Vaswani et al., 2017)

• Query heads: H
• Key heads: H
• Value heads: H
• Ratio: H_q = H_kv = H

Every Query head has its own dedicated Key head and Value head. Head i computes attention using its own Q_i, K_i, V_i.

Pros: Maximum expressivity. Each head can learn completely independent subspace relationships (e.g., one head tracks syntax, another tracks coreference, another tracks positional patterns).

Cons: During inference, every single Key and Value head across all layers must be kept in GPU VRAM (the KV cache). For a model with 32 heads and 32 layers, that is 32 × 32 = 1,024 separate Key/Value streams to cache.

---

2. Multi-Query Attention (MQA — Shazeer, 2019)

• Query heads: H
• Key heads: 1
• Value heads: 1
• Ratio: H_q = H, H_kv = 1

All H Query heads share the exact same Key head and Value head.

Pros: Drastically reduces KV cache size by H× (e.g., 8× or 32× smaller cache!). The Key and Value projections are computed once and broadcast to all Query heads.

Cons: Slight reduction in model capacity / perplexity on complex reasoning tasks. All query heads must agree on what is “important” in the context because they share the same keys and values.

---

3. Grouped-Query Attention (GQA — Ainslie et al., 2023)

The modern standard in Llama 2 (70B), Llama 3 (8B & 70B), Mistral 7B, Qwen, and many others.

• Query heads: H_q
• Key heads: G (number of groups)
• Value heads: G
• Ratio: H_kv = G, with H_q = G × (H_q / G)

Query heads are divided into G groups. Each group shares 1 Key head and 1 Value head.

Example: 32 Query heads grouped into 8 KV heads (H_q = 32, H_kv = 8). Each KV head serves 32 / 8 = 4 Query heads.

Pros: Delivers 4× to 8× KV cache memory reduction with quality virtually indistinguishable from full MHA. It is the sweet spot for production serving.

---

Visual Comparison: Head Topologies

MHA (H=8):       Q1 Q2 Q3 Q4 Q5 Q6 Q7 Q8
                  |  |  |  |  |  |  |  |
                 K1 K2 K3 K4 K5 K6 K7 K8
                  |  |  |  |  |  |  |  |
                 V1 V2 V3 V4 V5 V6 V7 V8

MQA (H=8):       Q1 Q2 Q3 Q4 Q5 Q6 Q7 Q8
                  \  \  \  \  \  \  \  \
                   \  \  \  \  \  \  \  \
                     K1 (shared)
                      |
                     V1 (shared)

GQA (H_q=8, G=2):  Q1 Q2 Q3 Q4 | Q5 Q6 Q7 Q8
                    \  \  \  \    \  \  \  \
                     \  \  \  \    \  \  \  \
                     K1 (group 0)  K2 (group 1)
                      |             |
                     V1 (group 0)  V2 (group 1)

In GQA, the Key and Value heads are repeated (broadcast) along the head dimension to match the Query head count during attention computation.

---

Hands-On Exercise 2: KV Cache Memory Comparison Across MHA, GQA, and MQA

Problem: Model specs:

• Sequence length S = 4096
• Layers L = 32
• Head dimension d_head = 128
• Precision = 2 bytes (FP16/BF16)
• Query heads H_q = 32

Formula: KV Cache Bytes = 2 × L × H_kv × d_head × S × bytes_per_elem (The leading 2 accounts for Keys and Values.)

Compute exact memory for:

1. MHA (H_kv = 32)

Bytes = 2 × 32 × 32 × 128 × 4096 × 2
      = 2,147,483,648 bytes
      = 2.0 GiB (2.15 GB)

2. GQA (H_kv = 8)

Bytes = 2 × 32 × 8 × 128 × 4096 × 2
      = 536,870,912 bytes
      = 0.5 GiB (536.9 MB)   ← 4× reduction!

3. MQA (H_kv = 1)

Bytes = 2 × 32 × 1 × 128 × 4096 × 2
      = 67,108,864 bytes
      = 0.0625 GiB (67.1 MB)   ← 32× reduction!

Takeaway: Moving from MHA to GQA (8 groups) saves 1.5 GiB of VRAM per request at this context length. Moving to MQA saves ~2.0 GiB. For batched inference serving many concurrent requests, this directly translates to higher throughput or lower GPU count.

---

Hands-On Exercise 3: Head Grouping and Attention Tensor Shapes

Problem: Given batch size B=1, sequence length T=4, H_q = 8, H_kv = 2, d_head = 16.

Questions:

1. How many query heads per KV head?
2. Shape of Query tensor?
3. Shape of Key tensor?
4. How is attention evaluated? (Describe the broadcasting step.)

Solution:

1. Query heads per KV head:

H_q / H_kv = 8 / 2 = 4 query heads per KV head.

2. Shape of Query tensor:

[B, H_q, T, d_head] = [1, 8, 4, 16]

3. Shape of Key tensor:

[B, H_kv, T, d_head] = [1, 2, 4, 16]

4. How attention is evaluated:

The Key tensor has shape [1, 2, 4, 16] but needs to match Query’s head dimension of 8 for the matrix multiplication Q @ K^T.

Broadcasting step: The H_kv = 2 Key heads are repeated (broadcast) 4 times along the head dimension:

K_broadcast = repeat(K, heads=4)  → shape [1, 8, 4, 16]

Now shapes align:

Q: [1, 8, 4, 16]   @   K_broadcast^T: [1, 8, 16, 4]
                                        → scores: [1, 8, 4, 4]

Each of the 4 Query heads in group 0 attends to the same K_0, V_0; each of the 4 Query heads in group 1 attends to the same K_1, V_1.

The same broadcasting applies to the Value tensor for the final A @ V multiplication.

---

Summary Reference Table

Property	MHA	GQA (Example)	MQA
Query heads (H_q)	32	32	32
KV heads (H_kv)	32	8	1
Query heads per KV head	1	4	32
KV cache ratio (vs MHA)	1× (baseline)	4× smaller	32× smaller
KV cache (L=32, S=4096, d=128, FP16)	2.0 GiB	0.5 GiB	67 MB
Quality retention	Maximum	~MHA	Slight drop on complex reasoning
Primary model examples	GPT-1/2/3, BERT	Llama 2/3, Mistral, Qwen	PaLM, early T5 variants
Key paper	Vaswani et al. (2017)	Ainslie et al. (2023)	Shazeer (2019)

---

Why Modern LLMs Standardly Adopt GQA

GQA is now the default choice for new open-weight models because it hits the Pareto frontier:

1. Memory efficiency: 4×–8× KV cache reduction enables longer contexts, larger batches, or smaller GPUs.
2. Quality preservation: Unlike MQA, GQA retains nearly all of MHA’s modeling capacity. The grouped heads still provide enough diversity for the query heads to specialize.
3. Implementation simplicity: GQA is a drop-in replacement for MHA in most codebases — only the projection matrices and cache shapes change. The attention computation logic (with broadcasting) is identical.
4. Training compatibility: GQA models can be initialized from MHA checkpoints (Ainslie et al., 2023 showed how to merge MHA heads into groups), enabling efficient conversion of existing models.

For production serving, GQA is the architecture that makes long-context, high-throughput inference economically viable on current hardware.

---

Primary Source Citations

• Vaswani et al. (2017), Attention Is All You Need, Section 3.2.3 — arXiv:1706.03762
• Shazeer (2019), Fast Transformer Decoding: One Write-Head is All You Need — arXiv:1911.02150
• Ainslie et al. (2023), GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints — arXiv:2305.13245

---

What Comes Next

• Next lesson: A Decoder Block from Input to Logits — Trace residuals, RMSNorm, SwiGLU, stacked blocks, output projection, weight tying, and architecture variations.
• Previous lesson: Attention from Numbers: Queries, Keys, and Values
• Track overview: Start Here: LLM Foundations