Design

How the engine is put together, the decisions worth defending, and an honest account of what is and is not verified.

Back to the README

Repository layout

proofinfer/
  Cargo.toml                 no [dependencies]; release profile: opt-level 3, lto, codegen-units 1
  LICENSE                    MIT
  CONTEXT.md                 design notes: format, forward pass, gotchas, TCB
  src/
    lib.rs                   module map
    ops.rs                   kernels: rmsnorm, softmax, silu, matmul, RoPE, GQA
    model.rs                 config, hardened checkpoint loader, State, forward, generate
    tokenizer.rs             BPE vocabulary reader, encoder, decoder
    main.rs                  CLI
  tests/
    engine.rs                forward/generation properties
    loader.rs                hostile-input totality
    tokenizer.rs             conformance to Meta's token vectors
  reference/
    diff_test.py             differential harness vs PyTorch
    mutation_check.py        proves the differential harness can fail
    mutation_check_tests.py  proves the Rust test suite can fail
    bench.py                 benchmark + byte-identity check
    llama2c/                 vendored UNMODIFIED from karpathy/llama2.c (MIT)
  verus/                     Verus proof of KV-cache index safety (24 conditions,
                             machine-checked; see its README for the scope)
                             plus check_citations.py, which keeps the proof
                             from silently drifting away from src/
  .github/workflows/ci.yml   four jobs, no `|| true`, no continue-on-error

Design decisions worth defending

Zero dependencies, seriously. Cargo.lock contains exactly one package, and CI asserts it before anything else runs. Arg parsing, file IO, the f32 kernels, the BPE tokenizer and the CLI are all built on std. For a project about machine-verified inference, the smallest possible dependency surface is the point: every crate in the lock file widens the trusted base by another transitive unsafe blob.

The loader treats the file as hostile. The checkpoint header dictates how much memory gets allocated, so it cannot be trusted to be right. Three properties, each with a test that fails if it breaks:

1. Totality. For any byte string, from_bytes returns Ok or Err, never a panic. tests/loader.rs feeds it 2000 pseudo-random buffers, half prefixed with a plausible header so the cursor gets past validation into the size arithmetic, which is where the bugs actually live.
2. No allocation from an unvalidated number. A header claiming 4096 dim × 32000 vocab on a 64 KiB file is a perfectly legal set of numbers, so the only thing preventing a multi-gigabyte allocation is checking the bytes are present before creating the Vec.
3. Overflow is an error, not a wraparound. Every size product goes through checked_mul. This forced a design change: computing each tensor's length inline at the point of reading meant a header that overflowed at wq reported a confusing "token_embedding needs 8589934584 bytes" instead. TensorSizes now computes the whole shape table up front.

forward allocates nothing. Every buffer lives in State and is written in place, so the memory footprint of decoding is a sum you can write down: scratch is a handful of vectors, and the cache dominates at 2 * n_layers * seq_len * kv_dim.

Greedy only, no sampling. Sampling needs an RNG, and an RNG would make the byte-identical comparison against run.c impossible unless both sides consumed the identical stream. Greedy decoding is a pure function of the weights and the prompt, so two implementations that agree on the maths produce byte-identical output — a far stronger claim than "the samples look similar".

Bugs found by the byte-identity check that no unit test caught. Generation originally stopped on EOS; llama2.c stops on BOS, which is the document delimiter these models are trained with, so with stories15M our output ran straight past the end of the first story. Output was also printed by decoding the whole sequence at once, whereas run.c prints per token through safe_printf, which drops single non-printable bytes. Both were found by bench.py and neither was findable from the differential test, which does not look at text at all.

Trusted computing base

Stated plainly, because it is the honest answer to "how much do you actually trust this?": the Rust compiler and std, the f32 semantics of the CPU, the checkpoint parsing code, and the reference implementation being tested against. Everything above those is tested. Those are assumed, and the assumption is written down rather than implied.

What is formally verified, and what is not

The first thing worth proving mechanically is index safety of the KV-cache slicing in forward: that for any pos < seq_len and any layer, every cache index computed in ops::attention and State::forward is in bounds.

verus/ contains a Verus development for exactly that, and it is machine-checked: Verus 0.2026.09.20.aef82ed reports 24 verified, 0 errors under --no-cheating, which rejects assume, admit and external_body outright. That is a real run, not a claim.

What is proved: the bounds of every KV-cache index expression for all shapes satisfying the loader's structural preconditions and all pos < seq_len; that those preconditions follow from the checks the loader actually performs; that the GQA mapping h / kv_mul lands in 0..n_kv_heads; and that the index arithmetic cannot overflow usize. The real slice expressions are transcribed into exec fns operating on actual Vec<f32>, so Verus discharges genuine slice-indexing obligations rather than assertions about symbolic expressions.

