Transformers & attention

Every current language model is a transformer, and the part of a transformer that does the work is attention. This section explains it in plain language with no equations, because being able to say what attention does in five minutes is a genuinely separate skill from being able to implement it — the first is asked in interviews constantly and the second is asked rarely. The implementation, with the shapes and the arithmetic, is the from-scratch section later in this topic.

The problem#

To understand a word you often need another word far away. In the sentence about the trophy not fitting in the suitcase because it was too big, working out what "it" refers to means looking back at two candidate nouns and deciding which one the word "big" makes sensible. Change "big" to "small" and the answer flips to the other noun. Nothing local to the word "it" contains that information.

Older sequence models processed one step at a time and carried a summary forward, so distant information degraded. They also could not be parallelised across the sequence, because step ten could not begin until step nine finished. The sequence-models topic measures exactly how badly the first problem bites: in a plain recurrent network the gradient reaching fifty steps back is around 5e-19, which is not small but absent.

Every position looks at every position#

Attention lets each position look directly at all the others in the same sequence — self-attention — and decide how much each one matters for it. Distance costs nothing: the fiftieth word and the first are one step apart, exactly as the fiftieth and the forty-ninth are. And because every position can do this at the same time, the whole sequence is processed in parallel, which is what made training at scale possible.

The output for each position is a blend: a weighted average of information drawn from every position, where the weights are how relevant each one was judged to be. Those weights are not fixed. They are computed fresh for every position on every input, which is why the same word contributes differently in different sentences, and why a model can resolve "it" one way in one sentence and the other way when a single adjective changes.

Queries, keys and values#

A useful analogy is a lookup. Each position produces a query describing what it is looking for, a key describing what it offers, and a value which is the information it contributes if selected. Every query is compared against every key, the comparisons become the weights, and the weights blend the values.

query   "I am a pronoun, what noun do I refer to?"
key     "I am a singular concrete noun"
value   the information contributed if selected

match query against every key   ->  a weight per position
blend the values by those weights  ->  this position's output

That last point is the one worth holding on to. If a position had only one representation, what it offers to others and what it seeks from them would be forced to be the same thing, and they are not. The pronoun is seeking a noun and offering the fact that it is a pronoun. Separating query from key lets those be different, and separating value from both lets what gets passed on differ again from what made it match.

Multiple heads#

One set of weights captures one kind of relationship. Several heads run in parallel with different learned parameters, so one can track grammatical agreement while another tracks which noun a pronoun refers to and another tracks position in a list. Their outputs are combined, so the layer can express several relationships at once rather than being forced to choose.

It is tempting to think each head learns one clean, nameable job, and the honest position is that some do and most do not. Interpretability research has found heads with crisp roles — copying a token from earlier, tracking the previous token — and also found that heads share work, change role with context, and rarely divide up as neatly as the explanation suggests. The interpretability topic measures this directly: ablating heads one at a time leaves accuracy at 1.000 for six of the eight, and above 0.99 for the other two, while the logit margin moves by between 1.02 and 5.18, so the ones that look unused are doing something.

The cost#

Every position compares against every position, so doubling the length quadruples the comparisons. This is the single most important consequence of the design, and it is easier to believe as a table than as a sentence.

  tokens     comparisons    vs 1k    KV cache, 8B model
   1,000       1,000,000       1x               0.12 GB
   2,000       4,000,000       4x               0.24 GB
   8,000      64,000,000      64x               0.98 GB
  32,000   1,024,000,000    1024x               3.91 GB
 128,000  16,384,000,000   16384x              15.62 GB
1,000,000 1,000,000,000,000 1000000x           122.07 GB

Two things follow that come up constantly in practice. Long context is expensive in a way that is not linear, so a prompt twice as long is not twice the cost; and the memory that grows is the cache of keys and values for everything already read, which is why generation slows as a conversation lengthens even though each new token is one token. Most of the work in modern architectures — the variants in the next section — is about making one or both of those grow more slowly.

How deep to go#

For applied work this level is sufficient: what attention does, why it replaced recurrence, what the three roles are for, and what it costs. If you can say those in five minutes without writing anything down, you can hold your end of any conversation an AI engineering interview will have about transformers.

