We're Teaching Your Browser to Paint: A 4B Image Model on WebGPU, Proven to Six Decimal Places

There's a 1.37 GB file sitting in our artifact store tonight that shouldn't be possible. It's FLUX.2 Klein — a 4-billion-parameter image generation model — packed into 2-bit ternary weights that our browser kernels already know how to read. Not a demo build. Not a quantized-into-mush approximation. The same model our fleet serves, byte-for-byte in the places that matter, headed for a <canvas> near you.

This is the story of how we're getting there, and more importantly, how we know each step is right. Because the defining property of GPU numerics work is that wrong code doesn't crash — it produces a plausible image. A conv with a flipped padding still renders something. Attention with the wrong scale still makes shapes. The only defense is a reference you trust and the discipline to diff against it at every single stage.

The bet: local-first means images too

If you've used aitherium.com lately, you know the deal: the text brain runs on your hardware. Your browser's GPU, your node, your machine — the cloud is the fallback, not the default. We shipped a pile of fixes this week to make that promise real on desktops (that's its own story, involving a circuit breaker that outlived the bug it was protecting against, and an RTX 5090 that was 94% full of our own vLLM fleet when the browser asked it for room — we'll write that postmortem too).

But images were the asterisk. Image generation ran on our fleet, behind a sign-in, metered — because a diffusion model was "obviously" too heavy for a browser. This week we decided to delete the asterisk. Two fronts:

