How it works
The state engine
Inference engines treat a context as memory inside a process. That is the right model for a chatbot and the wrong one for an agent, because agents do not run forward in a line — they branch, backtrack and resume. Three operations fix it, and they build on each other.
Paging
Positions live in fixed-size pages rather than one flat allocation
(kv.DefaultPageTokens = 64). A context is an ordered list of page
pointers. This buys two things: a context grows without reallocating, and —
far more importantly — two contexts can point at the same
page. Everything below depends on that one fact.
The prefix tree
Agent traffic is overwhelmingly repetitive: the same system prompt, the same
tool definitions, the same conversation re-sent with one more turn on the
end. Cache.Match(tokens) walks a radix tree over token sequences
and returns a context holding the longest prefix already computed, along with
how many tokens matched. Only the new suffix is prefilled.
Copy-on-write fork
Context.Fork() copies the page-pointer slice and increments
refcounts. Nothing else. A page is cloned only when a branch writes to a page
whose refcount is above one — in Context.writable, which is the
only place in the package where copying happens.
Two properties are asserted in tests rather than assumed:
- Isolation. A branch that writes is not visible to its parent. Getting this wrong produces cross-talk between agents that reads like a model problem, not a memory one.
- Exactness. A forked context continues identically — the test asserts a logit difference of zero, not a small one. A branch that drifts by a rounding error produces agents that quietly disagree about their own history.
The fork path, measured on a laptop:
Transcript — the whole session as text
~/llamay $ llamay ctx build -m qwen2.5:0.5b -o build/demo/work/base.llamayctx -p "You are re
viewing a database migration. The table is users, the column is email, and the index is uniq
ue."
prefilled 23 tokens in 175ms
saved build/demo/work/base.llamayctx (552.14 KiB) in 7ms
~/llamay $ ls -lh build/demo/work/base.llamayctx | awk '{print $5, $9}'
552K build/demo/work/base.llamayctx
~/llamay $ # eight branches off that one prefill
~/llamay $ llamay ctx fork -m qwen2.5:0.5b -i build/demo/work/base.llamayctx -n 8
source build/demo/work/base.llamayctx, 23 tokens, 1 pages
branches 8 in 1µs (200ns each)
pages allocated 0
pages a copy would need 8
memory held 1.50 MiB (a copy would be 12.00 MiB)
~/llamay $ # pages allocated: 0. The branches share every page until one writes.
Snapshots
A context serialises position-major rather than page-major, so a snapshot is independent of the page size that produced it. Two checks are non-negotiable and both are in the format itself.
At a prompt size worth keeping, with thirty-two branches off it:
Transcript — the whole session as text
~/llamay $ llamay bench -m qwen2.5:0.5b -gpu -forks 32 -prompt 2048 -gen 32 -reps 1 llamay v0.1.58-1-gb14f550 · Qwen2.5 0.5B Instruct · qwen2 · 7 threads · kernels neon+i8mm · matmul metal:Apple M4 · gemm own kernels · kv f32 weights 373.71 MiB · repack on · context 32768 · prompt 2048 · generate 32 phase rate best note prefill 650.4 tok/s 3.149s one command buffer per 256-token tile decode 41.7 tok/s 768ms 506 dispatches per token, one submission cpu encoding 40ms · gpu execution 3.854s (1% encoding) phase rate best note fork x32 65µs 0 pages allocated (a copy would be 1024) snapshot 129ms 48.01 MiB for 2048 tokens, at f32 restore 74ms vs 3.149s to re-prefill gpu time: staging 149ms · kernels 9.77s · copy-back 191ms Restoring a 2048-token context is 42.6x faster than recomputing it. ~/llamay $ # 0 pages allocated where a copy would have been 1024 — that is the fork
How it is proved
The package is pointer arithmetic, and the way it is checked is a reference implementation run beside it in lockstep. The reference is deliberately stupid: every context holds its own full copy of every position, nothing is shared at all, and the prefix tree is a flat list scanned linearly for the longest key that is a prefix of the query. It is obviously correct because it does nothing clever, so a sequence of operations where the two disagree is a bug in the clever one.
TestModelledOperations drives random sequences of every operation
the package offers — allocate, extend, overwrite, fork, truncate, close,
retain, match, under an arena and without one, at each of the three
precisions — and asserts the whole invariant set after every single
step:
- A page's refcount is exactly the number of contexts that point at it. A refcount of one has to mean precisely one holder, because that is the fact
writableacts on when it writes in place. - No page is held by a context after the cache freed it, none sits in the cache at or below zero, and
allocated − freedequals the number of live pages. - Every arena slot the arena believes is out is backing a live page, and no two pages claim the same slot.
- A position returns the same value forever unless this context writes it — fork isolation stated as an invariant rather than as a test case.
- The tree's entry set and eviction order match the flat-list model exactly, and its structure is a well-formed radix tree.
- Closing every context and clearing the tree returns the cache and the arena to empty.
At 500 seeds of 400 operations against each of the six configurations, all of
it holds; the committed run is smaller so the suite stays fast. Around that:
forks at a page boundary and one position either side, twelve-deep forks of
branches that have themselves been forked and written, every truncation point
of a shared context, every single-byte corruption and every truncation of a
snapshot, and the concurrency the scheduler actually creates — sixteen
goroutines forking one parent, snapshot-and-restore racing a decode that is
rewriting the shared prefix, eviction racing allocation — under
-race.
The two rows to read are fork is exact and snapshot resumes:
Transcript — the whole session as text
~/llamay $ llamay verify -m testdata/azmx-tiny.gguf check result detail prefill == decode pass max logit difference 0.00e+00 vq8 kv cache pass moves 0.3% of the sampler's mass over 64 positions (worst 0.3% at position 55), greedy agrees at all 64, max logit shift 0.028 — values quantised, key s exact q8 kv cache note moves 0.3% of the sampler's mass over 64 positions (worst 0.4% at position 30), greedy differs at 1 of 64, max logit shift 0.034 — keys quantised too, use vq8 instead vq8 bound has teeth pass a one-bit cache moves 37.2% of the sampler's mass, over th e 5% bound — so the row above is not vacuous on this file tokenizer round trip pass 7 samples exact fork is exact pass difference 0.00e+00 (must be zero) snapshot resumes pass difference 0.00e+00 thread-count invariant pass difference 0.00e+00 across 1, 2 and 4 threads All invariants hold. ~/llamay $ # fork is exact and snapshot resumes: 0.00e+00, which is the only passing value
Four things that were wrong
Kept here because each one is the reason an invariant above is phrased the way it is.
| The bug | How it presented | The fix |
|---|---|---|
| The prefix tree leaked arena slots | Nothing looked wrong: counters balanced, pages were collected, the entry bound held. Only the arena went on believing slots were occupied, so a server doing exactly what the tree is for drained its device arena over hours and silently stopped using the accelerated attention path. | One release path instead of four hand-written copies. Forty retains against a four-entry bound had left eighty slots out for eight live pages. |
| Recycled arena slots were handed out dirty | The heap path gets zeroed storage from make and the arena path got whatever the last page left, so the same program produced different bytes depending on whether a GPU arena happened to be attached — and an unzeroed slot wrote another conversation's KV into a snapshot. |
Slots are zeroed on issue. The snapshot is a file whose entire purpose is to be shipped elsewhere, which is what made this more than an aesthetic problem. |
| Eviction left the node standing | Five hundred twelve-token inserts against a four-entry bound left five hundred and one nodes, each holding a full copy of its token edge — tens of megabytes retained for prefixes evicted precisely because nobody wanted them. | Eviction prunes upward from the victim, and a node with one child and no context — a split that no longer splits anything — collapses into its child. |
Restore believed the length it was told |
The position count is four bytes; the positions it promises are not. A one-megabyte body could claim several gigabytes before the lie ran into the end of the stream and the checksum that would have caught it. | Token ids are read into a plain slice before anything is allocated for them, and a declared length beyond sixteen million positions is refused by name. |
The same weights, by content
A snapshot used to name its model by layout: the architecture and the shape
of every tensor. A fine-tune has the same layout as its base, so a context
computed by one could be restored into the other and produce confident
nonsense. Since 0.3.0 a snapshot (format version 3) also records the
SHA-256 of the weights file — combined with the adapter's,
for a LoRA model — and a restore checks it after the layout. A mismatch is
409 Conflict.
Older snapshots, which carry no digest, still restore by layout, with a
warning in the log. Under -profile enterprise, or on a request
that carries a classification marking, they are refused unless the server
was started with -allow-legacy-snapshots.
Spill: memory pressure without losing a context
-ctx-mem-budget 8G gives each model a budget for resident KV.
Past it, the least recently used idle context with pages of its own
is written to -ctx-spill-dir and its pages freed. The next
request that names it reads it back, exactly — the file is at the cache's
own precision — and deletes the file. A context in use is never spilled.
- Spill files are always encrypted: under
-context-key, or under a random key that exists only for the life of the process. - A spilled context lists as
"spilled": truewith zero pages; thellamay_context_spill*metrics count spills, restores, bytes and time. - If a read-back fails, the context is dropped and the request gets a 404 that says so.
What this composes into
- Best-of-N sampling falls out for free: N branches from one position, sharing every page below it.
- Agents exploring alternatives from a shared history pay for that history once.
- A session survives a closed laptop — snapshot on exit, restore on start, no re-prefill.
- A context moves to another server with its owner and marking, which is the local half of a disaggregated prefill/decode split. Moving a context.
- Many more idle conversations than fit in memory, with spill doing the paging.
What is not done
RetainandSnapshotread the context they are given without the cache lock, so the single-owner rule for aContextcovers them too. This is documented rather than enforced.- The prefix tree is bounded by entry count rather than by bytes, so a handful of very long prefixes can pin far more memory than many short ones.
- Snapshots are written at full precision by default — a snapshot outlives the settings that produced it — but the payload precision is a choice. On a 2,202-token context of a 2B model: 447.3 MB at
f32, 293.5 MB atvq8, 139.8 MB atq8. - Spill evicts whole contexts, not pages, and only idle ones. The shared prefix tree is not spilled.
- There is no object-storage leg: a snapshot goes where the client puts it, and a move is server to server.