ONNX export, quantization and the browser runtime #

This describes the inherited current baseline and pinned WebAssembly runtime. Parameter/tensor counts, 4-bit/8-bit layout, context budgets and download, memory and timing figures are baseline facts or historical receipts, not universal constraints or measurements of a replacement. The rebrand changes neither weights nor engine. Private experiments now train and export genuine GPT-2, GPT-NeoX/Pythia and standard Llama/GQA candidates; none is selected or published. Their own receipts are in training.md.

The tools inspect actual config, token tables, projection topology and separate rank-four KV tensors through train/onnx_utils.py. Tied tables are shared; untied Pythia heads remain independent. These are exercised family layouts, not a numerical-parity pass for every candidate or universal support: unfamiliar/merged cache contracts fail rather than receiving baseline-shaped dummy tensors. A new family still needs real export, cache, quality and browser proof (extending.md).

Export #

train/export.sh CHECKPOINT OUT [q4|q8|fp32] (default q4) produces a fresh, atomically staged directory for transformers.js. Nonempty destinations are refused. Tokenizer files are copied from the prepared checkpoint, not an exporter-resaved tokenizer:

OUT/config.json  generation_config.json  tokenizer.json  tokenizer_config.json
    special_tokens_map.json
OUT/onnx/model_q4.onnx           q4: 4-bit body, int8 token table/head
    # or model_quantized.onnx   q8: int8 weights (not model_q8.onnx)
    # or model.onnx             fp32: unquantized weights
OUT/SHA256SUMS export.json quantization.json
OUT/eval/model_<dtype>_full.onnx same weights, logits for every position
                                (quant_eval.py only; never published)

It runs optimum-cli export onnx --task text-generation-with-past on CPU in a dedicated export environment ($EXP, or the tools on PATH if EXP is unset), then quantize.py. The exercised export environment uses PyTorch 2.14.1 CPU, transformers 4.57.6, optimum 2.1.0, optimum-onnx 0.1.0, ONNX 1.23.2, ONNX Runtime 1.30.0 and accelerate 1.15.0. Training also uses transformers 4.57.6 but has a separate ROCm Torch wheel; keep environments isolated. The exporter needs accelerate for real tied-table de-duplication. Receipts record checkpoint/graph hashes, versions, parameter dimensions and quantization decisions; an unfamiliar architecture is not an exporter promise.

Quantization (train/quantize.py) #

The published baseline uses transformers.js dtype q4 (model_q4.onnx). Private experiments additionally support q8 and fp32. The current procedure is:

1. gemm_to_matmul() converts GPT-2 Conv1D Gemms into MatMul + Add. NeoX and Llama's actual linear projections are inspected directly.
2. quantize_model() stores q4 body matrices as asymmetric packed UINT4, block size 32 along K, fp32 scales and uint4 zero points. Blocked DequantizeLinear uses opset 21. A K dimension not divisible by 32 receives an explicitly recorded int8 fallback, never a silently padded matrix.
3. Actual token embeddings and output projections use per-axis int8 in q4/q8; q8 body weights also use int8. No activation quantization is introduced. A real tied table is stored once; independent input/head tables remain independent, with their actual dimensions rather than baseline names.
4. last_token_logits() slices the actual head input to the last position. The full-logits evaluation twin retains every position for canonical masked loss, including diary EOS. fp32 exports skip weight quantization.

The runtime must set session.disable_quant_qdq = 1 (web/brain.ts passes it, with session.disable_prepacking = 1, through transformers.js session_options.extra.session). With it, ONNX Runtime constant-folds the DequantizeLinear nodes at load and runs plain fp32 kernels. Without it, its QDQ fusion turns them into MatMulNBits, the 4-bit matmul, which in the WebAssembly build is several times slower (the benchmark below did not finish loading in 2 minutes with fusion on).

Measurements (headless Chromium on a 44-core x86 server, 4 threads, stock GPT-2 small, test/bench.html: 32 greedy tokens after a 58-token prompt; memory from performance.measureUserAgentSpecificMemory(), which includes the WebAssembly heap at its high-water mark):

