Skip to content
LLM Inference
S21The Capstone·495 lines

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

A NumPy matmul per weight matrix per token cannot compete with fused int4 kernels; expect roughly two orders of magnitude. What it reproduces is picoLM's outputs, not its performance. Where the two differ, the runnable file measures the gap rather than glossing over it.

Core idea

§1 — GGUF, and why it is mmap-able

DiagramGGUF on disk
GGUF ON DISK — ONE MMAP, NO COPIESheadermagic · version · n_tensors · n_kvmetadataarch · dims · heads · rope · tokenizer vocabtensor infoname · shape · type · offsetpad to general.alignmentTENSOR DATA — ONE OPAQUE BLOBtoken_embdQ4_Kblk.0.attn_qQ4_Kblk.0.ffn_downQ6_Koffset 0offset 0x2A0000offset 0x2C4000matmul(x, W_raw, …) — reads these bytes in placemmap() maps the whole file. A weight is read only when a matmul touches it, and that is how 1.1B params fit in 512 MB of RAM.
Every tensor is a byte range at a known offset. Nothing else is in the format, and that is what makes the weights mappable rather than loadable.

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

Diagramthe Q4_K super-block
Q4_K — 256 WEIGHTS IN 144 BYTESd (fp16)2 Bdmin (fp16)2 Bscales · 12 Beight 6-bit scales and eight 6-bit mins, bit-packedqs · 128 B — 256 nibblessub 032 × 4-bit·32 × 4-bit·32 × 4-bit·32 × 4-bit·32 × 4-bit·32 × 4-bit·32 × 4-bitsub 732 × 4-bitscale 0 · min 0scale 1 · min 1scale 2 · min 2scale 3 · min 3scale 4 · min 4scale 5 · min 5scale 6 · min 6scale 7 · min 7w = d × scale[sub] × q − dmin × min[sub]The nesting is the trick: one fp16 scale for the super-block, one 6-bit scale per 32-weight sub-block. 4.5 bits per weight, scales included.
One fp16 scale for 256 weights, with eight 6-bit scales nested inside it. Nesting is what buys 4.5 bits per weight without losing the outliers.

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.

code/picolm/quant.py (excerpt)python
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

A dequantiser that is self-consistent is worth nothing. The runnable file checks all eight formats against blocks written by hand from the ggml spec, and that check caught a real bug. Compiled against picoLM's own 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

Diagramfused dequant + dot
FUSED DEQUANT + DOT — THE S08 LESSON, IN CDEQUANTIZE, THEN MATMULint4 weightson disk, mmappedfp32 copywritten to RAM, then read backmatmul4× the trafficof never quantizingFUSED: DECODE IN REGISTERS, ACCUMULATEint4 weightson disk, mmappedvec_dot(q, x)never leaves registersaccumulator1× the trafficThe unfused path moves the int4 weights AND writes and re-reads an fp32 tensor: strictly more traffic than never quantizing. That is 'my quantized model got slower'.
The unfused path reads the int4 weights, then writes and re-reads an fp32 copy: strictly more traffic than never quantizing at all.

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

S08's point was that quantization only pays if the dequantisation is fused. Here is the counterexample, running in the same repository: this port reads 4-bit weights and is slower than fp32 NumPy would be, for the reason that chapter gave.

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.

the layout that is not S03'spython
# 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] * sin

Both 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

Diagramone token, one position
MODEL_FORWARD(TOKEN, POS) — ONE TOKEN, ONE POSITIONembedding lookupone dequantized row× N_LAYERSRMSNormQ/K/V projRoPE (interleaved)write K,V → cache (fp16)online softmax over cache[0..pos]no score array is ever allocatedout proj + residualSwiGLU FFN + residualKV CACHE · FP16 · [LAYERS, SEQ, KV_DIM]t0t1t2t3t4posfinal RMSNorm → logitsThere is no batch dimension anywhere. One user, one core. S06, S09, S10 and S11 are absent for that reason, and the engine is 2,944 lines.
One token per call is S01; `pos` indexing a persistent cache is S05; fp16 storage is S08; the online softmax is S12; n_kv_heads is S03. Five chapters, one signature.

Here is the whole engine, condensed. Read the signature first, model_forward(m, token, pos), and notice how much of this course it implies.

picolm/model.c (the whole engine, condensed)c
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:

code/picolm/model.py (the same thing)python
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

There is no batch dimension anywhere, no block manager, no scheduler, no continuous batching and no chunked prefill. S06 and S09 through S11 are absent altogether. picoLM serves one user on one core, so none of those mechanisms would pay for themselves, and that is why it is 2,944 lines instead of 200,000. An inference engine's size is set by its concurrency requirements, not by the model.

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.

Simulatorwill it fit?
will it fit?
weight format
device
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

DiagramSentencePiece BPE
SENTENCEPIECE BPE — HIGHEST SCORE WINSnormalize' ' → ▁, and prepend oneseed: one token per characterbyte fallback: <0xHH>merge loopjoin the adjacent pair with the best score"the cat"▁the▁catthecat▁thescore −4▁catscore −5S02's byte-level BPE picks the LOWEST merge rank. SentencePiece picks the HIGHEST score. Same loop, different key, different token ids.
Normalise, seed one token per character, then merge greedily. It is S02's loop with a different sort key.

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

A grammar mask makes malformed JSON unreachable. It does not make the model finish. At a 40-token cap the runnable file shows 30% of runs hitting the budget mid-object, and a truncated {"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

Diagramthe loop, after nineteen chapters
PICOLM.C — THE LOOP, AFTER NINETEEN CHAPTERSload + mmaptokenizeprefill: forward(t, pos) for each prompt tokenDECODE: FORWARD(TOK, POS++) UNTIL EOSgrammar mask (optional)samplestream bytes outforward()It is still the S01 while-loop. Every chapter since has been an optimisation of this shape, and picoLM keeps only the ones that pay off for a single user.
Load, tokenize, prefill, then the same while-loop S01 opened with. Everything in between became a module.
code/picolm/run.py (excerpt)python
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 += 1

That 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.

Run it locally
Builds a real GGUF file from scratch, maps it back and runs the full engine on it: hand-computed quantised blocks checked against the format spec, the interleaved-vs-NeoX RoPE difference measured, tokenizer round-trips, a forward pass with cache-immutability checks, sampler and grammar behaviour, and a line count against the C.
$ python code/s21_picolm.py
Expect: Every dequantiser exact against hand-written bytes, RoPE layouts differing by more than 0.1, a bit-identical prefix replay, 200/200 grammar-valid samples with the truncation rate reported honestly, and a line count printed against picoLM's 2,944.

Only 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:

running a real modelbash
# 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" --json

Step 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

  1. 1
    Implement 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.
  2. 2
    Add Q5_K and Q5_1 to quant.py. Both are in picoLM's type table but unimplemented here; the block layouts are in quant.h. Verify against hand-written bytes the way the other formats are.
  3. 3
    Give the port a block manager from S06 and a scheduler from S10, and serve two concurrent requests from one mapped model. Measure where the crossover is: at what concurrency does paging start paying for its own complexity?
  4. 4
    Point 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.

Check yourselfQuestion 1 of 6

picoLM maps a 638 MB model on a device with 512 MB of RAM and runs. Why is that not a contradiction?

Score 0/6