microgpt in MLPL
A literate walk through the sw-MLPL port of Karpathy's microgpt.py, compared with the Python original and the Rust port

Andrej Karpathy's microgpt.py trains a GPT and samples from it in ~200 lines of dependency-free Python. Its docstring sets the rule:

This file is the complete algorithm. Everything else is just efficiency.

Two ports follow that rule in different directions:

• microgpt-rs (python-to-rust.org) keeps the algorithm's shape: a scalar autograd engine, now a flat tape of nodes, and the same per-scalar loops, compiled.
• microgpt.mlpl (this repo) changes the shape. sw-MLPL is an array language with reverse-mode autograd, adam, cross_entropy and softmax built in. The Value class disappears, every per-element loop becomes a whole-array expression, and the per-token KV-cache loop becomes one masked pass over a whole name.

This document is the MLPL program itself, told section by section. Every MLPL block below runs in one shared :session, so state carries from block to block like lines typed into the REPL, and the output under each block was produced by running it. The blocks that make up the program (the ones without :tangle no) also tangle, in order, into one standalone script; scripts/check-literate.sh runs that script and requires its output to be byte-identical to microgpt.mlpl's. The prose therefore describes code that really is the implementation. The production program splits the same definitions into lib/*.mlpl so the test suites can include them.

Each section states its math as display equations (rendered with MathJax in the HTML). The model's functions also carry their equation as data: an @formula annotation (Unicode) plus an @ascii form, following the companion repos' convention (see idiomatic-mlpl.md). Every summation states its limits and every symbol is defined.

This is the faithful variant: it follows microgpt.py's structure line for line. The idiomatic and compact variants trade that fidelity for shorter code built from MLPL's Model DSL.

Two kinds of helper that Python gets for free: indexing a vector by a vector of positions, and f"{x:4d}" / f"{x:.4f}" formatting. MLPL's at and take accept a single index only, so u:gather1 reshapes to a one-column table and uses the row gather gather_rows.

def u:gather1(v, idx) {
  "Elements of rank-1 v at the (rank-1) integer positions idx.";
  reshape(gather_rows(reshape(v, [len(v), 1]), idx), [len(idx)])
}

def u:pad_left(s, width) {
  "Right-align string s in width columns, width <= 16 (like Python's {:>width}).";
  str_concat(str_slice("                ", 0, relu(width - str_len(s))), s)
}

def u:fmt_int(x, width) {
  "f'{x:widthd}' for a non-negative integer x.";
  u:pad_left(to_string(round(x)), width)
}

def u:fmt_fixed4(x) {
  "f'{x:.4f}' for x >= 0 (round half up at the 4th decimal).";
  scaled = floor(x * 10000 + 0.5);
  whole = floor(scaled / 10000);
  frac = to_string(scaled - whole * 10000);
  str_concat(str_concat(to_string(whole), "."), str_concat(str_slice("0000", 0, 4 - str_len(frac)), frac))
}

def u:write(s) {
  "Write s to stdout with no trailing newline (for Python's end='\\r').";
  unwrap(write_stdout(tokenize_bytes(s)))
}

[u:fmt_int(7, 4), u:fmt_fixed4(1.91456)]

["   7", "1.9146"]

docs = [line.strip() for line in open('input.txt') if line.strip()]
random.shuffle(docs)

Python holds the 32,033 names as a list of strings; microgpt-rs as a Vec<&str>. MLPL is happiest with one flat array, so the corpus stays a single byte array (tokenize_bytes). Each name is a span between two newlines: find the newline positions, add virtual ones at -1 and n, and take differences. Blank lines give empty spans and are dropped, which is Python's if line.strip(). The names file has nothing to strip, so instead of implementing strip() the parser rejects any byte that would need it.

The same function also finds the vocabulary, sorted unique characters, the array way: sort every byte (grade_up), keep the first byte of each run, drop the newline. The result is inverted into a 256-entry lookup table from byte to token id. The shuffle waits until training (section 7).

