TransformerCPP: Building a Transformer from Scratch in C++

I wrote transformer.cpp to understand how much code sits behind a Transformer layer. I wanted to be able to start at an attention calculation and follow it through the matrix multiplication, the array indexing, and the backward pass. That meant writing the tensor library as well as the model.

The resulting implementation runs on the CPU and includes an encoder-decoder Transformer, automatic differentiation, and a thread pool for tensor operations. Most of this article concerns the tensor library. To explain why it is organized the way it is, I'll first go through the model that uses it.

1. The model

I used the encoder-decoder architecture from Attention Is All You Need. 1 The encoder reads a source sequence; the decoder processes a target sequence with access to the encoder's output. For a batch of $B$ examples, the two inputs contain $T_s$ and $T_t$ token IDs per example. Each side has its own embedding layer and positional encoding, after which every token is represented by $D$ floating-point values.

An encoder layer applies self-attention and a feed-forward network. A decoder layer adds cross-attention between its self-attention and feed-forward sublayers. In the cross-attention call, I pass the decoder representation as the query input and the encoder output as both the key and value inputs. The attention implementation can therefore handle both kinds of attention through the same interface.

Each sublayer returns a tensor with the same shape as its input. After dropout, I add the input to the result and apply layer normalization:

$$y=\operatorname{LayerNorm}\left(x+\operatorname{Dropout}(F(x))\right).$$

Normalization is performed over a token's $D$ features, with learned scale and bias parameters. 2 This follows the post-normalization layout of the original Transformer. Preserving the feature width across sublayers makes the residual addition possible without another projection.

At the end of the decoder, a linear layer converts each token's $D$ features into $V$ vocabulary scores. For example, with a vocabulary of 100 tokens, there are 100 scores at every target position, regardless of the internal feature width. forward returns the resulting $(B,T_t,V)$ tensor to the caller, which computes the loss during training.

2. Attention, one row at a time

For a query $q_i$, attention compares it with every key $k_j$ by taking their dot product. A larger score gives the corresponding value vector more influence over the output. Softmax turns the scores into coefficients that sum to one, so a row with coefficients $0.7$, $0.2$, and $0.1$ would produce

$$o_i=0.7v_1+0.2v_2+0.1v_3.$$

The same calculation for all queries can be expressed with two matrix products: 1

$$S=\frac{QK^\top}{\sqrt d}+M,\qquad P=\operatorname{softmax}(S),\qquad O=PV.$$

The first product computes the scores; the second combines the values using their probabilities. Here $d$ is the width of one head and $M$ is the attention mask. In softmax, the reduction runs along the last axis of $S$, which contains the keys. Each row then describes how one query divides its attention among them.

Multi-head attention requires some rearrangement before these products. The query, key, and value projections each produce a tensor of shape $(B,T,D)$. With $H$ heads and $d=D/H$, I reshape and transpose each projection as follows:

$$(B,T,D) \xrightarrow{\text{reshape}}(B,T,H,d) \xrightarrow{\text{transpose}}(B,H,T,d).$$

My matrix multiplication kernel treats the final two axes as a matrix and everything before them as batch dimensions. Putting the heads ahead of the sequence axis lets the kernel process them as additional batches. If the query and key lengths are $T_q$ and $T_k$, respectively, the matrix dimensions within each head are

$$\underbrace{Q}_{T_q\times d} \underbrace{K^\top}_{d\times T_k} \longrightarrow \underbrace{S}_{T_q\times T_k}, \qquad \underbrace{P}_{T_q\times T_k} \underbrace{V}_{T_k\times d} \longrightarrow \underbrace{O}_{T_q\times d}.$$

These lengths can differ in cross-attention, so the code tracks them separately. Once the heads have been evaluated, I transpose their outputs to $(B,T_q,H,d)$, reshape to $(B,T_q,D)$, and apply the output projection.

Computing softmax

