{"course":{"slug":"how-llms-work","title":"How a Language Model Actually Works","url":"https://addaly.com/learn/how-llms-work","level":"intermediate","track":"start"},"lesson":{"slug":"attention-in-forty-lines","title":"Attention in forty lines of NumPy","url":"https://addaly.com/learn/how-llms-work/attention-in-forty-lines","position":11,"of":78,"minutes":10,"answer":"Writing attention once, with real numbers, converts it from an analogy into an operation you can predict the behaviour of.","keyIdea":"Writing attention once, with real numbers, converts it from an analogy into an operation you can predict the behaviour of."},"body":{"format":"markdown","text":"## The whole operation, runnable\n\nThe previous lesson described attention without matrices. Here it is with them, in code that runs on any laptop with NumPy installed and nothing else. Fifteen minutes with this file is worth more than any diagram.\n\n```python\nimport numpy as np\n\ndef softmax(x, axis=-1):\n    x = x - x.max(axis=axis, keepdims=True)      # subtract max for numerical safety\n    e = np.exp(x)\n    return e / e.sum(axis=axis, keepdims=True)\n\ndef attention(X, Wq, Wk, Wv, causal=True):\n    \"\"\"X: (T, d) — one row per token position.\"\"\"\n    Q = X @ Wq          # (T, dk)   what each position is looking for\n    K = X @ Wk          # (T, dk)   what each position advertises\n    V = X @ Wv          # (T, dv)   what each position will hand over\n\n    scores = Q @ K.T / np.sqrt(Q.shape[-1])      # (T, T)\n\n    if causal:\n        mask = np.triu(np.ones_like(scores), k=1).astype(bool)\n        scores = np.where(mask, -np.inf, scores)\n\n    W = softmax(scores)                          # (T, T) rows sum to 1\n    return W @ V, W                              # (T, dv)\n\nT, d, dk = 6, 8, 4\nrng = np.random.default_rng(0)\nX = rng.normal(size=(T, d))\nout, W = attention(X, rng.normal(size=(d, dk)),\n                      rng.normal(size=(d, dk)),\n                      rng.normal(size=(d, dk)))\nprint(np.round(W, 2))\n```\n\nRun it. Print `W`. You will see a lower-triangular matrix whose every row sums to 1. That single object is the entire \"which positions matter to which\" story, and it is 6 × 6 for six tokens.\n\n## Three details that matter\n\n**The division by √dk.** Dot products of vectors with dimension 64 have a standard deviation roughly eight times larger than those of dimension 1. Without the scaling, the scores spread wide, softmax saturates, and gradients vanish. Delete that division and re-run: the weight matrix becomes almost one-hot, each position attending to exactly one other. That is a live demonstration of why the term is there.\n\n**Subtracting the max in softmax.** Purely numerical. `exp(1000)` overflows; `exp(1000 - 1000)` does not. Every production implementation does this and it changes no result.\n\n**`-inf` before the softmax, not zero after.** Masking with negative infinity makes those entries exactly zero *after* normalisation, so the remaining weights still sum to one. Zeroing afterwards would leave rows summing to less than one, which quietly rescales every value.\n\n## Multiple heads\n\nA head is one run of the above with its own three matrices. A layer runs several in parallel and concatenates the results, then applies one more projection:\n\n```python\ndef multihead(X, heads, Wo):\n    outs = [attention(X, *h)[0] for h in heads]\n    return np.concatenate(outs, axis=-1) @ Wo\n```\n\nIn real implementations the heads are packed into single large matrices and split by reshaping, purely for speed. Conceptually it is the loop above.\n\n## Now watch the cost\n\nLook at `Q @ K.T`. Its shape is (T, T). Set T to 1,000 and that is a million entries per head per layer. Set T to 100,000 and it is ten billion, per head, per layer. Time it yourself:\n\n```python\nfor T in (256, 512, 1024, 2048):\n    X = rng.normal(size=(T, 512))\n    import time; t = time.time()\n    attention(X, *(rng.normal(size=(512, 64)) for _ in range(3)))\n    print(T, round(time.time() - t, 3), \"s\")\n```\n\nEach doubling roughly quadruples the time. You have just measured, on your own machine, the quadratic cost that shapes the entire economics of long context.\n\n## What FlashAttention changes, and what it does not\n\nProduction systems do not build that (T, T) matrix in memory. FlashAttention computes the same result in tiles, keeping partial softmax statistics as it goes, so memory use grows linearly rather than quadratically and far less data moves between GPU memory and the compute units.\n\nBe precise about the claim: it is an **exact** algorithm, not an approximation. The output is identical to the naive version up to floating-point rounding. The *arithmetic* is still quadratic; what became linear is the memory traffic, and on modern hardware memory traffic is usually the thing that costs you time. This is one of the clearest examples in the field of a large practical speed-up from changing nothing about the mathematics.\n\n## The thing worth keeping\n\nAttention is a weighted average of value vectors, where the weights come from comparing a query against keys and pushing the results through a softmax. Everything else (heads, masks, scaling, caching, tiling) is engineering around that sentence. When you next read a paper claiming a new attention variant, you will be able to see immediately which part of these forty lines it is replacing."},"check":{"question":"In the code above, the causal mask sets future positions to negative infinity before the softmax rather than zeroing their weights afterwards. Why does the order matter?","options":[{"text":"Negative infinity is faster on a GPU than a multiplication by zero.","correct":false,"tempting":"Performance framing is plausible for a low-level detail, and masking is indeed implemented for speed in production kernels."},{"text":"Softmax normalises the row to sum to one, so masking first excludes those positions from the normalisation; zeroing afterwards would leave the row summing to less than one and silently shrink the output.","correct":true},{"text":"Zeroing afterwards would leak future information through the gradient during training.","correct":false,"tempting":"Information leakage is exactly the right worry for a causal mask, which makes this a well-aimed guess about the wrong step."},{"text":"The softmax cannot handle exact zeros in its input without producing NaN values.","correct":false,"tempting":"Numerical stability is a real concern in softmax, so a NaN explanation sounds like inside knowledge."}]},"prev":{"slug":"attention","title":"Attention, without a single matrix","url":"https://addaly.com/learn/how-llms-work/attention"},"next":{"slug":"the-causal-mask-and-parallel-training","title":"Why training is parallel and generation is not","url":"https://addaly.com/learn/how-llms-work/the-causal-mask-and-parallel-training"},"authorship":"Founder and series editor: Dr Mustafa Mun. Written by AI models to a written standard. Exam keys checked blind by a second model: 3,665 of 3,666 agreed. Academic review invited.","review":{"status":"Draft, open for academic review"},"cite":"Addaly (2026) 'Attention in forty lines of NumPy', in How a Language Model Actually Works. Edited by M. Mun. Addaly. Available at: https://addaly.com/learn/how-llms-work/attention-in-forty-lines (Accessed: 13 October 2026).","licence":{"name":"CC BY-NC-SA 4.0","url":"https://creativecommons.org/licenses/by-nc-sa/4.0/"},"reader":{"schema":1,"title":"Attention in forty lines of NumPy","courseTitle":"How a Language Model Actually Works","module":{"index":2,"slug":"inside-the-block","title":"Inside the block"},"answer":null,"keyIdea":"Writing attention once, with real numbers, converts it from an analogy into an operation you can predict the behaviour of.","hero":{"kind":"plate","image":null,"alt":null,"caption":null},"blocks":[{"t":"h2","id":"the-whole-operation-runnable","runs":[{"text":"The whole operation, runnable"}]},{"t":"p","runs":[{"text":"The previous lesson described attention without matrices. Here it is with them, in code that runs on any laptop with NumPy installed and nothing else. Fifteen minutes with this file is worth more than any diagram."}]},{"t":"code","lang":"python","label":"Python","text":"import numpy as np\n\ndef softmax(x, axis=-1):\n    x = x - x.max(axis=axis, keepdims=True)      # subtract max for numerical safety\n    e = np.exp(x)\n    return e / e.sum(axis=axis, keepdims=True)\n\ndef attention(X, Wq, Wk, Wv, causal=True):\n    \"\"\"X: (T, d) — one row per token position.\"\"\"\n    Q = X @ Wq          # (T, dk)   what each position is looking for\n    K = X @ Wk          # (T, dk)   what each position advertises\n    V = X @ Wv          # (T, dv)   what each position will hand over\n\n    scores = Q @ K.T / np.sqrt(Q.shape[-1])      # (T, T)\n\n    if causal:\n        mask = np.triu(np.ones_like(scores), k=1).astype(bool)\n        scores = np.where(mask, -np.inf, scores)\n\n    W = softmax(scores)                          # (T, T) rows sum to 1\n    return W @ V, W                              # (T, dv)\n\nT, d, dk = 6, 8, 4\nrng = np.random.default_rng(0)\nX = rng.normal(size=(T, d))\nout, W = attention(X, rng.normal(size=(d, dk)),\n                      rng.normal(size=(d, dk)),\n                      rng.normal(size=(d, dk)))\nprint(np.round(W, 2))"},{"t":"p","runs":[{"text":"Run it. Print "},{"text":"W","code":true},{"text":". You will see a lower-triangular matrix whose every row sums to 1. That single object is the entire “which positions matter to which” story, and it is 6 × 6 for six tokens."}]},{"t":"h2","id":"three-details-that-matter","runs":[{"text":"Three details that matter"}]},{"t":"p","runs":[{"text":"The division by √dk.","b":true},{"text":" Dot products of vectors with dimension 64 have a standard deviation roughly eight times larger than those of dimension 1. Without the scaling, the scores spread wide, softmax saturates, and gradients vanish. Delete that division and re-run: the weight matrix becomes almost one-hot, each position attending to exactly one other. That is a live demonstration of why the term is there."}]},{"t":"p","runs":[{"text":"Subtracting the max in softmax.","b":true},{"text":" Purely numerical. "},{"text":"exp(1000)","code":true},{"text":" overflows; "},{"text":"exp(1000 - 1000)","code":true},{"text":" does not. Every production implementation does this and it changes no result."}]},{"t":"p","runs":[{"text":"-inf","b":true,"code":true},{"text":" before the softmax, not zero after.","b":true},{"text":" Masking with negative infinity makes those entries exactly zero "},{"text":"after","i":true},{"text":" normalisation, so the remaining weights still sum to one. Zeroing afterwards would leave rows summing to less than one, which quietly rescales every value."}]},{"t":"h2","id":"multiple-heads","runs":[{"text":"Multiple heads"}]},{"t":"p","runs":[{"text":"A head is one run of the above with its own three matrices. A layer runs several in parallel and concatenates the results, then applies one more projection:"}]},{"t":"code","lang":"python","label":"Python","text":"def multihead(X, heads, Wo):\n    outs = [attention(X, *h)[0] for h in heads]\n    return np.concatenate(outs, axis=-1) @ Wo"},{"t":"p","runs":[{"text":"In real implementations the heads are packed into single large matrices and split by reshaping, purely for speed. Conceptually it is the loop above."}]},{"t":"h2","id":"now-watch-the-cost","runs":[{"text":"Now watch the cost"}]},{"t":"p","runs":[{"text":"Look at "},{"text":"Q @ K.T","code":true},{"text":". Its shape is (T, T). Set T to 1,000 and that is a million entries per head per layer. Set T to 100,000 and it is ten billion, per head, per layer. Time it yourself:"}]},{"t":"code","lang":"python","label":"Python","text":"for T in (256, 512, 1024, 2048):\n    X = rng.normal(size=(T, 512))\n    import time; t = time.time()\n    attention(X, *(rng.normal(size=(512, 64)) for _ in range(3)))\n    print(T, round(time.time() - t, 3), \"s\")"},{"t":"p","runs":[{"text":"Each doubling roughly quadruples the time. You have just measured, on your own machine, the quadratic cost that shapes the entire economics of long context."}]},{"t":"h2","id":"what-flashattention-changes-and-what-it-does-not","runs":[{"text":"What FlashAttention changes, and what it does not"}]},{"t":"p","runs":[{"text":"Production systems do not build that (T, T) matrix in memory. FlashAttention computes the same result in tiles, keeping partial softmax statistics as it goes, so memory use grows linearly rather than quadratically and far less data moves between GPU memory and the compute units."}]},{"t":"p","runs":[{"text":"Be precise about the claim: it is an "},{"text":"exact","b":true},{"text":" algorithm, not an approximation. The output is identical to the naive version up to floating-point rounding. The "},{"text":"arithmetic","i":true},{"text":" is still quadratic; what became linear is the memory traffic, and on modern hardware memory traffic is usually the thing that costs you time. This is one of the clearest examples in the field of a large practical speed-up from changing nothing about the mathematics."}]},{"t":"h2","id":"the-thing-worth-keeping","runs":[{"text":"The thing worth keeping"}]},{"t":"p","runs":[{"text":"Attention is a weighted average of value vectors, where the weights come from comparing a query against keys and pushing the results through a softmax. Everything else (heads, masks, scaling, caching, tiling) is engineering around that sentence. When you next read a paper claiming a new attention variant, you will be able to see immediately which part of these forty lines it is replacing."}]}],"check":{"question":[{"text":"In the code above, the causal mask sets future positions to negative infinity before the softmax rather than zeroing their weights afterwards. Why does the order matter?"}],"options":[{"text":[{"text":"Negative infinity is faster on a GPU than a multiplication by zero."}],"correct":false,"tempting":[{"text":"Performance framing is plausible for a low-level detail, and masking is indeed implemented for speed in production kernels."}]},{"text":[{"text":"Softmax normalises the row to sum to one, so masking first excludes those positions from the normalisation; zeroing afterwards would leave the row summing to less than one and silently shrink the output."}],"correct":true,"tempting":null},{"text":[{"text":"Zeroing afterwards would leak future information through the gradient during training."}],"correct":false,"tempting":[{"text":"Information leakage is exactly the right worry for a causal mask, which makes this a well-aimed guess about the wrong step."}]},{"text":[{"text":"The softmax cannot handle exact zeros in its input without producing NaN values."}],"correct":false,"tempting":[{"text":"Numerical stability is a real concern in softmax, so a NaN explanation sounds like inside knowledge."}]}],"order":[0,3,2,1],"why":null},"prev":{"slug":"attention","title":"Attention, without a single matrix","minutes":8},"next":{"slug":"the-causal-mask-and-parallel-training","title":"Why training is parallel and generation is not","minutes":8}}}