def u:dataset(txt) {
  "Parse input.txt: one doc per non-empty line, plus the char vocabulary.";
  bytes = tokenize_bytes(txt);
  n = len(bytes);
  bad = reduce_add(lt(bytes, 33) * ne(bytes, 10)) + reduce_add(gt(bytes, 126));
  if bad { err(str_concat("corpus has non-name bytes: ", to_string(bad))) } else {
    nl = compress(eq(bytes, 10), range(n));
    bounds = concat(concat([0 - 1], nl), [n]);
    m = len(bounds) - 1;
    starts_all = u:gather1(bounds, range(m)) + 1;
    lens_all = u:gather1(bounds, range(m) + 1) - starts_all;
    keep = gt(lens_all, 0);
    sorted = u:gather1(bytes, grade_up(bytes));
    first = concat([1], ne(u:gather1(sorted, range(n - 1) + 1), u:gather1(sorted, range(n - 1))));
    uniq = compress(first, sorted);
    uchars = compress(ne(uniq, 10), uniq);
    vocab = len(uchars);
    lut = zeros(256) - 1;
    i = 0;
    repeat vocab { lut = scatter(lut, at(uchars, i), i); i = i + 1 };
    ok({
      bytes: bytes,
      starts: compress(keep, starts_all),
      lens: compress(keep, lens_all),
      n_docs: reduce_add(keep),
      uchars: uchars,
      lut: lut,
      bos: vocab,
      vocab_size: vocab + 1
    })
  }
}

d = unwrap(u:dataset(load("input.txt")));
print("num docs:", to_string(d.n_docs));

num docs: 32033

The whole parse, vocabulary included, takes about 19 ms. An earlier version found the vocabulary with 256 passes of eq over the corpus and took 254 ms; see benchmarks.md.

uchars = sorted(set(''.join(docs)))
BOS = len(uchars)
tokens = [BOS] + [uchars.index(ch) for ch in doc] + [BOS]

Python searches uchars once per character; microgpt-rs uses a HashMap. MLPL encodes a name with one gather through the lookup table. The same trick encodes all the names training will visit in one pass (u:doc_batch, used in section 7). That matters in this interpreter because reading a large array copies it, so per-name access to the corpus costs ~1.8 ms, against 1.8 ms for 1000 names at once.

def u:encode_bytes(d, bs) {
  "Token ids for a rank-1 byte array (every byte must be in the vocab).";
  u:gather1(d.lut, bs)
}

def u:encode(d, s) {
  "Token ids for a string (no BOS).";
  u:encode_bytes(d, tokenize_bytes(s))
}

def u:decode(d, ids) {
  "The string for a rank-1 array of non-BOS token ids.";
  decode_bytes(u:gather1(d.uchars, ids))
}

def u:doc_bytes(d, i) {
  "Raw bytes of doc i (file order).";
  u:gather1(d.bytes, at(d.starts, i) + range(at(d.lens, i)))
}

def u:doc_tokens(d, i) {
  "microgpt.py's tokens for doc i: [BOS] + ids + [BOS].";
  concat(concat([d.bos], u:encode_bytes(d, u:doc_bytes(d, i))), [d.bos])
}

def u:doc_batch(d, doc_ids, block_size) {
  "Token rows for many docs at once: {tokens: [m, block_size + 2] = [BOS] + ids + [BOS] padded with BOS, n: [m] positions to train}.";
  m = len(doc_ids);
  lens = u:gather1(d.lens, doc_ids);
  pos = reshape(range(block_size), [1, block_size]);
  valid = lt(pos, reshape(lens, [m, 1]));
  idx = (reshape(u:gather1(d.starts, doc_ids), [m, 1]) + pos) * valid;
  bytes = reshape(gather_rows(reshape(d.bytes, [len(d.bytes), 1]), idx), [m, block_size]);
  ids = reshape(gather_rows(reshape(d.lut, [256, 1]), bytes), [m, block_size]);
  body = valid * ids + (1 - valid) * d.bos;
  bos = zeros([m, 1]) + d.bos;
  {tokens: concat(bos, concat(body, bos, 1), 1), n: lens + 1 - relu(lens + 1 - block_size)}
}

