picoLM, in Python
picoLM runs a 1.1B model on a $10 board with 256 MB of RAM in 2,944 lines of C. Port it to NumPy one module at a time and you find every chapter of this course inside it, including the two it leaves out on purpose.
- GGUF & mmap
- K-quants
- fused dequant
- SentencePiece BPE
- grammar masking
- porting
Motivation
The problem
You have spent twenty chapters building parts. Whether that adds up to being able to read somebody else's engine is a separate question, and the only way to answer it is to pick an engine and rebuild it.
picoLM is small, which is what makes it a good choice: 2,944 lines of C11 across seven modules, no dependencies, and a 1.1B model running on a $10 board with 256 MB of RAM. You can read all of it in an afternoon, and it loads the same GGUF files llama.cpp does.
So this chapter is a port. Seven Python modules mirror its seven C ones, and every one of them is checked against the original rather than against our own intuition about what it probably does.
- lines of C11
- 2,944
- seven modules, zero dependencies
- of that, in Python
- 43%
- the rest is SIMD, threads and memory managed by hand
- dequantisers bit-exact
- 8/8
- against bytes written from the ggml spec
This port is readable, not fast
Core idea
§1 — GGUF, and why it is mmap-able
picoLM opens the model with mmap and never copies a weight. The weights stay on flash; the kernel pages in only what a matmul touches, and several processes can share one copy. This is what makes 1.1B parameters fit on a board with 512 MB of RAM.
The format is dull by design: a header, self-describing metadata, a tensor index of (name, shape, type, offset), padding to general.alignment, then one opaque blob. There is nothing to decompress and no pointers to fix up. Add the offset to the mapped base and you are looking at the weights.
Loading is not reading
weights 638 MB mapped does not mean 638 MB of RAM. Mapped pages are clean and file-backed, so the kernel can evict them under pressure and fetch them again. The KV cache is dirty anonymous memory and cannot be evicted at all. Because of that asymmetry, what exhausts memory on a small device is context length rather than parameter count.Mechanics
§2 — K-quants
S08 argued that the whole game is one scale per group, not per tensor. A K-quant nests that idea. A Q4_K super-block holds 256 weights in 144 bytes: one fp16 scale and one fp16 minimum for the block, then eight 6-bit scales and eight 6-bit minimums, one pair per 32-weight sub-block, bit-packed into 12 bytes, and finally 128 bytes of nibbles.
def dequantize_q4_K(raw, n: int) -> np.ndarray:
"""144 bytes -> 256 weights (4.5 bpw)."""
blk, nb = _blocks(raw, Q4_K, n)
d, dmin = _fp16_field(blk, 0), _fp16_field(blk, 2)
packed, qs = blk[:, 4:16], blk[:, 16:144]
sc = np.empty((nb, 8), dtype=np.float32)
mn = np.empty((nb, 8), dtype=np.float32)
for j in range(8):
if j < 4:
sc[:, j] = packed[:, j] & 63
mn[:, j] = packed[:, j + 4] & 63
else: # six-bit fields straddle byte boundaries
sc[:, j] = (packed[:, j + 4] & 0xF) | ((packed[:, j - 4] >> 6) << 4)
mn[:, j] = (packed[:, j + 4] >> 4) | ((packed[:, j] >> 6) << 4)
q = qs.reshape(nb, 4, 32)
out = np.empty((nb, 4, 64), dtype=np.float32)
d_, dm_ = d[:, :, None], dmin[:, :, None]
out[:, :, :32] = d_ * sc[:, 0::2, None] * (q & 0xF) - dm_ * mn[:, 0::2, None]
out[:, :, 32:] = d_ * sc[:, 1::2, None] * (q >> 4) - dm_ * mn[:, 1::2, None]
return out.reshape(-1)The awkward loop over j comes from the format rather than the port: in the first four sub-blocks the 6-bit fields sit in their own bytes, and in the last four they straddle byte boundaries. Get it wrong and you have plausible weights and a model that runs, with output that is quietly worse.
Verify against the spec, because even the C can be wrong
quant.c and run on identical random bytes, six formats came out bit-identical; Q2_K and Q3_K did not. picoLM reads those two with the wrong field offsets and bit order, so it silently mis-decodes any Q2_K or Q3_K file llama.cpp produces. The port follows ggml, the format's source of truth, rather than the C it started from.- Q2_K bits/weight
- 2.62
- 84 bytes per 256 weights
- Q4_K bits/weight
- 4.50
- not 4.0; the scales are not free
- Q6_K bits/weight
- 6.56
- where ffn_down and output usually land
Explore
§3 — Tensor ops, and the one thing this port gets wrong
picolm/tensor.c is eleven functions: a threaded matmul, plus rmsnorm, softmax, rope and a handful of one-liners. matmul carries the performance story. It walks the quantised bytes and accumulates the dot product as it decodes, so the only thing crossing the memory bus is the 4-bit data.
This port cannot do that. Expressing a fused loop in NumPy means going element-by-element in Python, and that is slower than the thing it is trying to avoid. So matmul here dequantises per call. That is the unfused path S08 warns about, chosen for readability and measured in the runnable file rather than hidden.
The lesson survives the compromise
The RoPE layout is not the one you learned
S03 taught the GPT-NeoX convention, pairing dimension i with i + D/2, and warned that a second convention also goes by the name RoPE and is not compatible with it. picoLM uses the second one.
# picoLM / llama.cpp: rotate the pair (x[2i], x[2i+1])
for i in range(half):
q0, q1 = qh[i * 2], qh[i * 2 + 1]
qh[i * 2] = q0 * cos[i] - q1 * sin[i]
qh[i * 2 + 1] = q0 * sin[i] + q1 * cos[i]
# S03 / GPT-NeoX: rotate the pair (x[i], x[i + D/2])
out[:half] = x[:half] * cos - x[half:] * sin
out[half:] = x[half:] * cos + x[:half] * sinBoth are correct, and S03 is not wrong either. The two conventions differ by a permutation of the Q and K weight matrices, which convert_hf_to_gguf.py applies at conversion time. A HuggingFace checkpoint stores NeoX-ordered weights; the GGUF stores them pre-permuted, so llama.cpp's interleaved rotation lands on exactly the same attention scores. Which layout you need is therefore a property of the file you loaded rather than of the architecture. Llama GGUF files want the interleaved layout, and a model run with the wrong one still produces grammatical, confident, subtly wrong text. When a hand-written loader produces garbage that almost works, this is the first thing to check.
Build it
§4 — The forward pass
Here is the whole engine, condensed. Read the signature first, model_forward(m, token, pos), and notice how much of this course it implies.
float *model_forward(model_t *m, int token, int pos) {
dequantize_row(embd_row, s->x, dim, w->type_token_embd);
for (int l = 0; l < c->n_layers; l++) {
rmsnorm(s->xb, s->x, s->attn_norm_w[l], dim);
matmul(s->q, s->xb, lw->attn_q, dim, dim, lw->type_attn_q);
/* … K, V … */
rope(s->q, k_tmp, head_dim, n_heads, n_kv_heads, cos_pos, sin_pos);
for (int d = 0; d < kv_dim; d++)
key_pos_fp16[d] = fp32_to_fp16(k_tmp[d]); /* the cache is fp16 */
for (int h = 0; h < n_heads; h++) {
/* online softmax — no att[] buffer is ever allocated */
for (int t = 0; t <= pos; t++) { /* … */ }
}
/* … output projection, SwiGLU, residuals … */
}
rmsnorm(s->x, s->x, s->output_norm_w, dim);
matmul(s->logits, s->x, w->output, dim, c->vocab_size, w->type_output);
return s->logits;
}And the port, the same thing with the buffers taken away:
def forward(self, token: int, pos: int) -> np.ndarray:
c = self.cfg
cos_pos, sin_pos = self.cos[pos], self.sin[pos]
x = quant.dequantize_row(raw_row, c.n_embd, emb.qtype)
for l in range(c.n_layers):
xb = tensor.rmsnorm(x, self.attn_norm[l])
q = tensor.matmul(xb, *self._w(f"blk.{l}.attn_q.weight"), c.n_embd, c.n_embd)
k = tensor.matmul(xb, *self._w(f"blk.{l}.attn_k.weight"), c.n_embd, c.kv_dim)
q = tensor.rope_interleaved(q, c.n_heads, c.head_dim, cos_pos, sin_pos)
k = tensor.rope_interleaved(k, c.n_kv_heads, c.head_dim, cos_pos, sin_pos)
self.key_cache[l, pos] = k.astype(np.float16) # the cache is fp16
self.val_cache[l, pos] = v.astype(np.float16)
for h in range(c.n_heads):
g = h // c.kv_mul # the KV head this query shares
scores = (kh[:, g, :] @ qh[h]) * inv
out[h] = tensor.softmax(scores) @ vh[:, g, :]
# … output projection, SwiGLU, residuals …
x = tensor.rmsnorm(x, self.output_norm)
return tensor.matmul(x, *self._w(self.output_name), c.n_embd, c.vocab_size)- One token per call: the S01 loop, unchanged after twenty chapters.
- `pos` indexes a persistent cache: S05, with K stored after RoPE because a cached token's position never changes.
- The cache is fp16: S08, applied to the cache instead of the weights. At TinyLlama's full 2,048 context that halves ~88 MB to ~44 MB, which is headroom a small board notices.
- Attention in the C is an online softmax: S12. picoLM's comment says it plainly: the
att[]buffer was removed, and no score array is ever allocated. (The port keeps the plain two-pass softmax so the arithmetic stays readable.) - `n_kv_heads < n_heads`: S03's GQA. The cache is small enough to be fp16 in the first place because of it.
What is missing is the point
Production notes
Try it
picoLM's headline claim is a 1.1B model on a $10 board with 256 MB of RAM; its hardware table also lists the Pi Zero 2 W at $15 and 512 MB. Check that row. Set the parameters to 1.1B, the format to Q4_K and the device to the Pi Zero, then push the context length up and watch what breaks.
- weights (mmapped)
- 590 MB
- KV cache (resident)
- 44 MB
- bits / weight
- 4.50
- KV bytes / token
- 22 KB
- device RAM
- 512 MB
- headroom
- 289 MB
- fits — picoLM runs this
- tight — swapping likely
- does not fit
Weights are mmapped, so they do not all have to be resident. The KV cache does. That asymmetry is why picoLM stores it in fp16, and why context length rather than parameter count is what runs you out of memory.
layers 22 · kv heads 4 · head dim 64
Two things worth noticing. Switching the KV cache from fp16 to fp32 costs you more headroom than switching the weights from Q4_K to Q6_K, because the cache is resident and the weights are not. At 16k context nothing fits whichever weight format you pick: the cache alone has taken two-thirds of the board's memory. picoLM ships a --ctx flag for that reason, and S05's arithmetic is the one you carry to every new device.
Continue
§5–§6 — Tokenizer, sampler, grammar
S02 built a byte-level BPE from the GPT lineage: merges are ranked, and the lowest rank wins. Llama GGUF files carry the other family, SentencePiece vocabularies where every token has a score and the highest score wins. Same loop, different ordering key, different token ids.
- Spaces become
▁(U+2581) and one is prepended, so"Once"and" Once"tokenize identically. - Anything not in the vocabulary falls back to a
<0xHH>byte token, so the vocabulary is total. - The merge loop is greedy and quadratic. That is fine: it runs once per request, on a short string.
The sampler is 105 lines and is exactly S04: greedy at temperature 0, otherwise temperature → softmax → top-p. It accumulates until cum >= top_p and keeps the token that crossed. That is S04's + 1, in production code.
In 175 lines the grammar module does what S15 describes at a quarter of the size: a character-level pushdown automaton for JSON, lifted to the token level by asking whether every character a token would emit keeps the machine alive. Illegal tokens are set to −∞ before the sampler runs.
Masking guarantees valid, not complete
{"name": "ad is just as unusable to the caller as invalid JSON. Check accepting before you return, or budget for the closing brackets.Step 8
§7 — The CLI, and what it all cost
for pos, t in enumerate(ids): # prefill: one position at a time
logits = model.forward(t, pos)
out_ids, pos = [], len(ids)
while len(out_ids) < max_tokens and pos < model.cfg.max_seq_len:
row = grammar.mask_logits(logits, state, pieces) if grammar else logits
nxt = sampler(row)
if nxt == tok.eos_id:
break
out_ids.append(nxt)
stream.write(dec.push(nxt))
logits = model.forward(nxt, pos)
pos += 1That is the whole of picolm.c worth quoting. After twenty chapters of optimisation it is still a while-loop that calls a model, picks a token and appends it. That was the first claim this course made.
$ python code/s21_picolm.pyOnly NumPy is required — setup instructions.
To run a real model, download any Llama-architecture GGUF and point the CLI at it. The loader under test in the suite is the same one:
# any Llama-architecture GGUF works — TinyLlama is the one picoLM targets
cd code
python -m picolm.run --model tinyllama-1.1b-chat.Q4_K_M.gguf \
--prompt "Once upon a time" -n 40
# grammar-constrained, so the output is JSON or nothing
python -m picolm.run --model tinyllama-1.1b-chat.Q4_K_M.gguf \
--prompt "Describe a cat as JSON" --jsonStep 9
What's next
Nothing. That was the course.
You have written a naive loop and made it fast; paged, shared and compressed the memory; batched, scheduled and preempted; broken the one-token-per-pass rule and constrained what the model may say; split a model across GPUs and prefill across machines. Then you took somebody else's engine and rebuilt it, and it turned out to be made of the same parts.
Two obvious next moves. Port a second engine and see what its authors chose differently, or take this one somewhere picoLM does not go: put a block manager and a scheduler on top of this forward pass and you have the beginnings of S20 on a Raspberry Pi.
Exercises
- 1Implement the fused path,
vec_dot_q4_K_f32, as a C extension or with Numba, and measure it against the NumPy version in this port. You should reproduce S08's traffic argument as a wall-clock number. - 2Add Q5_K and Q5_1 to
quant.py. Both are in picoLM's type table but unimplemented here; the block layouts are inquant.h. Verify against hand-written bytes the way the other formats are. - 3
- 4Point the CLI at a real TinyLlama GGUF, then at the same model quantized to Q2_K, and compare the outputs at temperature 0. This is your own version of S08's quality-versus-bits table, on a model you can run yourself.
Check yourself
Questions
Answer before you look. A wrong answer tells you more than a right one, because the explanation points at the part of the chapter to reread.
picoLM maps a 638 MB model on a device with 512 MB of RAM and runs. Why is that not a contradiction?