spec⟶model

The way a language model works, roughly, is that you have a big pile of numbers, and you multiply a vector by some of the numbers, and then by some more of the numbers, and you do this a few dozen times, and then a word comes out. Then you do it again for the next word.

And you might think: well, if I know exactly which numbers are in the pile, and I know exactly which multiplications happen in which order, couldn't I just write that program down? Like, the actual program, the one that does those multiplications and no others? And the answer is yes, obviously, and also that is the whole idea of this book. It is a compiler. You give it a description of a model, and it gives you back a directory of kernels, a driver and a Makefile, and you type make and you get ./model, and ./model decodes.

I want to be careful here, because people get very excited about this sort of thing, and then they get disappointed. The compiler does not make your model smarter. It does not make the pile of numbers smaller.² What it does is remove everything between you and the multiplications that is not a multiplication. That turns out to be a lot.

1.1Five tools in a trench coat

The compiler is not one program. It is five programs standing on each other's shoulders, and each of them reads a text file and writes a text file. This is either a deeply principled architectural decision or what happens when you like grep a lot. It is both. You can stop the chain at any arrow: ./compile --emit-bir, --emit-tir or --emit-kir writes that stage and goes home. The front end, spec, reads your description and flattens it. The middle, tir, turns that into a graph of tensor operations and tidies it up. kir decides how that graph becomes kernels on one particular piece of hardware. And emit writes the kernels out as source code, which some other compiler, written by people with much larger budgets, finishes off.

Strictly, the coat holds more than five. emit is one binary, but it's really a set of back ends: one that writes C for a Mac, one that writes CUDA, one that writes HIP. Which one you get depends on a sixth input that isn't a program at all: the target, a small file naming the device (m2_max, ada, mi300x, …). kir reads it to decide the kernels; emit reads it to decide the language.

The nice thing about putting a text file between every stage is that when something goes wrong, there is a place to stand and look at it. “The output is wrong” becomes “the .tir is fine and the .kir is wrong,” which is a much smaller and more pleasant problem. You can also edit the files by hand. You should not do this. But you can, and knowing that you can is calming.³

1.2You describe the model; you do not program it

Here is what a model looks like going in. You say how big things are. You say what the layers are. You say “do this twice” and the front end does it twice. And then, importantly, you say where every single weight comes from in the checkpoint, because the compiler is not going to go rummaging through a 30-gigabyte file hoping something has the right name.

-- Moonlight-mini: same architecture as Moonlight-16B-A3B at exercise size.model Moonlight_Mini {  dim d = 128    dim L = 2    dim H = 4    dim V = 512    dim T = 256  const eps = 1e-5   const theta = 1e4   const scale = 2.446  embed tok : Embed(V, d)  spine PreNorm(d)  for i in 1..L {    layer attn{i} : MLA(d, H, dq, dc, dn, dr, dv; norm = RMSNorm(d), ...)    layer ffn{i}  : if i == 1 then GLU(d, Fd; norm = RMSNorm(d))                    else MoE(d, F, E, k, Es; norm = RMSNorm(d), act = silu, ...)  }  norm final : RMSNorm(d)  head lm    : Linear(d, V)}

You will notice that line 10 says the first feed-forward layer is a plain GLU and the rest are a mixture of experts. This is a real thing that real models do, and the language just lets you say it. The grammar for this part is short enough to print, so here it is:⁴

1.3Follow one normalization all the way down

One way to think about a compiler is that it is a machine for making things more specific. You start with “normalize this,” which is the kind of thing you say at a meeting, and you end with “for each of these 128 numbers, do this exact arithmetic,” which is the kind of thing a GPU will actually do. Each step in between is a reduction: a phrase in one language is replaced by phrases in a lower one, until nothing from the higher language is left. Here is line 13 of the listing, on its way down. These are the real lines out of the real files; I did not tidy them.

So: the front end replaces d with 128, so nobody downstream ever has to ask what d is. The middle end gives the result a name, s5, and makes the weight explicit. And the kernel stage takes “rmsnorm,” which is a word, and turns it into a reduce and a map, which are loops. Loops are what hardware wants. Hardware has never once wanted a word. If you took a compilers course you learned “reduce” as the thing a parser does on the way up. Here we are going down. Same word, opposite direction. Compiler people are like this.