print("vocab size:", to_string(d.vocab_size));
vocab_size = d.vocab_size;

vocab size: 27

The first name in the file, emma, encodes exactly as in Python (e is 4, m is 12, BOS is 26):

In symbols: a name of \(L\) characters \(c_1 \dots c_L\) becomes the token sequence below, with \(\mathrm{BOS} = 26\). The model is trained on its first \(n\) transitions, where \(T = 16\) is the block size:

\[ \tau = (\mathrm{BOS},\ \mathrm{id}(c_1),\ \dots,\ \mathrm{id}(c_L),\ \mathrm{BOS}), \qquad n = \min(T,\ L + 1) \]

print(decode_bytes(d.uchars));
u:doc_tokens(d, 0)

abcdefghijklmnopqrstuvwxyz
26 4 12 12 0 26

class Value:
    def __init__(self, data, children=(), local_grads=()): ...
    def __add__(self, other): ...      # +, *, **, log, exp, relu, ...
    def backward(self): ...            # topological sort, then the chain rule

This is the largest difference between the three programs:

• Python: ~40 lines. Each Value is a heap object pointing at its children; backward() sorts the graph topologically, then applies the chain rule. The garbage collector frees each step's graph.
• Rust: ~170 lines. Nodes are appended to a flat Vec (the tape) in forward order, which is already a topological order, so backward is one reverse loop; truncate(num_params) frees a step's graph.
• MLPL: 0 lines. param[shape] declares a trainable array leaf. grad(expr, W) records expr on the interpreter's tape, whole-array ops (matmul, softmax, gather_rows, cross_entropy) rather than scalars, and runs it backward. u: functions called inside are inlined onto the tape.

A tiny example, d/dW of sum(W^2) = 2W:

W_demo = param[2, 2];
W_demo = [[1, 2], [3, 4]];
grad(reduce_add(W_demo * W_demo), W_demo)

2 4
6 8

For the real model, tests/test_gradcheck.mlpl compares grad with central finite differences for all nine matrices; they agree to ~3e-11.

matrix = lambda nout, nin, std=0.08: [[Value(random.gauss(0, std)) for _ in range(nin)] for _ in range(nout)]
state_dict = {'wte': matrix(vocab_size, n_embd), 'wpe': matrix(block_size, n_embd), 'lm_head': matrix(vocab_size, n_embd)}
for i in range(n_layer):
    state_dict[f'layer{i}.attn_wq'] = matrix(n_embd, n_embd)    # ... wk, wv, wo
    state_dict[f'layer{i}.mlp_fc1'] = matrix(4 * n_embd, n_embd)
    state_dict[f'layer{i}.mlp_fc2'] = matrix(n_embd, 4 * n_embd)

The same nine matrices, with the same names (layer0_attn_wq for layer0.attn_wq) and Python's [nout, nin] shapes. Python draws every weight from one Mersenne Twister stream; here each matrix gets its own seeded randn. (In --rs-parity mode, not shown here, the weights come from microgpt-rs's own SplitMix64 stream instead; see section 9.)

Python and Rust keep a flat params list for Adam. MLPL cannot store a list of param leaves in a variable (an array literal of matrices is an array), so the nine names are written inline in the adam call.

With vocabulary \(V = 27\), width \(d = 16\), block size \(T = 16\) and MLP width \(4d\), the parameter count is

\[

• \underbrace{4d^2}_{W_q,W_k,W_v,W_o} + \underbrace{2 \cdot 4d^2}_{W_1,\,W_2}

\;=\; 864 + 256 + 1024 + 2048 \;=\; 4192 . \]

n_layer = 1;      # depth of the transformer (number of layers)
n_embd = 16;      # width of the network (embedding dimension)
block_size = 16;  # maximum context length (the longest name is 15 chars)
n_head = 4;       # number of attention heads
head_dim = n_embd / n_head;
init_std = 0.08;