Front one shipped already: if you self-host our image backend (curl -fsSL https://aitherium.com/install-bonsai-image.sh | sh — one command, NVIDIA GPU required, uninstalls as cleanly as it installs), aitherium.com now detects it and renders on your card. No account, no meter. We found, while wiring this, that our own CORS allowlist didn't include our own apex domain — the browser was being blocked from your machine by our policy. That's the kind of bug you only find by walking the whole path.

Front two is the fun one: the model itself, in the tab, on WebGPU.

The lucky break that wasn't luck

FLUX.2 Klein's transformer ships as gemlite-int2 — ternary weights (every value is -1, 0, or +1, times a scale) in a CUDA-specific packed format. Unloadable in a browser. Everyone's first instinct is "write a converter and pray the formats roughly align."

We measured instead. Pulled one layer apart byte by byte, in the container, before writing a line of converter:

• gemlite codes: {0, 1, 2}, packed LSB-first, four per byte, along the input axis
• gemlite dequant: (q − zero) × scale
• and the measurement that made the whole project: zeros == −scales, exactly, every layer

Which means gemlite's dequantization is (q − 1) × scale. Now look at our browser kernel format — Q2_0, the 2-bit ternary type in our PrismML llama.cpp fork, built for the text models months ago: 128-weight blocks, an f16 scale, 2-bit codes packed LSB-first, dequantized as (q − 1) × d.

Identical. Same number system, same bit order, same formula. The "conversion" is a transpose and regroup — zero numeric change. The only rounding anywhere in the pipeline is scale f32→f16, so the converter asserts that roundtrip is exact and aborts the run if it ever isn't. It never has. The 100 ternary linears of a 4B diffusion transformer now load with the same kernels that run our text models. We didn't write a single new quantization kernel.

Golden vectors, or: how to port a model without lying to yourself

Here's the discipline, and it's the actual point of this post.

Before porting anything, we ran the real pipeline — the exact diffusers + PrismML code our fleet serves — and dumped ground truth at every seam: the text encoder's output for fixed prompts, the noise schedule's exact sigmas, every intermediate latent of a seeded 4-step generation, the final image tensor. Those files are the contract. Every ported component gets held to them, mechanically, in tests.

Then we ported in layers, each proven before the next started:

The VAE decoder (latents → pixels). CPU reference first, scalar loops written to be obviously correct. Validated against the real 168 MB checkpoint — every weight name resolves, or the loader throws by name. Then the GPU version, diffed against the CPU on identical inputs. Measured on an RTX 5090: max difference 4.5e-8. Float noise. Not "close enough" — indistinguishable.

The transformer forward (the MMDiT — 5 joint blocks, 20 single blocks, 24 heads). Ported scalar from the pipeline source, then replayed the golden inputs through our converted 2-bit file. Max difference vs the real model's output: 6.75e-6, on outputs of magnitude ~0.8. One number validating four things simultaneously — the architecture port, the weight conversion, the file format, the tensor naming — because getting any of them wrong moves outputs by whole units, not millionths. Conventions are where ports die: text tokens concatenate first; block modulations split as (shift, scale, gate) but the output norm splits (scale, shift); RoPE runs on four position axes with interleaved pairs, after the per-head RMSNorm. Each of those, wrong, renders a plausible image. The golden diff catches all of them at once.

The scheduler (4-step flow matching). Proven without ever running the transformer in the test: we replay the dumped model outputs through our Euler step and demand the dumped next-step latents come back. Scheduler bugs can't hide behind model noise. Batch-norm latent scaling round-trips at 7.45e-8.

The GPU driver, and the bug that looked like broken math

Tonight the GPU forward pass came alive. The driver mirrors the CPU reference stage for stage on our existing kernel family — the 2-bit matmuls, a new token LayerNorm, a new f16 matmul (no shader-f16 extension needed; we unpack half-floats manually, a trick the codebase already used for scales), and bidirectional attention.

Two engineering choices worth stealing:

Slice weights, not activations. The model fuses projections — one matrix computes Q, K, V, and the MLP input in a single multiply. On GPU, un-fusing the output means strided reads or thousands of per-token copies. But 2-bit quantized rows are independent byte runs — so we slice the weight at load time, once, on CPU, into its logical parts. The fused output projection (which consumes attention-output and MLP-output concatenated per token) splits along its input axis instead: two matmuls, summed. The math is identical; the memory traffic isn't.

Never trust a claimed convention — reconstruct it. Our text-model RoPE kernel was audited as "matches the reference." The audit contradicted itself in the same paragraph about which pairing convention it used. Rather than adjudicate, the driver computes its rotation tables on the CPU with the same function the golden-proven reference uses, and the GPU kernel only rotates. Right by construction beats right by review.

And the confession, because our best bugs are the instructive ones: the first GPU differential run failed at 0.097 absolute — catastrophic, convention-error territory. The per-op bisection layer pointed at exactly the two new kernels. Cause? Our WGSL sources ship inside a generated TypeScript file, and we'd edited the shaders without regenerating it. The new GPU entry points didn't exist; pipeline creation failed asynchronously; the dispatches silently no-opped; and the "results" were stale memory from a buffer pool. A build-step omission, cosplaying as broken numerics. It cost one cycle and taught the permanent lesson: when the GPU disagrees with the CPU, first prove the GPU ran your code at all.

With the module actually regenerated: the full miniature-model differential on the 5090 lands at 0.59% relative — and our activations run through 8-bit quantization on the GPU path, which predicts a few-percent floor. Convention errors sit at 100%. We're two orders of magnitude from the alarm line, exactly where the physics says an honest implementation should be. The full-scale run against the real 1.37 GB file is grinding away as this post is written.

What's left, and what it means

Assembly, not research: the text encoder (whose 4-bit format we've already cracked and dequantized — (q − zero) × scale again, exactness verified), the worker that strings encoder → 4 denoise steps → VAE decode behind a progress bar, and one end-to-end 256px render compared against the golden final latent. The flag that lists "Bonsai Image" in your model picker stays off until that render matches. That's the rule that made everything above trustworthy: nothing ships on vibes.

When it lands, the stack on a machine with a decent GPU looks like this: a 4B language model and a 4B image model, both in browser tabs or on your own node, both in ternary weights we can prove against their reference implementations, with the cloud as a fallback you can see and choose. Your prompts render where you're standing.

The models are on HuggingFace. The kernels are clean-room, in our fork, diffed against scalar references you can read in an afternoon. The golden vectors don't care about anyone's marketing — including ours.

That's the whole cool of it: not that a browser can run a diffusion transformer — but that we can hand you the numbers proving it runs correctly.

— Aitherium Engineering, 2 a.m., watching a differential harness converge