Attention Mechanisms

Attention lets every token build its representation from a weighted mix of the other tokens' information, with weights computed from how well its query matches their keys. This chapter derives scaled dot-product and multi-head attention, the causal mask, the quadratic cost, and the variants that make attention affordable at scale - FlashAttention, grouped-query attention and sliding windows.

Scaled Dot-Product Attention

Concept

Attention is the core operation that allows each token to "look at" every other token and decide what information to incorporate. The mechanism is elegant: the more similar a query is to a key, the more the corresponding value contributes to the output.

Formal definition:

Attention(Q, K, V) = softmax( Q·Kᵀ / √d_k ) · V

Where:

• Q (Queries): "What am I looking for?" - derived from the current token
• K (Keys): "What do I have to offer?" - derived from all tokens
• V (Values): "What do I actually provide?" - derived from all tokens
• d_k: dimension of the key/query vectors (used for scaling)

Step-by-step forward pass:

Input X: [seq_len, d_model]

1. Project to Q, K, V:
   Q = X · W_Q    [seq_len, d_k]
   K = X · W_K    [seq_len, d_k]
   V = X · W_V    [seq_len, d_v]

2. Compute attention scores (raw):
   scores = Q · Kᵀ / √d_k     [seq_len, seq_len]
   # Each entry scores[i,j] = how much token i attends to token j

3. (Optional) Apply causal mask:
   scores[i,j] = -inf  for all j > i   (future masking)

4. Normalize with softmax (per row):
   weights = softmax(scores)     [seq_len, seq_len]
   # Each row sums to 1; weights[i] is token i's attention distribution

5. Weighted sum of values:
   output = weights · V          [seq_len, d_v]
   # Each token's output = weighted average of all values

Why scale by √d_k?

If the components of q and k are independent with mean 0 and variance 1, their dot product has variance d_k (standard deviation √d_k). Without scaling, large dot products push softmax into regions of extremely small gradients (near-zero everywhere except the argmax). Dividing by √d_k keeps the dot product variance at ~1, keeping softmax in its useful gradient range.

Tricky Q: Why does large Q·K variance cause softmax to become "spiky"?
Softmax is exp(x_i) / Σ exp(x_j). With large magnitude differences, one term dominates the denominator exponentially - the distribution collapses to a one-hot. The gradient of softmax at this extreme is near-zero → training stalls.

---

Multi-Head Attention

Concept

A single attention computation attends to the sequence with one "perspective." Multi-head attention runs h parallel attention operations with different learned projections, then concatenates the results.

Why multiple heads?

Each head learns to attend based on different relationships:

• Head 1 might focus on syntactic dependencies (subject-verb agreement)
• Head 2 might track coreference (pronoun → noun)
• Head 3 might focus on local context (adjacent words)
• Head 4 might attend to the beginning of the sequence (topic sentence)

No single set of Q/K/V projections can capture all these simultaneously - different projections reveal different structure.

Formal definition:

MultiHead(Q, K, V) = Concat(head_1, ..., head_h) · W_O

where head_i = Attention(Q · W_Qi, K · W_Ki, V · W_Vi)

• Each head has its own W_Q, W_K, W_V of shape [d_model, d_k]
• d_k = d_model / h (so total compute is similar to single large attention)
• Output projection W_O: [h × d_v, d_model] recombines head outputs

Parameter count per attention layer:

W_Q, W_K, W_V each: d_model × d_model = 4096² ≈ 16.8M (for d_model=4096)
W_O:                 d_model × d_model = 4096² ≈ 16.8M
Total per layer:     4 × d_model² ≈ 67M (for 4096)

Tricky Q: Why use separate projection matrices per head rather than just splitting the embedding dimension?

Splitting the raw embedding would restrict each head to d_model/h of the input features and give it no learned choice of what to compare. A learned projection lets each head compute its own linear transformation of the full embedding - its own notion of "similar" for queries and keys. (In code this is one d_model × d_model matrix per Q, K and V whose output is reshaped into h heads, which is mathematically the same as h separate projections of the full input.)

---

Self-Attention vs Cross-Attention

Concept

Self-attention: Q, K, and V all come from the same input sequence. Used in both encoder and decoder (for attending within a sequence).

# Self-attention in a decoder block
Q = K = V = hidden_states  (same source)
output = MultiHead(Q, K, V)

Cross-attention: Q comes from one sequence (e.g., the decoder), while K and V come from another (e.g., encoder output). This is how encoder-decoder models condition the decoder on the encoded input.