wte = param[vocab_size, n_embd];
wte = randn(101, [vocab_size, n_embd]) * init_std;
wpe = param[block_size, n_embd];
wpe = randn(102, [block_size, n_embd]) * init_std;
lm_head = param[vocab_size, n_embd];
lm_head = randn(103, [vocab_size, n_embd]) * init_std;
layer0_attn_wq = param[n_embd, n_embd];
layer0_attn_wq = randn(104, [n_embd, n_embd]) * init_std;
layer0_attn_wk = param[n_embd, n_embd];
layer0_attn_wk = randn(105, [n_embd, n_embd]) * init_std;
layer0_attn_wv = param[n_embd, n_embd];
layer0_attn_wv = randn(106, [n_embd, n_embd]) * init_std;
layer0_attn_wo = param[n_embd, n_embd];
layer0_attn_wo = randn(107, [n_embd, n_embd]) * init_std;
layer0_mlp_fc1 = param[4 * n_embd, n_embd];
layer0_mlp_fc1 = randn(108, [4 * n_embd, n_embd]) * init_std;
layer0_mlp_fc2 = param[n_embd, 4 * n_embd];
layer0_mlp_fc2 = randn(109, [n_embd, 4 * n_embd]) * init_std;

num_params = size(wte) + size(wpe) + size(lm_head) + size(layer0_attn_wq) + size(layer0_attn_wk) + size(layer0_attn_wv) + size(layer0_attn_wo) + size(layer0_mlp_fc1) + size(layer0_mlp_fc2);
print("num params:", to_string(num_params));

num params: 4192

microgpt.py follows GPT-2 with three simplifications: RMSNorm instead of LayerNorm, no biases, and ReLU instead of GeLU. The MLPL functions keep Python's names and structure, but each works on an [n, 16] array, one row per position of the name, instead of a list of scalars.

m[i] = beta1 * m[i] + (1 - beta1) * p.grad
v[i] = beta2 * v[i] + (1 - beta2) * p.grad ** 2
m_hat = m[i] / (1 - beta1 ** (step + 1))
v_hat = v[i] / (1 - beta2 ** (step + 1))
p.data -= lr_t * m_hat / (v_hat ** 0.5 + eps_adam)

Python and Rust keep the moment buffers m and v themselves. MLPL's adam(loss, [params], lr, b1, b2, eps) keeps them per param across calls and applies exactly this bias-corrected rule; tests/test_training.mlpl checks two steps against the formula to 1e-12. One adam call records a single tape for all nine gradients (0.74 ms, against 5.8 ms for nine separate grad calls), and it returns the loss before the update, which is Python's loss.data.

For each parameter, with gradient \(g_k\) at step \(k = 1 \dots K\) (\(K = 1000\)), \(\beta_1 = 0.85\), \(\beta_2 = 0.99\), \(\epsilon = 10^{-8}\), \(\eta = 0.01\):

\[ m_k = \beta_1 m_{k-1} + (1-\beta_1)\, g_k, \qquad v_k = \beta_2 v_{k-1} + (1-\beta_2)\, g_k^{2}, \] \[ \theta_k = \theta_{k-1} - \eta\Big(1 - \frac{k-1}{K}\Big)\, \frac{m_k / (1-\beta_1^{k})}{\sqrt{v_k / (1-\beta_2^{k})} + \epsilon}. \]

learning_rate = 0.01;
beta1 = 0.85;
beta2 = 0.99;
eps_adam = 0.00000001;

for step in range(num_steps):
    doc = docs[step % len(docs)]
    ... forward, loss.backward(), Adam ...
    print(f"step {step+1:4d} / {num_steps:4d} | loss {loss.data:.4f}", end='\r')

First, the training data. The shuffled visiting order picks the 1000 names the run will see, and u:doc_batch encodes all of them in one pass. Then come two performance moves particular to this interpreter:

• Keep only what inference needs, and expunge the rest. Every u: call copies the global environment, so leaving the 228k-byte corpus in scope made the loop 2.6x slower (2.09 s against 0.8 s).
• Encode up front, because reading a large array copies it.

