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:

~/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:

~/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 writable acts 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 − freed equals 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:

~/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 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": true with zero pages; the llamay_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

• Retain and Snapshot read the context they are given without the cache lock, so the single-owner rule for a Context covers 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 at vq8, 139.8 MB at q8.
• 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.