Go further in two cases. If you are interviewing for research or for a role that trains models, implement causal multi-head attention from scratch — the from-scratch section in this topic does it in NumPy, and the exercise is an afternoon rather than a week. And if you are debugging inference cost or latency, read the modern-variants section instead, because what you need is not the mechanism but what recent architectures changed about it and why.

You know what attention does. This section fills in everything around it: how text becomes vectors, what one block contains, what stops a model reading ahead when it is trained to predict the next word, and how the final vector becomes a token. Having the whole path in your head is worth the hours because most of the practical questions, such as why generation is slow, why streaming exists, or why long context costs what it does, are answered at some specific point along it.

token ids        [9906, 1917, 0]
embedding        one learned vector per id, width d_model
position         rotate or add, because attention cannot see order

repeat N times:
  x = x + attention(norm(x))        mix information across positions
  x = x + feed-forward(norm(x))     transform each position on its own

final norm
unembedding      d_model -> one score per vocabulary entry
sample           scores become probabilities, one is chosen

Embeddings#

Tokenization turns text into integers, and each integer is a row index into a learned table. Look up the row and you have a vector of width d_model, which for a model of this size is 4,096 numbers. Nothing about that lookup is clever: the table is a plain matrix of vocabulary size by d_model, trained like everything else, and the row for a token is the only thing the model knows about it before any layer has run.

Two consequences follow and both come up in interviews. The embedding table is large because vocabulary is large, so on a small model it is a noticeable share of the parameters and on a large one it is not. And because every token becomes the same width, the rest of the network never has to know how many tokens there are; d_model is fixed and the sequence length is whatever it is.

Why position has to be supplied at all#

This is the part that sounds like a detail and is not. Attention compares every position with every other and blends the results, and none of those operations knows where anything sits. Shuffle the input tokens and you get the same outputs back in a different order. That is not an approximation or a weakness of the explanation; it is exactly true, and you can watch it happen in twelve lines.

import numpy as np

rng = np.random.default_rng(0)
d = 8
Wq, Wk, Wv = (rng.standard_normal((d, d)) / np.sqrt(d) for _ in range(3))

def attention(X):
    Q, K, V = X @ Wq, X @ Wk, X @ Wv
    s = Q @ K.T / np.sqrt(d)
    p = np.exp(s - s.max(axis=1, keepdims=True))
    return (p / p.sum(axis=1, keepdims=True)) @ V

X = rng.standard_normal((4, d))        # four token embeddings, no position anywhere
order = [3, 0, 2, 1]                   # shuffle the tokens

same = np.abs(attention(X)[order] - attention(X[order])).max()
print(f"no position:  shuffling the input just shuffles the output, gap {same:.1e}")

P = rng.standard_normal((4, d)) * 0.5  # one learned vector per slot in the sequence
moved = np.abs(attention(X + P)[order] - attention(X[order] + P)).max()
print(f"with position: the same shuffle changes every value, gap {moved:.3f}")

# no position:  shuffling the input just shuffles the output, gap 4.4e-16
# with position: the same shuffle changes every value, gap 1.378

So position is injected from outside. The original models added a learned vector per slot, which is what the second half of the code above does. Current models rotate the query and key vectors by an angle proportional to position instead, which makes the comparison between two positions depend on the gap between them rather than on where each one sits, and removes the fixed-size table that capped how far a model could read. That is rotary position embedding, and the modern-variants section in this topic covers it and the tricks used to stretch it.

Residual connections and normalisation#

Each block adds its output to its input rather than replacing it. Written out, x becomes x plus attention of x, then x plus feed-forward of x. The running vector that every block reads from and writes back into is usually called the residual stream, and it is the closest thing the architecture has to a working memory: a layer that has nothing to contribute at a position can contribute nearly nothing and the information already there passes through untouched.

def cos(a, b):
    return float(a @ b / (np.linalg.norm(a) * np.linalg.norm(b)))

def rms_norm(v):
    return v / np.sqrt(np.mean(v * v) + 1e-6)

D = 256
g = np.random.default_rng(1)
A = [g.standard_normal((D, D)) / np.sqrt(D) for _ in range(32)]   # 32 blocks
h0 = g.standard_normal(D)
plain = res = h0.copy()