The look-ahead mask contains zero on and below the diagonal and negative infinity above it. Adding it to the scores makes future positions contribute zero after exponentiation. The remaining entries need to be normalized without overflowing the exponential, so I subtract the largest score in each row:

$$p_i=\frac{e^{s_i-m}}{\sum_j e^{s_j-m}},\qquad m=\max_j s_j.$$

Subtracting $m$ multiplies every exponential by the same factor, which cancels in the fraction. All exponent arguments are now at most zero. Blanchard, Higham, and Higham examine the numerical accuracy of this shifted formulation. 3

There is an unresolved edge case when every entry in a row is masked. The maximum is then negative infinity, and the subtraction produces NaNs. A caller using this kernel needs at least one finite score in each row, or a separate policy for fully masked rows.

3. Reshaping and transposing the arrays

The reshapes above share data with their inputs. The transposes copy it. This follows from a restriction in my tensor representation: every tensor is interpreted as a contiguous, row-major array.

Tensor stores its dimensions in a std::vector<int> and its values in a shared std::vector<float>. Consider an array with shape $(2,3,4)$. Each row contains four values, and each of the two outer slices contains three rows. To reach element $(i,j,k)$, skip $i$ slices, $j$ rows, and $k$ values:

$$\operatorname{offset}(i,j,k)=12i+4j+k.$$

The coefficients $(12,4,1)$ are the strides. I calculate them from the shape whenever indexing requires them. A reshape changes the dimensions used to interpret this array while preserving the order of its values. Its result shares the original data buffer and allocates a separate gradient buffer.

Now consider transposing a matrix:

$$A=\begin{bmatrix} 1&2&3\\ 4&5&6 \end{bmatrix}, \qquad A^\top=\begin{bmatrix} 1&4\\ 2&5\\ 3&6 \end{bmatrix}.$$

Both matrices must be stored row by row under the contiguous-storage convention:

$$\begin{aligned} \operatorname{buffer}(A)&=[1,2,3,4,5,6],\\ \operatorname{buffer}(A^\top)&=[1,4,2,5,3,6]. \end{aligned}$$

Reinterpreting the first buffer with shape $(3,2)$ would put 1 and 2 in the first row. The transpose needs 1 and 4 there. My implementation allocates a second buffer and writes each input value at its transposed offset, including when attention splits and recombines heads.

I could avoid the copy by storing arbitrary strides on the tensor. The transposed matrix would use strides $(1,3)$ against the original buffer, and reading its first row would visit offsets 0 and 3. Every kernel consuming that tensor would then have to support this access pattern. Keeping inputs contiguous simplified the kernels at the expense of copying transposed data.

Reshape has a different consequence: its input and output alias the same storage. Calling set through either tensor changes the values visible through both. If backward needs one of those values, modifying it after the forward pass changes the derivative being computed. The API currently leaves that constraint to the caller.

4. Building the backward pass

A layer's local variables disappear when its forward method returns. Some of their values are still needed to differentiate the result. For an elementwise product $y=a\odot b$, for instance, the derivative with respect to $a$ requires the forward value of $b$:

$$\bar a=\bar y\odot b,\qquad \bar b=\bar y\odot a.$$

A bar denotes the derivative of the loss with respect to that variable. To evaluate these expressions later, the result tensor stores references to both operands and records OperationType::Mul. Other operations save whatever their derivatives require, such as the permutation used by a transpose.

This is an operator-overloading implementation of reverse-mode automatic differentiation. 4 Calling backward on a result accumulates the incoming gradient and selects a derivative routine using the operation tag. That routine computes the parent gradients and recursively calls backward on each parent.

For a matrix product $Y=XW$, the routines compute

$$\bar X=\bar YW^\top,\qquad \bar W=X^\top\bar Y.$$

I use the existing transpose and matrix multiplication operations for these calculations. As a result, the backward pass also incurs their allocations and copies. When $W$ is shared across a batch, the contributions from different batch elements are reduced to a single gradient with the shape of $W$.