Equation (1.1) is just the previous paragraph again, but with the kind of brackets that make it look like it was reviewed.

1.4The compiler will not guess

There is a certain kind of software that, when you give it something incomplete, fills in the gap with whatever seems plausible and carries on cheerfully. This is great until you find out, three weeks later, that your embedding table has been random numbers the whole time. This compiler does the other thing. If a weight has nowhere to come from, it tells you which one, on which line, and stops. Formally there are two cases, and exactly one of them applies:

The rule on the right is my favourite rule in the book. It is not an edge case. It is a feature, it has a test, and the test checks the exact wording of the error message, because a vague error message is just a guess with extra steps.

1.5Same seed, same tokens, every time

Here is a thing that is surprisingly hard to get from a GPU: the same answer twice.

Proof (sketch). Floating-point addition is not associative, which is the kind of fact that sounds pedantic until it ruins your week. If you add up a thousand numbers in a different order, you can get a slightly different total. On a GPU, the order is often “whichever thread got there first,” which is to say, it varies. So the compiler simply does not let it vary. There are no floating-point atomics in the emitted kernels, and every sum is added up in an order that was decided at compile time. Every output is then a function of M, B and s and nothing else. To check, run it twice and cmp the outputs. It says nothing. That is the proof.⁵ ∎

1.6Where it runs

The compiler is written in C11 and has no dependencies beyond the C library, which these days is almost an act of protest. To get it running, you hand it to the C compiler you already have, and a native compiler comes out the other side. Textbooks draw this as one T sitting in the arm of another, which looks like a small construction site and is in fact the entire story of how compilers have been built for sixty years.

Then you pick where the model should run, which is one argument. The back end picks what language the kernels are written in. The target picks the specific chip, because “a GPU” is about as useful a description as “a car” when you are trying to decide how many threads to launch.

The models that ship with a description include GPT-2 (124M), Mixtral 8×7B, Moonlight-16B-A3B and some of the GLM family, each with a miniature version you can compile in a few seconds and break on purpose. Weights can stay in bf16, or be squeezed into 4 or 8 bits in groups of 128, if you are the sort of person who looks at a memory bandwidth number and takes it personally.⁶

1.7Exercises

An asterisk means you will need a checkpoint. Two asterisks means you will need a free afternoon and a tolerance for reading C.

make && make -j test

./compile specs/moonlight_mini.spec && ./out/model

./compile specs/moonlight_mini.spec --emit-kir

./compile specs/gpt2_124m.spec --ckpt path/to/gpt2 --vocab path/to/vocab

1.8Things happen

Bake can now read a data checkpoint, and bake export learned to cut fused concatenations back into their operands. Sequence RoPE arrived for multi-head latent attention. Part three of the long, quiet campaign to remove t0. Somebody, somewhere, is asking whether this could be done in Python. It could not.⁷

In one sentence. Parse the model, evaluate every number, expand every loop, and hand back a flat list where nothing is left to work out.

You know what a front end does. You have written one, probably in a course, probably for a language with while loops and a sad little type checker. This one is easier and stricter. The language has no runtime: every dim is a constant, every for has constant bounds, and every if is decidable at compile time. So the front end just runs the whole program, and what's left is the trace.

That trace is the .bir, the block IR: one op per line, in six kinds (dim, const, block, share, step, bind). It is deliberately boring. Boring is what you want from the file every other stage will read.

The loop is more fun. Listing 1.1 says “for i in 1..L, the feed-forward is a GLU if i = 1, otherwise a MoE.” After the front end has run it, the if has been decided, the {i} has been spliced into every name, and the bind block's {i-1} prefix has turned into real checkpoint keys, slices and all:

That last line is worth a second look. The checkpoint stores every head's query projection in one tensor, 24 rows per head: 16 for the content part, 8 for the rotary part. The spec says so once, with a slice, and the front end writes it out per layer. Nobody downstream ever does index arithmetic on checkpoint keys again.

Notice what the front end does not do: optimize. It has no opinions. It doesn't know what a GPU is. That's on purpose, and the next chapter is about why.

./compile specs/moonlight_mini.spec --emit-bir