print("layer    no residual          with residual")
for layer, Ai in enumerate(A, start=1):
    plain = np.tanh(plain @ Ai)
    res = res + np.tanh(rms_norm(res) @ Ai)
    if layer in (1, 2, 4, 8, 16, 32):
        print(f"{layer:5d}    cos {cos(plain, h0):+.3f}  |h| {np.linalg.norm(plain):6.3f}"
              f"    cos {cos(res, h0):+.3f}  |h| {np.linalg.norm(res):6.2f}")

# layer    no residual          with residual
#     1    cos -0.026  |h| 10.339    cos +0.843  |h|  19.04
#     2    cos -0.033  |h|  7.728    cos +0.760  |h|  22.35
#     4    cos -0.100  |h|  5.283    cos +0.653  |h|  25.06
#     8    cos +0.003  |h|  3.820    cos +0.479  |h|  33.14
#    16    cos +0.017  |h|  2.723    cos +0.409  |h|  39.62
#    32    cos -0.027  |h|  2.209    cos +0.256  |h|  59.01

That is the case for residuals in one table, and it is the same argument in the forward direction as the usual one about gradients flowing backwards. The cost is visible in the right-hand column: the norm of the residual stream grows as every layer adds to it, reaching 59 after thirty-two layers from a start of about 16. Normalisation is what keeps that under control. Each sub-layer normalises its input before using it, so a block sees a vector of predictable scale whatever the stream has accumulated, and a final normalisation runs before the unembedding for the same reason.

Layer normalisation subtracts the mean, divides by the standard deviation, then applies a learned scale and shift. RMSNorm, which almost everything current uses, drops the mean and the shift: divide by the root mean square and apply a scale. Two learned vectors of width d_model per block, so a few thousandths of one per cent of the parameters, which is worth knowing because it is a common guess that normalisation is expensive. It is not; it is a rounding error in the parameter count and a real cost only in memory traffic.

The feed-forward block, where the parameters actually are#

After attention has moved information between positions, the feed-forward block transforms each position independently: project up to a wider hidden dimension, apply a non-linearity, project back down. Older models used two matrices with a GELU between them. Current ones use a gated variant, SwiGLU, which does the up-projection twice and multiplies one copy through a swish by the other, so there are three matrices instead of two and the hidden width is cut to compensate.

People consistently expect attention to dominate the parameter count, because attention is the part everyone talks about. Work it out for a real configuration and it is not close.

# Shapes taken from a published 8B open-weights config.
V, d_model, layers = 128256, 4096, 32
heads, kv_heads, head_dim, ffn = 32, 8, 128, 14336

parts = [
    ("token embedding", V * d_model),
    ("attention Q", layers * d_model * heads * head_dim),
    ("attention K and V", layers * 2 * d_model * kv_heads * head_dim),
    ("attention output", layers * heads * head_dim * d_model),
    ("feed-forward (3 matrices)", layers * 3 * d_model * ffn),
    ("RMSNorm scales", layers * 2 * d_model + d_model),
    ("unembedding", V * d_model),
]
total = sum(n for _, n in parts)
print(f"{'component':26}{'parameters':>14}{'share':>8}")
for name, n in parts:
    print(f"{name:26}{n:14,}{n/total:8.1%}")
print(f"{'TOTAL':26}{total:14,}{1.0:8.1%}")

attn = sum(n for name, n in parts if name.startswith("attention"))
ffnp = dict(parts)["feed-forward (3 matrices)"]
print(f"\nfeed-forward is {ffnp/attn:.2f}x the whole attention mechanism")
print(f"per layer: attention {attn//layers:,}  feed-forward {ffnp//layers:,}")
for bits in (16, 8, 4):
    print(f"weights at {bits:2d}-bit: {total * bits / 8 / 1024**3:6.2f} GiB")

# component                     parameters   share
# token embedding              525,336,576    6.5%
# attention Q                  536,870,912    6.7%
# attention K and V            268,435,456    3.3%
# attention output             536,870,912    6.7%
# feed-forward (3 matrices)  5,637,144,576   70.2%
# RMSNorm scales                   266,240    0.0%
# unembedding                  525,336,576    6.5%
# TOTAL                      8,030,261,248  100.0%
#
# feed-forward is 4.20x the whole attention mechanism
# per layer: attention 41,943,040  feed-forward 176,160,768
# weights at 16-bit:  14.96 GiB
# weights at  8-bit:   7.48 GiB
# weights at  4-bit:   3.74 GiB

