llm-fog-fhe: fhe-llm-demo — encrypted tiny QA on the Niobium Fog

A client asks a small memorized question-answering model a natural-language question and gets the answer back without the compute party (the Fog worker) ever seeing the question — and without the weights ever existing in the clear server-side. The whole 4-slot answer window is produced in ONE forward pass inside ONE Fog job (no autoregression, so the 1-hour Fog wall has huge margin).

record/replay pipeline

(source: design/flow.dot; re-render with Graphviz: dot -Tsvg flow.dot -o flow.svg)

Scope honesty: this is a capability demo of fully-encrypted inference, not a general-purpose LLM. The model has D=16, one FFN block, 97-token vocabulary (61 corpus words + 4 specials + 32 filler tokens that pad the vocab to 4 groups of 32 lanes), and answers MEMORIZED questions (20 training pairs). Zero-shot generalization is deliberately out of scope and is not designed against.

Architecture — Rust-first, one C++ kernel

Everything plaintext (training, tokenizing, the client-side pipeline, slot packing, argmax, the two-home security layout, file formats) is Rust. The ONLY C++ is src/kern, because the wire format is OpenFHE serialization: ciphertexts, keys, and the recorded trace must be byte-compatible with the OpenFHE / libnbfhetch stack the Fog worker runs, and niobium::compiler() (the record/replay API) is a C++-only template API with no C boundary.

build.rs          compiles src/kern to libllmkern.a, links OpenFHE +
                  libnbfhetch from the niobium-client checkout
src/kern/
  llm_kernel.h    extern "C" ABI (keygen, encrypt, server run, decrypt)
  llm_kernel.cpp  the only NEW C++ (src/run_circuit.* + llm_common.* are
                  kept kernel-side C++ too; see below)
src/              Rust crate (lib + one `llm` binary)
  consts.rs       one source of truth for D/FF/SEQ/V/slot geometry
  corpus.rs       20 QA pairs + vocab (port of twin/corpus.py)
  artifacts.rs    data/ readers + writer (frozen float32/text artifacts)
  twin.rs         candle twin: Chebyshev-GELU fit, MIX, train (candle-nn
                  AdamW), demo, range-check, export, differential
  enc.rs          builds the packed slot vectors the kernel encrypts
                  (8 per answer position + 7 shared + ANSWER_WINDOW x
                  V_GROUPS heads; 31 at V=97)
  backend.rs      FFI bridge: every kernel call + error channel
  main.rs         `llm` CLI (keygen/encrypt/server/decrypt/check, twin)
src/run_circuit.* C++ encrypted forward pass + llm_common.* (slot geometry
                  and load helpers) — kernel-side code, unchanged from the
                  C++ era, shared by local record, local sim, and Fog replay

cargo build --release produces ONE binary, target/release/llm.

Status (local gates, verified on RUST-trained weights)

Gate	Mode	Result
`llm twin diff`	candle forward vs plain f64 recomputation	MATCH (worst diff ~5e-8)
`llm twin train`	candle twin, banned penalty, AdamW	perfect train+eval (ep 1325)
`make cpu`	real-math record + fhetch_sim replay diff	MATCH, decrypted `paris`
`make sim`	hollow record + fhetch_sim replay (Fog recording path)	PASS
`make sweep`	6 prompts end-to-end via record/replay	6/6 PASS + leak tripwire clean
`make check`	pure-OpenFHE circuit vs plaintext math	MATCH (~1e-8 on logits)

Fog submission: keygen at ring 2^16, encrypt, then

fogrs submit ./target/release/llm server run <server_home> --hollow --target=FOG

(fogrs provisions the Fog job, exports NBCC_FHETCH_SERVER/TOKEN/REPLAY, and spawns the binary with everything after the app path passed through; the kernel wraps niobium::compiler() for the record/replay dance). The worker-side contract is byte-identical to the C++ era (same circuit, same artifact formats), and the Rust build has a live Fog PASS under it: job b01e500b, fog-worker-batch-1, "what do cows drink" -> "milk" (1.97 GB trace in, 12.6 MB logits out, 50 s worker wall).

Layout

data/          frozen artifacts exported from the trained twin checkpoint
               (model_weights.bin [private], tok.bin/pos.bin [client-only],
                mix.csv [public structure], gelu.txt [public math],
                vocab.txt, params.json). `llm twin export` regenerates all of
               these; the C++ never touches the checkpoint.
twin/          model-rs.safetensors (the candle checkpoint)

Security model (two-home split, enforced in Rust)

The client and server sides use SEPARATE directories, and secret material lives only in the client home:

client_home/ (0700): cc.bin, pk.bin, sk.bin (0600), mk.bin, rk.bin, plus prompt.txt, expected.txt, answers.txt (all 0600). This never ships.
server_home/ (0700): cc.bin, pk.bin, mk.bin, rk.bin (public and evaluation keys) + ct_*.bin encrypted inputs + ct_logits*.bin results. This is exactly what the worker receives — no secret key, no plaintext prompt, no plaintext answer.

The Rust CLI refuses to run encrypt/decrypt/check when the two paths resolve to the same directory (symlinks canonicalized), and the kernel hard-fails at startup any time an sk.bin is present in server_home — local or worker-bound — since the two-home split means no legitimate run needs the secret key there. All binaries umask(0077) before the first write, so secret files are born 0600 with no world-readable window; the sweep target carries a leak tripwire on every server dir.

Circuit geometry (see src/consts.rs + src/run_circuit.h for the full writeup)