The same reduction is necessary for a broadcast bias. Suppose a projection adds one bias vector to every token in a batch. Increasing its $d$th entry changes the $d$th feature at every one of those positions, so the derivative includes all of them:

$$\bar b_d=\sum_{n=1}^{B}\sum_{t=1}^{T}\bar Y_{n,t,d}.$$

reduce_gradient sums over the dimensions introduced by broadcasting. Without that sum, the backward result would still have separate entries for each use of the parameter.

For softmax, the probability of one element depends on every score in its row through the denominator. After differentiating and collecting terms, the gradient is

$$\bar s_i=p_i\left(\bar p_i-\sum_j\bar p_jp_j\right).$$

The sum can be evaluated once per row and reused for every element. My backward kernel assumes softmax ran along the final axis, as it does in attention. The forward API accepts other axes too; supporting them in backward requires recording the selected axis alongside the operation.

Multiple uses of a tensor

Backward must add contributions whenever a tensor is used more than once. If $z=x+x$ and the gradient arriving at $z$ is 1, each use of $x$ contributes 1. The addition routine calls backward on $x$ twice, leaving 2 in its gradient buffer.

On that second call, only the newly received contribution should continue to $x$'s parents. Sending the accumulated value of 2 would propagate the first contribution again. The implementation therefore keeps accumulation and propagation separate: it stores the total gradient, then passes the incoming contribution to the local derivative routine.

This recursive traversal can evaluate the derivative of a shared ancestor repeatedly. A traversal in reverse topological order would allow each node to collect all its contributions before running its derivative once. That is one change I would make before using the library with larger, heavily branched graphs.

The graph also determines when memory can be released. A result holds shared_ptr references to its parents, keeping their forward values alive until the result is released. Tensor::create establishes shared ownership before an operation calls shared_from_this() to obtain a reference to itself. Parameters remain alive through a separate static registry populated by the factory. Moving that registry into the model would give each model control over its parameter lifetime.

5. CPU execution

For matrix multiplication, I assign each task a range of output elements. A task computes the full dot product for each element in its range, so workers never need to combine partial results or write to the same output location. The operation submits these tasks to a persistent thread pool and waits for the batch to complete.

This introduces synchronization at every operation boundary. On small tensors, the queue, worker notification, and completion barrier may account for a substantial part of the work. Increasing the worker count can only help when the tasks contain enough computation to cover that overhead.

The matrix multiplication loop also has poor locality in one operand: it reads across a row of the left matrix and down a column of the right. Column reads are separated by the row width in memory. Goto and van de Geijn describe how packing operands and blocking matrix multiplication improves reuse through the memory hierarchy. 5 Those changes concern the loop within each task, so they remain relevant even after the computation has been divided among workers.

Attention has a further storage cost. My implementation allocates the entire score matrix and another tensor for its probabilities. In self-attention, each contains $BHT^2$ floats. At long sequence lengths, avoiding those intermediates becomes as important as speeding up individual matrix products. FlashAttention addresses this with tiled exact attention designed around GPU memory transfers. 6 A blocked CPU implementation would need its own measurements to establish the benefit on the hardware used here.

References

1. Ashish Vaswani et al. Attention Is All You Need. NeurIPS, 2017.
2. Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer Normalization. arXiv:1607.06450, 2016.
3. Pierre Blanchard, Desmond J. Higham, and Nicholas J. Higham. Accurately computing the log-sum-exp and softmax functions. IMA Journal of Numerical Analysis, 41(4):2311–2330, 2021.
4. Atilim Gunes Baydin, Barak A. Pearlmutter, Alexey Andreyevich Radul, and Jeffrey Mark Siskind. Automatic differentiation in machine learning: a survey. Journal of Machine Learning Research, 18(153):1–43, 2018.
5. Kazushige Goto and Robert A. van de Geijn. Anatomy of High-Performance Matrix Multiplication. ACM Transactions on Mathematical Software, 34(3), Article 12, 2008.
6. Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS, 2022.