Seventy per cent in the feed-forward layers, seventeen across all of attention, thirteen in the two vocabulary matrices, and effectively nothing in the norms. The gap is wider still per layer: 176 million parameters of feed-forward against 42 million of attention, a ratio of 4.2 to 1. Attention gets the attention because it is the interesting mechanism and because its cost grows with the square of sequence length, but in sheer weight the model is mostly a stack of position-wise transformations, which is the usual reading of where a model keeps what it knows.

The same table answers a practical question. Weights at 16-bit are 14.96 GiB, which does not leave much of a 24 GiB consumer card for anything else, and at 4-bit they are 3.74 GiB, which does. The inference-cost section in the quantization topic takes that further, including what the cache on top of the weights costs.

Causal masking#

Language models predict the next token, so each position may only attend to earlier positions. A mask blocks attention to the future. This is what lets the whole sequence train in parallel while still learning left-to-right prediction. Drop the mask and you have an encoder-only model, such as BERT, where every position sees the whole input at once: useful for classification and embedding, and unable to generate text left to right.

The mask is also why the same architecture serves two different jobs. With it you have a decoder-only model that generates. Without it you have an encoder, where every position is conditioned on the whole input at once, which is what you want for classification and for producing a single embedding of a document and useless for writing the next word. Nothing else about the block changes.

Output: from the last vector to a token#

The final vector at the last position is projected to a score for every token in the vocabulary. Sampling from those scores produces the next token, which is appended to the input, and the process repeats. That autoregressive loop is generation.

The projection is the unembedding, a matrix of d_model by vocabulary size, and it is the largest single matrix in the model at 525 million parameters in the table above. Some models tie it to the embedding table, reusing the same weights transposed, which saves those parameters and costs a little quality; the 8B configuration above does not, which is why the two rows are counted separately. Its output is one raw number per vocabulary entry, called logits, and logits are not probabilities. Turning them into probabilities is a softmax, and everything a sampling parameter does happens between those two steps.

vocab = ["the", "a", "this", "my", "an", "our", "some", "no"]
logits = np.array([8.2, 6.9, 5.4, 4.8, 4.1, 3.7, 2.2, 1.0])

def softmax(z, temperature=1.0):
    z = z / temperature
    e = np.exp(z - z.max())
    return e / e.sum()

print(f"{'token':8}" + "".join(f"{f'T={t}':>9}" for t in (0.5, 1.0, 1.5)))
cols = [softmax(logits, t) for t in (0.5, 1.0, 1.5)]
for i, w in enumerate(vocab):
    print(f"{w:8}" + "".join(f"{c[i]:9.4f}" for c in cols))

p = softmax(logits)
order = np.argsort(-p)
cum = np.cumsum(p[order])
keep = order[: int(np.searchsorted(cum, 0.9)) + 1]
print(f"\ntop-p 0.9 at T=1 keeps {len(keep)} of {len(vocab)} tokens: "
      f"{[vocab[i] for i in keep]}")
print(f"the other {len(vocab) - len(keep)} share {1 - cum[len(keep)-1]:.4f} of the mass "
      f"and are dropped before sampling")
print(f"greedy would always pick {vocab[order[0]]!r} at p={p[order[0]]:.4f}")

# token       T=0.5    T=1.0    T=1.5
# the        0.9264   0.7155   0.5495
# a          0.0688   0.1950   0.2310
# this       0.0034   0.0435   0.0850
# my         0.0010   0.0239   0.0570
# an         0.0003   0.0119   0.0357
# our        0.0001   0.0079   0.0274
# some       0.0000   0.0018   0.0101
# no         0.0000   0.0005   0.0045
#
# top-p 0.9 at T=1 keeps 2 of 8 tokens: ['the', 'a']
# the other 6 share 0.0895 of the mass and are dropped before sampling
# greedy would always pick 'the' at p=0.7155

Three things are worth taking from that. Temperature is a division applied before the softmax, not a scaling of probabilities afterwards, which is why its effect is not linear and why the difference between 0.7 and 1.0 is larger than it looks. Top-p truncates the sorted distribution at a cumulative mass and renormalises, so on a confident step it keeps two tokens and on an uncertain one it keeps dozens; that is the point of it, and the reason it is usually preferred to a fixed top-k. And temperature zero is not a temperature at all, since dividing by zero is undefined: implementations treat it as take the highest logit, which is why it is deterministic in a way no small temperature is.

