Running LLMs in the Browser with Three.js
Three-LLM enables local LLM inference in the browser by compiling transformer architectures into Three.js TSL compute shaders.
Deep dive
- Compute Stack: Uses Three.js TSL (Three Shading Language) to generate WGSL kernels at runtime.
- Architecture Support: Runs GPT-2, Llama-style, Phi, and Qwen3.5 (including Gated DeltaNet variants).
- Performance Gains: Moving from 400+ separate command submissions to one single compute call improved throughput by up to 4.7x.
- Memory Management: Currently limited by FP32 overhead; future support for FP16 and 8-bit quantization is planned via WGSL language extensions.
- Implementation: The library directly consumes standard Hugging Face SafeTensors checkpoints rather than requiring custom model conversion.
Decoder
- Inference: The process of using a trained model to make predictions or generate text.
- WGSL (WebGPU Shading Language): The native shader language for WebGPU, used to program the GPU for graphics and general-purpose compute.
- Compute Shader: A program written for a graphics API to perform arbitrary calculations on the GPU rather than traditional rendering tasks.
Original article
TL;DR
Three-LLM runs GPT-2, SmolLM2, Qwen, and Phi models locally in the browser by turning their inference graphs into Three.js TSL compute shaders on WebGPU. Try the demo at https://three-llm.ben3d.ca.
I started this project to test two things: how far I could push the general compute capabilities of Three.js and WebGPU, and what today's small language models can do when they run in a browser. The models range from a 3-million parameter story generator to a modern 800-million parameter chat model.
Three.js is known as a rendering library, but its WebGPU renderer also exposes storage buffers, compute dispatches, workgroup memory, atomics, and GPU readback. A language model puts all of those capabilities to work. It needs hundreds of ordered compute dispatches, large matrix-vector products, reductions, persistent attention or recurrent state, and a tight loop between the GPU and JavaScript as each new token appears.
I built Three-LLM as both a test of that compute stack and a usable inference library. It loads ordinary Hugging Face configurations, tokenizers, and SafeTensors checkpoints in the browser, builds the model from reusable Three.js Shading Language (TSL) compute kernels, and runs inference on the user's GPU through WebGPU. It does not need a server-side inference runtime or model-specific WebAssembly binary.
The package includes CPU reference implementations for validation and GPU implementations for actual use. The library currently recognizes GPT-2, Llama-style, Gemma 3, Phi, and Qwen3.5 architectures. I also made a quick chat application so you can try five checkpoints without writing an application first:
- TinyStories GPT-2 3M
- GPT-2 124M
- SmolLM2 135M
- Qwen3.5 0.8B
- Phi-1.5 1.3B
I made the demo to test the library without first writing an application. Model files range from 15 MB to 2.8 GB before Three-LLM expands their weights to 32-bit floats. On a phone, stick with TinyStories or SmolLM2. Qwen and Phi need enough memory that many mobile devices will fail to load them.
Using the library itself starts with a regular Three.js WebGPURenderer:
import { createTSLRunner } from 'three-llm';
import { WebGPURenderer } from 'three/webgpu';
const renderer = new WebGPURenderer();
await renderer.init();
const runner = await createTSLRunner(
'https://huggingface.co/HuggingFaceTB/SmolLM2-135M/resolve/main/',
);
const result = await runner.generate(renderer, 'Once upon a time,', {
maxNewTokens: 64,
temperature: 0.7,
topK: 10,
onToken: (text) => console.log(text),
});
Before the transformer
The central operation remains: y = Wx + b
An LLM owns many matrices containing learned weights. Inference repeatedly multiplies vectors by those matrices, applies a few nonlinear functions, and moves information between token positions. A transformer gives those operations a particular structure.
What happens when an LLM generates one token
The full loop has six parts:
- A tokenizer converts text into a sequence of integer token IDs.
- An embedding table turns each ID into a vector.
- A stack of transformer blocks updates that vector using the current token and the preceding context.
- A final matrix projection produces one score, called a logit, for every token in the vocabulary.
- A sampling rule chooses one token from those scores.
- The model appends that token and runs again.
The model generates a sentence one token at a time. It predicts a token, adds it to the context, and repeats until it selects a stop token or reaches a limit.
Processing the supplied prompt is called prefill. Generating subsequent tokens one at a time is called decode. These two phases run the same model but have different performance characteristics. Prefill has many known input tokens that a specialized implementation can process together. Decode has only one new token per step and must repeatedly read most of the model's weights.
GPT-2, one block at a time
GPT-2 first applies byte-level Byte Pair Encoding to the input. BPE repeatedly combines frequent adjacent symbols, giving the model a fixed vocabulary that can still represent any input text.
Each resulting token ID selects a row from a learned token-embedding table. GPT-2 adds a second learned embedding for the token's position in the sequence. The sum is the first vector sent through the transformer.
Every GPT-2 block then runs this sequence:
- Layer normalization
- Causal multi-head self-attention
- A residual addition
- A second layer normalization
- A two-layer MLP using the GELU activation
- Another residual addition
Residual connections add each block's result back to its input. They let later blocks refine a representation without requiring every block to rebuild it.
After the last block, GPT-2 applies one more normalization and projects the result to the vocabulary. If the vocabulary contains 50,257 tokens, that projection returns 50,257 logits. The sampler turns those scores into the next token ID.
Attention is a lookup built from the current context
Self-attention lets the current token retrieve information from earlier tokens. Each block projects its input into three vectors for every attention head:
- A query describes what the current position is looking for.
- A key describes what each position offers.
- A value contains the information retrieved from that position.
The query takes a dot product with each preceding key. Dividing by the square root of the head dimension keeps the scale stable, and softmax converts the scores into weights: Attention(Q,K,V) = softmax(QKᵀ / √d)V
A causal mask excludes future positions. During generation there are no future tokens yet, but the same rule lets prompt tokens attend only to text that precedes them.
Recomputing every earlier key and value for every generated token would waste most of the work. A KV cache stores them after their first computation. Each decode step writes one new key and value, then compares the new query against the cached keys. The cache grows with context length and becomes an important memory and bandwidth cost.
Expressing the model as TSL compute
Three-LLM represents each model operation as a small TSL compute kernel. A runner connects those kernels in the order required by the model recipe.
One GPU invocation owns one output element and loops across the input vector:
return Fn(() => {
const outputIndex = instanceIndex.toVar('outputIndex');
If(outputIndex.lessThan(uint(outputSize)), () => {
const sum = biasNode.element(outputIndex).toVar('sum');
Loop(
{ start: uint(0), end: uint(inputSize), type: 'uint', condition: '<' },
({ i }) => {
const weightIndex = i.mul(uint(outputSize)).add(outputIndex);
sum.addAssign(inputNode.element(i).mul(weightNode.element(weightIndex)));
},
);
outputNode.element(outputIndex).assign(sum);
});
})().compute(outputSize, [workgroupSize]);
Loading Hugging Face checkpoints at runtime
Three-LLM reads a Hugging Face config.json, tokenizer files, and SafeTensors weights. The configuration selects a recipe. The recipe describes the graph family, normalization, MLP type, residual layout, position encoding, head counts, and checkpoint layout. Tensor-name aliases map names into the common fields expected by the runners.
From GPT-2 to SmolLM2
SmolLM2 135M is close to GPT-2 in parameter count, but it uses the newer Llama-style decoder pattern. Four architectural changes matter to the inference implementation: RMSNorm, Rotary Position Embeddings, Grouped-Query Attention, and SwiGLU activation.
Phi changes the block layout
Phi-1.5 is a 1.3-billion-parameter model. Its inference graph differs from GPT-2 in two ways: it uses partial RoPE, and its attention and MLP branches run in parallel from the same normalized input.
Qwen3.5 needs a different kind of memory
Qwen3.5 0.8B uses a hybrid stack: three Gated DeltaNet layers for each gated full-attention layer. Full-attention layers retain a KV cache. Gated DeltaNet layers replace that growing history with fixed-size recurrent state.
Optimizations that paid off
I tested kernel rewrites, command submission, prefill behavior, readback, and model-loading changes while building the library. A few changes accounted for most of the measured gains:
- Submit a forward pass once: Recording ordered nodes into one Three.js compute call reduced 427 command submissions to one.
- Do not compute unused prompt logits: Only the final prompt token needs vocabulary logits.
- Reuse prompt prefixes: Three-LLM compares the new token sequence with the previous one and reuses the matching prefix.
- Reduce sampling on the GPU: Three-LLM can run a hierarchical maximum or small top-k reduction on the GPU to avoid synchronization points.
- Free copies once the GPU owns the weights: After Three.js creates the GPU bindings, Three-LLM releases static CPU weight arrays.
The FP32 cost
Three-LLM currently runs every GPU kernel with 32-bit floating-point storage and arithmetic. While this provides a simple baseline, it imposes a large memory and bandwidth cost compared to native FP16 or quantized weights. Future development will focus on adding FP16 support and quantized dot products.
Try it
You can run the live Three-LLM chat demo, install three-llm from npm, or read the source on GitHub.
Fresh dev news, AI-digested daily. Never miss a trend at devoured.fyi