The inspiration for the title came from the popular online course, From Nand to Tetris: Building a Modern Computer From First Principles.
Background
So, I make challenges for NUS Greyhats, a cybersecurity interest group in Singapore. Every year around June, we run our flagship academic CTF, Grey Cat the Flag. As far as I can remember, this CTF (2022 to be exact, when I was still in high school) was what got me absolutely hooked on the scene; thus I felt it was a great honour to be able to contribute back to the community, in this case through challenge making.
The Qualifiers for the CTF took place around late May. In line with fully online CTFs that took place at this period of time (Q2 2026, think Opus 4.7/4.8), the Conversation was still mostly about LLMs’ disruption of the CTF scene, along with a hint of beauty within the authors’ struggle in writing LLM-resistant, human-friendly challenges (that was when AGI hadn’t arrived yet, anyway). My prime contender was 3D Maze, which, despite receiving some valid criticism, I was extremely proud of, almost on the level of my Magnum Opus.
So of course, I felt the imperative to write a challenge that could one-up this experience, for the Finals. Finals was on-site, which allowed us to ban LLM usage, which in turn allowed us to focus primarily on the human solving experience. Lying in bed as I fell asleep across multiple days, I tried to come up with a challenge that would exemplify some aspect of human intelligence for solving a Reverse Engineering challenge; eventually I came across the idea of abusing the heck out of abstraction. Deliberately looking for something dramatic and thematic for the Finals and for the surrounding mood, I thought it would be fun to end the challenge off with some emulation of an LLM, and to require the participants to repeatedly lift the challenge in order to uncover it. I instantly fell in love with the idea upon landing on it, and all subsequent ideas did not give me as much of a dopamine hit as that one. So after a while I just cut brainstorming and went straight towards ideating the implementation.
The agenda I had in mind was:
Lambda Calculus / Turing Machine
-> Subroutines
-> Register VM
-> Procedures
-> Transformer ArchitectureUnderstandably, Lambda Calculus and Turing Machines are comfortably interconvertible with each other, so I initially had an additional layer of Lambda Calculus -> Turing Machine. But after realizing the challenge was complex enough and interpreting Lambda Calculus is probably worth its own challenge by itself, I decided to drop it. From the get-go I expected most of the difficulty to lie within lifting the Turing Machine into a Register VM; the remaining stages would probably take at most an hour for participants of expected calibre.
Well true enough, said remaining stages took me perhaps just a few days (spread out effort) to implement, with no prior knowledge of transformers; same was true for perhaps 80% of the Turing machine construction. The remaining 20% however ended up just ballooning and ballooning and ballooning that I even spent the first few hours of the challenge desperately pumping lines and lines of the remaining transition rules out thinking I was basically there… except I wasn’t and I eventually gave up and left the challenge unreleased.
Well, thankfully the participants seemed to struggle at a little and enjoy quite a lot my other two “minion” challenges, so the whole challenge making was not a waste of effort. No teams ended up fully clearing Rev, mainly due to another boss challenge, mad props to the author. So even if I had managed to pump the challenge out, it would have been unpolished and underappreciated anyway. Overall, the event ended up being quite a satisfying success.
What I felt was left unspoken was the lingering effect of LLMs’ disruption, even on high calibre players in a no-LLM setting. Personally I have gotten the feeling that humans have gotten slower at deep, purely technical problems; this could either be attributed to us gradually drifting our specializations (subconsciously or not) towards more economically rewarding metrics, or it could just be that the baseline unit of productivity has risen so fast around me that I have acquired the illusion of such. Maybe I should write a separate post for itself, but for now I have digressed for far too long.
Anyway, over the past few weeks, I gradually squeezed out pockets of time from my already busy schedule to finish up v1 of the challenge and to dedicate a writeup to it. The code is still visibly messy, but what I felt was more interesting were the concepts and ideas behind it.
I already have big ideas for a refactor, so expect v2 to be significantly more ergonomic. Well, if it comes out anytime soon anyway.
Introduction
This writeup will be split into two parts:
- (This post) Implementing a simple transformer from an even simpler bytecode VM.
- (Next post) Implementing said bytecode VM from a Turing Machine.
This first part sounds reasonably simple, or at least imaginable, for a reversing challenge, hence we are tackling it first. Essentially, if you could understand the architecture well enough to implement it in the language of your choice (without abusing fancy features, probably), then you could probably translate it to C, and you could probably translate that to Assembly.
While simple, our bytecode VM of course needs to be expressive enough: control flow instructions need to be present where appropriate. Just a handful of them covers most of the common computations, anyway. Instead, our VM would be simple in terms of the operations on data, because imaginably that would be a bigger headache to implement on our future Turing Machine.
The Transformer
Going by The Memes, LLMs are just linear algebra anyway. And everyone knows linear algebra just involves addition and multiplication. Sounds easy!
Well, turns out things were a little more involved than I thought. But to realize that, I needed to study the architecture anyway. I’m still not an expert on this topic, and I only learnt enough to emulate something functional, so I will not be on some high horse here and larp and yap. I just found these materials insanely helpful for acquiring a satisfactory picture:
- 3Blue1Brown video series, for a comprehensive outlook.
- Obligatory Wikipedia article, for high-level skimming.
- This visualizer, especially around the QKV attention mechanism as well as the matrix multiplication convention.
- That which shall not be named, which also brought me to this PyTorch tutorial.
For the purposes of a simple challenge, it would probably be overkill to literally adapt a whole LLM, no matter how tiny the model is. Instead, the idea is to specifically emulate the autoregressive aspect while keeping the spirit of the architecture intact as much as possible. I am picturing the participant reaching the climax of the solve, running their custom VM which sequentially spits out the final flag piece by piece:
$ python3 bootstrap-model.py
b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00grey' -> b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00rey{'
b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00grey{' -> b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00rey{h'
b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00grey{h' -> b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00rey{hu'
b'\x00\x00\x00\x00\x00\x00\x00\x00\x00grey{hu' -> b'\x00\x00\x00\x00\x00\x00\x00\x00\x00rey{hum'
b'\x00\x00\x00\x00\x00\x00\x00\x00grey{hum' -> b'\x00\x00\x00\x00\x00\x00\x00\x00rey{huma'
b'\x00\x00\x00\x00\x00\x00\x00grey{huma' -> b'\x00\x00\x00\x00\x00\x00\x00rey{human'
b'\x00\x00\x00\x00\x00\x00grey{human' -> b'\x00\x00\x00\x00\x00\x00rey{humani'
b'\x00\x00\x00\x00\x00grey{humani' -> b'\x00\x00\x00\x00\x00rey{humanit'
b'\x00\x00\x00\x00grey{humanit' -> b'\x00\x00\x00\x00rey{humanity'
b'\x00\x00\x00grey{humanity' -> b'\x00\x00\x00rey{humanity_'
b'\x00\x00grey{humanity_' -> b'\x00\x00rey{humanity_c'
b'\x00grey{humanity_c' -> b'\x00rey{humanity_cl'
b'grey{humanity_cl' -> b'rey{humanity_cli'
b'rey{humanity_cli' -> b'ey{humanity_clim'
b'ey{humanity_clim' -> b' {humanity_climb'
b'y{humanity_climb' -> b'{humanity_climbs'
b'{humanity_climbs' -> b'humanity_climbs_'
b'humanity_climbs_' -> b'emanity_climbs_t'
b'umanity_climbs_t' -> b'banity_climbs_th'
b'manity_climbs_th' -> b'atity_climbs_the'
b'anity_climbs_the' -> b'tity_climbs_the_'
b'nity_climbs_the_' -> b'(ty_climbs_the_l'
b'ity_climbs_the_l' -> b'cy_climbs_the_la'
b'ty_climbs_the_la' -> b'h_climbs_the_lad'
b'y_climbs_the_lad' -> b'{climbs_the_ladd'
b'_climbs_the_ladd' -> b'climbs_the_ladde'
b'climbs_the_ladde' -> b'limbs_the_ladder'
b'limbs_the_ladder' -> b'icbs_the_ladeer}'
b'imbs_the_ladder}' -> b'cbs_the_ladder}\x00'
grey{humanity_climbs_the_ladder}Above: Output of the toy LM simulated using Python. Left of arrow is the current context sequence, right of arrow is the model prediction. At each step, only the final token (in this case, a character) is appended to the context for the following step.
The goal here is clearly to conceal and then subsequently reliably retrieve a string of characters called the flag. To simulate this process, the design choice here would be to use individual characters as our tokens.
Architecture
For a simplistic, high-level overview, we can view our model architecture as being broken down into 4 parts:
---> [Encoder] ---> [Attention] ---> [FFN/MLP] ---> [Decoder] --->- Our input sequence (“context”; here a string of characters) gets parsed by the encoder into an array of vectors (so, a matrix).
- This array of tokens gets passed into the attention mechanism to attend to each other, that is, get their values adjusted based on other tokens.
- This array of modified vectors gets passed into the MLP mechanism. Note that this step is per vector, that is, they no longer affect each other (that solely happens in the attention step).
- This final array gets converted by the decoder, each into a probability distribution of their corresponding next tokens.
One could think of the MLP step as a function independently applied to each vector
where
On real systems, there are many more steps involved, including (but not limited to) having multiple Attention -> MLP passes, hence every single vector in the array needs to get processed through the MLP step. Otherwise, solely for next-token generation, for our simplified model with only one pass, one could realize that only the last vector needs to be processed through the MLP step. The rest of the output tokens are discarded at the end, anyway.
From what I know, another benefit of running the MLP step for every vector is to improve the training process, where we get even more instances to compute the training loss and apply to our model, rather than just one instance per sample of (context, output token).
Otherwise, for this challenge, doing so also allows us to reuse existing matrix multiplication procedures, instead of having to specifically craft a matrix-vector multiplication; ultimately this makes our code cleaner. It also has the side effect of slightly increasing obscurity for the participant.
Altogether, our toy transformer shall look like this:
import torch
import torch.nn as nn
CONTEXT_SZ = 16
VOCAB_SZ = 256
MODEL_DIM = 16
FEATURE_DIM = 16
FNN_DIM = 16
class LLM(nn.Module):
def __init__(self):
super().__init__()
self.tok_embed = nn.Embedding(VOCAB_SZ, MODEL_DIM, padding_idx=0)
self.pos_embed = nn.Embedding(CONTEXT_SZ, MODEL_DIM)
self.q = nn.Linear(MODEL_DIM, FEATURE_DIM, bias=False)
self.k = nn.Linear(MODEL_DIM, FEATURE_DIM, bias=False)
self.v = nn.Linear(MODEL_DIM, FEATURE_DIM, bias=False)
self.attn_out = nn.Linear(FEATURE_DIM, MODEL_DIM, bias=False)
self.fnn = nn.Sequential(
nn.Linear(MODEL_DIM, FNN_DIM),
nn.ReLU(),
nn.Linear(FNN_DIM, MODEL_DIM)
)
self.unembed = nn.Linear(MODEL_DIM, VOCAB_SZ, bias=False)
def forward(self, inp):
_, T = inp.shape
x = self.tok_embed(inp) + self.pos_embed(torch.arange(T)[None, :])
qk = (self.q(x) @ self.k(x).transpose(-2, -1)) / (FEATURE_DIM**0.5)
mask = torch.triu(torch.full((T, T), -torch.inf), diagonal=1)
x = x + self.attn_out(torch.softmax(qk + mask, dim=-1) @ self.v(x))
x = x + self.fnn(x)
return self.unembed(x)While crafting the model, the major tension is of course trying to balance model size (for challenge ergonomics) and performance (in this case, whether it overfits by too easily spitting out the flag, or it underfits by not spitting out the flag even with the correct prefix). Ironically enough, we specifically want to avoid the probabilistic aspect of LLMs; that is done by always picking the token with the highest probability as the canonical next token, also known as having a temperature of 0.
The numbers were selected through some trial-and-error, but essentially the idea is that I just kept decreasing the constants (and also cutting out redundant mechanisms (not elaborated here) along the way) while maintaining zero or extremely low training loss (as long as the flag sequence itself is memorized). Overfitting on the flag is hopefully mitigated by also training on “junk text” (so that hopefully fewer sequences converge on the flag).
Another thing to consider was that to be able to comfortably reuse the same procedure, we would of course want the matrix representations of each layer to be identical in shape to each other as much as possible. 16x16 matrices happened to land on the sweet spot for me, hence also the context window of 16 characters.
Training
The training script looks like this:
vs = (
b'Show this to the challenge author!\0',
b'grey{humanity_climbs_the_ladder}\0',
b'The quick brown fox jumps over the lazy dog.\0',
b'public static void main(String[] args) {}\0',
)
vs = tuple(map(lambda x: b'\0'*(CONTEXT_SZ-1) + x, vs))
batch_x = []
batch_y = []
for v in vs:
for end in range(CONTEXT_SZ, len(v)):
batch_x.append(list(v[end-CONTEXT_SZ:end]))
batch_y.append(list(v[end-CONTEXT_SZ+1:end+1]))
batch_x, batch_y = torch.tensor(batch_x), torch.tensor(batch_y)
print(batch_x)
print(batch_x.shape)
print(batch_y)
print(batch_y.shape)
model = LLM()
loss = nn.CrossEntropyLoss()
opt = torch.optim.AdamW(model.parameters(), lr=3e-3, weight_decay=0.0)
def epoch():
opt.zero_grad()
y_pred = model(batch_x)
l = loss(y_pred.transpose(1, 2), batch_y)
l.backward()
opt.step()
acc = (y_pred.argmax(dim=-1) == batch_y).float()[:, -1].mean()
return l.item(), acc.item()
import numpy as np
def save():
for k, v in model.state_dict().items():
np.save(f'weights/{k.replace(".", "_")}.npy', v.detach().cpu().numpy())
for i in range(1000):
print(f'{i:03} / 1000: {epoch()}')This part is more-or-less hacked together from random snippets including from the PyTorch tutorial. Apparently which way to pad tokens (left or right) for a partially filled context is a design decision as well; I went with left padding because that was supposedly simpler both conceptually and implementation-wise.
Notice the four training strings I used. The first two need to be completely memorized; the idea for the challenge is to have the challenge LLM output the first string by default, then following the instructions, the author (i.e., me) will come over and let the participant know to replace the seed with the flag header (i.e., grey).
The script was run in interactive mode (python3 -i), so when I felt the training loss was satisfactory I just interrupted the script and invoked save() manually. Could have had a more sophisticated mechanism, but it was meant to be quick and dirty anyway.
Code Lowering
Of course, my little Turing Machine is not going to be able to import an arbitrary torch library. We still need one more step, which is to extract just the essence of the model workings, and re-materialize it on a much simpler interface. For my sanity, I followed what I laid out at the start, and gradually lowered the program to (raw) Python, C, then (pseudo-)Assembly.
Here, it is already good to start keeping some low-level details in mind. For example, to what lengths am I willing to implement floating-point operations? Of course, for my sanity, more complex operations will not exist as machine instructions, but rather as entire procedures.
Very quickly, the Python version looks something like this:
import os
import numpy as np
CONTEXT_SZ = 16
VOCAB_SZ = 256
MODEL_DIM = 16
FEATURE_DIM = 16
FNN_DIM = 16
class LLM():
def __init__(self):
# just a helper function
def rec_map(f, xs):
return [
rec_map(f, x) if isinstance(x, list) else f(x)
for x in xs
]
# loading the weights
for fn in os.listdir('weights'):
if not fn.endswith('.npy'): continue
x = fn[:-4]
if x.endswith('_weight'):
x = x[:-len('_weight')]
setattr(self, x, rec_map(
lambda _x: Float(_x),
np.load(f'weights/{fn}').tolist()
))
def run(self, seq):
assert len(seq) == CONTEXT_SZ
x = list(seq)
# embedding
...
# attention
...
# fnn
...
# unembedding
out = []
...
return bytes(out)
# FLAG = b'grey{humanity_climbs_the_ladder}'
FLAG = b'grey'.rjust(CONTEXT_SZ, b'\0')
a = LLM()
while True:
print(FLAG[-CONTEXT_SZ:], end=' -> ')
x = a.run(FLAG[-CONTEXT_SZ:])
print(x)
if x[-1] == 0:
break
FLAG += x[-1:]
print(FLAG.strip().decode())BTW, the full version of this script (bootstrap-model.py) was what produced the earlier “intended solve” demo shown at the start of the Transformer section.
As we get to the C version, we start “dressing up” the code to flesh out the participant-facing UI. Thus, the corresponding C version looks something like this:
#include ...
#define DIM 16
#define SQRT_FEAT_DIM 4
#define CONTEXT_SZ DIM
#define VOCAB_SZ 256
#define MODEL_DIM DIM
#define FEATURE_DIM DIM
#define FNN_DIM DIM
struct weights {
...
};
_Static_assert(sizeof(struct weights) == 40064, "unaligned");
struct weights M;
void setup(void) {
int fd = open("weights.bin", O_RDONLY);
if (fd == -1) exit(1);
read(fd, &M, 40064);
}
char run(char *seq) {
float a[DIM][DIM];
float b[DIM][DIM];
float c[DIM][DIM];
float d[DIM][DIM];
/* embed */
embed(a, seq);
...
/* attn */
...
/* fnn */
...
/* unembed */
return extract_vocab(a[CONTEXT_SZ-1]);
}
struct __attribute__((packed)) {
char head[CONTEXT_SZ];
char tail[48];
} buf = {
.head = "\0\0\0\0\0\0\0\0\0\0\0\0Show",
.tail = { 0 }
};
int main(void) {
setup();
char *head_ptr = buf.head;
char *tail_ptr = buf.tail;
for (int i = 0; i < 48; ++i) {
*tail_ptr = run(head_ptr++);
if (!*tail_ptr++) break;
}
puts(buf.tail-4);
return 0;
}The idea is simple:
- Maintain a buffer of
CONTEXT_SZcharacters. The initial seed is left-padded with padding tokens. - In a loop, send the buffer through the LLM churn machine.
- Once we detect an end-of-sequence token (in this case the null byte), we stop.
- Otherwise, append the newly generated token to the existing sequence, and shift the context window correspondingly.
- For this implementation specifically, we limit the program to generate only 48 new tokens.
Just realized, I guess I have a potential OOB infoleak from the missing null-termination? LOL
Now, let’s look at each step in more detail:
def forward(self, inp):
_, T = inp.shape
x = self.tok_embed(inp) + self.pos_embed(torch.arange(T)[None, :])
qk = (self.q(x) @ self.k(x).transpose(-2, -1)) / (FEATURE_DIM**0.5)
mask = torch.triu(torch.full((T, T), -torch.inf), diagonal=1)
x = x + self.attn_out(torch.softmax(qk + mask, dim=-1) @ self.v(x))
x = x + self.fnn(x)
return self.unembed(x)Embedding
Corresponding PyTorch code:
self.tok_embed = nn.Embedding(VOCAB_SZ, MODEL_DIM, padding_idx=0)
self.pos_embed = nn.Embedding(CONTEXT_SZ, MODEL_DIM)
# ...
def forward(self, inp):
_, T = inp.shape
x = self.tok_embed(inp) + self.pos_embed(torch.arange(T)[None, :])
# ...One convention I had to get used to while understanding the mechanism was the direction of matrix multiplication. Consider a typical matrix-vector multiplication
One might expect a transformation over an array of data vectors to be represented as such:
But nope, data vectors are transposed into rows instead, and the transformation matrix correspondingly gets transposed and applied from the right:
Anyway. For our embedding step, per standard literature, each of our tokens gets converted into its corresponding vector (ready for the rollercoaster of transformations!) through a combination (simple addition) of two components:
- Token embedding, where each token literally maps to one “default” vector value.
- Positional embedding, in order to also encode the positional information of the token (i.e., where the token is within the context).
As code:
# embedding
# x is a list of bytes (our tokens), left-padded with zeros
for i, k in enumerate(x):
x[i] = [0 for _ in range(MODEL_DIM)]
for j in range(MODEL_DIM):
x[i][j] = self.tok_embed[k][j] + self.pos_embed[i][j]As literal as it gets!
One fun simplification is to notice how the same values for pos_embed always get written to the input matrix regardless of x. That is, we simply need to construct the intermediate matrix from mapping just the token embedding vectors, then one-shot perform matrix addition between this intermediate matrix and the full positional embedding matrix.
C representation below:
/* embedding */
// seq is an array of bytes, left-padded with zeros
// a is an uninitialized scratch matrix
embed(a, seq);
add(a, M.pos_embed);
/* definitions */
#define DIM 16
#define CONTEXT_SZ DIM
void embed(float x[][DIM], char *seq) {
for (int i = 0; i < CONTEXT_SZ; ++i) {
int k = (unsigned char)seq[i];
// exact same shenanigans as the Python counterpart
for (int j = 0; j < DIM; ++j) {
x[i][j] = M.tok_embed[k][j];
}
}
}Attention
If you haven’t yet browsed through any of the previously listed explainers, here is an extremely brief high-level overview (think architectural motivation), the way I understand it.
Essentially, the goal of this step is to enrich each vector with additional information that only arises because of the presence of some other token in the current context. Recall that after this step, in the MLP step the vectors do not interact with each other at all, so all interactions must arise from this step.
In the pure LLM sense, even in linguistics, a “word” can have multiple meanings based on context, hence by itself we can only attach some abstract idea to it. Only when certain other “key words” are present can we have a more concrete, specialized idea of what our original word meant. One random example:
- “Kingdom” — fantasy, monarchy, Europe, knights, medieval?
- United “Kingdom” — country, state, geopolitics, Britain, constitutional monarchy
To produce such a cool effect using Matrix Magic, the Attention Mechanism uses each vector
- To score how much
should affect the meaning of some other vector . For example, a “cloverleaf” later down the sentence probably shouldn’t be too affected by the earlier “United”. - If it does affect the meaning of
, which direction to nudge towards. In a sense, this forms the axis of -ness.
Edit: The author later realizes that the above UK example is probably pretty bad for justifying the “axis of
Specifically, for a pair of vectors
- dot the query representation of
, against the key representation of , , to compute a scalar; - add to
the value representation of , scaled accordingly.
It thus becomes a textbook exercise in a linear algebra course to derive the formula
where the
Three important caveats —
One; especially for next-token generation, we cannot reasonably see in the future, so it does not make sense for a token to attend to some other token at a later position. As such, we mask off the strict upper triangle of the
Two; right now each dot product can reasonably return any arbitrary real value (like a score). The Literature generally applies some normalization to have a more well-rounded effect on which tokens and to what extent the current token should focus its attention on. The details are currently beyond my knowledge.
Three; according to 3Blue1Brown here, it seems common to split the value calculation into two matrices: the first to denote the actual representations in a lower-rank space (
If you’re curious, our full attention step now looks like this:
But at this point, it’s probably clearer to look at the implementation:
# attention
pre_x = [[value for value in row] for row in x]
Q = matmul(x, transpose(self.q))
K = matmul(x, transpose(self.k))
V = matmul(x, transpose(self.v))
cur = matmul(Q, transpose(K)) # QK^T
cur = masked_softmax(cur) # softmax(...)
cur = matmul(matmul(cur, V), transpose(self.attn_out)) # V_down V_up
for i in range(len(x)):
for j in range(len(x[0])):
x[i][j] += cur[i][j]Similarly in C. We expect to be a bit more register-starved, so the goal is to be able to reuse as many variables as possible.
/* attention */
// a, b, c, d are scratch matrices (2D arrays)
// a contains the vectors from the embedding step (i.e., X)
// matmul_t multiplies by the transpose of the third argument
matmul_t(b, a, M.q); // Q = X @ q^T
matmul_t(c, a, M.k); // K = X @ k^T
matmul_t(d, b, c); // attn = Q @ K^T
matmul_t(b, a, M.v); // Vd = X @ v^T
masked_softmax(d); // (attn)
matmul(c, d, b); // attn = attn @ Vd
matmul_t(b, c, M.attn_out); // _X = attn @ Vu
add(a, b); // X += _XWhat a surprise, they consist of the same steps. We are still one abstraction level too high, so we still need to iron out the specifics of each procedure. For brevity I will just show the C version, since they are roughly the same anyway:
/* matrix operations all assume a fixed size of DIM * DIM */
// matrix addition; straightforwardly element-wise
void add(float a[][DIM], float b[][DIM]) {
for (int i = 0; i < DIM; ++i) {
for (int j = 0; j < DIM; ++j) {
a[i][j] += b[i][j];
}
}
}
// matrix multiplication
// matmul_t multiplies `a` by the _transpose_ of `b`
void matmul_t(float r[][DIM], float a[][DIM], float b[][DIM]) {
for (int i = 0; i < DIM; ++i) {
for (int j = 0; j < DIM; ++j) {
r[i][j] = 0;
for (int k = 0; k < DIM; ++k) {
// note flipped access for b (b[j][k] instead of b[k][j])
r[i][j] += a[i][k] * b[j][k];
}
}
}
}
void matmul(float r[][DIM], float a[][DIM], float b[][DIM]) {
for (int i = 0; i < DIM; ++i) {
for (int j = 0; j < DIM; ++j) {
r[i][j] = 0;
for (int k = 0; k < DIM; ++k) {
// standard i-k k-j construct
r[i][j] += a[i][k] * b[k][j];
}
}
}
}
// sqrt(16) == 4
#define SQRT_FEAT_DIM 4
// performs 3 tasks:
// 1. masks off the strict upper triangle of the matrix
// 2. divides each row by sqrt(DIM)
// 3. computes row-wise softmax over the lower triangle of the matrix
void masked_softmax(float x[][DIM]) {
for (int i = 0; i < DIM; ++i) {
float total = 0;
for (int j = 0; j < DIM; ++j) {
if (i >= j) {
x[i][j] = _exp(x[i][j] / SQRT_FEAT_DIM); // (2) (3)
total += x[i][j]; // (3)
} else {
x[i][j] = 0; // (1)
}
}
for (int j = 0; j < DIM; ++j) {
if (i >= j) {
x[i][j] /= total; // (3)
}
}
}
}Feedforward
As mentioned a few times above, this step operates entirely on individual vectors, just that we can “batch-process” them as a matrix for efficiency? and convenience.
This step seems to be quite battle-tested in the realms of deep learning. Again a short rundown as I best understand it: Consider a certain number of “abstract metrics” you want to score (as numerical values) a particular vector
If we do the same for all our metrics, then congratulations, we have just rederived matrix-vector multiplication!
These metrics are still quite “abstract”, that is, it is hard to immediately assign some semantic meaning to these numbers at this stage. We can attempt to introduce complexity by stacking even more layers, that is, evaluating new metrics on the previous array of metrics itself:
But as you might know in linear algebra, all these weights are collapsible into just one matrix, i.e., directly stacking more layers does not add any complexity.
Unless we introduce some nonlinearity to our metrics. The most common modification these days (at least here) is ReLU (or some variant of it), where it forms sort of an activation function (the signal only “passes through” above a certain threshold):
To adjust the threshold, we add a bias term to each metric:
Thus forming one full layer that we will be implementing for our model.
Again, since this whole computation can be done in parallel, we can batch all the vectors into one large matrix
where
/* fnn */
matmul_t(b, a, M.fnn_0_w); // _X = X @ fnn_0_w
add_bc(b, M.fnn_0_b); // _X += fnn_0_b
relu(b); // (_X)
matmul_t(c, b, M.fnn_2_w); // _X = _X @ fnn_2_w
add_bc(c, M.fnn_2_b); // _X += fnn_2_b
add(a, c); // X += _X
// element-wise ReLU (if <0, then set to 0)
void relu(float x[][DIM]) {
for (int i = 0; i < DIM; ++i) {
for (int j = 0; j < DIM; ++j) {
if (x[i][j] < 0) {
x[i][j] = 0;
}
}
}
}
// adds a vector to a matrix via broadcast.
// per convention, each row of a matrix is a target vector
void add_bc(float a[][DIM], float *b) {
for (int i = 0; i < DIM; ++i) {
for (int j = 0; j < DIM; ++j) {
a[i][j] += b[j];
}
}
}In the context of LLMs, obviously there would be way more layers in order to “predict the next token”. Specifically for next-token generation, it could be seen as looking at the current last token in the context (which by now has been enriched by surrounding context), and “asking questions” about its encoded information, then using those questions to vote out a judgment as to what the next token is most likely going to be.
For our case, we don’t really care, we just need the model to memorize that some sequences of characters have to produce a specific next character.
You may have also noticed an additional add(a, c); operation at the end of the step. The Literature calls this residual connection; I did not dive too deep into why this is the practice, but a quick Chat told me that this is done to “stabilize training by improving gradient flow”. If you were attentive, we actually did this during the Attention step as well! Conceptually, each step outputs an offset to “correct” the input vector towards the intended embedding.
Unembedding
Conceptually the reverse process of our first step (embedding). In practice this is a bit more involved, with approximations and probability distributions involved.
With greedy decoding (top-1 sampling) though, life becomes much simpler: we pick the “best” token.
What is “best” still has to be computed however, and that is not as simple as reverse table lookup. Well, almost —
For a typical transformer, the probability distribution is formed by computing one final metric for each possible token, roughly along the intuition for how strongly the model “favours” that particular token. We then normalize (using softmax) this tuple of metrics (also known as a vector) such that they form a probability distribution. Of course, this would mean the dimension of this resulting vector shall equal the vocabulary size.
Well, another property of softmax is called order preservation: the largest component of the vector remains the largest after transformation. If we are going to deterministically pick the component with the largest probability anyway, we might as well skip the softmax step and just find the largest value!
/* unembedding */
// as mentioned, we only care about the last vector in the context
return extract_vocab(a[CONTEXT_SZ-1]);
/*
* over one vector (output of ffn),
* transform into another VOCAB_SZ-dimension vector,
* then finds the index with the highest score.
*/
#define VOCAB_SZ 256
#define MODEL_DIM DIM
char extract_vocab(float *x) {
int best_idx = -1;
float best_score = 0;
for (int i = 0; i < VOCAB_SZ; ++i) {
float cur_score = 0;
for (int k = 0; k < MODEL_DIM; ++k) {
cur_score += x[k]*M.unembed[i][k];
}
if (best_idx == -1 || best_score < cur_score) {
best_idx = i;
best_score = cur_score;
}
}
return (char)best_idx;
}And we’re done!
Compilation
Our C version is low-level enough to mostly have a smooth mapping over to the machine level.
The easiest starting point to demonstrate is our run function:
char run(char *seq) {
float a[DIM][DIM];
float b[DIM][DIM];
float c[DIM][DIM];
float d[DIM][DIM];
/* embed */
embed(a, seq);
add(a, M.pos_embed);
/* attn */
matmul_t(b, a, M.q);
matmul_t(c, a, M.k);
matmul_t(d, b, c);
matmul_t(b, a, M.v);
masked_softmax(d);
matmul(c, d, b);
matmul_t(b, c, M.attn_out);
add(a, b);
/* fnn */
matmul_t(b, a, M.fnn_0_w);
add_bc(b, M.fnn_0_b);
relu(b);
matmul_t(c, b, M.fnn_2_w);
add_bc(c, M.fnn_2_b);
add(a, c);
/* unembed */
return extract_vocab(a[CONTEXT_SZ-1]);
}Architecture
Here is a good place to start thinking harder about our target architecture. During the implementation I of course ideated first with how even to carry out an execution cycle for an instruction. A few key decisions that came out from the prototype were:
- There shall be 16 registers (4 bits of representation)
- The word size shall be 16 bits
- Naturally from above, the memory can contain 2^16 (65536) words
That is, we had to also make sure that all the weights (as 16-bit reals) of our model can fit within memory. Hence also motivating the small size of mostly dimensions of 16.
For ease of compilation, we will be adapting System-V’s (or just any sensible system in general) calling convention of
- Passing arguments in registers, spilling over to the stack as necessary
- Passing return values in registers
- Splitting into callee-saved (local) and caller-saved (scratch) registers, as appropriate.
Though unnecessary, my ip and sp registers are also bunched into the 16 for the prototype in order to cut some complexity. Dedicated fp is redundant as usual. This means we will have 14 GP registers from r0 to r13; the higher numbers will be for locals and the lower numbers will be for volatile computation (scratch + info passing).
The callee-caller split can be more dynamic where basically I adjusted the split based on my own program in order to minimize spillovers (to the stack). The basic idea for the tradeoff is:
- For callee-saved, we don’t have to protect them everytime we call a function. However, if a function wants to use those registers, it will have to save them in the prologue and load them back in the epilogue.
- For caller-saved, it’s the opposite problem: we don’t have to deal with them in the prologue and epilogue, but everytime we want them to persist through a function call, we have to save them on the stack.
For our run function, the split is basically obvious: we only have 4 local variables (a, b, c, d) for the entire function, and 4 is a reasonable number, so we can go with a 10-4 split for now (r10 to r13 callee-saved):
_run:
;;; prologue
push r13 ; save callee-saved,
push r12 ; so we can then freely use r10-r13
push r11
push r10
mov r10, NEG_DIM ; r10 = -16
lea r13, [r10*DIM+rsp] ; (d) r13 = rsp - 16*16
lea r12, [r10*DIM+r13] ; (c) r12 = r13 - 16*16
lea r11, [r10*DIM+r12] ; (b) r11 = r12 - 16*16
lea r10, [r10*DIM+r11] ; (a) r10 = r11 - 16*16
mov r9, rsp ; r9 is temporary fp
mov rsp, r10
push r9 ; push old sp location onto stack,
; so we know how to unwind at the end
;;; ...
;;; epilogue
pop rsp ; unwind, getting rid of the local matrices
pop r10 ; restore callee-saved
pop r11
pop r12
pop r13
retComfortable!
Everything so far follows quite closely to simple x86 assembly. We do have to be mindful of the more complex ones such as lea, that is, how extensive the operands can get. But again, we can easily break a complex instruction down into simpler ones, if required later. I could already foresee memory operations potentially being a hassle, so it would definitely help to bring our code closer in spirit to using a load-store architecture.
In the meantime, we have our series of function calls:
;;; embedding
mov r1, r0
mov r0, r10
call _embed ; _embed(r10, r0)
mov r0, r10
lea r1, _m_pos_embed
call _add ; _add(r10, _m_pos_embed)As described earlier, calling convention is straightforward enough: specifically we shall just dump our arguments into r0, r1, and so on, and similarly expect return values to be in r0, r1, and so on. Based on our C representation, we don’t expect to have too many parameters anyway, so there is no worry about spillage onto the stack.
Thus, recall for run that we will have just one parameter, being our left-padded input sequence. This parameter goes straight into the second argument of embed, where the first argument will be a.
This a then gets added with _m_pos_embed (some constant address to .rodata) to round off the embedding step. The original input sequence is no longer necessary.
The rest is literally rinse-and-repeat, so maybe I will show just a few more:
;;; feedforward
mov r0, r11
mov r1, r10
lea r2, _m_fnn_0_w
call _matmul_t ; matmul_t(b, a, M.fnn_0_w)
mov r0, r11
lea r1, _m_fnn_0_b
call _add_bc ; add_bc(b, M.fnn_0_b)
mov r0, r11
call _relu ; relu(b)
mov r0, r12
mov r1, r11
lea r2, _m_fnn_2_w
call _matmul_t ; matmul_t(c, b, M.fnn_2_w)
mov r0, r12
lea r1, _m_fnn_2_b
call _add_bc ; add_bc(c, M.fnn_2_b)
mov r0, r10
mov r1, r12
call _add ; add(a, c)
;;; unembedding
mov r0, DIM_MINUS_ONE
lea r0, [r0*DIM+r10]
call _extract_vocab ; extract_vocab(a[DIM-1])Control Flow
Let’s look at a simple yet representative function:
void relu(float x[][DIM]) {
for (int i = 0; i < DIM; ++i) {
for (int j = 0; j < DIM; ++j) {
if (x[i][j] < 0) {
x[i][j] = 0;
}
}
}
}Two for-loops, one if-branch, both represented as branches / jumps in assembly:
_relu:
mov r5, F_ZERO
mov r1, DIM_MINUS_ONE
_relu_loop_i:
mov r2, DIM_MINUS_ONE
_relu_loop_j:
lea r3, [r1*DIM+r2]
mov r4, [r0+r3]
jlf r5, r4, _relu_cont ; if (!(0 < x[i][j])) {
mov [r0+r3], r5 ; x[i][j] = 0;
_relu_cont: ; }
loop r2, _relu_loop_j
loop r1, _relu_loop_i
retWait, the for-loops look a little different from what I expected?
Some slight design decisions were made. Of course I could have used the standard x86 idiom:
mov r1, 0 ; int i = 0;
_relu_loop_i:
cmp r1, DIM ; i < DIM;
jge _relu_loop_i_exit
;;; ...
inc r1 ; ++i;
jmp _relu_loop_i
_relu_loop_i_exit:
retBut look at how unwieldy that is! Just a few of my complaints:
- Having to implement a separate flag register is too unnecessary, so
cmpis already out. - In general, look at just how many instructions we need specifically for our for-loop!
- The structure also looks super bloated, especially when compared to the surrounding code.
Knowing that our main loop construct is going to just be iterating through a row or column of a matrix, it’s quite reasonable to optimize our instruction set specifically for that construct. We aren’t even using any whiles or gotos.
- There is no practical need for us to separate the compare from the branch, so a simple
blt(as in MIPS / RISC-V) would suffice. - Though that changes the logic a little, we can iterate downwards instead of iterating upwards, so that we can compress the increment-and-loop idiom into a literal decrement-then-loop instruction.
Thus it becomes:
mov r1, DIM_MINUS_ONE ; int i = DIM-1;
_relu_loop_i:
;;; ...
loop r1, _relu_loop_i ; i >= 0; --i;I guess we are cheating a bit if we omit the very first check of the loop, but that’s not a big deal for us. DIM is a constant. (If really going for the uniformity, then jumping to the loop instruction after initialization would suffice). But I love my code being concise.
As it turns out, the almighty x86 basically has a loop instruction as well, though with a very slightly different implementation detail. Cool!
But anyway, genuinely this is nice for the sanity of the reader as well, because if you just squint a little bit:
for (...) // mov r1, DIM_MINUS_ONE
{ // _relu_loop_i:
/* ... */
} // loop r1, _relu_loop_iThey basically have the same structure :D
One more detail: the blt instruction here is named jlf. b vs j was purely cosmetic; the difference here is that our instruction works for both ints and floats. More on that in Part II.
Armed with this knowledge, let’s look at something slightly more involved:
void matmul(float r[][DIM], float a[][DIM], float b[][DIM]) {
for (int i = 0; i < DIM; ++i) {
for (int j = 0; j < DIM; ++j) {
r[i][j] = 0;
for (int k = 0; k < DIM; ++k) {
r[i][j] += a[i][k] * b[k][j];
}
}
}
}That is, the standard matrix multiplication (no transpose involved). Since we are all experts now, here is the corresponding Assembly dump:
_matmul:
mov r3, DIM_MINUS_ONE ; for (int i = DIM-1; i >= 0; --i)
_matmul_loop_i: ; {
mov r4, DIM_MINUS_ONE ; for (int j = DIM-1; j >= 0; --j)
_matmul_loop_j: ; {
mov r6, F_ZERO ; float sum = 0;
mov r5, DIM_MINUS_ONE ; for (int k = DIM-1; k >= 0; --k)
_matmul_loop_k: ; {
lea r7, [r3*DIM+r5] ;
mov r7, [r1+r7] ; float tmp1 = a[i][k];
lea r8, [r5*DIM+r4] ;
mov r8, [r2+r8] ; float tmp2 = b[k][j];
fmul r7, r8 ; tmp1 *= tmp2;
fadd r6, r7 ; sum += tmp1;
loop r5, _matmul_loop_k ; }
lea r7, [r3*DIM+r4] ;
mov [r0+r7], r6 ; r[i][j] = sum;
loop r4, _matmul_loop_j ; }
loop r3, _matmul_loop_i ; }
retAgain minor difference in behaviour, but functionally equivalent for the sake of our program.
One thing to note: with that, we have already made use of 9 scratch registers. Only r9 is left to be classified as either type, but any further and we will need either optimizations or concessions.
Softmax
This is perhaps the most complicated function in the whole program, rejoice:
void masked_softmax(float x[][DIM]) {
for (int i = 0; i < DIM; ++i) {
float total = 0;
for (int j = 0; j < DIM; ++j) {
if (i >= j) {
x[i][j] = _exp(x[i][j] / SQRT_FEAT_DIM);
total += x[i][j];
} else {
x[i][j] = 0;
}
}
for (int j = 0; j < DIM; ++j) {
if (i >= j) {
x[i][j] /= total;
}
}
}
}Before we proceed though, there is one more detail to consider. Per Wikipedia (and Chat, for that matter):
The standard softmax is numerically unstable because of large exponentiations. The safe softmax method calculates instead
where
is the largest score involved. Subtracting by it guarantees that the exponentiations result in at most 1, see translation invariance in § Mathematical properties.
This is especially relevant to us, due to a limited number of bits (just 16!) that we are using to represent reals. Casual careless exponentiation potentially overblows intermediate values, even though our softmax pushes everything between
Anywho, we only need a simple tweak to the algorithm, to have a separate loop to find the row-wise maximum, then subtracting that off each value before throwing it into the exponentiation. With a few shuffles we end up with:
void masked_softmax(float x[][DIM]) {
for (int i = 0; i < DIM; ++i) {
float hi = -FLT_MAX;
for (int j = 0; j <= i; ++j) {
x[i][j] /= SQRT_FEAT_DIM;
if (hi < x[i][j]) {
hi = x[i][j];
}
}
float total = 0;
for (int j = 0; j < DIM; ++j) {
if (i >= j) {
x[i][j] = _exp(x[i][j] - hi);
total += x[i][j];
} else {
x[i][j] = 0;
}
}
for (int j = 0; j <= i; ++j) {
x[i][j] /= total;
}
}
}Now, we are ready to tackle this step by step. Firstly, one annoyance comes from the _exp, which we will implement as a function call. (i.e., reasonably not as an instruction).
With this function call, we cannot just dump our local variables in the scratch registers, at risk of clobbering important information. That is, we will be maxing out our callee-saved registers. We can first get a sense of how many pieces of information we need to keep alive through the _exp call:
x, the target matrix;i, the current row index;j, the current column index (specifically for the_exploop);total, the running total;hi, the offset applied to each value before_exp.
That’s not a lot (quite close to 4)! We can easily come up with two mitigations:
- Store the 5 variables above in
r9tor13, that is, designatingr9as callee-saved; - Use a separate loop (prior to the
_exploop) to subtract offhifor each value. That way we don’t have to makehipersist through the call.
Well, both sound reasonable enough! But too bad, I wanted to give myself (and the participants) a tiny bit of challenge, and introduce some friction: we are going to save one important variable in r9, while still having r9 caller-saved, so that we are forced to employ some stack action >:) (Partly also motivated by “ooh single digit registers are caller-saved how clean” :P)
_masked_softmax:
push r13 ; preserve callee-saved registers for use
push r12
push r11
push r10
mov r9, r0 ; move argument (x) to r9
mov r10, DIM_MINUS_ONE ; for (int i = DIM-1; i >= 0; --i)
_masked_softmax_loop_i: ; {
mov r12, F_MIN ; float hi = -FLT_MAX;
mov r11, r10 ; for (int j = i; j >= 0; --j)
_masked_softmax_loop_j0: ; {
;;; ...
loop r11, _masked_softmax_loop_j0
; }
mov r2, F_NEG_ONE ;
fmul r12, r2 ; hi *= -1;
mov r13, F_ZERO ; float total = 0;
mov r11, DIM_MINUS_ONE ; for (int j = DIM-1; j >= 0; --j)
_masked_softmax_loop_j1: ; {
;;; ...
loop r11, _masked_softmax_loop_j1
; }
mov r11, r10 ; for (int j = i; j >= 0; --j)
_masked_softmax_loop_j2: ; {
;;; ...
loop r11, _masked_softmax_loop_j2
; }
loop r10, _masked_softmax_loop_i
; }
pop r10 ; unwind; restore callee-saved registers
pop r11
pop r12
pop r13
retSo we have r9 = x, r10 = i, r11 = j, r12 = hi, r13 = total.
Now for each of the inner loop blocks:
(1) Divide every relevant element by sqrt(DIM), while finding max
mov r11, r10
_masked_softmax_loop_j0:
lea r1, [r10*DIM+r11] ;
mov r0, [r9+r1] ; float cur = x[i][j];
mov r2, SQRT_DIM ; float tmp = SQRT_DIM;
fdiv r0, r2 ; cur /= tmp;
mov [r9+r1], r0 ; x[i][j] = cur;
jlf r0, r12, _masked_softmax_loop_j0_cont
; if (!(cur < hi)) {
mov r12, r0 ; hi = cur;
_masked_softmax_loop_j0_cont: ; }
loop r11, _masked_softmax_loop_j0(2) Running _exp on each element added with the negated hi. Remember the additional challenge of having to save r9 on the stack before calling _exp.
mov r11, DIM_MINUS_ONE
_masked_softmax_loop_j1:
push r9 ; save x on stack
mov r0, F_ZERO ; float result = 0;
jlf r10, r11, _masked_softmax_loop_j1_cont
; if (!(i < j)) {
lea r1, [r10*DIM+r11] ;
mov r0, [r9+r1] ; result = x[i][j];
fadd r0, r12 ; result += (-hi);
call _exp ; result = _exp(result);
fadd r13, r0 ; total += result;
_masked_softmax_loop_j1_cont: ; }
pop r9 ; restore x from stack
lea r1, [r10*DIM+r11] ;
mov [r9+r1], r0 ; x[i][j] = result;
loop r11, _masked_softmax_loop_j1Secretly, there is another challenge involved: at the expense of some spoilers, jlf only supports jumping to 4-bit signed offset, that is, the furthest instruction it can jump forward to is 7 instructions down the line. Well, we managed to meet the requirement quite comfortably here :”) I can’t remember precisely, but this requirement probably influenced me to structure the three inner loops for this function in the way I did.
(3) Divide every relevant element by total. As you could expect, this is quite similar in spirit to (1):
mov r11, r10
_masked_softmax_loop_j2:
lea r1, [r10*DIM+r11] ;
mov r0, [r9+r1] ; float cur = x[i][j];
fdiv r0, r13 ; cur /= total;
mov [r9+r1], r0 ; x[i][j] = cur;
loop r11, _masked_softmax_loop_j2And that’s it! That wasn’t that complicated after all, right?
Exponentiation
As mentioned, it is highly unlikely that we will be implementing exponentiation as an instruction, so we need to express it as a procedure.
The “raising to a power” operator, pow(a, b) = exp(a) =
The infinite expression here is explicitly denoted as a limit as a nod to its nature as a convergent series, that is, the chain of partial sums approach arbitrarily close to the true value with an increasing number of steps. Theoretically, we can attempt to define our function through the least fixed point of the step operator
Note that computationally this formulation does not terminate, making the function undefined for
Thus approximating
Anyway, all that is to say we can implement exponentiation as such:
#define EPS ...
float exp(float x) {
float res = 0.0;
float term = 1.0;
// float x = x;
float i = 0.0;
while (term > EPS) {
res += term;
i += 1.0;
term *= x;
term /= i;
}
return res;
}Correspondingly in Assembly:
_exp:
mov r1, F_ONE ; increment for i
mov r2, F_ZERO ; float i = 0.0;
mov r3, F_ONE ; float term = 1.0;
mov r4, F_ZERO ; float res = 0.0;
mov r5, F_ZERO ; float EPS = 0.0; (abusing rounding)
_exp_loop: ; do {
fadd r4, r3 ; res += term;
fmul r3, r0 ; term *= x;
fadd r2, r1 ; i += 1.0;
fdiv r3, r2 ; term /= i;
jlf r5, r3, _exp_loop ; } while (term > EPS);
mov r0, r4 ;
ret ; return res;Quite later on, when I was proofreading my work with That Thing, I learned about catastrophic cancellation: consider
where successive terms alternate between positive and negative, potentially ballooning near our maximum representation, all while the final result (through the cancellations) turns out to be less than
Fortunately, fixed point approximations of real numbers (spoilers!) seem to handle this issue slightly better than typical floating point ones. But the “standard mitigation” to this issue is so simple that we might as well just do that and not have to worry about it.
You might have guessed it already: the solution is to just compute
; implicitly x <= 0
_exp:
mov r1, F_NEG_ONE
fmul r0, r1 ; now x >= 0
; ...
mov r0, F_ONE
fdiv r0, r4 ; take the reciprocal
retThough potentially we may cause additional overflows where the original negative value would not have underflowed, but we just pray that those values are insignificant anyway.
Interfacing
The rest of the mathematical functions are quite cookie-cutter. Hence let’s round this off with the actual program entrypoint:
struct __attribute__((packed)) {
char head[CONTEXT_SZ];
char tail[48];
} buf = {
.head = "\0\0\0\0\0\0\0\0\0\0\0\0Show",
.tail = { 0 }
};
int main(void) {
char *head_ptr = buf.head;
char *tail_ptr = buf.tail;
for (int i = 0; i < 48; ++i) {
*tail_ptr = run(head_ptr++);
if (!*tail_ptr++) break;
}
return 0;
}The idea:
- It is a loop invariant for
tail_ptrto be 16 (CONTEXT_SZ) bytes away fromhead_ptr. The demarcated buffer in between forms the 16-byte sequence passed intorun, for one round of inference. - No pushdown of new tokens into context is needed; we simply attach every new token to the end of the current sequence, then move the window (
head_ptrandtail_ptr) 1 byte rightward. - The program kicks off with the seed left-padded in the window. The program halts when the model predicts a null byte (marking the end of sequence).
Thus we have:
_main:
lea r10, _buf_head ; char *head_ptr = buf.head;
mov r11, 47 ; for (int i = 47; i >= 0; --i)
_main_loop: ; {
mov r0, r10 ;
call _run ; char next_token = run(head_ptr);
mov r1, DIM ;
mov [r10+r1], r0 ; head_ptr[DIM] = next_token;
mov r1, rsp ; ┐
mov rsp, r10 ; │
pop r10 ; ├── ++head_ptr;
mov r10, rsp ; │
mov rsp, r1 ; ┘
mov r1, 1 ;
jlf r0, r1, _main_end ; if (next_token < 1) break;
loop r11, _main_loop ; }
_main_end:
hltI don’t exactly remember why I made the increment so convoluted; I do know though that using lea (in the sense of lea r10, [r10+1]) would have complications in my architecture (more in part II), and I just either didn’t think of using fadd or felt that using real addition instead of discrete increment was a bit “unclean”.
This was also the only location in the whole program where increment appears (even decrement was folded into the loop instruction), so it felt overkill to dedicate one whole instruction to it.
Instead, I could flaunt my ingenuity (saying this ironically) by abusing the pop instruction, which does effectively increment a register! This only came at a cost of requiring 1 scratch register (along with ballooning the pseudo-instruction into 5 real instructions), tradeoffs I was perfectly willing to accept for handling this instruction only ever appearing once.
Intermission
…and we’re done (for now)! We’ve descended from our toy LLM to the machine code level.
With the Assembly fully in place, we could hack up a quick n’ dirty Assembler + VM to sanity check that it works. The Python script is a bit long to fit here (and into an already overblown writeup), so I might potentially cover it in Part II. But essentially, you can have it work like any other interpreter for a reverse engineering challenge.
As a summary to lead into Part II, here are the instruction forms we have utilized:
push rmov r, imm16lea r, [r*DIM+r]mov r, rpop rretcall labellea r, labelmov r, [r+r]jlf r, r, labelmov [r+r], rloop r, labelfmul r, rfadd r, rfdiv r, rhlt
Purely by mnemonic alone, there are only 12 unique instructions! Though of course we know in x86 that every instruction basically does like everything. Mostly we will be deconflicting the movs and leas into 5 different forms of data moving:
li r, imm16/labelmov r, rlea r, [r*DIM+r]ld r, [r+r]st r, [r+r]
Trading 2 mnemonics for 5, that bumps our numbers up to 15 instructions total. We also know that ret is practically identical to pop rip, ending with a final count of 14, comfortably below a heuristic threshold of 16.
You could of course probably have told from the Assembly listings that some creativity was involved and not explicitly shown here in terms of fitting the instructions into these 14, but nothing was too drastic that would significantly affect readability or intuitive understanding (which are ergonomically important for participant experience).
Anyway, thanks for reaching the end of this long-winded post! See you in Part II (when it is ready), where we effectively reinvent a virtual CPU :)