Let’s make GPT-2 in assembly

What GPT-2 does

Before we write any code, we need the full list of jobs. GPT-2 makes text one token at a time. A token is a word or a part of a word. For each new token, the model does the same loop:

So the jobs are:

1. Get 124 million weights into memory.
2. Read and write without libc.
3. Turn text into token ids.
4. Multiply matrices quickly.
5. Compute exp, square roots and GELU without a math library.
6. Run attention over all earlier tokens.
7. Pick a token from the scores.

We take them in that order.

Step 1: weights

We use the files that Hugging Face publishes, as they are: model.safetensors for the weights, and vocab.json and merges.txt for the tokenizer. There is no conversion step. No Python runs before or during a run.

A safetensors file is simple. The first 8 bytes give the length of a JSON header. The header lists every tensor by name, with its shape and its byte offset. The raw numbers come after the header. At startup, the assembly reads the header, finds the 148 tensors that it needs by name, and keeps a pointer to each one.

We do not load this file. We ask the kernel to map it into memory with mmap. The file appears as one block of memory, and the kernel reads the parts from disk when the code first touches them.

One problem. The checkpoint stores the large matrices as [input][output]. Our fastest kernel wants [output][input]. Then every output value is one dot product of one row of the matrix with x, and the CPU reads that row from start to end. A kernel for the checkpoint’s own layout gave correct output, but it made only about 60 tokens per second. So at startup, the program copies the 48 large matrices one time, transposed, into memory that it gets from the kernel. This uses 340 MB of RAM. The full startup, with the tokenizer, takes about a tenth of a second.

Step 2: Talking to the kernel without libc

Normally, a program calls open() or write() in libc, and libc asks the kernel. We skip libc and ask the kernel directly. On macOS for ARM, a system call works like this:

1. Put the arguments in registers x0 to x5.
2. Put the system call number in x16.
3. Run svc #0x80. The CPU jumps into the kernel.
4. The result comes back in x0. If it fails, the carry flag is set.

This is how we print a token:

mov     x0, #1          // file 1 is the terminal (stdout)
                                // x1 = address of the bytes, x2 = length
mov     x16, #4         // system call 4 is write
svc     #0x80           // ask the kernel

The whole program uses five system calls:

Memory without malloc. The program needs space for its working values and for a 75 MB cache (more about the cache in Step 6), and for the tokenizer tables. We declare all of it as zero-filled sections in the binary. The loader gives us the memory when the program starts. The only memory that we ask for is the one block for the transposed matrices.

One surprise. A fully static binary, with no libraries at all, is the purest form. But macOS on Apple Silicon stops such a program at launch with SIGKILL. So the binary must name libSystem for the loader. It imports no functions from it, and nm -u shows an empty list.

Step 3: Text to tokens

The model does not read letters. It reads token ids, from a vocabulary of 50,257 pieces of text. GPT-2 uses byte-pair encoding (BPE) to split text into these pieces. It works in two parts.

Part 1: split the text into chunks. GPT-2 uses a regular expression for this. We have no regex library, so we write its rules as plain branches. A chunk is one of these:

• an ending like 's, 're or 'll
• a run of letters, with one optional space in front
• a run of digits, with one optional space in front
• a run of other symbols, with one optional space in front
• a run of spaces

Part 2: merge. Each chunk starts as single bytes. GPT-2 has a list of 50,000 merges, in order of rank. At each step, we find the two neighbours with the lowest rank and join them. We stop when no neighbours can join. Try it:

␣ marks a space. GPT-2 keeps the space with the word after it, so “␣assembly” is one token, number 10474. Some words do not join completely. “␣tokenizer” ends as two tokens.

At startup, the program reads vocab.json and merges.txt and builds all the tokenizer tables itself. It decodes the \uXXXX escapes in the JSON and the UTF-8 in the merge list. To find a pair’s rank quickly, it builds a hash table with 131,072 slots. The key is the two token ids packed into one 32-bit number: (a << 16) | b. We multiply the key by a large odd constant and keep the top 17 bits. That number is the slot. If the slot holds a different key, we try the next slot.

Step 4: The matrix-vector product

This one function does almost all of the work. For each token, GPT-2 multiplies its 768 numbers by large weight matrices: about 124 million multiply-adds per token. Each output is one dot product:

out[i] = w[i][0]*x[0] + w[i][1]*x[1] + ... + w[i][767]*x[767]

A normal instruction multiplies one pair of numbers. The M1 also has NEON registers that hold four floats each. One fmla instruction does four multiply-adds at the same time. Press run to compare:

The real loop goes further. Each pass loads 32 weights and 32 inputs and runs eight fmla instructions into eight separate totals. With eight totals, no instruction must wait for the result of the one before it. This is the core of the loop:

loop:
    ld1     {v4.4s-v7.4s}, [x2], #64      // load 16 weights
    ld1     {v16.4s-v19.4s}, [x6], #64    // load 16 inputs
    ld1     {v24.4s-v27.4s}, [x2], #64    // 16 more weights
    ld1     {v28.4s-v31.4s}, [x6], #64    // 16 more inputs
    fmla    v0.4s, v4.4s, v16.4s          // 4 multiply-adds
    fmla    v1.4s, v5.4s, v17.4s          // 4 more
    ...                                   // 8 fmla in total
    subs    x7, x7, #32                   // 32 fewer to go
    b.ne    loop

The real limit is memory. Each token reads all 498 MB of weights one time. At 100 tokens per second, that is about 50 GB of reads per second. That is close to the most that one M1 core can pull from memory. So after this point, a faster loop gives very little. The speed of memory sets the speed of the model.

Step 5: Math without a math library

We also need a few functions that C programs get from libm:

• Square root is easy. ARM has an instruction for it: fsqrt. LayerNorm uses it.
• exp(x) has no instruction. Softmax needs it, and so does GELU.
• GELU uses tanh, and we can write tanh with exp: tanh(z) = 1 − 2 / (e2z + 1).

So everything depends on exp. We build it in three steps:

1. Write x = k·ln 2 + r, where k is a whole number and r is small. Then ex = 2k · er.
2. r is always between −0.35 and 0.35. For such a small number, a short polynomial gives er almost exactly: 1 + r + r²/2 + … + r⁶/720.
3. 2k needs no math at all. A float stores its power of two in 8 exponent bits. We write k + 127 into those bits, and the result is exactly 2k.

The bits in orange are the exponent field. Move the slider and watch k change them. The error stays below one part in ten million in most places, which is the precision of a 32-bit float.

Step 6: ATTENTION(with cache)

Attention lets each token look back at the tokens before it. For each of 12 heads, the new token makes a query q. Every earlier token has a key k and a value v:

1. Score each earlier token: q · k / 8.
2. Turn the scores into weights with softmax. The weights add up to 1.
3. The output is the weighted sum of the values.

The cache. The keys and values of earlier tokens never change. So we compute them one time and keep them: 12 layers × 1,024 positions × 768 numbers, for keys and for values. That is 75 MB, and it is the zero-filled memory from Step 2. Each new token adds one row and reads the rows before it.

A head looks at 64 numbers. NEON has 32 registers of 4 floats, so the query of one head fits in 16 registers. It stays there while we score all earlier tokens.

Step 7: Pick the next token

At the end, the model gives one score for each of the 50,257 tokens. There are two common ways to choose:

• Greedy: take the highest score. The output is the same on every run.
• Sampling: divide the scores by a temperature, apply softmax, and draw at random. A low temperature makes safe choices. A high temperature makes surprising ones.

The bars show the 10 most likely tokens, scaled so that they add up to 100%. Look at the top two. The scores for “Java” and “Python” differ by only 0.00005. That is smaller than the rounding noise of 32-bit floats, so the order of the additions decides the winner. PyTorch’s generate() and our assembly both pick Python.

For random numbers we use xorshift64*: three shifts, three XORs and one multiply. The seed comes from cntvct_el0, the CPU’s hardware timer. A program can read it directly, with no system call. We read the same timer to measure tokens per second.

Does it work?

A script compares the program with Hugging Face GPT-2 in PyTorch. Both use 32-bit floats and greedy choice, so the token ids must be identical.

And the speed, on one M1 core:

The result is 1,834 lines of source and 6,324 bytes of machine code. The weights are nearly 80,000 times larger than the code that runs them.

Limits

• macOS on Apple Silicon only. The system call numbers are specific to macOS.
• ASCII text is exact. For other text, the tokenizer treats every byte above 127 as a letter. This is usually right, but not always.
• One thread. More cores would help only a little, because memory is the limit.
• 1,024 tokens of context, the same as GPT-2.

Try it

You need an Apple Silicon Mac and the Xcode command line tools.

git clone https://github.com/ArneshBanerjee/gpt2.s
cd gpt2.s
make weights      # download the 3 original files, about 550 MB
make
./gpt2 -n 100 "The meaning of life is"

Credits

GPT-2 is the work of OpenAI. The weights come from the Hugging Face checkpoint. The idea to run GPT-2 somewhere unusual came from Shubham Patil’s GPT-2 in CSS.

• Code: github.com/ArneshBanerjee/gpt2.s

The merge steps and the token scores in this post come from the real GPT-2 vocabulary and model. The attention weights in the cache figure are an example.