Then the loop. The chosen token is appended and the whole thing runs again, which is the fact that explains most of what you notice about using a model: output is produced strictly one token at a time so latency scales with answer length, streaming is possible because each token is finished when it appears, and a model cannot revise a token it has already emitted because that token is now part of its own input.

Open the config of any open-weights model released recently and you will find a dozen names that no introduction to transformers mentions: grouped-query attention, RoPE scaling, SwiGLU, a router and an expert count. This section goes through them one at a time and says what each changes about memory, speed or quality, with the arithmetic where there is arithmetic. It is the difference between knowing how a transformer works and being able to read what someone is actually shipping.

The model in the paper is not the model you will serve#

The 2017 architecture is the one most explanations still teach: multi-head attention, learned absolute position embeddings, a GELU feed-forward layer, layer normalisation. Every one of those has been replaced in the models being served today. Interviewers ask for the replacements by name, because each is a decision with a cost attached rather than a detail.

Grouped-query attention#

Multi-head attention gives every query head its own key and value heads. Keys and values are what the KV cache stores, so the cache grows with the number of heads. Grouped-query attention keeps all the query heads and lets several of them share one key and value head. Multi-query attention is the same idea taken to a single shared pair for the whole layer.

32 query heads, 32 KV heads   multi-head        full cache
32 query heads, 8 KV heads    grouped-query     cache / 4
32 query heads, 1 KV head     multi-query       cache / 32

KV cache per token = 2 x layers x kv_heads x head_dim x bytes

32 layers, 8 KV heads, head_dim 128, 16-bit   ->  128 KiB per token
the same model with 32 KV heads               ->  512 KiB per token

On a served model the cache, not the weights, is usually what caps how many requests fit on a card. Dividing it by four divides the memory cost of concurrency by four, or buys four times the context length for the same memory. This is the single change that made long context affordable, and it is why an 8B model can serve dozens of 8k conversations on one GPU.

The quality cost is small and real, and it grows as sharing gets more aggressive. Multi-query attention loses more and was reported to be less stable to train, which is why grouping, commonly eight key and value heads, is the compromise nearly everyone settled on.

Sliding-window and local attention#

Restrict each position to the last w tokens instead of the whole prefix. Cost per layer becomes linear in sequence length rather than quadratic, and the cache only has to hold w tokens per layer no matter how long the conversation gets.

What it costs is direct reach. Information still travels further than the window, because each layer moves it another w positions, so a stack of L layers reaches about L times w. What is lost is precise retrieval: a single token far outside the window cannot be attended to directly, which is exactly what a long-context lookup needs. Models that use it usually interleave, putting a few full-attention layers among many windowed ones so at least some layers can see everything.

Rotary position embeddings#

Attention treats its input as a set, so position has to be supplied from outside. Early models added a learned vector per position. That has two problems: the table has a fixed number of rows, so there is no vector at all for a position beyond the training length, and it encodes where a token sits rather than how far apart two tokens are.

RoPE rotates each query and key vector by an angle proportional to its position index. The dot product of a rotated query with a rotated key then depends only on the difference between the two positions, so relative distance falls out of the arithmetic instead of being learned. Nothing is added to the values and no table has to be sized in advance.

Extending context afterwards means changing the rotation frequencies so the angles a longer sequence produces stay inside the range the model saw in training. Position interpolation scales them down uniformly; the NTK-aware and YaRN variants scale them unevenly across frequency bands, leaving the high frequencies that carry local order alone. All of them work to a point, and all of them cost some quality, most visibly on precise retrieval inside the stretched range.

SwiGLU and RMSNorm#

The feed-forward layer used to be an up-projection, an activation such as GELU, and a down-projection. SwiGLU does the up-projection twice, passes one copy through a swish activation and multiplies the two elementwise, so one path gates the other. That is three matrices where there were two, so the hidden width is cut to roughly two thirds to keep the parameter count level. The gate lets the layer suppress channels per token, and it measures slightly better for the same budget. There is no deeper story than that.