In one sentence. Turn each block into tensor operations from a closed set of 30 builtins, then rewrite the graph wherever the arithmetic says it's cheaper.

A layer like MLA(…) is a word. The middle end replaces it with a small graph of builtins: einsum, rmsnorm, rope, sdpa and friends. There are exactly 30 of them, and a _Static_assert makes sure there are never more than 32, because each one gets a bit in a mask.Is a 32-bit mask a reason to keep a language small? It is now.

Then it optimizes. Here is the rule that keeps the whole compiler sane, and it applies to every stage:

3.1Worked example: absorbing MLA

The poster child is multi-head latent attention, the attention in DeepSeek-style models and in Moonlight. The pitch of MLA is that each token is squeezed into one small latent vector c, and every head's key and value are recovered from it by an up-projection:

The lowering writes exactly that. It's the obvious program: build every head's k and v for every token, cache them, run sdpa. It is also the expensive one, because the cache now holds the expanded keys and values and not the small latent they came from. You compressed the token and then uncompressed it straight into the most expensive memory you have.

The fix is two lines of linear algebra that you could do on a napkin. Matrix multiplication is associative, so the up-projections can move to the other side of the dot product and the other side of the sum:

Read left to right: fold W_uk into the query once per step, attend over the latents directly, and apply W_uv once to the weighted sum. Nothing per head is ever cached. That's the rewrite called absorb_mla, and Figure 3.1 is the whole idea.

Here is what tir actually writes for layer 1. Note the state lines: those are the caches, and they are now 32 and 8 wide. mla_q takes Wuk; mla_attn takes Wuv. Equation (3.2), in the wild:

On a toy model that's a 4× smaller cache. On the real thing it is the whole reason MLA exists:

So why isn't absorption just always on? Because it isn't always cheaper. Folding W_uk into the query costs flops on every step. If the context is short and the heads are few, the plain form can win. So the rewrite is priced: bytes cached per position plus flops, and it fires only when the sum goes down. That is Principle 3.1 doing its job. The translation stays dumb; the optimizer makes the call, out loud, where you can turn it off.Full context = T positions × L layers: 256 × 2 for the mini, 8 192 × 27 for the 16B. Plain counts each head's key at dn + dr and value at dv.

Want the plain version back? tir --no-opt, or --opt-disable=absorb_mla to switch off just this rule. That's also the oracle: optimized and unoptimized have to agree. And notice what tir did not price: kernel launches. It can't; launches aren't in its language. So it doesn't get a vote on fusion. That's the next chapter's job.

build/bin/tir --no-opt specs/moonlight_mini.bir

In one sentence. Write every tensor op as an explicit loop nest, then decide which loops share a kernel, for exactly one device.

Here the builtins disappear. Every node becomes one of five loop-nest forms, map, reduce, contract, gather and scatter, each with an explicit index map for every array it touches. You saw a reduce and a map back in §1.3. Here is a contract, which is a matrix-vector product admitting what it is:

Now kir has the one thing tir didn't: it can see launches, barriers and memory traffic. So kernel boundaries get decided here, once, for every back end. That includes the fun one, horizontal fusion. When several matvecs read the same input, kir concatenates their weights into one derived tensor (that's the hf0 you will find in the file) and does one big matvec in place of several small ones.The concatenated weight never exists in the checkpoint. bake builds it. See Chapter 6.

4.1Worked example: two projections, one matvec

Go back to MLA for a second. Every step, the normalized input x is multiplied by Wdkv (128 → 32, the latent) and, separately, by Wkr (128 → 8, the rope key). Same x, two skinny matvecs, two launches, and x read from memory twice. kir sees that and writes this instead:

The GLU in layer 1 gets the same treatment, and it's a bigger win: gate and up projection, 256 columns each, become hf1.W : [128, 512] = concat(ffn1.Wg, ffn1.Wu). Note the = where you'd expect a <-. A <- means “load this from the checkpoint”; an = means “compute this from other weights.” Somebody has to do that computing, and it's not the GPU at runtime.

Raw kir, before any of this scheduling, is still runnable. That's the oracle again: whatever the scheduler does, raw kir is the answer it has to match.

build/bin/kir --check specs/moonlight_mini.kir

In one sentence. Print the kernels as C, CUDA or HIP, add a driver and a Makefile, and let the vendor's compiler do the last mile.

