A tour of how language models generate text, how GPT-2 and newer Phi and Qwen architectures differ, and how Three-LLM implements their inference graphs with Three.js TSL compute shaders.
Ben Houston • • 17 min read
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# #
I define model, weight, activation, MLP, and forward pass in Neural network basics, for graphics people. The inference path here does not require the training details.
The central operation remains the one from that primer:
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 is a useful starting point because its architecture contains the pieces that still define most decoder language models without many later variations. It builds on the decoder half of the transformer introduced in Attention Is All You Need.
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. The idea began as a data-compression algorithm and was later adapted to subword tokenization for neural models.
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:
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. The paper Fast Transformer Decoding: One Write-Head is All You Need describes this incremental-inference bottleneck and introduces one of the first widely used ways to reduce it.
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. The main building blocks include:
TSLLinear
for matrix-vector multiplicationTSLNormalize
andTSLRMSNorm
for normalizationTSLAttention
for full causal attentionTSLMLP
andTSLGatedMLP
for feed-forward blocksTSLAdd
for residual connectionsTSLLogitSampler
for reducing vocabulary logits on the GPU
The linear layer shows the basic pattern. 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]);
I kept this kernel plain so you can see the operation it performs. Three-LLM constructs readable operators in the browser, and TSL generates WGSL for the WebGPU backend. It does this work at runtime instead of asking a model compiler to select tiled matrix kernels ahead of time.
For GPT-2, the general decoder runner connects those operators as:
embedding
→ LayerNorm
→ packed QKV projection
→ attention
→ output projection
→ residual add
→ LayerNorm
→ GELU MLP
→ residual add
→ final LayerNorm
→ vocabulary projection
The GPU attention implementation uses four compute stages for each layer:
- Write the new key and value into the cache.
- Prepare the query.
- Compute one scaled query-key score per head and cached token.
- Run stable softmax and mix the cached values.
The CPU runner performs the same operations in TypeScript. Browser checkpoint tests compare greedy GPU output with the CPU reference for all five catalog models. That reference path has been essential because a shader can run without errors while still producing subtly wrong text.
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 then map names such as GPT-2's
transformer.h.0.attn.c_attn.weight
or Llama's
model.layers.0.self_attn.q_proj.weight
into the common fields expected by the runners.
The extra code lets the browser consume the checkpoint published by the model author. Three-LLM builds it without an offline conversion step or a separately compiled runtime for each model.
From GPT-2 to SmolLM2# #
SmolLM2 135M is close to GPT-2 in parameter count, but it uses the newer Llama-style decoder pattern. The SmolLM2 report covers the model family and its training.
Four architectural changes matter to the inference implementation:
RMSNormuses the root mean square without subtracting the mean.Rotary Position Embeddingsrotate query and key components according to position instead of adding a learned position vector.Grouped-Query Attentionlets groups of query heads share fewer key and value heads, shrinking the KV cache.SwiGLUreplaces GPT-2's dense GELU MLP with a gated MLP.
The gated MLP computes two projections from the same input. One passes through SiLU and gates the other before a final down projection:
Three-LLM handles these changes with recipe flags and reusable kernels rather than a separate SmolLM2 runner.
Phi changes the block layout# #
Phi-1.5 is a 1.3-billion-parameter model introduced in Textbooks Are All You Need II. Its inference graph still uses LayerNorm and a dense GELU MLP, but it differs from GPT-2 in two important ways.
First, Phi uses partial RoPE. It rotates part of each query and key while leaving the remaining dimensions unchanged.
Second, its attention and MLP branches run in parallel from the same normalized input:
normalized = LayerNorm(x)
attentionOutput = Attention(normalized)
mlpOutput = MLP(normalized)
output = x + attentionOutput + mlpOutput
GPT-2 instead finishes its attention branch and residual addition before
normalizing again for the MLP. The parameterized decoder runner supports both
layouts. When the recipe says residual: 'parallel'
, it constructs Phi's branching graph.
Qwen3.5 needs a different kind of memory# #
Qwen3.5 0.8B is the most complex model in the demo. Three-LLM runs its text backbone and skips the checkpoint's vision tensors. Qwen's official introduction and the Transformers architecture documentation describe a hybrid stack: three Gated DeltaNet layers for each gated full-attention layer.
Full-attention layers retain a KV cache and can retrieve a specific earlier token. Gated DeltaNet layers replace that growing history with fixed-size recurrent state. Their memory use stays fixed as the model processes more tokens.
Gated Delta Networks combine two ideas. A learned decay gate controls how much old state survives, while a delta update changes the association for the current key toward its new value. Earlier work describes this as a fast-weight programming rule. The relationship between linear attention and recurrent networks is developed in Transformers are RNNs.
For one decode step, Three-LLM's TSLGatedDeltaNet
performs:
- Query, key, value, gate, decay, and update projections
- A short causal convolution over recent state
- Query and key normalization
- Decay of the recurrent matrix
- A delta-rule correction using the new key and value
- A query against the updated matrix
- Gated normalization and an output projection
The Qwen runner chooses the mixer from the checkpoint's per-layer type:
if (block.layerType === 'linear_attention') {
mixer = new TSLGatedDeltaNet(ln1.outputNode, block.delta, {
hiddenSize: this.hiddenSize,
numKHeads: weights.linearKeyHeads,
numVHeads: weights.linearValueHeads,
keyDim: weights.linearKeyDim,
valueDim: weights.linearValueDim,
kernelSize: weights.linearConvKernel,
});
} else {
// Build gated full attention with a conventional KV cache.
}
Qwen's architecture warrants a dedicated runner, recurrent state management, and a full set of DeltaNet TSL kernels. The surrounding residual and gated MLP structure reuses the same operators as the other model families.
Optimizations that paid off# #
I tested kernel rewrites, command submission, prefill behavior, readback, and model- changes while building the library. A few changes accounted for most of the measured gains.
Submit a forward pass once#
SmolLM2 executes more than 400 compute dispatches for one token. The first implementation submitted each dispatch separately. Recording the ordered nodes into one Three.js compute call reduced 427 command submissions to one.
That change raised TinyStories decode from about 120 to 566 tokens per second, a 4.7× improvement. SmolLM2 rose from about 34 to 75 tokens per second, a 2.2× improvement. JavaScript-to-WebGPU submission overhead had dominated these small models.
The GPU still runs each dispatch, switches pipelines, and reads intermediate buffers. One command submission removes CPU and queue overhead. Fusing the model into fewer shaders would require a separate optimization.
Do not compute unused prompt logits#
Only the final prompt token needs vocabulary logits. Earlier prompt tokens need to update attention caches or recurrent state, but projecting every one of them across a 49,152-token vocabulary throws the result away.
Three-LLM records a prefill path without the final normalization and vocabulary projection, then computes logits for the last prompt token. Together with the submission changes, warm SmolLM2 time to first token fell from roughly 760 ms to about 260 ms in the measured test.
Reuse prompt prefixes#
Chat prompts repeat the earlier conversation on every turn. Three-LLM compares the new token sequence with the previous one and reuses the matching prefix. The runner only prefills the appended text when its cache state can be reused.
Reduce sampling on the GPU#
Reading an entire vocabulary back to JavaScript creates a synchronization point after every generated token. Three-LLM can run a hierarchical maximum or small top-k reduction on the GPU. Greedy generation reads four bytes for one token ID. A supported top-k request reads a small candidate set and finishes sampling on the CPU.
The full-logit path remains available for sampling options that need every score. This keeps the API flexible while removing the common greedy readback bottleneck.
Free copies once the GPU owns the weights#
a large checkpoint can leave the source tensors, unpacked arrays, and GPU buffers alive at the same time. After Three.js creates the GPU bindings, Three-LLM releases static CPU weight arrays while retaining the embedding data needed to look up future token rows. The demo's model proxy also splits large files into concurrent requests so one multi-gigabyte response does not become a serial bottleneck.
Measurements rejected two plausible optimizations. Shared-workgroup
normalization and broad vec4
matrix-vector kernels failed to improve representative SmolLM2 throughput, so I left out the extra complexity.
The FP32 cost# #
Three-LLM currently runs every GPU kernel with 32-bit floating-point storage
and arithmetic. A checkpoint may arrive as FP16 or BF16, but the
expands non-embedding tensors into Float32Array
values. Embedding rows are converted as they are copied. KV caches, activations, logits, and Qwen's recurrent state are FP32 too.
This provides one simple baseline across WebGPU adapters, but it imposes a large memory and bandwidth cost. SmolLM2's 269 MB BF16 checkpoint becomes roughly twice that size for its GPU weights. Qwen's 1.7 GB and Phi's 2.8 GB checkpoint files make the problem much more visible. Autoregressive decode reads much of that weight data again for every token.
In a near-matched SmolLM2-135M test, Three-LLM reached about 78 tokens per second with the base checkpoint expanded to FP32. WebLLM reached about 80 with an unquantized FP32 Instruct checkpoint using the same architecture and parameter scale. The gap was much larger for prompt prefill and for larger quantized models. WebLLM compiles model-specific kernels ahead of time, supplies separate prefill and decode programs, and ships reduced-precision or quantized weights. Three-LLM accepts regular Hugging Face checkpoints and assembles general operators at runtime. The two designs optimize for different goals.
Where this can go next# #
Native FP16 is the clearest next precision step. The WGSL f16 extension and WebGPU's
shader-f16
device featureQuantized weights offer a larger capacity gain. The WGSL packed_4x8_integer_dot_product extension exposes hardware instructions that multiply four packed 8-bit integer pairs at once. I have also been working to expose this language extension through TSL. A Three-LLM path built around it would need a weight format, per-block scales, adapter feature detection, and kernels that account for dequantization cost. The extension supplies the dot-product primitive; a complete quantization design must provide everything around it.
Prefill also remains far behind a compiled engine. Three-LLM records several one-token passes together, but it does not yet turn prompt processing into true batched matrix-matrix operations. A dedicated batched prefill graph could reuse weights across prompt tokens and expose more parallel work.
Finally, graph fusion could remove dispatches and intermediate buffers.
Candidates include the gate and up projections, projection plus residual
addition, and parts of attention. The failed normalization and vec4
experiments are a useful warning: each fused path needs full-model measurements on representative GPUs.
Try it# #
You can run the live Three-LLM chat demo, install three-llm from npm, or read the
Start with TinyStories to see the complete GPT-2 path at a scale that loads almost anywhere. SmolLM2 shows the modern Llama-style changes without a huge checkpoint. Phi demonstrates a parallel transformer block, and Qwen exercises the hybrid Gated DeltaNet architecture. On mobile, stay with TinyStories or SmolLM2.
Building this test of Three.js compute support required tokenization, checkpoint , several transformer families, recurrent linear attention, GPU sampling, prompt caching, and complete autoregressive generation. The same WebGPU renderer can draw a scene and run the language model that describes it.