# Cross-attention in encoder-decoder
Q = decoder_hidden_states
K = V = encoder_output
output = MultiHead(Q, K, V)  # decoder attends to encoder

Cross-attention is absent in decoder-only models (LLaMA, GPT) - they have no separate encoder to cross-attend to.

---

Causal (Autoregressive) Masking

Concept

Decoder-only models generate tokens left-to-right. During training, the model sees the full sequence but must not be allowed to "cheat" by attending to future tokens when predicting position i.

The causal mask: A lower-triangular boolean matrix applied to attention scores before softmax:

For seq_len=4:
Mask:
  [1, 0, 0, 0]    token 0 can only attend to itself
  [1, 1, 0, 0]    token 1 can attend to 0 and 1
  [1, 1, 1, 0]    token 2 can attend to 0, 1, 2
  [1, 1, 1, 1]    token 3 can attend to all

Masked positions get score = -inf → softmax produces probability ≈ 0 for those positions.

Causal masking enables parallel training: Even though generation is sequential, training can process the entire sequence in parallel by applying the mask. The model simultaneously predicts token 1 from token 0, token 2 from tokens 0–1, token 3 from tokens 0–2, etc. - all in a single forward pass.

---

Attention Complexity: The Quadratic Bottleneck

Concept

Standard self-attention is O(n²) in both memory and compute for sequence length n:

• The attention score matrix is [n × n] - n² entries
• Computing Q·Kᵀ requires n × n × d_k multiply-adds

Practical implications:

Sequence length	Attention matrix entries (n²)	Memory per head (FP16)	Per layer with 32 heads
2K	~4.2M	~8 MB	~256 MB
32K	~1.1B	~2 GB	~64 GB
128K	~17B	~32 GB	~1 TB
1M	~1.1T	~2 TB	~64 TB - impossible to materialize

(Per sequence; multiply by batch size. FlashAttention never materializes this matrix, which is why long context is feasible at all.)

This is why long-context models need specialized attention implementations.

---

Efficient Attention Variants

Flash Attention

Concept

Flash Attention (Dao et al., 2022) achieves the same mathematical result as standard attention but avoids materializing the full n×n attention matrix in GPU HBM (High Bandwidth Memory). Instead, it tiles the computation to fit in fast on-chip SRAM.

Why this matters - GPU memory hierarchy:

GPU SRAM (on-chip): ~192 KB per SM, ~20 TB/s bandwidth
GPU HBM (off-chip):  80 GB on A100, ~2 TB/s bandwidth

Standard attention writes the full [n,n] attention matrix to HBM and reads it back → bottlenecked by the 10× slower HBM. Flash Attention fuses the attention operations into a single kernel that keeps everything in SRAM:

Standard: QKᵀ → write to HBM → read from HBM → softmax → write → read → ×V
Flash:    Tile(Q,K,V) → compute entire attention within SRAM → write only output to HBM

Results:

• 2–4× faster wall-clock time on long sequences
• Memory usage: O(n) instead of O(n²) - enables much longer contexts
• Numerically identical to standard attention (not an approximation)

Flash Attention 2 (2023): Improved parallelism across GPU thread blocks, better work partitioning. ~2× faster than Flash Attention 1.

Flash Attention 3 (2024): Exploits H100 tensor core pipelining, asynchronous execution. Another 1.5–2× speedup on H100.

---

Multi-Query Attention (MQA) and Grouped-Query Attention (GQA)

Concept

Both are modifications of multi-head attention that share key and value projections across heads, reducing KV cache memory (critical for inference - see KV Cache).

Multi-Head Attention (MHA): Each head has its own K and V projections.

h heads × (d_k + d_v) × d_model = full KV memory

Multi-Query Attention (MQA): All heads share a single K and V.

1 shared K + 1 shared V - KV memory reduced by h×
Quality trade-off: slightly worse, especially for complex tasks
Used by: PaLM, Falcon, Gemma 1 2B

Grouped-Query Attention (GQA): G groups of heads, each group shares one K/V.

G groups, h/G heads per group
KV memory reduced by h/G×
Best of both worlds: quality close to MHA, memory close to MQA
Used by: LLaMA-3 (h=32 heads, G=8 groups), Gemma, Mistral

MHA:  [Q1 K1 V1] [Q2 K2 V2] [Q3 K3 V3] [Q4 K4 V4]  (4 KV pairs)
GQA:  [Q1 Q2 K1 V1]  [Q3 Q4 K2 V2]                  (2 KV pairs, G=2)
MQA:  [Q1 Q2 Q3 Q4  K  V ]                           (1 KV pair)