RMSNorm drops the mean subtraction from layer normalisation: divide the vector by its root mean square and apply a learned scale, with no shift. One less reduction over the vector, marginally less arithmetic, and no measurable quality cost, which is why it has become the default.

Mixture of experts#

The feed-forward layer is split into many separate experts, and a small router sends each token to one or two of them. The others do nothing for that token. The model gets far more parameters without the compute per token rising to match.

Dense model       8B total parameters      8B active per token
Sparse model    120B total parameters      5B active per token

compute per token    follows the active count
memory to serve      follows the total count

The other costs are load balancing, which needs an auxiliary term in the loss to stop the router sending nearly everything to a few favoured experts, and communication: when experts are spread across devices, every token's routing decision becomes an all-to-all exchange.

FlashAttention#

Standard attention computes the full matrix of scores for every pair of positions, writes it to GPU memory, reads it back for the softmax, writes it again, and reads it once more to multiply by the values. On a long sequence that matrix is enormous, and moving it is what takes the time rather than the multiplying.

FlashAttention never writes it. It walks over blocks of queries and keys small enough to sit in on-chip memory, keeps a running maximum and sum so the softmax can be accumulated block by block, and builds the output as it goes. The backward pass recomputes the blocks it needs instead of storing them.

It is exact. The output matches standard attention to floating-point noise, it performs slightly more arithmetic rather than less, and the entire win is memory traffic. It also makes attention memory linear in sequence length rather than quadratic, which is what let context windows grow at all. If it appears in a list beside sparse and linear attention as another approximation, that is the thing to correct.

This is the single most commonly reported from-scratch task in ML engineer and research loops, and it is usually asked without a framework: a blank editor, NumPy, forty minutes. The mathematics is four lines. What actually goes wrong is shapes — a transpose on the wrong axis, a mask broadcast into the wrong dimension, a softmax over the wrong one — so this section states the shape of every intermediate and ends with two assertions that catch most of it.

One head, one sequence#

Start with a single sequence of T positions, each a vector of d_model features. Three learned matrices turn that one input into a query, a key and a value for every position. They are the same input three times, projected differently; there is no separate source for the keys in self-attention.

import numpy as np

rng = np.random.default_rng(0)

T, d_model, d_k = 5, 8, 8
x = rng.normal(size=(T, d_model))                 # (T, d_model)

W_q = rng.normal(size=(d_model, d_k)) * 0.1       # (d_model, d_k)
W_k = rng.normal(size=(d_model, d_k)) * 0.1
W_v = rng.normal(size=(d_model, d_k)) * 0.1

q = x @ W_q                                       # (T, d_k)
k = x @ W_k                                       # (T, d_k)
v = x @ W_v                                       # (T, d_k)

Each row of q asks a question; each row of k advertises what that position offers. Comparing every query with every key is one matrix multiplication, and it produces a square matrix of T by T: how much position i should attend to position j.

scores = q @ k.T / np.sqrt(d_k)                   # (T, d_k) @ (d_k, T) -> (T, T)

The division is not decoration. A dot product of two d_k-dimensional vectors with unit-variance entries has variance d_k, so without scaling the scores grow with the feature dimension, the softmax saturates into a near one-hot row, and the gradient through it goes to nothing. Dividing by the square root of d_k puts the variance back to roughly one.

The causal mask#

A decoder may not look ahead. Position i may attend to positions 0 through i and no further, which is exactly the lower triangle of the score matrix including the diagonal. The mask is applied to the scores, before the softmax, by setting the forbidden entries to negative infinity — so that after exponentiating they are exactly zero and the remaining weights still sum to one.

# True strictly above the diagonal: the future, which must not be attended to
mask = np.triu(np.ones((T, T), dtype=bool), k=1)
scores = np.where(mask, -np.inf, scores)          # (T, T), still

print(np.round(scores, 2))
# [[-0.03  -inf  -inf  -inf  -inf]
#  [ 0.11 -0.05  -inf  -inf  -inf]
#  [ 0.03  0.    0.    -inf  -inf]
#  [ 0.02  0.    0.02 -0.02  -inf]
#  [ 0.05 -0.1   0.02 -0.05  0.03]]

Softmax, stably#

Write it yourself rather than reaching for a library, because the stable form is part of what is being asked. Subtracting the row maximum before exponentiating changes nothing mathematically — the constant cancels between numerator and denominator — and it means the largest exponent is always exp(0), so nothing overflows.

