Resumo Técnico: Llamas na Web

Declaração do Problema

Executar Modelos de Linguagem de Grande Escala (LLMs) localmente em navegadores web oferece vantagens significativas para privacidade, latência e eficiência energética, especialmente à medida que desenvolvedores de modelos lançam modelos de pesos abertos projetados para implantação na borda. No entanto, a realização desse potencial é impedida por três restrições principais nos frameworks de inferência baseados em navegadores existentes (por exemplo, WebLLM, Transformers.js):

1. Ineficiências de Memória: Motores existentes frequentemente alocam memória GPU dinamicamente e carregam pesos de forma ineficiente. Isso leva ao crescimento da memória durante a execução, causando lentidão ou travamentos devido a limites rígidos de memória impostos pelos navegadores por aba.
2. Falta de Design Portátil em Desempenho: Embora o WebGPU forneça portabilidade funcional em hardware diversificado, os frameworks existentes carecem de especialização suficiente de kernels. Eles falham em ajustar para capacidades específicas do dispositivo (por exemplo, operações de subgrupo, núcleos tensor), resultando em desempenho subótimo na paisagem heterogênea de hardware dos navegadores.
3. Suporte Limitado a Quantização: Para suportar a grande variedade de formatos de quantização usados no ecossistema llama.cpp (variando de esquemas legados a K-quants e I-quants avançados), motores de inferência exigem designs de kernels flexíveis. Os frameworks existentes frequentemente carecem da extensibilidade para suportar esses formatos diversos de forma eficiente ou para integrar novos rapidamente.

Metodologia: LlamaWeb

Os autores apresentam o LlamaWeb, um backend WebGPU para llama.cpp projetado para abordar essas restrições. O sistema é construído dentro da base de código llama.cpp para aproveitar seu extenso suporte a modelos e ecossistema ativo. O design foca em três pilares centrais:

1. Inferência Local Eficiente em Memória

Para minimizar a sobrecarga de memória e garantir estabilidade em dispositivos com restrições de memória:

• Alocação Estática de Memória: Toda a memória necessária para a execução do modelo, incluindo estruturas de dados intermediárias para operações como FlashAttention e regiões de tamanho fixo com slots para parâmetros de kernel, é alocada estaticamente na inicialização. Isso evita a criação dinâmica de buffers e fragmentação.
• Carregamento Otimizado: O sistema utiliza a biblioteca wllama para transmitir pesos assincronamente do Sistema de Arquivos Privado de Origem (OPFS) diretamente para buffers WebGPU. Ele evita materializar pesos no heap WebAssembly, que é apenas crescente e não pode liberar memória durante o tempo de vida de uma aba.
• Arena de Buffer de Parâmetros: Em vez de manter pools dinâmicos de pequenos buffers para parâmetros de kernel (uma fonte comum de fragmentação), o LlamaWeb aloca um único buffer estático com slots configuráveis, rotacionando por eles para garantir que os parâmetros não sejam sobrescritos até que a execução do kernel seja concluída.

2. Biblioteca de Kernels Portátil em Desempenho

Para maximizar o desempenho em hardware heterogêneo:

• Geração de Shaders Templada: O backend usa um pré-processador WGSL personalizado (pre-wgsl) e uma biblioteca de kernels que gera shaders com base em recursos específicos do dispositivo (por exemplo, suporte a subgrupo) e tipos de dados.
• Parâmetros Ajustáveis: Parâmetros-chave, como tamanhos de blocos para multiplicação de matrizes e tamanhos de grupos de trabalho, são ajustáveis. Os autores derivaram padrões "portáteis em desempenho" por meio de um estudo empírico que varreu milhares de configurações em quatro GPUs de diferentes fornecedores.
• Agendamento de Execução: O tempo de execução agrupa eficientemente submissões de kernel e minimiza a sincronização CPU-GPU. Ele armazena em cache pipelines compilados para evitar custos de compilação redundantes após a primeira passagem direta.

3. Design de Shader Consciente de Quantização

Para suportar os diversos formatos de quantização do llama.cpp:

• Desquantização Integrada: Rotinas de desquantização são co-projetadas com kernels de computação. Para multiplicação de matrizes, blocos quantizados são carregados, desquantizados em memória compartilhada e reutilizados. Para multiplicação matriz-vetor (fase de decodificação), a desquantização ocorre diretamente em registradores para melhorar os padrões de acesso à memória.
• Extensibilidade: O sistema representa formatos de quantização como buffers planos de inteiros sem sinal de 32 bits, com lógica de desquantização selecionada no momento da compilação. Isso permite que novos formatos (por exemplo, o recém-introduzido q1_0 para modelos de 1 bit) sejam integrados adicionando uma rotina de desquantização sem reescrever implementações centrais de operações matriciais.