Keeping the proof honest

A green Verus run is not enough on its own, because the proof verifies a model. If someone changed h / kv_mul in ops.rs, the proof would keep verifying, faithfully, about code that no longer exists — and the green tick would be actively misleading.

So verus/check_citations.py checks the link. Every source line the proof and its README cite must still contain the construct it is cited for. It is bidirectional — an uncited citation and an unused expectation are both failures, so neither the prose nor the table can quietly stop being checked. Stdlib only, so it costs nothing in CI.

The check has been negative-tested, which matters more than it passing: it goes red when the GQA mapping is mutated in place, when an unrelated edit shifts the line numbers, when the attention window changes, and when a cache write is deleted from forward.

What it still does not check is the reasoning. An edit that preserves a cited expression's text while changing what it computes elsewhere would pass, as would a gap in which expressions were transcribed at all. It is a tripwire on the likely accident, not a proof of correspondence.

What is not proved

verus/README.md lists eleven items; these are the two that matter most:

1. It is a proof about a transcription, not about the crate. src/ops.rs and src/model.rs contain no verus! macro and are not compiled by Verus. The citation check above narrows how far this can go wrong; it does not close it.
2. Bounds safety does not catch wrong-but-in-bounds bugs. Mutant #3 (h / kv_mul -> h % n_kv_heads) is also in range and sails straight through. That is the boundary between what a verifier can say and what only the differential test can say, and the proof file says so itself.

Running the proof

The proof CI job fetches the prebuilt Verus release, verifies it against the sha256 GitHub publishes for that asset, installs the toolchain Verus asks for (it prints the exact rustup install line, which the job parses rather than hardcoding a version), and runs the verifier. Verus is not a dependency of this crate and is not vendored.

That job was executed, not just reasoned about. Its steps were run verbatim in an emulated x86_64 Linux container — the same architecture and OS family as GitHub's ubuntu-latest — and produced:

--- 2. Fetch and unpack Verus (x86-linux) ---
-rw-r--r-- 1 root root 485678212 verus.zip
verus.zip: OK                      # sha256 matched the pinned digest
unpacked: 1.6G
--- 3. Locate the Verus binary ---
Found Verus at verus-dist/verus-x86-linux/verus
--- 4. Install the toolchain Verus requires ---
verus: required rust toolchain 1.98.1-x86_64-unknown-linux-gnu not found
  rustup install 1.98.1-x86_64-unknown-linux-gnu
Verus requires: 1.98.1-x86_64-unknown-linux-gnu
--- 5. Verify the KV-cache index-safety proof ---
verification results:: 24 verified, 0 errors

That run is also what justified the ANSI-colour fix in the toolchain-discovery step: on Linux the parse still had to strip the escape sequence to recover 1.98.1-x86_64-unknown-linux-gnu from the same coloured error message.

What is still unexecuted, precisely: the GitHub-hosted actions themselves — actions/checkout, actions/cache, actions/setup-python and dtolnay/rust-toolchain. Those are standard, they are not written here, and the container provided its own curl/unzip/python3/rustup in their place. The remaining risk on a first CI run is therefore in the YAML and the runner environment, not in the commands this repository contributes. If the proof job does go red, the other three jobs are unaffected and the log will name the step.

Nothing else here is machine-checked. The loader's totality rests on 2000 fuzz iterations plus reasoning, not a proof; the differential test establishes agreement with a reference on six configurations and one trained checkpoint, not equivalence to the architecture in general.

Two other things were scoped and deliberately not done, rather than half-done: a differential test against vLLM (needs a GPU and a Hugging Face export path, and the harness is written so that adding a third oracle is a new config table rather than a new file), and Triton kernels for rmsnorm, softmax and matmul (a benchmarking exercise with no bearing on whether the engine is correct, which is the claim this repository makes). The measurement in this repository is CPU and single-threaded throughout; nothing here says anything about GPU inference.

Credits

The PyTorch reference model, the legacy export format, the tokenizer format, the C engine, and the expected token id vectors all come from karpathy/llama2.c, MIT licensed. The files under reference/llama2c/ are vendored byte-for-byte unmodified; reference/llama2c/PROVENANCE.md records the sha256 of each so a reviewer can verify that, and it matters more than it might look — a differential test against a reference that has been edited is worth nothing.

The token id expectations in tests/tokenizer.rs originate in Meta's example_text_completion.py and are checked by llama2.c's test.c.

Source: https://github.com/blackdragoon26/proofinfer/blob/main/docs/design.md in the repository.