Rust · WebGPU · no CUDA, no Python
How little does it take to run a language model?
Forge is the floor a dedicated, vendor-independent runtime has to stand on before it can carry real research — and eventually, local, safety-focused AI work — on hardware you control end to end.
Scroll down and it will run on yours, and tell you how fast.
verified — CPU and WGPU
logits agree to 8.4e-5, and GPT-2's
output matches HuggingFace
transformers token for token
(tests/gpt2_e2e.rs, on an RTX A5000).
Watch it think
A 10.77M-parameter character-level GPT — 6 blocks, 6 heads, 384 dimensions, a 65-character vocabulary — trained from scratch by Forge on Tiny Shakespeare and running here on your own GPU. Below is a single head of a single block: head 1 of block 1, doing the one calculation attention is made of.
Left to right: the Query for the
newest character, 64 numbers deep; the
Keys of the last 8 characters; the
score each key earns, which is nothing more exotic than
q · kᵢ / √64; the softmax of those
scores, drawn as the causal triangle; and the
Values that softmax weights, summed
into the head's output. Every cell is a number this model
computed on your GPU during the run beside it — not bucketed, not
averaged, not a stand-in. Hover one and it will tell you which.
A head is 64 of the model's 384 dimensions, which is small enough to show honestly; the whole layer at once is what made the earlier version of this section unreadable. The other five blocks are folded into the stack underneath — same code, same live attention, one slab each.
drag to rotate · scroll to zoom · hover any cell for its value
Nothing has been computed yet
Every cell here is filled in by the model as it runs. Start it and the numbers arrive one character at a time — 43 MB of weights, fetched once and then cached by your browser.
Next character
The model's own probabilities, over the whole vocabulary.
The calculation, in words
Your browser has no WebGL, so the 3D stage is not drawn. The model still ran on your GPU and the numbers beside it are real. What the stage would show is one head of one block: the newest position's Query, a vector of 64 numbers, scored against the Key of every position before it as q · kᵢ / √64; a softmax over those scores, which is where the causal mask makes the picture a triangle; and the Values of those same positions, weighted by that softmax and summed into the head's output. Six heads do that at once, their outputs are concatenated back to 384 dimensions, and an MLP that widens 384 → 1536, applies GELU and comes back writes into the same residual stream — six blocks deep.
Weights (43 MB) are fetched only when you press Run, then cached by your browser.
No server does the work — your GPU does. Forge needs no
SharedArrayBuffer, because
rayon is native-only, so this page is
plain static files.
What a token costs
10.77 million parameters is not a small language model by the industry's use of the phrase — it is three orders of magnitude below it. The research question is what that floor actually buys, and the only answer worth printing is the one measured on the machine in front of you.
Nothing here is filled in until you press Run above — every figure is produced by your own hardware, so none of it can go stale.
What is not in that budget is as much the point: no CUDA toolkit, no Python interpreter, no PyTorch, no ONNX runtime, and no server round trip. The cost of a token here is the cost of the arithmetic plus one WebGPU dispatch queue.
Runs where you are
One set of WGSL kernels, four places to run them. The browser tab and the workstation take the same code path — not an equivalent one, the same one.
| Where | Backend | What it needs |
|---|---|---|
| Linux workstation | wgpu → Vulkan | a Vulkan ICD (mesa or the vendor driver) |
| macOS, Apple silicon | wgpu → Metal | nothing beyond the OS |
| Windows | wgpu → D3D12 | nothing beyond the OS |
| This page | wasm32 → WebGPU | Chrome/Edge 113+ or Safari 26+ |
The wasm build is the library with a
wasm-bindgen facade over the async
inference API — not a port. Native-only dependencies
(rayon,
pollster, the
forge-top terminal UI) are behind
target and feature gates, which is why this page needs no
cross-origin isolation headers and ships as static files.
Why Forge is built this way
Three decisions, and what each one buys.
Rust, not Python
Explicit memory and dispatch, no interpreter in the loop, and one binary to run. Training and inference are the same code on the same tensors, so there is no second implementation to keep honest.
WebGPU, not CUDA
One set of WGSL kernels reaches every target above through wgpu, with no vendor toolchain to install — and no vendor to be locked to.
A CPU reference backend
Every kernel has a mathematically identical plain-Rust twin to check against. That is what makes the parity numbers mean something — and what makes the thing teachable, since the reference and the kernel can be read side by side.
Read the whole thing
Small enough to read in an afternoon, and laid out so you can. Each entry below is a link to the actual source.
- src/models/gpt2/ — config, the transformer blocks, KV-cache decode, generation, and the trace the section above is drawn from. Each block is LayerNorm → causal self-attention (6 heads × 64) → LayerNorm → MLP (384 → 1536 → 384, GELU).
- src/nn/, src/autograd/, src/optim/ — Linear, LayerNorm, Embedding with a row-chunked token table; a tape-based reverse mode; AdamW.
- src/ops.rs — shape-checked dispatch, one API over both backends.
-
src/backend/
— the WGPU backend and the CPU reference. Every
shaders/*.wgslkernel has a plain-Rust twin here computing the same thing the slow, obvious way; that pairing is what makes the parity claim checkable rather than assertable.
One compatibility note worth knowing before you load a checkpoint:
Forge stores Linear weights as
[in, out] (the HuggingFace Conv1D
convention) rather than nanoGPT's
[out, in], and reads safetensors
rather than PyTorch pickles — so files are not interchangeable with
nanoGPT .pt checkpoints.
Run this page locally
./scripts/build_site.sh ./scripts/serve_web.sh # plain static files — no COOP/COEP headers needed
Context for the section above
Context 256 characters, 65-token vocabulary, LM head weight-tied to the token embedding. GPT-2 124M is the same code with 12 blocks, 12 heads and n_embd 768 — 548 MB of weights, so it stays a local artifact rather than something this page downloads.
23 WGSL kernels in Rust
The complete compute surface, generated from
shaders/ at build time so this list
cannot drift from the code.
Forward — 15
- add
- dropout
- embedding
- gather_nll
- gelu
- kv_append
- layernorm
- matmul
- merge_heads
- scale
- softmax
- split_heads
- sum_rows
- unmerge_heads
- unsplit_heads
Backward — 6
- ce_bwd
- gelu_bwd
- layernorm_bwd_dp
- layernorm_bwd_dx
- scatter_add
- softmax_bwd
Optimizer — 2
- adamw
- sumsq
Train it yourself
The model in the section above is not a download — it is what this command produces. Training and inference are the same code on the same tensors, so there is no separate training stack to install and nothing that only exists at train time.
Tiny Shakespeare, from scratch
./scripts/download_shakespeare.sh # 6 blocks / 6 heads / 384 dims, 65-token vocab cargo run --release --example train_shakespeare -- \ --backend wgpu
What comes out
A safetensors checkpoint and its vocabulary, byte-identical in
layout to the one this page loads. Point
--model at it and the demo above
runs your weights instead — the forward pass, the KV cache and
the trace do not know the difference.
Quick start
Rust edition 2024 and a Vulkan, Metal, or D3D12 driver. No Python, no CUDA toolkit.
Generate with GPT-2 124M
# fetch weights + tokenizer ./scripts/download_gpt2.sh cargo run --release --example generate -- \ --backend wgpu --prompt "Hello Forge!"
Browse models in the terminal
# forge-top: model browser + live dashboard # tokens/s, VRAM, GPU util, temperature, power cargo run --release --features tui \ --bin forge-top -- --path models/