def softmax(z, axis=-1):
    # subtracting the max is a no-op mathematically and stops exp overflowing
    z = z - z.max(axis=axis, keepdims=True)
    e = np.exp(z)
    return e / e.sum(axis=axis, keepdims=True)

w = softmax(scores, axis=-1)                      # (T, T), each row sums to 1
out = w @ v                                       # (T, T) @ (T, d_k) -> (T, d_k)

The masked entries survive this correctly: a row's maximum is always a real number, because the diagonal is never masked, so a masked entry becomes negative infinity minus a finite number, and its exponential is zero. No row is entirely masked, so there is no division by zero to worry about here.

The two assertions to run#

Nearly every bug in an implementation of this shows up in one of two checks. Run both before claiming it works, and say out loud that you are running them: an interviewer watching you write your own test is watching the thing they are actually grading.

# every row is a probability distribution over the positions it can see
assert np.allclose(w.sum(axis=-1), 1.0)

# nothing attends to the future: the strict upper triangle is exactly zero
assert np.all(np.triu(w, k=1) == 0.0)

Several heads at once#

Multiple heads are not multiple loops. The projections stay one matrix each of shape (d_model, d_model); the head split is a reshape of the feature axis into n_heads chunks of d_head, followed by a transpose that moves heads in front of positions so the score matrix is computed per head by ordinary batched matrix multiplication.

def causal_mha(x, W_q, W_k, W_v, W_o, n_heads):
    B, T, d_model = x.shape                       # (B, T, d_model)
    d_head = d_model // n_heads

    q = x @ W_q                                   # (B, T, d_model)
    k = x @ W_k
    v = x @ W_v

    def split(t):
        # (B, T, d_model) -> (B, T, n_heads, d_head) -> (B, n_heads, T, d_head)
        return t.reshape(B, T, n_heads, d_head).transpose(0, 2, 1, 3)

    q, k, v = split(q), split(k), split(v)

    # transpose the last two axes only: (B, n_heads, d_head, T)
    scores = q @ k.transpose(0, 1, 3, 2) / np.sqrt(d_head)   # (B, n_heads, T, T)

    mask = np.triu(np.ones((T, T), dtype=bool), k=1)         # broadcasts over B and heads
    scores = np.where(mask, -np.inf, scores)
    w = softmax(scores, axis=-1)                             # (B, n_heads, T, T)

    ctx = w @ v                                              # (B, n_heads, T, d_head)
    # heads back behind positions, then flatten them into the feature axis
    ctx = ctx.transpose(0, 2, 1, 3).reshape(B, T, d_model)   # (B, T, d_model)
    return ctx @ W_o, w                                      # (B, T, d_model)

Two details in that function are where people lose marks. The transpose before the score matmul swaps the last two axes and leaves the batch and head axes alone; writing k.T would reverse all four. And the final reshape must follow the transpose back to (B, T, n_heads, d_head), because reshaping straight from (B, n_heads, T, d_head) interleaves positions and heads into nonsense that still has the right shape and will not raise an error.

B, T, d_model, n_heads = 2, 6, 16, 4
x = rng.normal(size=(B, T, d_model))
W_q, W_k, W_v, W_o = (rng.normal(size=(d_model, d_model)) * 0.1 for _ in range(4))

y, w = causal_mha(x, W_q, W_k, W_v, W_o, n_heads)
print(y.shape, w.shape)                           # (2, 6, 16) (2, 4, 6, 6)

assert np.allclose(w.sum(axis=-1), 1.0)
assert np.all(np.triu(w, k=1) == 0.0)

# a third check worth knowing: causality means a prefix is unaffected by what
# follows it, so truncating the input must not change the earlier outputs
y_short, _ = causal_mha(x[:, :4], W_q, W_k, W_v, W_o, n_heads)
assert np.allclose(y[:, :4], y_short)

What is deliberately not here#

This is attention, not a transformer block. A block wraps it in a residual connection and a normalisation, then a feed-forward layer with its own residual, and a real model adds positional information — rotary embeddings, applied to q and k before the scores, in every current model. Modern attention also shares keys and values across heads rather than giving each its own. Those are covered elsewhere in this topic; if an interviewer asks for attention, give them this and then say what would surround it.