Llama e Qwen rodando no navegador via Three.js e WebGPU

Um LLM rodando dentro do navegador com Three.js

Ben Houston construiu o Three-LLM, uma biblioteca que executa modelos de linguagem localmente no navegador convertendo seus grafos de inferência em compute shaders escritos na Three.js Shading Language (TSL), rodando sobre o WebGPU. O objetivo declarado do projeto foi duplo: testar até onde vai a capacidade de compute geral do Three.js e do WebGPU - tradicionalmente vistos como bibliotecas de renderização - e verificar o que os pequenos modelos de linguagem atuais conseguem fazer direto no navegador, sem servidor de inferência nem binário WebAssembly específico por modelo.

O Three.js, apesar da fama de biblioteca de renderização, expõe no seu renderizador WebGPU storage buffers, dispatches de compute, memória de workgroup, atomics e leitura de dados da GPU. Um modelo de linguagem usa tudo isso: centenas de dispatches ordenados, produtos matriz-vetor grandes, reduções, estado persistente de atenção e um ciclo apertado entre GPU e JavaScript a cada token gerado. A biblioteca carrega configurações, tokenizadores e checkpoints SafeTensors do Hugging Face em tempo de execução e monta o modelo a partir de kernels TSL reutilizáveis, gerando WGSL para o backend WebGPU no próprio navegador.

Arquiteturas suportadas

O catálogo reconhece GPT-2, arquiteturas estilo Llama, Gemma 3, Phi e Qwen3.5. O chat de demonstração inclui cinco checkpoints: TinyStories GPT-2 3M, GPT-2 124M, SmolLM2 135M, Qwen3.5 0.8B e Phi-1.5 1.3B - variando de 15 MB a 2,8 GB de arquivo. Em celulares, o autor recomenda ficar nos modelos menores, pois Qwen e Phi exigem memória que muitos dispositivos móveis não têm.

O texto explica o ciclo de geração de tokens (tokenização, embeddings, blocos de transformer, projeção para logits, amostragem) e distingue prefill, que processa o prompt, de decode, que gera um token por passo. O KV cache, que evita recomputar chaves e valores anteriores, cresce com o contexto e se torna um custo importante de memória e banda. Cada família exige adaptações: o SmolLM2 usa RMSNorm, RoPE, Grouped-Query Attention e SwiGLU; o Phi-1.5 usa RoPE parcial e blocos paralelos de atenção e MLP; o Qwen3.5 empilha camadas Gated DeltaNet - que trocam o KV cache crescente por estado recorrente de tamanho fixo - a cada três camadas de atenção total.

Otimizações e custos

Três mudanças concentraram a maior parte dos ganhos medidos. Agrupar todos os dispatches de um passo em uma única submissão ao WebGPU (de 427 para 1 no SmolLM2) elevou o decode do TinyStories de cerca de 120 para 566 tokens por segundo (4,7×) e o do SmolLM2 de 34 para 75 (2,2×). Evitar projetar logits para tokens não finais do prompt reduziu o tempo até o primeiro token do SmolLM2 de aproximadamente 760 ms para 260 ms. Reuso de prefixos de prompt, amostragem reduzida na GPU e liberação de cópias de pesos após a carga completam a lista. Duas otimizações plausíveis - normalização com workgroup compartilhado e kernels matriciais com vec4 - não melhoraram o desempenho e foram descartadas.

Tudo roda em FP32: checkpoints FP16 ou BF16 são expandidos, o que dobra o tamanho em memória (o checkpoint de 269 MB do SmolLM2 vira cerca de o dobro em pesos de GPU). Em um teste quase equiparado, o Three-LLM atingiu cerca de 78 tokens por segundo no SmolLM2-135M contra cerca de 80 do WebLLM em FP32, mas a diferença cresce no prefill e em modelos quantizados maiores - o WebLLM compila kernels específicos com antecedência e usa pesos quantizados, enquanto o Three-LLM aceita checkpoints comuns do Hugging Face e monta operadores gerais em tempo de execução; são objetivos de projeto diferentes.

Próximos passos

O autor aponta quatro frentes: suporte nativo a FP16 via extensões f16 do WGSL e do WebGPU; pesos quantizados usando a extensão packed_4x8_integer_dot_product, que já vem sendo exposta via TSL, embora exija formato de pesos, escalas por bloco e detecção de recursos do adaptador; prefill em lote verdadeiro, ainda muito atrás de motores compilados; e fusão de grafos para eliminar dispatches e buffers intermediários - com a ressalva de que cada caminho fundido precisa de medições em GPUs representativas.