Home / LLM Fundamentals

Transformers and attention

Self-attention step by step, multi-head attention, the decoder-only architecture and the KV cache.

12 min read

The Transformer (Vaswani et al., 2017) replaced recurrence with attention, which lets every token look at every other token in parallel. Nearly every modern LLM is a stack of Transformer blocks.

Self-attention in four steps

For a sequence of token vectors X (shape n × d):

1. Project each token into a query, key and value: Q = XW_Q, K = XW_K, V = XW_V.
2. Score every pair: S = QKᵀ / √d_k. Dividing by √d_k keeps the dot products from growing with dimension and saturating the softmax.
3. Normalize each row with softmax so the weights for one token sum to 1.
4. Mix: Attention(Q, K, V) = softmax(S) · V. Each output is a weighted average of value vectors.

Intuition: the query is "what am I looking for", the key is "what do I contain", and the value is "what I pass along if selected".

Causal masking

Decoder-only models (GPT-style) generate left to right, so position i must not see positions after i. A causal mask sets those scores to −∞ before the softmax.

Multi-head attention

Instead of one attention with dimension d, run h heads of dimension d/h in parallel and concatenate. Different heads can specialize (syntax, coreference, nearby tokens) at the same total cost.

Variants you should know:

• Multi-Query Attention (MQA): all heads share one K and V. Much smaller KV cache, slight quality loss.
• Grouped-Query Attention (GQA): groups of heads share K/V. The common middle ground in current models.

A Transformer block

x = x + Attention(LayerNorm(x))
x = x + FeedForward(LayerNorm(x))

• Residual connections let gradients flow through deep stacks.
• Pre-norm (normalizing before each sub-layer) trains more stably than the original post-norm. RMSNorm is a popular cheaper variant.
• The feed-forward network (often SwiGLU) holds most of the parameters and is thought to store much of the factual knowledge.
• Mixture-of-Experts (MoE) models replace the single FFN with many experts and a router that picks a few per token, so total parameters grow without growing per-token compute.

Complexity

Attention is O(n²) in sequence length for both compute and memory of the score matrix. That is why long context is expensive. Mitigations include FlashAttention (an IO-aware exact kernel that avoids materializing the full matrix), sliding-window attention, and sparse or linear attention variants.

The KV cache

During generation, the keys and values of earlier tokens never change. The model caches them so each new token only computes its own Q, K, V and attends to the cache.

• Prefill: process the whole prompt in parallel. Compute-bound. Determines time-to-first-token.
• Decode: generate one token at a time. Memory-bandwidth-bound. Determines tokens per second.

KV cache size ≈ 2 × layers × kv_heads × head_dim × seq_len × bytes_per_value per sequence. This is often what limits batch size and context length in serving, and it is why GQA, paged attention (vLLM) and KV-cache quantization exist.

Encoder, decoder, or both