Home / Fundamentals / Transformer Architecture

Self-attention from first principles

Scaled dot-product attention with shapes and code, why √d_k, masking, multi-head attention, FLOP accounting, FlashAttention's online softmax, and MQA/GQA/MLA.

16 min read

Attention lets every position build its representation as a weighted combination of other positions, with weights computed from the content itself. It replaced recurrence because it parallelizes across the sequence and gives any two positions a direct path to each other.

The computation

For one sequence of T token vectors X ∈ ℝ^{T×d} and one head of size d_k:

Q = X W_Q      K = X W_K      V = X W_V          (each T × d_k)
S = Q Kᵀ / √d_k                                  (T × T)
P = softmax(S + M)        row-wise; M is the mask (0 or −∞)
O = P V                                          (T × d_k)

Row i of P is a probability distribution over positions saying how much position i reads from each other position. The output is a convex combination of value vectors.

Why separate queries, keys and values? Relevance needn't be symmetric (position i may need j without j needing i), and what a position offers for matching (key) can differ from what it contributes when selected (value).

Why divide by √d_k

If the components of q and k are independent with mean 0 and variance 1, their dot product q·k = Σ q_j k_j has variance d_k. With d_k = 128, raw scores have a standard deviation of ~11. Softmax over numbers that spread out puts nearly all mass on the maximum, and its gradient, diag(p) − p pᵀ, goes to nearly zero. Dividing by √d_k restores unit variance so the softmax starts in a regime where it can learn.

Causal multi-head attention in NumPy

import numpy as np

def softmax(x, axis=-1):
    x = x - x.max(axis=axis, keepdims=True)          # numerical stability
    e = np.exp(x)
    return e / e.sum(axis=axis, keepdims=True)

def causal_mha(X, Wq, Wk, Wv, Wo, n_heads):
    T, d = X.shape
    dk = d // n_heads
    def split(Z):                                    # (T, d) -> (H, T, dk)
        return Z.reshape(T, n_heads, dk).transpose(1, 0, 2)
    Q, K, V = split(X @ Wq), split(X @ Wk), split(X @ Wv)
    S = Q @ K.transpose(0, 2, 1) / np.sqrt(dk)       # (H, T, T)
    mask = np.triu(np.ones((T, T), dtype=bool), k=1) # True above the diagonal
    S = np.where(mask, -np.inf, S)
    P = softmax(S, axis=-1)
    O = P @ V                                        # (H, T, dk)
    O = O.transpose(1, 0, 2).reshape(T, d)           # concat heads
    return O @ Wo

Points interviewers probe:

• Masking happens before softmax with −∞ (or a large negative number in fp16/bf16), so masked positions get exactly zero weight.
• Padding masks in batched inputs work the same way, masking key positions that are padding.
• Multi-head attention splits d into H heads of size d/H, runs attention independently, concatenates and projects with W_O. Same parameters and FLOPs as one wide head, but heads can attend to different things (local syntax, coreference, delimiters).
• Parameters: W_Q, W_K, W_V, W_O are each d × d, so 4d² per layer with standard multi-head attention.

Where the FLOPs go

Per token, per layer (forward pass), counting a multiply-add as 2 FLOPs:

The context-dependent attention terms (4·T·d) equal the parameter matmuls (24d²) when T ≈ 6d. For d = 4,096 that's ~24k tokens. Below that, attention's quadratic term isn't the main cost; beyond it, it dominates. Memory is a different story: naively materializing the T × T score matrix per head is what made long context expensive in practice.

FlashAttention: exact attention without the T × T matrix

GPUs have a small, fast on-chip memory (SRAM) and a large, slower HBM. Naive attention writes S and P (each T × T per head) to HBM and reads them back, so it is memory-bound. FlashAttention computes the same result in tiles that fit in SRAM and never stores the full matrix.

The trick is an online softmax. Process keys and values in blocks; for each query row keep a running maximum m, a running normalizer ℓ, and an unnormalized output o:

for each block of keys/values (K_b, V_b):
    s      = q · K_bᵀ / √d_k                 # scores for this block
    m_new  = max(m, max(s))
    ℓ      = e^(m − m_new) · ℓ + Σ e^(s − m_new)
    o      = e^(m − m_new) · o + e^(s − m_new) · V_b
    m      = m_new
output = o / ℓ

Rescaling by e^(m − m_new) corrects earlier partial sums whenever a larger score appears, so the final result equals the exact softmax. FLOPs are unchanged (slightly more in the backward pass, which recomputes attention instead of storing it), but HBM traffic drops dramatically and memory becomes linear in T. Later versions improve parallelism across the sequence and use newer hardware features.

Sharing keys and values: MQA, GQA, MLA

During generation, keys and values of past tokens are cached, and that KV cache dominates memory at long contexts. Variants reduce it:

GQA with 4–8 groups keeps most of MHA's quality with a several-fold smaller cache, which is why it's common. MLA goes further by caching a low-rank latent, trading extra compute for memory.

Other forms of attention

• Cross-attention: queries from one sequence (a decoder), keys and values from another (an encoder or image features).
• Sliding-window attention: each position attends to the last w positions; stacking layers still propagates information further.
• Linear attention and state-space models: replace the softmax kernel to get linear time; historically weaker at precise recall over long contexts, so many recent designs are hybrids with some full-attention layers.

Interview angles

• Write attention from memory with correct shapes and a causal mask.
• Give the variance argument for √d_k.
• Count parameters and FLOPs, and know when the quadratic term dominates.
• Explain FlashAttention as IO-awareness plus online softmax, and that it's exact.
• Relate MQA/GQA/MLA to KV-cache memory.