Slot(d,u) = 64*d + u, d in [0,16), u in [0,32); slots 0..1023 of 32768 live, the rest encrypted zero. The 64-wide zero-padded blocks make every d-sum / u-sum butterfly contamination-free (only block 0 / lane 0 is read downstream). MIX is applied client-side before encryption (public constants on the client's own data). The GELU polynomial coefficients fold into the client-ENCRYPTED W2 ciphertexts; its cubic coefficient is |c3| < 2e-18 for this checkpoint so the p^3 term is dropped (contribution ~2.4e-16, far below decode noise). The public structure masks travel as ciphertexts: creating plaintexts mid-trace leaves the replay read uninitialized towers, so everything the trace multiplies is ciphertext x ciphertext.

Depth 5 (mult depth 8), per position: 8 CTC mults, 21 rotations, 2 mask mults; circuit total 4 answer positions x that. Trace ~168k instructions at ring 4096 (~1.6 s local replay; V_GROUPS head groups scale it a few percent each).

Build and run

Requires a built niobium-client (Release) at $(NIOBIUM), default ~/code/niobium-client. Rust (cargo) is required; there is NO Python anywhere in this repo.

Version control is fossil (~/code/llm-fog-fhe.fossil). The cargo output dir (target/, ~1 GB of built binaries) is NOT tracked; it is excluded by the checked-in .fossil-settings/ignore-glob, so every checkout ignores it automatically.

make                # cargo build --release -> target/release/llm
make twin           # train the candle twin + activation-range check
make data           # re-export frozen data/ artifacts from the checkpoint
make cpu            # one-prompt end-to-end gate (PROMPT=... to override)
make sim            # hollow-record gate (Fog's recording path)
make sweep          # six memorized prompts + leak tripwire
make check          # pure-OpenFHE circuit vs plaintext diff diagnostic
make fog            # LIVE Fog roundtrip (ring 2^16; math runs on the worker)
make fog-sweep      # all six prompts as live Fog jobs + leak tripwire

Binaries behind the CLI:

./target/release/llm twin train        # train (GELU fit computed at runtime)
./target/release/llm twin demo --prompt "what do cows drink"
./target/release/llm twin range        # max |activation| after retraining
./target/release/llm twin export       # regenerate all data/ artifacts
./target/release/llm twin diff         # candle vs plain-f64 differential

./target/release/llm client keygen  <client_home> [ring_dim]
./target/release/llm client encrypt <client_home> <server_home> "<prompt>" [app_root]
./target/release/llm server run     <server_home> [app_root] [--hollow] [--no-ring-dim-check] ...
./target/release/llm client decrypt <client_home> <server_home> [app_root]
./target/release/llm server check   <client_home> <server_home> [app_root]

The local gates use ring 4096 for speed (--no-ring-dim-check); the Fog runs only N = 2^16, so make fog / make fog-sweep keygen at ring 65536 and submit via fogrs (pure-Rust Fog CLI; override with FOGSUBMIT=fog). Under the hood, the step the worker performs is exactly

fogrs submit target/release/llm server run build/fog_home/server . --hollow --target=FOG

— record happens with hollow (placeholder) values locally, and every ciphertext multiply/rotation executes on the worker during replay. Only build/fog_home/server ships; the client dir is the secret half (see the security-model section above). Everything that needs the secret key — keygen, encryption, decryption, argmax — stays client-side by design and can never move to the Fog.

Suggested steps (how to actually use it)

If you just want the demo working, do 1-3. If you want to run on the real Fog hardware, continue to 4. 5 is the loop for changing the model.

Build: make — produces target/release/llm.
Understand what goes where: skim the security-model section. client_home keeps the secret key, prompt, and answer; server_home carries only ciphertexts + encrypted weights and is the exact directory a Fog worker receives.
Local roundtrip (ring 4096, math replayed by fhetch_sim): make cpu for one prompt ("what is the capital of france" -> paris; override with make cpu PROMPT="what do cows drink"), then make sweep to exercise all six memorized prompts plus the leak tripwire. Everything green => the circuit and the trained twin agree.
Live Fog roundtrip (ring 2^16, math runs on a Mistic Fog worker): make fog for one prompt — keygen, hollow record, fogrs submit (~1-4 min per job incl. upload), decrypt, corpus assertion; then make fog-sweep for all six prompts as sequential Fog jobs. Needs valid Fog credentials (~/.fog/credentials); Fog quota is 4 in-flight jobs, so these deliberately run one at a time.
Change the model (the vocab ceiling is structural — raising V is now data + retrain, see design/VOCAB_SCALE.md): a. Edit src/corpus.rs: replace VOCAB_FILLERS entries with real words, and/or add QA pairs. make twin retrains from scratch (~1,300 epochs; checkpoint shapes do not migrate across a V change). b. make data re-exports the frozen artifacts. c. Bump V in src/consts.rs AND the V mirror in src/llm_common.h — together. The runtime guard rejects any run where the two disagree (silent one-sided bumps misalign the output ciphertext groups). d. Re-run the local gates in order: make cpu -> make check -> make sim -> make sweep, then make fog before claiming anything about the Fog.

Budget intuition: the head (the only vocab-dependent part) costs 5 CTC per answer position per vocab group (32 words), so even V=10,000 (313 groups, 1,252 probes) is ~2 min of worker wall out of the 1 h Fog budget — the binding constraints are probe plumbing, encrypted- weight upload, and model capacity, not compute.

License

BSD-3-Clause — see LICENSE.