Both are recorded as upstream issues.

num_steps = 1000;
train_ids = u:gather1(shuffle(range(d.n_docs), 42), mod(range(num_steps), d.n_docs));
train_docs = u:doc_batch(d, train_ids, block_size);
train_tokens = train_docs.tokens;
train_n = train_docs.n;
uchars = d.uchars;
BOS = d.bos;
n_docs = d.n_docs;
expunge(["d", "train_ids", "train_docs", "txt"]);

The loop itself. train N { } binds step and collects each iteration's final value into last_losses. Each step selects one name's row, runs one adam (forward, backward and update of all nine matrices, about 0.6 ms), and writes Python's progress line, carriage return included. Its output is not kept here; the next block summarizes it.

train num_steps {
  row = take(train_tokens, 0, step);
  n = at(train_n, step);
  inp = u:doc_in(row, n);
  tgt = u:doc_tgt(row, n);
  mask = u:causal_mask(n);
  lr_t = learning_rate * (1 - step / num_steps);
  loss = adam(u:loss(inp, tgt, mask), [wte, wpe, lm_head, layer0_attn_wq, layer0_attn_wk, layer0_attn_wv, layer0_attn_wo, layer0_mlp_fc1, layer0_mlp_fc2], lr_t, beta1, beta2, eps_adam);
  u:write(str_concat(str_concat(str_concat(str_concat(str_concat("step ", u:fmt_int(step + 1, 4)), " / "), u:fmt_int(num_steps, 4)), " | loss "), str_concat(u:fmt_fixed4(loss), "\r")));
  loss
}

The loss per 100-step window. A single name's loss is noisy, so compare windows:

w = 0;
repeat 10 {
  m = reduce_add(last_losses * ge(range(num_steps), 100 * w) * lt(range(num_steps), 100 * w + 100)) / 100;
  print(str_concat(str_concat(str_concat("steps ", u:fmt_int(100 * w + 1, 4)), str_concat("-", u:fmt_int(100 * w + 100, 4))), str_concat("  mean loss ", u:fmt_fixed4(m))));
  w = w + 1
};
print(str_concat("step 1000 loss ", u:fmt_fixed4(at(last_losses, num_steps - 1))));

steps    1- 100  mean loss 2.7501
steps  101- 200  mean loss 2.5411
steps  201- 300  mean loss 2.5331
steps  301- 400  mean loss 2.4659
steps  401- 500  mean loss 2.4630
steps  501- 600  mean loss 2.4239
steps  601- 700  mean loss 2.3984
steps  701- 800  mean loss 2.4216
steps  801- 900  mean loss 2.3898
steps  901-1000  mean loss 2.3630
step 1000 loss 2.2796

The three implementations track each other (CPython / microgpt-rs / MLPL): 2.77 / 2.69 / 2.71 for steps 1-100, and 2.28 / 2.36 / 2.37 for steps 901-1000. The small differences come only from their different random streams.

token_id = BOS
for pos_id in range(block_size):
    logits = gpt(token_id, pos_id, keys, values)
    probs = softmax([l / temperature for l in logits])
    token_id = random.choices(range(vocab_size), weights=[p.data for p in probs])[0]
    if token_id == BOS: break

Generation produces one token at a time, so here the KV cache would save work. MLPL's cache builtins (gen_state=/=gen_append) serve only Model DSL chains, so the port recomputes the prefix: at most 16 tokens, 0.24 ms per token. Each draw takes the next uniform from one sequential stream and picks the first token whose cumulative weight reaches u * total, which is what random.choices and microgpt-rs's choices do.

With temperature \(\tau_s = 0.5\), last-position logits \(z\) and uniform \(u \in [0, 1)\):

\[ p = \mathrm{softmax}(z / \tau_s), \qquad i^{*} = \min\Big\{\, i \;:\; \sum_{j=0}^{i} p_j \;\ge\; u \sum_{j=0}^{V-1} p_j \Big\}. \]