Variant	KV cache size	Quality	Used by
MHA	Full (h KV pairs)	Best	BERT, early GPT
GQA	h/G KV pairs	Near-MHA	LLaMA-3, Gemma, Mistral
MQA	1 KV pair	Slightly lower	PaLM, Falcon, Gemma 1 2B

---

Sliding Window Attention (Mistral)

Concept

Instead of attending to all past tokens, each token attends only to the W nearest tokens (window size W). Long-range information propagates through multiple layers.

Standard: token 100 attends to all 100 previous tokens
Sliding:  token 100 attends only to tokens 96–100 (W=5)
After L layers: receptive field = L × W

Mistral 7B v0.1 used W=4096 with an 8K training context; across 32 layers the theoretical receptive field is about 131K tokens. Per-layer attention compute falls from O(n²) to O(n×W), and a rolling KV cache of W entries caps memory. Gemma 2 and 3 interleave sliding-window layers with global-attention layers.

Trade-off: Cannot directly attend to distant context within a single layer. Works well for most generation tasks; less effective for tasks requiring precise recall of distant facts.

---

Code

import numpy as np
import torch
import torch.nn.functional as F

# Manual scaled dot-product attention (NumPy for clarity)
def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    Q: [batch, heads, seq, d_k]
    K: [batch, heads, seq, d_k]
    V: [batch, heads, seq, d_v]
    """
    d_k = Q.shape[-1]
    scores = np.matmul(Q, np.swapaxes(K, -2, -1)) / np.sqrt(d_k)

    if mask is not None:                    # mask is True where attention is NOT allowed
        scores = np.where(mask, -1e9, scores)

    weights = np.exp(scores - scores.max(axis=-1, keepdims=True))
    weights /= weights.sum(axis=-1, keepdims=True)  # softmax

    output = np.matmul(weights, V)
    return output, weights

# Demo
seq_len, d_k, d_v = 4, 8, 8
Q = np.random.randn(1, 1, seq_len, d_k)
K = np.random.randn(1, 1, seq_len, d_k)
V = np.random.randn(1, 1, seq_len, d_v)

# Causal mask: True above the diagonal = future positions, which are blocked
causal_mask = np.triu(np.ones((seq_len, seq_len)), k=1).astype(bool)

output, weights = scaled_dot_product_attention(Q, K, V, mask=causal_mask)
print("Attention weights (causal):")
print(weights.squeeze().round(3))
# Should show upper triangle ≈ 0 (masked future positions)

# PyTorch Flash Attention (requires PyTorch >= 2.0)
# torch.nn.functional.scaled_dot_product_attention uses Flash Attention when available
Q_t = torch.randn(2, 8, 512, 64)  # batch, heads, seq, d_k
K_t = torch.randn(2, 8, 512, 64)
V_t = torch.randn(2, 8, 512, 64)

# is_causal=True automatically applies causal masking with Flash Attention backend
# torch.nn.attention.sdpa_kernel replaces the deprecated torch.backends.cuda.sdp_kernel
from torch.nn.attention import sdpa_kernel, SDPBackend
with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
    output_flash = F.scaled_dot_product_attention(Q_t, K_t, V_t, is_causal=True)

print(f"\nFlash Attention output shape: {output_flash.shape}")

---

Study Notes

Must-know for interviews:

• Attention(Q,K,V) = softmax(QKᵀ/√d_k) · V - know this formula by heart
• Scale by √d_k to prevent softmax from collapsing to one-hot with large dot products
• Multi-head attention: different projection matrices per head give genuinely different "views," not just embedding slices
• Causal masking sets future positions to -inf before softmax → enables parallel training of autoregressive models
• Standard attention is O(n²) memory and compute - FlashAttention computes the same result in O(n) memory via SRAM tiling (compute stays O(n²))
• GQA (Grouped-Query) shares K/V across head groups - standard in LLaMA-3, Gemma, Mistral for KV cache efficiency
• Mixture-of-experts and multi-head latent attention (MLA), which compresses the KV cache further, are covered in Modern Architectures

Check Yourself

Exercises

References

• Vaswani et al., Attention Is All You Need (2017)
• Dao et al., FlashAttention (2022); Dao, FlashAttention-2 (2023); Shah et al., FlashAttention-3 (2024)
• Shazeer, Fast Transformer Decoding: One Write-Head is All You Need (MQA) (2019)
• Ainslie et al., GQA: Training Generalized Multi-Query Transformer Models (2023)
• Jiang et al., Mistral 7B (2023)

Last reviewed: 2026-09