AI in U — Inference, Training, and Model Adaptation
U compiles ML pipelines to native binaries and WebAssembly. No Python, no pip, no CUDA toolkit. One language for data loading, model inference, training, and deployment. The Safebox runner compiles to a 34KB static binary — that is pure executable code (C compiled with -O2), statically linked against libc and libm, with no bundled model weights, no embedded runtime interpreter, and no dynamically loaded libraries beyond what the OS provides.
The ML runtime is 26,000+ lines of C across 40 headers: tensor ops, neural-network layers, training with autograd, flash attention, SIMD-accelerated matmul, quantized inference (Q4/Q8/BitNet), WebGPU compute shaders, multimodal encoders, prefix caching, session pooling, and a full Safebox runner. Everything described below is implemented and tested.
What U Can Do Today
Load any major model format
import Safetensors
import GGUF
// HuggingFace models (safetensors):
weights = Safetensors.load("model.safetensors")
ww = Safetensors.get(weights, "model.layers.0.self_attn.q_proj.weight")
// Handles F16, BF16, F32, F64 — auto-converts to double.
// llama.cpp models (GGUF):
model = GGUF.parse(buffer)
ww = GGUF.get(model, "blk.0.attn_q.weight")
// Handles Q4_0, Q8_0, F16, F32, BF16 — auto-dequantizes.
// Metadata (layers, heads, vocab, context length) is in the header.
arch = GGUF.get_str(model, "general.architecture") // "llama"
Run transformer inference (Llama / Nemotron / Mistral / Qwen)
Every piece of a transformer forward pass is implemented and tested:
import NN
import Safetensors
import Tokenizer
import KVCache
f main()
// Load model:
sf = Safetensors.load("nemotron-nano-4b.safetensors")
model = llama_load(sf, 32, 32, 3072, 9216, 131072)
// Tokenize:
tok = Tokenizer.load(vocab_text)
tokens = Tokenizer.encode(tok, "What is the capital of France?")
// Generate (autoregressive with KV cache):
output = llama_generate(model, tokens, tokens.length, 128, 0.7, 0.9)
print(Tokenizer.decode(tok, output))
The forward pass: embedding → per-layer (RMSNorm → Q/K/V → RoPE → multi-head attention with causal mask → KV cache → output projection → residual → RMSNorm → SwiGLU FFN → residual) → final norm → LM head → top-p sampling.
Train models (XOR example — tested, trains to 4/4 accuracy)
import NN
import Tensor
import Autograd
f main()
xx = Tensor.from_data([0,0, 0,1, 1,0, 1,1], [4, 2])
yy = Tensor.from_data([0, 1, 1, 0], [4, 1])
w1 = Tensor.he_normal([8, 2])
b1 = Tensor.zeros([8])
w2 = Tensor.he_normal([1, 8])
b2 = Tensor.zeros([1])
adam = NN.adam_new(4, 0.01, 0.9, 0.999, 1e-8, 0.0)
[0..2000].on(epoch => (
hidden = Tensor.relu(NN.linear(xx, w1, b1))
pred = Tensor.sigmoid(NN.linear(hidden, w2, b2))
loss = NN.mse_loss(pred, yy)
// ... backward + adam_step
))
// Result: [0.016, 0.992, 0.993, 0.007] — 4/4 correct
Flash Attention
The runtime implements FlashAttention-2: tiled QKV in SRAM-sized blocks, online softmax, no materialisation of the full attention matrix. Two entry points cover the two cases that matter:
// Decode — single new token against the full KV cache:
u_flash_attention_decode(Q, K, V, n_heads, head_dim, seq_len, out)
// Prefill — multiple query tokens with a causal mask:
u_flash_attention_prefill(Q, K, V, n_heads, head_dim, n_queries, kv_len, out)
Tile size is tuned for L1 cache. Prefill is O(n·d) memory instead of O(n²), which is what makes long-context inference fit on machines that cannot allocate a 32K×32K attention matrix.
SIMD-accelerated inference
The SIMD layer (u_simd.h, 695 lines) dispatches at compile time to AVX-512, AVX2, or ARM NEON. Key operations: dot product, saxpy, sum-of-squares, vector add/scale, and max — all f32.
Quantized matmul is SIMD-accelerated too. The runtime does Q4_0, Q8_0, Q4_K, and Q6_K matrix multiplications directly on the compressed blocks — no full dequantisation pass, no temporary allocation:
// A [4096,4096] weight matrix:
// F32 = 64 MB Q8_0 = 18 MB (3.6x) Q4_0 = 9 MB (7.1x)
// Q4_K = 9.2 MB Q6_K = 13.4 MB
// Results match the dequant-then-multiply path exactly.
Int8 activation quantisation (u_int8_dot.h) accelerates prefill further — 2.3x over the f32 path — using VPDPBUSD on x86.
GPU dispatch — WebGPU / WGSL
U's +V modifier tags data for GPU residency. The runtime emits WGSL compute shaders for WebGPU, which runs on any GPU (not just NVIDIA). Nine kernels ship today:
| Kernel | What it does |
|---|---|
| u_gpu_linear_f32 | Dense linear layer (f32) |
| u_gpu_linear_q8_0 | Quantised linear (Q8_0 weights) |
| u_gpu_linear_q4_0 | Quantised linear (Q4_0 weights) |
| u_gpu_matmul | 16×16 tiled matrix multiply |
| u_gpu_rmsnorm | RMS normalisation |
| u_gpu_softmax_rows | Row-wise softmax |
| u_gpu_rope | Rotary position embeddings |
| u_gpu_silu_mul | SiLU activation × gate |
| u_gpu_attn_decode | Single-query attention decode |
Weight residency uses a pointer-keyed cache, so weights move to the GPU once and stay there. The async bridge (u_gpu_async.h) maps GPU readback into U's cooperative fiber scheduler — a GPU compute step becomes an +A pending value, no busy-waiting.
ONNX model loading
In addition to Safetensors and GGUF, the runtime loads ONNX initializer weights via u_onnx_loader.h. It maps both HuggingFace and GGUF naming conventions to a uniform F32 weight struct, and supports external data files. Load a model exported from PyTorch, TensorFlow, or JAX, and use the same forward-pass code:
import ONNX
weights = ONNX.load("model.onnx")
// Same ULlamaWeightsF32 struct — same forward pass code
Adapt models without retraining
LoRA hot-swap — load a 10-50MB adapter, apply per-request, no model reload:
import LoRA
import Safetensors
adapter = Safetensors.load("formal-voice-lora.safetensors")
lora = LoRA.load(adapter, 32, 8, 16.0)
logits = llama_forward_lora(tokens, n_tok, model, kv, lora)
// Switch voice by swapping adapter — same base model
Steering vectors — shift behavior with a ~1MB vector, adjustable strength:
import Steering
steer = Steering.load(sf, "formal_direction", 2.0, 15)
// Adds a direction to hidden states at layer 15
// "Be more formal" = one vector, "Be more creative" = another
Soft prompts — per-user tuning at ~1KB:
soft = SoftPrompt.load(sf, "user_alice_prefix")
combined = SoftPrompt.apply(token_embeddings, soft)
// Prepends 4 learned embedding vectors — invisible to the user
KV cache grafting — document memory without re-reading:
// Process a 50-page document once:
doc_kv = process_document(model, document_tokens)
KVCache.save(doc_kv, "/srv/kv-caches/doc-123.bin")
// Every future question about this document:
conv_kv = KVCache.new(32, 8192, 4096)
KVCache.graft(conv_kv, doc_kv, 0, 0, doc_kv.cur_len)
logits = llama_forward(question_tokens, n, model, conv_kv)
// Model "remembers" the document — no re-reading, no RAG chunking
BitNet — 10x cheaper inference
import BitNet
bw = BitNet.quantize(weight) // float → 2-bit ternary (-1, 0, +1)
result = BitNet.matmul(input, bw) // add/sub only, no float multiply
// 4B model: 8GB at F16 → 1GB at BitNet. CPU-viable.
Diffusion models
The diffusion pipeline is implemented in runtime/multimodal/diffusion.h (806 lines):
| Component | Status | Notes |
|---|---|---|
| Linear layers | ✅ | Dense projections |
| Conv2d / Conv2d transpose | ✅ | 2D convolution with stride/padding, plus transposed (learnable upsampling) |
| GroupNorm / LayerNorm | ✅ | Normalisation layers |
| GELU / SiLU | ✅ | Activations |
| Multi-head + cross attention | ✅ | Self-attention and cross-attention with tiled online softmax |
| Safetensors / GGUF loader | ✅ | Loads SDXL, FLUX, and multimodal GGUF weights |
| UNet architecture | ✅ | Residual blocks, time + text conditioning, down/up sampling |
| VAE decoder | ✅ | Residual blocks + upsample stages |
| Noise schedulers | ✅ | DDPM, DDIM, Euler, Euler-A, LCM |
| DiT backbone | ✅ | Transformer-based diffusion (SD3/FLUX architecture) |
| MMDiT | ✅ | Dual-stream transformer, adaLN-Zero, joint attention, rectified-flow sampler |
| Diffusion cache | ✅ | Caches text-encoder output, cross-attn K/V, timestep embeddings across steps |
| CLIP text encoder | 🔜 | Transformer encoder — ops are all there, architecture wiring remains |
The pipeline supports text-to-image, img2img, text-to-video, and 3D multi-view generation. The diffusion step cache avoids redundant text-encoder and cross-attention recomputation across denoising steps.
Server-side inference vs. client-side WASM
Server-side (the normal path): The u-runner binary loads the model, listens on a Unix socket, speaks the Safebox Protocol.LLM.Local wire format. The model stays on the server. Users send prompts, get responses. No model transfer.
// Server: load model, serve inference
./u-runner model.safetensors /run/safebox/services/u-llm.sock
// Client (Safebots/LLM.js, unchanged):
resp = await safebots.llm({
model: 'local/nemotron-nano',
system: systemPrompt,
userInput: userMessage
})
Client-side WASM (for small models): U compiles to WebAssembly. A classifier, sentiment model, or embedding model can run in the browser with zero server roundtrip. The WASM bridge (u_wasm_bridge.c) exports a clean API — init, load_model, generate_next, tokenize, detokenize, prefill, get_stats — and dispatches GPU matmul to JS-side WebGPU when available:
<!-- In the browser: -->
<script>
const wasm = await WebAssembly.instantiate(modelBytes, shim);
wasm.exports.u_wasm_init();
wasm.exports.u_wasm_load_model(modelPtr, modelLen);
// Runs inference locally — no network, no latency, no server cost
// GPU matmul dispatched to WebGPU when available
</script>
Use WASM for: classification (<10MB models), embeddings, tokenization, sentiment, small transformers. Use server for: LLMs (>1GB), diffusion, anything that needs a GPU.
With TensorFlow models
TensorFlow models export to safetensors via the safetensors Python package:
# Convert once (Python):
from safetensors.torch import save_file
# or from TF: convert TF SavedModel → ONNX → safetensors
Then load in U:
weights = Safetensors.load("converted-model.safetensors")
U doesn't run TensorFlow graphs natively (and doesn't need to — safetensors is the universal weight format). For TensorFlow.js models specifically, the recommended path is server-side U inference rather than transferring multi-GB models to browsers. The u-runner serves the same purpose as TF.js serving but without the JavaScript overhead and with the .u.meta capability proof.
U vs. vLLM vs. llama.cpp
| U | vLLM | llama.cpp | |
|---|---|---|---|
| Binary size | 34KB (code only; weights separate) | ~3GB (Python + PyTorch) | ~2MB |
| Cold start | <1s | 10-30s | 2-5s |
| Weight formats | Safetensors + GGUF + BitNet | Safetensors | GGUF |
| Quantization | Q4_0, Q8_0, BitNet (ternary) | AWQ, GPTQ, FP8 | Q2-Q8, IQ, BitNet |
| GPU | +V → WebGPU (any GPU) |
CUDA only (+ROCm) | CUDA, Metal, Vulkan |
| Multi-user | fork() CoW — OS-level isolation | Internal batching — shared process | Server mode — shared process |
| KV cache | ZFS-backed, persistent, forkable | In-process, volatile | In-process, volatile |
| Adapters | LoRA hot-swap per-request | LoRA (restart required) | LoRA (restart required) |
| Training | Autograd + Adam built-in | No | No |
| Capability proof | .u.meta — compiler-proven |
None | None |
| WASM | Compiles to WASM | No | Yes (limited) |
| Safebox protocol | Native (Unix socket, HMAC) | Needs wrapper | Needs wrapper |
Where U is better
Provably safe inference. The binary's .u.meta section proves its capability surface. No o { Network } = the binary provably cannot exfiltrate prompts. No sandbox to escape — the capability doesn't exist in the compiled code. Neither vLLM nor llama.cpp can make this claim.
Persistent KV caches. vLLM and llama.cpp lose all KV caches on restart. U saves them to ZFS. System prompts, document memory, conversation state — all survive reboots. Fork a KV cache for branching conversations. Snapshot for rollback.
Per-request adapter swap. Load a different LoRA adapter for each user without restarting the model. vLLM requires a restart for adapter changes. llama.cpp requires restart or pre-loading all adapters.
Native Safebox integration. The u-runner speaks the Infrastructure's HMAC-authenticated Unix-socket protocol directly. vLLM and llama.cpp need a Python/Node wrapper to translate.
Where U trades off
GPU throughput. vLLM's PagedAttention and continuous batching on CUDA are highly optimized for multi-tenant GPU serving. U's +V GPU dispatch emits WGSL compute shaders for WebGPU (any GPU), but does not yet match CUDA-tuned kernels on peak throughput. U's advantage is on CPU inference and mixed CPU+GPU deployments where provable safety, persistent KV caches, and zero-dependency binaries matter.
Quantized matmul. ✅ Now implemented. U does Q4_0 and Q8_0 matmul in-place — weights stay compressed, no dequantization allocation. A [4096,4096] weight matrix: F32 = 64MB, Q8_0 = 18MB (3.6x savings), Q4_0 = 9MB (7.1x savings). Results match the dequant-then-multiply path exactly.
Continuous batching. vLLM batches multiple requests into one GPU kernel call. U uses fork() for isolation instead. On GPU, batching wins. On CPU with many cores, fork() is competitive.
Model coverage. vLLM and llama.cpp have been tested against hundreds of model architectures. U has been tested against Llama-class (Llama, Nemotron, Mistral, Qwen). Other architectures (GPT-2, BERT, T5, Mamba) would need architecture-specific forward loops — the ops are all there, the composition differs.
The honest summary
U replaces vLLM for CPU inference on attested hardware where provable safety, persistent KV caches, per-request adapters, and zero-dependency deployment matter more than peak GPU throughput. It replaces llama.cpp for deployments that need the Safebox governance chain, ZFS integration, and .u.meta capability proofs. It doesn't yet replace either for maximum-throughput GPU serving.
Cache Breakpoint Protocol
U's runner implements a cache breakpoint protocol compatible with Anthropic's API and the Safebots/LLM.js five-block prompt discipline. The protocol gives unlimited breakpoints (Anthropic's API limits to 4).
How it works
The Safebots/LLM.js wrapper structures every prompt into four blocks:
Each breakpoint saves a KV cache snapshot. On the next call, the runner checks prefix hashes from deepest to shallowest. The deepest match gives the most cached tokens — only the suffix after that match runs through the model.
Wire format
How it maps to Anthropic / OpenAI
| Provider | Cache mechanism | U equivalent |
|---|---|---|
| Anthropic | on content blocks, max 4 breakpoints | Unlimited breakpoints, per-bot tenant isolation, LRU eviction |
| OpenAI | Automatic prefix detection (no explicit control) | Explicit prefix hashes — deterministic cache hits, no guessing |
| U runner | KV cache snapshots on ZFS, fork()-shared across sessions | Survives restarts, branchable, graftable |
Integration with Safebots/LLM.js
The existing wrapper already hashes the prefix and manages slots. The U runner replaces the slot backend:
The wrapper computes for each block boundary and sends them to the runner. The runner looks up the deepest matching KV cache snapshot, grafts it into the session, and processes only the suffix. On a typical multi-turn conversation, the system prompt (500-2000 tokens) and pinned context (500-1000 tokens) are cached from the first call — subsequent turns only process the new user message.
Multimodal — Vision, Audio, Diffusion
The runtime/multimodal/ directory (3,236 lines) implements the encoder side of multimodal models:
Vision (ViT encoder)
Patch embedding, learnable position embeddings, bidirectional self-attention, CLS or mean pooling. Connects to a language model via cross-attention or prefix fusion. The cross-attention module (u_cross_attn.h, 493 lines) implements Q-from-decoder / KV-from-encoder attention with tiled online softmax — the same mechanism LLaVA uses for image understanding.
Audio (Whisper-style encoder)
STFT → mel filterbank → log-mel spectrogram → Conv1d downsampling → sinusoidal position embeddings → transformer encoder. The pipeline matches Whisper's architecture, so Whisper weights load directly via the multimodal GGUF loader.
Cross-modal fusion
Two fusion modes: PREFIX (prepend encoder output to the token sequence — the LLaVA approach) and CROSS_ATTN (interleaved cross-attention layers — the Flamingo approach). Both are configurable per-model.
Diffusion pipeline
Full noise schedulers (DDPM, DDIM, Euler, Euler-A, LCM), UNet with time and text conditioning, DiT backbone, and VAE decoder. Supports text-to-image, img2img, text-to-video, and 3D multi-view generation. The diffusion cache (u_diffusion_cache.h, 584 lines) caches text-encoder output, cross-attention K/V, and timestep embeddings across denoising steps — so only the changed noise step recomputes.
MMDiT (SD3 / FLUX)
The u_mmdit.h header implements Stability AI's MMDiT block: dual-stream transformer with adaLN-Zero conditioning, joint attention, and GELU-tanh MLP. Paired with a rectified-flow Euler sampler with time shift — the architecture behind FLUX and SD3.
SGLang-compatible serving
The sglang_runner.h (571 lines) implements a serving layer compatible with SGLang's execution model: radix attention tree for prefix sharing, continuous batching, chunked prefill, and constrained decoding (JSON schema and regex constraints). This is separate from the Safebox runner — it targets workloads where SGLang's request scheduling is already in use.
Compact Multi-Context Storage
Serving inference has three memory costs: model weights (immutable, identical across sessions), KV cache for shared prefixes (identical until a session diverges), and KV cache for each session's unique tokens (always per-session). A naive server duplicates all three. U's runtime applies copy-on-write at every level — OS page tables for weights, a radix tree for shared prefix KV, refcounted pages for divergent KV — so each cost is paid once at the appropriate scope, sharing is exact, and the model sees no difference.
Radix-tree prefix cache — the killer feature
The prefix cache (u_prefix_cache.h, 848 lines) is U's most important data structure for multi-user serving. It indexes cached KV state by the tokens themselves — a compressed-edge radix tree where a 2,000-token system prompt is one edge, not 2,000 nodes. Storage is proportional to distinct tokens across all sessions, not total tokens:
// 1,000 users with the same system prompt:
//
// root
// └── [system prompt: 2,000 tokens] ──┬── [user A turn 1] ── ...
// ├── [user B turn 1] ── ...
// └── [user C turn 1] ── ...
//
// The system prompt's KV is stored ONCE. Each user adds only their
// unique suffix. 1,000 × 2,000 tokens of redundant KV → 0.
Each node stores KV for its own token range only, not the full prefix. Restoring a session walks root→node copying each segment into place — a memcpy per layer per segment, several orders of magnitude cheaper than re-prefilling those tokens. Thread-safe with mutex + pin/release reference counting. LRU eviction when the cache exceeds its memory budget. Serialises to disk, so the cache survives process restarts.
Why this matters more than it sounds: every inference engine copies the system prompt's KV cache per session, or re-computes it from scratch. With the radix tree, the system prompt is computed once and never again — not copied, not serialised, not re-prefilled. The memory saving is linear in the number of concurrent users (shared prefix KV is stored once; each user's cost is only the tokens where their session diverges). The latency saving is linear in the prompt length: a 4,000-token system prompt at 7B costs ~35ms to prefill on a modern GPU at typical throughput (~100K tokens/s). After a deploy, 1,000 users reconnecting would each re-prefill that prompt — 35 seconds of aggregate GPU work. With the radix cache, the prompt's KV is already in the tree; restoring it is a pointer walk and a few memcpys.
Page-level copy-on-write
The CoW cache (u_kv_cow.h, 546 lines) divides KV state into refcounted 64-token pages. When two sessions share a prefix, they share the same physical pages. Writing to a shared page copies it first — CoW semantics, the same as a virtual-memory page fault, but in userspace with atomic refcounts:
// Fork a session's KV state:
UKVCacheCow* child = u_kv_cow_snapshot(parent);
// Cost: bump refcounts on existing pages. No data copied.
// The child diverges only when it appends new tokens — and only
// the page being written to is copied, not the whole cache.
Named checkpoints map string IDs to snapshots, so the Safebox runner can save, branch, and graft conversation state by name.
The runtime supports two execution modes for concurrent inference, chosen at deployment time:
Thread-based serving
The Safebox runner (u_safebox_runner.c) spawns one pthread per connection. All sessions share a single process, and the userspace CoW cache above provides per-session KV isolation: writes to a shared page copy the page first, so sessions cannot corrupt each other's state. This is the portable mode — it works on Linux, macOS, and Windows, and avoids the per-request overhead of a process fork. The tradeoff is that isolation is logical, not physical: a crash or memory corruption in one thread can bring down the entire process.
Fork-based process isolation (Linux / macOS)
On Unix systems, fork() gives you copy-on-write at the hardware level — the OS duplicates the page table, not the pages themselves. The session pool (u_session_pool.h, 388 lines) exploits this: the parent process loads the model weights and pre-warms the system prompt's KV cache, then each incoming request fork()s. The child starts with the entire address space — model weights, KV cache, tokenizer, LoRA adapters — already resident. No serialization, no deserialization, no copies. The thaw latency is ~0.1ms (the fork syscall itself).
What this means concretely:
- Model weights are loaded once via
mmap()withMAP_PRIVATE. Every forked child shares the same physical pages for the entire model. A 4B model at Q4_0 (~2.5GB) is 2.5GB of RAM no matter how many children are running. - KV cache — the warmed-up system prompt's KV state is in the parent's heap. After fork(), all children share those pages until they write. A child that appends its own tokens copies only the pages it touches — typically a few MB for a single turn, not the full cache.
- Isolation is physical, not logical. Each child is a separate process with its own address space. A crash, a timeout, or a hostile input in one session cannot corrupt another. The OS enforces this, not the runtime.
The tradeoff is platform support (no fork() on Windows) and per-request overhead (~0.1ms for the fork syscall plus page-table duplication). For high-throughput deployments on Linux where fault isolation matters — multi-tenant hosting, untrusted inputs — fork-based serving is the stronger guarantee.
Disk tier — ZFS snapshots and mmap
fork() is fast but volatile — everything is lost on restart. ZFS extends copy-on-write to disk. The runner saves KV state to a file on a ZFS dataset; the adapter takes a snapshot. The on-disk format matches the in-memory layout exactly (flat arrays of float or double, one per layer), so mmap() restores it zero-copy — no deserialization pass, no parsing, no allocation.
| Tier | Mechanism | Thaw latency | Survives restart | Isolation |
|---|---|---|---|---|
| Thread | Userspace CoW pages (u_kv_cow.h) | ~0.01ms | No | Logical (in-process) |
| Fork | fork() — OS page-table CoW | ~0.1ms | No | Process-level (OS-enforced) |
| ZFS | Snapshot → mmap | ~1ms | Yes | Any process on the same ZFS pool |
On reboot, the pool calls u_session_pool_load_warmup(pool, path). That mmap()s the saved KV file and registers it as the "system-prompt" checkpoint. The data is page-faulted in on demand — a 2GB KV cache "loads" in ~1ms; only the pages actually read during inference are brought into RAM. Branching a conversation (fork one checkpoint into two) is a ZFS clone, which is metadata-only regardless of size.
The same mmap path handles model weights. A GGUF or safetensors file on a ZFS dataset is mmap()'d read-only (u_mmap_open). The OS pages it in as needed. On a machine with 32GB of RAM serving a 70B model at Q4_0 (~40GB), the OS keeps the hot layers resident and pages cold ones. ZFS's ARC (adaptive replacement cache) acts as a second tier — frequently accessed blocks stay in the ARC even after the page cache evicts them.
Why this architecture fits mixture-of-experts
MoE models (Mixtral, OLMoE, Qwen2-MoE) activate only a subset of their parameters per token — typically 2 of 8 experts per layer. The mmap architecture is a natural fit: all expert weights are mapped into the address space, but only the activated experts' pages are faulted into RAM. A 47B Mixtral (8 experts, 2 active) pages in roughly 12B parameters' worth of pages per forward pass, with the OS and ZFS ARC keeping the popular experts hot and paging cold experts on demand. The GGUF loader parses MoE metadata (expert count, top-k, per-expert FFN size), loads the stacked expert weight tensors (ffn_gate_exps, ffn_up_exps, ffn_down_exps) and the router gate (ffn_gate_inp), and the forward pass branches at the FFN block: dense models run the existing SwiGLU path, MoE models route through softmax top-k selection and per-expert SwiGLU with renormalized gating weights. An optional expert LRU cache issues madvise(MADV_WILLNEED) for selected experts before the FFN fires, with hot-set pinning for experts that exceed a usage threshold and router lookahead prefetch that overlaps I/O with the attention computation of the current layer.
Named checkpoints — branch, graft, rollback
The checkpoint store (up to 256 named entries) supports three operations beyond save/load:
// Branch: two users continue from the same document context
u_checkpoint_fork(store, "doc-review-v3", "alice-branch");
u_checkpoint_fork(store, "doc-review-v3", "bob-branch");
// Both branches share the document's KV pages until they diverge.
// Graft: attach document memory to a new conversation
u_checkpoint_graft(store, "long-document", "alice-chat", "alice-with-doc");
// alice-with-doc = alice's conversation state + the document's KV suffix.
// The model "remembers" the document without re-reading it.
// Rollback: restore a conversation to an earlier point
u_checkpoint_load(store, "alice-turn-3");
// Named checkpoints make undo trivial.
Int8 KV quantization — 3.56x memory reduction
For long-context inference, the Q8 KV cache (u_kv_q8.h, 198 lines) quantises cached keys and values to int8 with a group size of 32, reducing per-value storage from 4 bytes to 1.125 bytes. Combined with prefix sharing, the savings compound: a shared 2,000-token system prompt stored once at int8 costs ~0.28x what one copy at f32 would. Opt-in per session:
u_inference_ctx_enable_kv_q8(ctx)
// From this point, new KV entries are stored at int8.
// Scratch buffer is reused across layers — no per-layer allocation.
// Attention code is unchanged — it sees dequantised f32, so the
// quantised path cannot alter attention semantics.
Diffusion cache — reusing previous images and prompts
The diffusion cache (u_diffusion_cache.h, 584 lines) applies the same principle to image generation. A diffusion sample runs the network 20–50 times over the same conditioning. Three things do not change from step to step:
| Cached | What it avoids | Saving at 50 steps |
|---|---|---|
| Text-encoder output | Re-encoding the prompt every step | 49 full encoder passes (98%) |
| Cross-attention K/V projections | Re-projecting K = enc @ Wk, V = enc @ Wv per layer per step | 49 × n_layers projection passes |
| Timestep embeddings | Recomputing sinusoidal embeddings from the schedule | Precomputed once for the whole schedule |
All three are exact — identical inputs produce identical outputs, no approximation anywhere.
The fourth reuse is trajectory resume: caching the latent at each denoising step so variations can start from a checkpoint instead of from noise. Two modes:
// EXACT resume — continue an interrupted sample, or raise the step count.
// The key covers seed + conditioning + schedule + guidance scale.
// A hit is the latent a full run would have produced.
latent = u_traj_resume_exact(cache, key, max_step, &found_step, &n);
// APPROXIMATE resume — "same image, adjusted prompt".
// Returns a latent from a DIFFERENT conditioning that shares the seed.
// Early steps fix composition; the new prompt drives the rest.
// Handoff at 30–50% of the schedule keeps layout, changes details.
latent = u_traj_resume_similar(cache, seed, handoff_step, &found_step, &n);
For iterative workflows — "same scene, make the sky darker" — the conditioning cache skips the text encoder entirely when the prompt repeats (hash + string compare, collision-safe), and the trajectory cache lets img2img variations start from a stored intermediate rather than re-denoising from scratch.
For MMDiT architectures (SD3, FLUX), the joint-attention cache tracks text and image halves separately. When only the prompt changes, the image-stream K/V stay valid across all layers and only the text rows are re-projected — halving the per-step projection cost on prompt edits.
What you can build with this
The caching layers above are not just optimisations — they enable architectures that would be impractical without them:
- Chatbot with 1,000+ concurrent users on a single GPU. Model weights: one mmap'd region, shared across all forked children (OS CoW). Prompt KV: stored once in the radix tree, shared by every session with that prefix. Each user's marginal cost is only the KV for their unique tokens — typically a few MB per turn, not the ~2GB a full 8K cache would cost.
- Document Q&A over a 100-page PDF. Prefill the document's KV once. Save it as a named checkpoint. Every subsequent question grafts the document checkpoint onto the user's conversation — the model "remembers" the document without re-reading it. At 32K tokens, skipping the re-prefill saves ~200ms per query on GPU, more on CPU.
- Interactive coding assistant with undo. Named checkpoints make conversation branching cheap. Roll back to any saved point, fork two alternatives, compare them. Each branch shares all the KV state that preceded the fork — the cost of exploring is proportional to the tokens that differ, not the total context.
- Zero-downtime deployment on a single machine. Save the warmed-up state to ZFS. Restart the process. The new process mmap's the checkpoint in ~1ms. No user re-prefills their system prompt. Compare to restarting a vLLM server, where every connection must re-warm from scratch.
- Image generation with iterative editing. The diffusion trajectory cache lets "same scene, darker sky" resume from step 15 instead of step 0. The conditioning cache skips the text encoder when the prompt repeats. The cross-attention cache skips K/V projection when only the latent changes. Each of these is exact — no approximate caching, no quality loss.
- Multi-tenant serving with process isolation. Each forked child is a separate OS process. A crash, OOM, or runaway generation in one session cannot touch another. The OS enforces this. In single-process servers (vLLM, llama.cpp's server mode), a bug in one request's KV management can corrupt every other session's state.
What this means in practice
| Scenario | Without U's caching | With U's caching |
|---|---|---|
| 1,000 users, same system prompt (2K tokens), 7B model | ~2TB KV (1,000 × 2GB) | ~2GB shared + per-user deltas |
| Model weights (physical RAM), 50 concurrent sessions, 4B Q4_0 | ~125GB if each process loads its own copy | ~2.5GB physical RAM — one mmap'd region, shared across all forked children via OS page-table CoW. Per-session KV still grows with each user's unique tokens. |
| Cold start after reboot | Re-prefill system prompt (~50ms per user) | mmap from ZFS (~1ms), then fork (~0.1ms per user) |
| New request, steady state | Load model + prefill prompt | fork() from parent (~0.1ms, model + KV already resident) |
| Branch a conversation | Deep-copy entire KV cache | Refcount bump (CoW pages) or ZFS clone (metadata-only) |
| Attach document memory to a new chat | Re-read the document | Graft the document's KV checkpoint (~0.1ms) |
| Long context (8K) with Q8 KV | ~2GB per session | ~560MB per session (3.56x reduction) |
| 50-step diffusion, same prompt | 50 encoder + projection passes | 1 encoder pass + 1 projection; 49 cached (98% saved) |
| Diffusion prompt edit ("darker sky") | Full 50-step re-denoise | Resume from cached step 15–25; denoise only the remainder |
The Runtime — Full Inventory
26,000+ lines of C across 40 headers in runtime/ml/ and runtime/multimodal/:
| Category | Count | Functions |
|---|---|---|
| Tensor ops | 42 | zeros, ones, full, eye, from_data, arange, linspace, rand, randn, xavier, he, add, sub, mul, div, scale, neg, abs, sqrt, clamp, where, sum, mean, max, min, argmax, argmin, dot, norm, matmul, transpose, reshape, cat, stack, relu, sigmoid, tanh, softmax, gelu, swish/silu, leaky_relu, elu, relu6 |
| NN layers | 11 | linear, conv2d, maxpool2d, layer_norm, rms_norm, batch_norm, dropout, embedding, attention, rope, causal_mask |
| Training | 15 | autograd tape, leaf, backward, grad, sgd, sgd_momentum, adam, adam_step, zero_grad, lr_cosine, lr_warmup, lr_step_decay, grad_clip_norm, grad_clip_value, mse_loss, cross_entropy |
| Gradients | 10 | sigmoid, tanh, gelu, softmax, relu, mse, cross_entropy, layer_norm, linear (input+weight+bias), tape (add, mul, matmul, scale) |
| Safetensors | 3 | parse (F16/BF16/F32/F64), get, info |
| GGUF | 6 | parse, get (Q4_0/Q8_0/F16/F32/BF16), get_u32, get_f32, get_str, info |
| BitNet | 2 | quantize (float→ternary), matmul (add/sub only) |
| Transformer | 5 | llama_forward, llama_forward_lora, llama_load, llama_generate, sample_top_p |
| LoRA | 3 | lora_load, linear_lora, forward_lora |
| Steering | 3 | load, apply, compute |
| Soft prompts | 2 | load, apply |
| KV cache | 8 | new, append, advance, get_k, get_v, reset, save, load, graft |
| Session pool | 6 | mmap_open, mmap_close, pool_new, pool_warmup, session_handle, pool_reap |
| Tokenizer | 4 | load, encode, decode, encode_bytes |
| Flash attention | 2 | flash_attention_decode (single query, online softmax), flash_attention_prefill (multi-query, causal) |
| SIMD | 11 | dot_f32, saxpy, sscale, sumsq, vadd, max (AVX-512/AVX2/NEON), matmul_q8_0_simd, matmul_q4_0_simd, matmul_q4_k_simd, matmul_q6_k_simd, matmul_q8_0_i8 |
| GPU (WGSL) | 9 | gpu_linear_f32/q8_0/q4_0, gpu_matmul, gpu_rmsnorm, gpu_softmax_rows, gpu_rope, gpu_silu_mul, gpu_attn_decode |
| Prefix cache | 6 | prefix_cache_new, lookup, insert, restore, save, load (radix tree, LRU, disk-persistent) |
| Session pool | 6 | mmap_open, mmap_close, pool_new, pool_warmup, session_handle, pool_reap |
| CoW KV cache | 5 | kv_cow_new, kv_cow_write, kv_cow_fork, kv_cow_checkpoint, kv_cow_restore (page-level, refcounted) |
| ONNX | 3 | onnx_open, onnx_close, onnx_find (initializer → F32, HuggingFace + GGUF naming) |
| Multimodal | 12 | vit_encoder, audio_encoder (Whisper), cross_attention, patch_embed, sinusoidal_pos_2d, vision_encoder_block, diffusion pipeline, MMDiT block, fusion (prefix + cross-attn) |
| Dataframe | 15 | from_csv_string, head, sort, select, drop, filter_mask, unique, sample, count, groupby_agg, with_column, join, get_column, to_table |
| Column | 12 | sum, mean, min, max, gt, lt, eq, str_len, str_upper, str_lower, str_contains, to_tensor |
Quick Start for the Safebox Infrastructure
The Safebox runner (u_safebox_runner.c, 1,317 lines) is a full HTTP/1.1 inference server. It speaks both Unix sockets and TCP, streams responses via SSE, and exposes internal endpoints for KV cache save/load/fork/graft — everything the Safebox infrastructure needs to manage model instances.
# Build the runner (one command):
gcc -O2 -o u-runner u_safebox_runner.c -I runtime/ -lm
# Result: 34KB binary (executable code only — model weights are mmap'd at runtime)
# Deploy alongside existing vLLM runners:
cp u-runner /opt/safebox/runners/
./u-runner /srv/safebox/models/<hash>/model.safetensors \
/run/safebox/services/u-llm.sock
# The runner speaks the same protocol as vLLM:
# Safebots/LLM.js → Protocol.LLM.Local → u-runner
# No code changes in Safebots or Safebox.