# What a transformer computes · Prerequisites Everything this book makes faster is, underneath, a transformer forward pass. This section describes that computation and nothing else: not why the architecture works, not how it is trained, and not why one variant beats another. What matters here is the shape of the arithmetic, because the shape is what the hardware responds to. Text does not enter the model as text. It is first split into **tokens** (The unit of input a language model actually processes: a chunk of text, often a word fragment, drawn from a fixed vocabulary produced by a tokenizer.), chunks drawn from a fixed vocabulary. The transformer paper encoded its sentences with byte-pair encoding over “a shared source-target vocabulary of about 37000 tokens,” and the principle has not changed even as vocabularies have grown: a token is usually smaller than a word and larger than a letter. This is the unit everything downstream is counted in. When a later chapter says a request has 8,000 tokens of context, or that decode produces one token per step, this is the thing being counted. Each token is then looked up in an embedding table, turning it into a vector of some fixed width. The paper calls that width d\_model and sets it to 512; production models are far wider, but the role is unchanged. From here on the model is manipulating vectors of that width, one per position, and the text is gone. One thing has to be added before the layers begin. Attention, as defined below, has no notion of order: it computes over a set of positions, and shuffling them would produce the same answer reordered. The paper is explicit that this is a problem to be fixed rather than a property to keep. “Since our model contains no recurrence and no convolution, in order for the model to make use of the order of the sequence, we must inject some information about the relative or absolute position of the tokens in the sequence,” which it does by adding **positional encodings** (A vector added to each token's embedding that encodes its position in the sequence, so that an otherwise order-blind attention mechanism can distinguish first from last.) to the embeddings, chosen to have the same dimension “so that the two can be summed.” How position is encoded has been revised many times since, and one of those revisions, rotary encoding, is load-bearing enough that it constrains a KV-cache optimization in [KV cache systems](https://learn-kernels.com/chapters/inference-engines/kv-cache-systems). ### What one layer does The model is then a stack of identical layers, six of them in the original paper and many more in a modern one. Each layer has two sub-layers: a multi-head self-attention mechanism and, in the paper’s words, “a simple, position-wise fully connected feed-forward network.” Around each of the two sits a residual connection followed by layer normalization, so that the output of a sub-layer is `LayerNorm(x + Sublayer(x))`: the sub-layer’s result is added to its own input rather than replacing it. The feed-forward half is the simpler and, in parameter count, usually the larger. It is two matrix multiplications with a nonlinearity between them, applied to each position independently, and it widens before it narrows: “the dimensionality of input and output is d\_model = 512, and the inner-layer has dimensionality d\_ff = 2048.” Four times wider in the middle, in that configuration. Because it treats every position separately, it is embarrassingly parallel across positions and maps onto the hardware without difficulty. It is also where mixture-of-experts models make their intervention, by keeping many such networks and routing each token to a few of them, which is the subject of [Mixture-of-experts serving](https://learn-kernels.com/chapters/distributed-inference/mixture-of-experts-serving). Attention is the half that makes the model interesting and, for this book’s purposes, expensive. It lets each position look at other positions and decide how much of each to take. Every token produces three vectors by multiplication with learned weight matrices: a **query** (The vector representing what a given position is looking for when it attends to the rest of the sequence.), a **key** (The vector a position advertises so that other positions can decide how much attention to pay to it.), and a **value** (The vector a position contributes to the output of every position that attends to it.). The paper defines the operation over them precisely: the input “consists of queries and keys of dimension d\_k, and values of dimension d\_v,” and then “we compute the dot products of the query with all keys, divide each by” the square root of d\_k, “and apply a softmax function to obtain the weights on the values.” > Figure. One attention head, end to end. Scaled dot-product attention as the paper defines it. Every query is scored against every key, which is the step whose cost grows with the square of the sequence length; the scores are scaled by the square root of the key dimension, normalized by softmax into weights that sum to one, and used to mix the values. Section 3.4 fuses this entire chain into a single kernel. Take the three steps in turn, because each one becomes a performance problem later. The first multiplication scores every query against every key, producing one number per pair of positions. That matrix is square in the sequence length, so doubling the context quadruples both the arithmetic and, if you store it, the memory. That single fact is why long context is hard, why [FlashAttention](https://learn-kernels.com/chapters/kernel-optimization/attention) is built around never writing that matrix to memory at all, and why [long context](https://learn-kernels.com/chapters/inference-engines/long-context-and-multimodal) has a section to itself. The division by the square root of d\_k is a scaling step, and it is cheap. The softmax that follows is not cheap in the way that matters: it must see an entire row before it can normalize any of it, because every output depends on the sum across the row. That dependency is what makes attention awkward to compute in pieces, and the technique for doing it anyway, one block at a time while carrying a running normalizer, is the mechanism underneath FlashAttention. The third step, weighting the values by those normalized scores, is another matrix multiplication. Finally, this runs several times in parallel rather than once. Multi-head attention “allows the model to jointly attend to information from different representation subspaces at different positions,” because “with a single attention head, averaging inhibits this.” The paper uses “h = 8 parallel attention layers, or heads,” each with d\_k = d\_v = d\_model / h = 64, and notes the consequence that makes this affordable: “due to the reduced dimension of each head, the total computational cost is similar to that of single-head attention with full dimensionality.” The heads divide the width rather than multiplying the work. Their count and width are the `n_heads` and `d_head` that appear throughout the later chapters, and the fact that heads are independent is what makes them a natural axis to split across GPUs in [Distributed inference](https://learn-kernels.com/chapters/distributed-inference). ### Nearly all of it is matmul Since the answer to what a transformer computes turns out to be mostly one operation, it is worth stating that operation exactly. Multiplying an M by K matrix with a K by N matrix produces an M by N result, in which every output element is the dot product of one row of the first and one column of the second: K multiplications and K additions. That gives a flop count of 2 M N K for the whole operation, and it is the rule behind every flop figure in this book. The per-token derivation in *Transformer Inference Arithmetic* is this rule applied repeatedly: multiplying a token’s 1 by d\_model vector through a d\_model by d\_model weight matrix costs 2 d\_model squared flops, and doing that for the query, key, and value projections costs three times as much. What makes this operation the one hardware is built around is not the flop count but the reuse hiding inside it. Each element of the first matrix is read once and then participates in N separate output elements; each element of the second participates in M. The arithmetic grows with the product of all three dimensions while the data grows only with their sums, so the larger the matrices, the more arithmetic the hardware extracts from each byte it fetches. That ratio is the thing the two budgets of the first section are competing over, and matmul is close to the best case for it. Contrast that with an elementwise operation, adding two vectors or applying a nonlinearity. It reads its inputs, performs one operation per element, and writes the result. There is no reuse available at all, so it spends almost entirely from the bandwidth budget and leaves the arithmetic units idle no matter how it is written. This is why such operations are worth folding into a neighbouring matmul rather than running on their own, and why fusion recurs throughout the later chapters as an optimization that saves no arithmetic whatsoever and is still worth doing. Step back now and count what the layer actually spends its time on. Projecting each token to a query, a key, and a value is three matrix multiplications. Scoring queries against keys is a fourth. Weighting the values is a fifth. Projecting the joined heads back to the model width is a sixth. Both halves of the feed-forward network are a seventh and an eighth. The nonlinearities, the normalizations and the softmax sit between them and account for a small fraction of the arithmetic. A transformer is, to a first approximation, a very large number of matrix multiplications with a little glue. This is why [Matrix multiplication](https://learn-kernels.com/chapters/kernel-optimization/matrix-multiplication) gets an entire section to itself, why the hardware chapters care so much about units that do nothing but multiply matrices, and why the glue matters more than its share of the arithmetic suggests: an operation that does little arithmetic but still has to read and write its operands is spending from the bandwidth budget, not the arithmetic one, which is what makes fusing it into a neighbouring matmul worth the trouble. After the last layer, the vector for a position is multiplied by one more matrix to produce a score for every token in the vocabulary. Those scores are the model’s output. Turning them into text is the subject of the next section. One notational habit is worth adopting before then, because the later chapters use it without explanation. The data flowing through the model is described by its shape, a list of dimensions. Activations inside a layer are usually written as three: how many independent sequences are being processed at once, how many tokens are in each, and how wide the model is. Weights carry no sequence or batch dimension at all, because the same matrix is applied to every token of every sequence, and that asymmetry is the one to hold on to. Activations grow with the work; weights do not. Reading a shape tells you immediately which budget an operation will spend from. An operation whose output shape is much smaller than the weights it reads will be limited by fetching those weights. An operation over a long sequence with a large batch has enough arithmetic in it to keep the units busy. The scaling book’s own worked example gives the sense of the quantities: prefilling 8,192 tokens, “a single activation only uses around 8,192 x 5,120 x 2 bytes = 80MB of memory,” three dimensions multiplied out, the last of them being the two bytes per value from the previous section. Source The byte-pair encoding and vocabulary size, the positional encoding rationale, the layer structure, the residual and layer-normalization arrangement, the feed-forward dimensionality, the definition of scaled dot-product attention, and the multi-head configuration and its cost argument are all from [Attention Is All You Need](https://arxiv.org/abs/1706.03762) (Vaswani et al., 2017), sections 3.1 through 3.5 and 5.1. The observation that the layer is mostly matrix multiplication is this book’s framing rather than the paper’s. The 2 M N K flop rule is the definition of the operation; the per-token application of it quoted here, and the three projections costing 2 times 3 times d\_model squared, are from [Transformer Inference Arithmetic](https://kipply.github.io/blog/transformer-inference-arithmetic/), which walks the whole layer that way. The 80 MB activation example is from [How to Scale Your Model: Inference](https://jax-ml.github.io/scaling-book/inference/). Rotary encoding is not from the transformer paper and is cited where it is used.