Contribuições Principais

1. Um Motor de Inferência Eficiente em Memória: Uma implementação WebGPU que aloca memória estaticamente e otimiza o movimento de dados, permitindo a execução de modelos úteis em dispositivos de borda com restrições de memória sem travamentos inesperados.
2. Uma Biblioteca de Shaders Flexível: Uma biblioteca para inferência de LLM que suporta ajuste externo, especializa shaders com base em recursos de hardware disponíveis (por exemplo, operações matriciais de subgrupo) e utiliza um tempo de execução de agendamento otimizado.
3. Kernels de Quantização Co-projetados: Um conjunto de kernels centrais (operações matriciais) projetados para funcionar com os extensos formatos de quantização do llama.cpp, permitindo a implantação de modelos otimizados para diferentes tamanhos e facilitando a integração de métodos de quantização emergentes.

Resultados Experimentais

Os autores avaliaram o LlamaWeb em 16 dispositivos de 8 fornecedores (incluindo NVIDIA, AMD, Intel, Apple, Qualcomm, Samsung, ARM e Imagination Technologies) executando 10 modelos de linguagem em quatro formatos de pesos.

• Eficiência de Memória: Comparado ao WebLLM e Transformers.js, o LlamaWeb reduziu o uso máximo de memória do navegador em 29–33% em média em várias combinações de dispositivo/navegador/SO. Em testes específicos (por exemplo, Safari em Apple M4), ele usou até 59% menos memória que o Transformers.js, que sofreu com vazamentos de memória nessas configurações.
• Taxa de Transferência de Decodificação: O LlamaWeb aumentou a taxa de transferência de decodificação em 45–69% comparado ao WebLLM e Transformers.js em quatro GPUs de fornecedores separados.
• Taxa de Transferência de Preenchimento (Prefill): Durante a fase de preenchimento, o LlamaWeb inicialmente teve desempenho inferior comparado ao WebLLM (em 21–51%) e Transformers.js. Os autores atribuem isso ao uso pelo WebLLM da fusão de kernels do TVM e às implementações otimizadas para subgrupo do Transformers.js. No entanto, ao executar o LlamaWeb nativamente (fora do navegador) com verificações de segurança desativadas, ele superou outros frameworks, sugerindo que verificações de segurança específicas do navegador e a falta de suporte estável a operações matriciais de subgrupo nos navegadores são os principais gargalos.
• Desempenho Nativo vs. Navegador: Em alguns dispositivos (por exemplo, Intel Arc B580), o backend WebGPU do LlamaWeb foi competitivo ou até mais rápido que backends nativos (por exemplo, SYCL) para tarefas específicas. No entanto, backends nativos (CUDA, Metal) geralmente superaram o WebGPU em até 10x no preenchimento e 2,5x na decodificação, em grande parte devido ao melhor aproveitamento de recursos específicos do fornecedor, como núcleos tensor e carregamentos de memória assíncronos.
• Impacto da Quantização: O estudo constatou que, embora a quantização reduza os requisitos de memória, ela nem sempre melhora o desempenho linearmente devido ao custo computacional da desquantização. Por exemplo, a transição de q8_0 para q4_k_m melhorou a taxa de transferência, mas a transição para q2_k a diminuiu em clusters de nível intermediário.

Significado e Alegações

O artigo afirma que o LlamaWeb representa um passo significativo para tornar a inferência de LLM baseada em navegador acessível, privada e performática. Ao integrar-se com o ecossistema llama.cpp, ele suporta uma vasta biblioteca de modelos (177.691 modelos compatíveis no formato GGUF) que outros frameworks de navegador não conseguem igualar.

Os autores enfatizam que seu trabalho demonstra que a portabilidade de desempenho é alcançável no navegador por meio de ajuste cuidadoso de kernels e planejamento estático de memória. Embora reconheçam que backends nativos atualmente detêm uma vantagem de desempenho, o LlamaWeb fornece uma alternativa competitiva multiplataforma que evita a necessidade de instalações nativas específicas de plataforma. O trabalho também destaca a necessidade de melhorias adicionais nas especificações do WebGPU, como melhor suporte para tipos inteiros menores e a remoção de verificações de segurança em tempo de execução que prejudicam o desempenho, para realizar plenamente o potencial da IA baseada em navegador.

O código é de código aberto e foi integrado à base de código llama.cpp por meio de mais de 100 pull requests, suportando 72 operadores centrais de ML e 23 formatos de dados.