By the time a program reaches emit, every interesting decision has been made. Which ops share a kernel? kir said. Which form of attention? tir said. emit's job is to be a very good typist. It writes kernels.*, main.*, bench.*, a manifest.txt and a Makefile, and then it's done.

“Typist” undersells it a little. The back end still owns the instruction level: tile shapes, vector widths, microkernels. On CUDA a 4-bit weight is read 32 nibbles per 16-byte load, one warp per output. On ROCm it's one wave per R outputs, with the dtype as a template constant. What a back end may not do is scan the graph for its own fusions. That would be the same optimization in two stages, and Principle 3.1 calls that a bug.

./compile specs/moonlight_mini.spec --backend mac && make -C out bench

In one sentence. Read a Hugging Face checkpoint, follow each tensor's recipe from the .tir, and write one flat file in exactly the order the program will read it.

bake is optional. Without it, ./model fills its weights deterministically and prints token ids, which is all you need to test a compiler. With it, you get words. Weirdly, it reads the .tir and not the .kir, because that's where each tensor's recipe lives: which checkpoint key, which slice, which dtype.

Three rules do most of the work. Order is the manifest's, not the file's. Orientation is the declared shape: if the source is stored the other way round, bake transposes it. Derived tensors, like kir's horizontally fused weights, are computed in fp32 from their operands' real values, never from already-quantized bytes, and then packed like anything else. Quantized weights go in contraction-major, so a dot product walks one contiguous run of nibbles and one scale per 128.

6.1Worked example: packing a 4-bit weight

Say a spec binds a weight with dtype q4g128. bake quantizes it to symmetric int4 in groups of 128 consecutive elements, each group with its own fp32 scale. That's half a byte per weight plus 4 bytes per 128 of them, so about 4.25 bits a weight, down from 16 in bf16.

The interesting decision is which 128 elements make a group. The kernel will compute y_j = Σ_i x_iW_ij, walking down a column. So bake lays each column out as one contiguous run, element (i, j) at offset j*n_in + i, and groups run down the column too. One 16-byte load is then 32 consecutive weights of the same dot product under one scale, which is exactly what a GPU lane wants to eat.

Proof. A square matrix looks the same either way round. Guessing wrong produces a model that runs, prints plausible numbers and is quietly wrong, which is the worst kind of wrong. See §1.4. ∎

./compile specs/gpt2_124m.spec --ckpt path/to/gpt2 --vocab path/to/vocab

In one sentence. Fifty-nine architectures written down as specs, and a ratchet that records how far through the compiler each one gets.

Writing a model down is cheap. You read its config.json and its modeling file, and an afternoon later you have a page of spec. Compiling it is the expensive part, because every architecture lab has had at least one idea that nobody else has had, and the compiler has to learn each one before it will stop complaining. So we write the specs first, sometimes well ahead of the compiler, and keep score.

The scorekeeping is a ratchet. Every make runs each spec as far down the pipeline as it will go: does it parse, does it flatten into .bir, does tir accept it, does it schedule into kernels for a target. tests/graph_coverage.txt records the furthest stage per spec, and the test fails if any spec falls short of its row. A spec can climb but never slip, which is the kind of guarantee you want from a ratchet and almost never get from a roadmap.

“Supported,” honestly, means the right-hand end of the table: compiles to a .target. The ones in specs/ (GPT-2, Moonlight, Mixtral, GLM-4.7-Flash and GLM-5.3) go further still. They have byte-exact fixtures at every stage, so a change that moves one of their files moves a diff you have to explain. A dash means the spec was written ahead of the compiler: it uses something the front end does not parse yet. That is a to-do list, not an embarrassment, although we would not object if it got shorter.

Two things stand out. First, the dense transformer is no longer the typical case: two thirds of these architectures route each token to a few of many experts, compress their keys, or replace attention with something linear for most of their layers. Second, the newer the idea, the shorter the bar: nearly half of the latent-attention models and the hybrids still stop at a dash, against about a quarter of the dense ones. The compiler learns ideas in roughly the order the labs had them, which is either reassuring or a little sad.

./compile models/llama3-8b.spec --emit-kir --target targets/m2_max.target

Bibliography