About this series

Transformers power modern language models, but explanations often jump from a colorful attention diagram to dense matrix notation. This series takes the middle route: clear intuition, enough mathematics to verify the mechanism, and practical implications for model architecture and GPU inference. It assumes familiarity with programming, arrays and basic algebra, not previous experience building Transformers.

Reading path. Parts 1–4 build one attention calculation. Parts 5–8 assemble a Transformer block, including positional information and feed-forward networks (FFNs). Parts 9–13 connect the math to memory-efficient GPU execution and text generation. Parts 14–15 close with training, implementation and evidence-based performance analysis.

Editorial conventions.  $T$ denotes token positions, $d_{model}$ the number of hidden features per token, $d_k$ Query/Key width, $B$ batch size and $H$ attention-head count. Dimensions are usually written as rows × columns. Approximate costs are analytic estimates, not measurements of a particular GPU.

1. Part 1 – Why a Word Needs Its Neighbors
2. Part 2 – Q, K and V: The Matching Engine
3. Part 3 – Why Divide Attention Scores by $\sqrt{d_k}$?
4. Part 4 – Softmax, Causal Masks and a Three-Token Walkthrough
5. Part 5 – Multiple Heads and Positional Information
6. Part 6 – Normalization and Residual Connections
7. Part 7 – Trace One Complete Transformer Block
8. Part 8 – The Feed-Forward Network Is Not Just Another Matrix
9. Part 9 – The GPU Cost of Attention
10. Part 10 – FlashAttention Without the Mystery
11. Part 11 – Why Predicting the Next Token Is a Learning Problem
12. Part 12 – Prefill, Decode and the First Generated Token
13. Part 13 – KV Cache, Grouped-Query Attention and PagedAttention
14. Part 14 – What Gradients and Distributed GPUs Really Change
15. Part 15 – Build, Verify and Profile Your Own Attention ← You are here

Overview

A mathematical explanation becomes much more convincing when a small reference program reproduces it. But correctness checks and performance measurements are separate tasks. A NumPy calculation that matches hand-worked attention weights proves something useful about the algorithm; it does not demonstrate NVIDIA Tensor Core utilization or reveal the fastest fused GPU kernel.

This final post shows an efficient validation path from a three-token example to a careful benchmarking plan. It deliberately avoids fabricated GPU timings.

Step 1: build a transparent NumPy reference

Start with three projected token positions, one two-feature head, and a causal mask. The code follows the same calculation order as Part 4. It is meant to run on a CPU so that the small arrays can be inspected easily.

import numpy as np
Q = np.array([[1., 0.], [0., 1.], [1., 1.]])
K = Q.copy()
V = np.array([[1., 0.], [0., 2.], [1., 1.]])
scores = (Q @ K.T) / np.sqrt(Q.shape[-1])
future = np.triu(np.ones(scores.shape, dtype=bool), 1)
scores = np.where(future, -np.inf, scores)
shifted = scores - scores.max(axis=-1, keepdims=True)
p = np.exp(shifted)
p /= p.sum(axis=-1, keepdims=True)
out = p @ V
assert p.shape == (3, 3)
assert np.allclose(p.sum(axis=-1), 1)
assert np.all(p[future] == 0)
np.testing.assert_allclose(p[1],
    [0.330238, 0.669762, 0.], atol=1e-6)
np.testing.assert_allclose(out[1],
    [0.330238, 1.339523], atol=1e-6)
print(out)

The assertions are more important than the printout. They check shape, row-wise normalization, exact zero weights on forbidden future positions and agreement with independently calculated reference values. If you change Q/K/V or the mask, update your hand-computed reference values rather than weakening the tests.

Step 2: test causal information flow

A valid causal decoder must not let an earlier position change solely because a later token was modified. For a complete deterministic Transformer-block reference, run the block twice: once with the original three-token input and once with only token 3 changed. The outputs at positions 1 and 2 should match within the numerical tolerance you declared.

This test checks more than a triangular mask shape. It checks that the entire block, including normalization, projections and the FFN, preserves the intended causal dependency. The source masterclass includes a complete NumPy implementation and a passing causal-prefix assertion for its toy block.

Step 3: compare with PyTorch scaled-dot-product attention

PyTorch provides torch.nn.functional.scaled_dot_product_attention (SDPA), which can select different implementations according to device, dtype, tensor shape, masks, dropout and software support. A useful experiment compares explicit attention with SDPA on the same tensors, checks the output within a declared tolerance, and confirms the intended causal setting.

Remember that attention APIs do not always share the same Boolean-mask convention. Also set dropout_p=0.0explicitly when checking inference equivalence. A successful result verifies numerical agreement for the tested case. It does not identify the actual selected GPU kernel or establish a speedup.

Step 4: benchmark the workload you care about

Before timing, write down the experimental configuration: GPU model, driver/runtime and framework versions, dtype, batch and concurrency, Query and Key lengths, Query/KV head counts, head dimensions, causal mode, actual backend, warmup procedure and timing mechanism.

A microbenchmark isolates one kernel. A serving benchmark includes all the components involved in generating tokens. They answer different questions. During prefill, use realistic long Query blocks; during decode, represent one new Query per active sequence and the relevant KV history. Record TTFT, inter-token latency and sustained throughput separately where applicable.

GPU launches are asynchronous, so a CPU-side timer can return before the GPU finishes. For isolated CUDA operations, use proper synchronization or CUDA events and multiple warmed-up repetitions. Report variability, not just the best run.

Step 5: profile before choosing an optimization

NVIDIA Nsight Systems is suited to examining the timeline: CPU launch gaps, GPU kernels, copies and communication overlap. NVIDIA Nsight Compute investigates individual kernels, including memory transactions, achieved arithmetic work, occupancy and stall reasons, subject to hardware support. PyTorch’s profiler can provide another layer of operator-level evidence.

Illustration · A measurement-driven optimization loop: define workload, profile, identify the bottleneck, change one variable and re-test.

A trace with heavy attention-intermediate HBM traffic suggests checking fused or tiled backends. A decode trace with large K/V reads suggests investigating cache layout, KV-head count and batching. A multi-GPU trace dominated by collective communication suggests revisiting partitioning and fabric placement. These are hypotheses to test, not conclusions from labels alone.

Do not mistake an operator name containing “flash” in a CPU trace for evidence that a CUDA FlashAttention kernel executed. Likewise, seeing low kernel time does not automatically establish a better end-to-end user experience.

Five closing questions to test your understanding

1. If $Q$ is $5\times64$ and $K$ is $7\times64$, what is the score shape? Answer: $5\times7$, because the product is $QK^\mathsf{T}$.

2. If a Query may consult only two Keys, when should future Keys be masked? Answer: Before softmax, so valid probabilities renormalize to one.

3. What does doubling context length do to full dense attention pair count? Answer: It multiplies the $T^2$ term by four.

4. Do 32 Query heads imply 32 KV heads? Answer: No. GQA and MQA share fewer KV heads across Queries.

5. Does FlashAttention reduce the same memory problem as PagedAttention? Answer: No. FlashAttention targets intermediate attention I/O; PagedAttention manages long-lived KV-cache allocation and address translation.

Series conclusion. The useful mental model is now complete: learned token vectors become Q/K/V projections; scaled and masked softmax mixes Value features; residuals, normalization and FFNs turn that mixture into updated hidden states; training adjusts the parameters through gradients; and inference engines execute the same mathematics under very different compute and memory conditions. When you can explain both the equations and the implementation boundaries, Transformer architecture becomes a system you can reason about rather than a collection of terms to memorize.

End of series. See you next time.