File	Download	Decode	Page memory
fp32	498 MB	34 tokens/s	1,753 MB
8-bit (quantize_dynamic: MatMulInteger, v1's WASM format)	125 MB	25–38 tokens/s	831 MB
4-bit MatMulNBits body (ORT's own 4-bit kernel)	213 MB	7.2 tokens/s	not measured
4-bit unpacked with BitShift/Cast/Mul ops	95 MB	37.7 tokens/s	3,118 MB (folding keeps every intermediate)
4-bit DequantizeLinear, fp16 scales, separate head copy	91 MB	34.6 tokens/s	1,483 MB
… with fp32 scales, no prepacking	97 MB	34.8 tokens/s	1,129 MB
… and one shared embedding/head table (shipped)	97 MB	35.0 tokens/s	804 MB

Phones were not measured.

Why the embedding stays 8-bit. Fake-quantizing the base model on pet data (loss in nats): fp32 4.229; body 4-bit asymmetric block 32 4.255; tied embedding 4-bit (block 32 or 16) 5.6–7.0, broken; embedding int8 per row 4.293. quantize_dynamic's 8-bit head also hurt, because it quantizes the head's input activations per tensor and GPT-2 has large activation outliers; for the same reason MatMulNBits with accuracy_level=4 (int8 activations) must not be used here.

ORT quirk. Without disable_quant_qdq, a per-axis int8 DequantizeLinear on the embedding made ONNX Runtime 1.30's graph optimizer abort with a C++ assertion (_M_is_engaged()). With the option set it loads.

In the inherited base-model receipt, the exported q4 scores 4.280 against 4.195 for the fp32 ONNX export on the same rows (+2.0%). Fine-tuned baseline numbers are in training.md. Run check_onnx.py EXPORT_DIR PROMPT [CHECKPOINT] after any change to quantize.py, passing the actual checkpoint and --dtype. The checker compares cached decoding against fresh prefill at each prefix, requires fp32 numerical and greedy parity, and records (rather than demands identity of) quantized greedy/logit drift. Tolerances are explicit in receipts; a failed numerical check is not publication evidence. quant_eval.py uses the shared independent held-out split and strict fully covered loss masks, reporting both token NLL and bits per supervised UTF-8 byte. Default sessions disable QDQ/prepacking like the app; --qdq/--prepack are separate experiments.

Optional native hosted runtime #

The Rust host loads a configured q4/q8/fp32 export with native ONNX Runtime and a native tokenizer. SUBTOKENLING_MODEL_DTYPE defaults to q4. Supported GPT-2, GPT-NeoX/Pythia and Llama/GQA layouts use real config/graph cache shapes, logit vocabulary and EOS IDs, not GPT-2-specific constants. Unfamiliar layouts are rejected; load/generation smoke does not qualify the model's behavior. It preserves the QDQ/prepacking settings and empty per-call KV, and takes the context fit, sampling settings and line/lookup stop rules from the shared rules in core/, as the browser worker does. One mutable session belongs to each serialized worker. Local mode remains WebAssembly only. Remote mode does not download a browser model; returning local loads the pinned compatible brain from cache or downloads it. The server cannot silently upgrade the selected brain.

See server setup and library inputs. Native CPU correctness and browser transfers are checked separately from speed probes. The native CUDA path uses I/O binding to keep KV tensors on the GPU between tokens and returns only logits for CPU sampling, with TF32 disabled. Caches are per generation and never carry across pets. Real model/HTTP checks and limited warm-burst measurements are described in testing; they do not establish sustained capacity or public deployment readiness. Published model/runtime files remain immutable.

Browser runtime #

transformers.js 4.3.1 with ONNX Runtime Web 1.31.0-dev.20260914-8d85527a0, vendored by train/vendor.sh into runtime/transformers-4.3.1/:

• transformers.min.js, the self-contained browser bundle. The transformers.web.js build imports bare onnxruntime-web and needs a bundler; do not use it.
• the ort-wasm-simd-threaded*.{mjs,wasm} builds (plain, jsep, jspi, asyncify). A browser fetches only the one it needs. The app no longer uses WebGPU, but the runtime directory is published and immutable, so it is kept as is.

The ORT version must match the one the transformers.js bundle was built against; vendor.sh checks the version string and refuses otherwise. Upgrading means a new runtime directory name, never an in-place update.

web/brain.ts #

The TypeScript worker is built as standalone brain.js; its typed messages are defined in web/types.ts. The main-thread client and loading fallbacks are in web/app/brain.ts. Bundling does not change the vendored runtime or model session options.

The worker imports the bundle from CFG.runtime, sets allowLocalModels = false, allowRemoteModels = true, remoteHost = CFG.models and remotePathTemplate = "{model}/" (so modelId resolves to the configured static model directory), useBrowserCache = true, and wasmPaths = <runtime>/. The HTTP metadata path reads the model's actual tokenizer class rather than assuming GPT-2. This only fetches configured static assets; it does not enable Hugging Face fallback or server inference. It then loads AutoTokenizer and AutoModelForCausalLM with device: "wasm" and the selected release dtype (default q4; q8 uses onnx/model_quantized.onnx, fp32 onnx/model.onnx), runs a one-token warmup and posts ready. Download progress is aggregated across files and drives the egg's cracks.

Before loading the runtime the worker loads the shared rules module (beside brain.js). Each generate call carries the request the shared turn built: a context budget (ctx; 1024 for react turns, 512 for idle turns), the token budget, temperature, top-k and top-p. The worker fits the prompt with the shared fit (counting with the tokenizer), stops with the shared stop rule, and reports prefillMs (until the first new token), tps (decoding speed after it) and prompt, the exact text it ran after trimming old lines to fit (shown in the debug panel, docs/testing.md).

Generation currently prefills the whole prompt on every call, including each lookup continuation. There is no KV reuse between calls. Reusing a cache within one lookup turn is a possible later optimization; reusing one across turns is harder because the clock, needs and log change near the front of the header and old chat lines are trimmed. Neither optimization has been implemented or benchmarked; measure on a phone before prioritizing it.

Files are stored in Cache Storage under transformers-cache, keyed by their absolute URLs (the selected graph under <models>/<id>/onnx/ and its JSON files). The page also calls navigator.storage.persist() without waiting for its optional permission prompt before starting a download or upgrade. The model files then come from Cache Storage even if they are gone from the server; the runtime's .wasm comes from the service worker's copy, the HTTP cache or the network. When the brain is ready the worker also reports the runtime files it fetched, so the page can ask the service worker to keep them for offline use.

Which brain runs, and upgrades #

On a visit with a hatched pet, web/app/brain.ts looks for cached brains among modelId and the page config's older list and loads the newest one it has. An egg always downloads modelId.

If that is an older brain, the pet runs on it and the page offers the upgrade: a dialog once per visit (what's new from news, features the new brain adds, the download size, and a mobile-data warning when navigator.connection reports cellular or Save-Data) plus a banner above the device and tabs and a button in settings that stay until the visitor upgrades. Upgrading fetches the new model's files into transformers-cache with progress while the old brain keeps running, waits until the old brain is between turns, then swaps workers; the new brain loads from the cache without a second download (checked in the preview: one request for the new .onnx). Nothing about the pet changes. If the download fails the old brain keeps running and the banner offers to try again.

Once a brain is ready, the page deletes cached files of every other brain under the models URL (keeping modelId too while an older brain runs, so a half-finished upgrade can resume). "Delete brain" in settings removes them all.

Fallback #

A 90-second watchdog, re-armed by each progress message, retries a stalled load once with numThreads = 1 (no SharedArrayBuffer needed), then reports "the model runtime did not start". A load error is reported directly. Not exercised against a real stall. Reported load errors, worker crashes and synchronous failures constructing the worker or sending its initial load message all cancel the watchdog and release any created worker, so a manual retry starts a fresh worker with fresh progress. The egg shows a retry button; settings also retain the download button. Retrying keeps saved pet/egg state and completed cached files. Stream and ONNX protobuf-parse errors evict only the failing brain's cached ONNX, so retry can replace an incomplete/corrupt download. Unrelated runtime failures retain it. A failed stream starts that file over; downloads do not resume by range.

The service worker leaves .onnx requests to the browser, including at custom model URLs. Its cache-first fetch proxy used to cut off a slow download in Linux Firefox with Error in input stream after about 30 seconds. The model still uses transformers-cache; this bypass does not disable offline brains. See testing.md for the slow-download regression recipe.

Why no WebGPU #

v1 shipped a 4-bit fp16 WebGPU model next to the WASM one, and in a desktop browser with a GPU it was slower than WASM. That was reported by the owner, not measured by our scripts, so the cause below is an explanation by reasoning, not a profile. Generation runs one token at a time with batch 1: each step dispatches every kernel of the graph (hundreds of small GPU dispatches for 24 layers), copies the logits back for sampling in JavaScript, and synchronises with the GPU before the next token. For a model of this size that fixed per-step overhead is larger than the arithmetic the GPU saves, and ONNX Runtime's WebGPU backend may also compile new shaders when prompt lengths change. A CPU with SIMD and a few threads does the same small matrix-vector work without that round trip. With GPT-2 small the CPU path is fast enough, so v2 removed WebGPU: one model file, one engine, no adapter probing, no fp16 export, and no engine setting.

Threads and cross-origin isolation #

WASM threading needs SharedArrayBuffer, which needs cross-origin isolation. The web server should send Cross-Origin-Opener-Policy: same-origin and Cross-Origin-Embedder-Policy: require-corp for the app's base path (see deployment.md). Everything the page loads is same-origin (the font is vendored), so nothing breaks, and crossOriginIsolated is true. Adding any third-party script, font, image or iframe to the page will be blocked by COEP; host it under <base>/ instead.

Observed once on the dev server with v1, after flipping isolation on for a browser profile that had cached the runtime without those headers: ort-wasm-simd-threaded.asyncify.mjs failed with net::ERR_BLOCKED_BY_RESPONSE until the HTTP cache was cleared. It was not root-caused. Enable the headers before first publishing.

Speed #

tokenling-v2 (v3 and v4 have the same architecture and file size; not re-measured), measured with test/bench.html in headless Chromium on a 44-core x86 server, model served from the same machine, cross-origin isolated:

Threads	Load (from server)	Decode, greedy / sampled	Prefill, 58 tokens	Page memory
4 (ORT default)	1.6 s	34.4 / 32.6 tokens/s	153 ms	805 MB
1	1.6 s	14.6 / 14.5 tokens/s	546 ms	803 MB

In the app (same machine, 4 threads) a react turn with a few hundred tokens of history took 240–290 ms of prefill and decoded at 29–34 tokens/s; most pet lines are 5–20 tokens. On the live site a fresh profile went from "get an egg" to a ready brain in 3.5 s over localhost networking. Without cross-origin isolation WebAssembly runs on one thread (the second row). Phones were not measured; expect several times slower, and use bench.html on a device to find out.