@formula "i* = min{ i : ∑(j=0..i) p_j ≥ u · ∑(j=0..V−1) p_j },   p = softmax(z / τ)"
@ascii "i* = min{ i : SUM(j=0..i) p_j >= u * SUM(j=0..V-1) p_j },   p = softmax(z / tau)"
def u:sample_token(logits, temperature, u) {
  "First token id whose cumulative softmax(logits / temperature) weight reaches u * total.";
  cdf = running_sum(softmax(logits / temperature, 0));
  k = reduce_add(lt(cdf, u * at(cdf, len(cdf) - 1)));
  k - relu(k - (len(cdf) - 1))
}

def u:generate(uniforms, start, temperature) {
  "Token ids (BOS excluded) of one sample, drawing uniforms[start], uniforms[start + 1], ...; it consumes u:consumed(ids) of them.";
  seq = [BOS];
  pos = 0;
  done = 0;
  while lt(pos, block_size) * (1 - done) {
    n = pos + 1;
    logits = last_row(u:gpt(seq, range(n), u:causal_mask(n)));
    tok = u:sample_token(logits, temperature, at(uniforms, start + pos));
    if eq(tok, BOS) { done = 1 } else { seq = concat(seq, [tok]) };
    pos = pos + 1
  };
  u:gather1(seq, range(len(seq) - 1) + 1)
}

def u:consumed(ids) {
  "Uniforms a sample used: one per token plus the closing BOS draw, unless it ran to block_size.";
  len(ids) + 1 - eq(len(ids), block_size)
}

def u:ids_to_text(ids) {
  "Characters for token ids (empty string for no ids).";
  if len(ids) { decode_bytes(u:gather1(uchars, ids)) } else { "" }
}

The loop over samples is a plain repeat. (Until sw-mlpl e6ee2203, a string-valued statement inside a repeat body failed the loop, and this block had to use a while; upstream issue j.)

temperature = 0.5;
num_samples = 20;
sample_uniforms = random(4242, [num_samples * block_size]);
cursor = 0;
print("");
print("--- inference (new, hallucinated names) ---");
sample_idx = 0;
repeat num_samples {
  ids = u:generate(sample_uniforms, cursor, temperature);
  cursor = cursor + u:consumed(ids);
  name = u:ids_to_text(ids);
  print(str_concat(str_concat(str_concat("sample ", u:fmt_int(sample_idx + 1, 2)), ": "), name));
  sample_idx = sample_idx + 1
};
sample_idx

--- inference (new, hallucinated names) ---
sample  1: aline
sample  2: erish
sample  3: jani
sample  4: meei
sample  5: jade
sample  6: kayish
sample  7: zayan
sample  8: aniai
sample  9: arari
sample 10: hayli
sample 11: arina
sample 12: arasth
sample 13: zuyee
sample 14: barey
sample 15: meina
sample 16: rion
sample 17: azien
sample 18: onila
sample 19: renan
sample 20: rikia

• Speed. The MLPL interpreter lands within 1.2x of compiled Rust and ~90x ahead of CPython. Each MLPL operation processes a whole array, while Python pays interpreter overhead per scalar operation. This is "everything else is just efficiency" from the other side: the efficiency comes from the language's array ops, not from a hand-written tape.
• Exact parity. Run as microgpt.mlpl -- --rs-parity, the port draws every random number from microgpt-rs's SplitMix64 stream, reimplemented in pure MLPL with 64-bit words as four 16-bit pieces. Its output (1000 loss lines ending 1.9146 and the same 20 names, amanion, alik, zarani, …) is byte-identical to microgpt-rs's. just parity checks this live.
• What MLPL made harder. String handling and number formatting; interpreter costs (global copies, array-read copies) that shaped the data layout; no vector indexing; a loop-body string bug. See upstream-issues.md for all of them, with reproducers.

Further reading: python-vs-mlpl.md (the same comparison as a plain walkthrough), benchmarks.md (speed log), plan.md (design decisions).