Resumen Técnico: Llamas en la Web

Declaración del Problema

Ejecutar Modelos de Lenguaje Grande (LLM) localmente en navegadores web ofrece ventajas significativas en privacidad, latencia y eficiencia energética, especialmente a medida que los desarrolladores de modelos lanzan modelos de pesos abiertos diseñados para despliegue en el borde. Sin embargo, la materialización de este potencial se ve obstaculizada por tres limitaciones principales en los marcos de inferencia basados en navegadores existentes (por ejemplo, WebLLM, Transformers.js):

1. Ineficiencias de Memoria: Los motores existentes a menudo asignan memoria GPU dinámicamente y cargan pesos de manera ineficiente. Esto conduce a un crecimiento de la memoria durante la ejecución, causando ralentizaciones o fallos debido a los límites estrictos de memoria impuestos por los navegadores por pestaña.
2. Falta de Diseño Portátil en Rendimiento: Aunque WebGPU proporciona portabilidad funcional en hardware diverso, los marcos existentes carecen de especialización de kernels suficiente. No logran ajustar las capacidades específicas del dispositivo (por ejemplo, operaciones de subgrupos, núcleos tensoriales), lo que resulta en un rendimiento subóptimo en el panorama heterogéneo de hardware de los navegadores.
3. Soporte Limitado de Cuantización: Para soportar la amplia variedad de formatos de cuantización utilizados en el ecosistema llama.cpp (desde esquemas heredados hasta K-quants e I-quants avanzados), los motores de inferencia requieren diseños de kernels flexibles. Los marcos existentes a menudo carecen de la extensibilidad para soportar estos formatos diversos de manera eficiente o para integrar nuevos rápidamente.

Metodología: LlamaWeb

Los autores presentan LlamaWeb, un backend WebGPU para llama.cpp diseñado para abordar estas limitaciones. El sistema se construye dentro del código de llama.cpp para aprovechar su extenso soporte de modelos y su ecosistema activo. El diseño se centra en tres pilares fundamentales:

1. Inferencia Local Eficiente en Memoria

Para minimizar la sobrecarga de memoria y garantizar la estabilidad en dispositivos con restricciones de memoria:

• Asignación Estática de Memoria: Toda la memoria requerida para la ejecución del modelo, incluidas las estructuras de datos intermedias para operaciones como FlashAttention y regiones de ranuras de tamaño fijo para parámetros de kernels, se asigna estáticamente al inicio. Esto evita la creación dinámica de búferes y la fragmentación.
• Carga Optimizada: El sistema utiliza la biblioteca wllama para transmitir pesos de forma asíncrona desde el Sistema de Archivos Privado del Origen (OPFS) directamente a búferes WebGPU. Evita materializar pesos en el montón de WebAssembly, que solo puede crecer y no puede liberar memoria durante la vida útil de una pestaña.
• Arena de Búfer de Parámetros: En lugar de mantener pools dinámicos de búferes pequeños para parámetros de kernels (una fuente común de fragmentación), LlamaWeb asigna un único búfer estático con ranuras configurables, rotando a través de ellas para garantizar que los parámetros no se sobrescriban hasta que la ejecución del kernel se complete.

2. Biblioteca de Kernels Portátil en Rendimiento

Para maximizar el rendimiento en hardware heterogéneo:

• Generación de Shaders Templados: El backend utiliza un preprocesador WGSL personalizado (pre-wgsl) y una biblioteca de kernels que genera shaders basados en características específicas del dispositivo (por ejemplo, soporte de subgrupos) y tipos de datos.
• Parámetros Ajustables: Parámetros clave, como los tamaños de baldosa para la multiplicación de matrices y los tamaños de grupo de trabajo, son ajustables. Los autores derivaron valores predeterminados "portátiles en rendimiento" mediante un estudio empírico que barrió miles de configuraciones en cuatro GPUs de diferentes fabricantes.
• Programación de Ejecución: El tiempo de ejecución agrupa eficientemente las envíos de kernels y minimiza la sincronización CPU-GPU. Almacena en caché las pipelines compiladas para evitar costos de compilación redundantes después del pase forward inicial.

3. Diseño de Shaders Consciente de la Cuantización

Para soportar los diversos formatos de cuantización de llama.cpp:

• Descuantización Integrada: Las rutinas de descuantización se diseñan conjuntamente con los kernels de cómputo. Para la multiplicación de matrices, los bloques cuantizados se cargan, se descuantizan en memoria compartida y se reutilizan. Para la multiplicación matriz-vector (fase de decodificación), la descuantización ocurre directamente en registros para mejorar los patrones de acceso a memoria.
• Extensibilidad: El sistema representa los formatos de cuantización como búferes planos de enteros sin signo de 32 bits, con la lógica de descuantización seleccionada en tiempo de compilación. Esto permite que nuevos formatos (por ejemplo, el recientemente introducido q1_0 para modelos de 1 bit) se integren añadiendo una rutina de descuantización sin reescribir las implementaciones centrales de operaciones matriciales.

Contribuciones Clave

1. Un Motor de Inferencia Eficiente en Memoria: Una implementación WebGPU que asigna memoria estáticamente y optimiza el movimiento de datos, permitiendo la ejecución de modelos útiles en dispositivos de borde con restricciones de memoria sin fallos inesperados.
2. Una Biblioteca de Shaders Flexible: Una biblioteca para inferencia de LLM que soporta ajuste externo, especializa shaders basados en características de hardware disponibles (por ejemplo, operaciones matriciales de subgrupos) y utiliza un tiempo de ejecución de programación optimizado.
3. Kernels de Cuantización Diseñados Conjuntamente: Un conjunto de kernels centrales (operaciones matriciales) diseñados para funcionar con los extensos formatos de cuantización de llama.cpp, permitiendo el despliegue de modelos optimizados para diferentes tamaños y facilitando la integración de métodos de cuantización emergentes.

Resultados Experimentales

Los autores evaluaron LlamaWeb en 16 dispositivos de 8 fabricantes (incluyendo NVIDIA, AMD, Intel, Apple, Qualcomm, Samsung, ARM e Imagination Technologies) ejecutando 10 modelos de lenguaje en cuatro formatos de pesos.

• Eficiencia de Memoria: En comparación con WebLLM y Transformers.js, LlamaWeb redujo el uso máximo de memoria del navegador en un 29–33% en promedio en diversas combinaciones de dispositivo/navegador/sistema operativo. En pruebas específicas (por ejemplo, Safari en Apple M4), utilizó hasta un 59% menos de memoria que Transformers.js, que sufría de fugas de memoria en esas configuraciones.
• Rendimiento de Decodificación: LlamaWeb aumentó el rendimiento de decodificación en un 45–69% en comparación con WebLLM y Transformers.js en cuatro GPUs de fabricantes separados.
• Rendimiento de Prefill: Durante la fase de prefill, LlamaWeb inicialmente tuvo un rendimiento inferior en comparación con WebLLM (en un 21–51%) y Transformers.js. Los autores atribuyen esto al uso por parte de WebLLM de fusión de kernels de TVM y a las implementaciones optimizadas por subgrupos de Transformers.js. Sin embargo, al ejecutar LlamaWeb nativamente (fuera del navegador) con las comprobaciones de seguridad deshabilitadas, superó a otros marcos, lo que sugiere que las comprobaciones de seguridad específicas del navegador y la falta de soporte estable de matrices de subgrupos en los navegadores son los principales cuellos de botella.
• Rendimiento Nativo vs. Navegador: En algunos dispositivos (por ejemplo, Intel Arc B580), el backend WebGPU de LlamaWeb fue competitivo o incluso más rápido que los backends nativos (por ejemplo, SYCL) para tareas específicas. Sin embargo, los backends nativos (CUDA, Metal) generalmente superaron a WebGPU hasta en 10 veces en prefill y en 2.5 veces en decodificación, en gran parte debido a una mejor utilización de características específicas del fabricante como núcleos tensoriales y cargas de memoria asíncronas.
• Impacto de la Cuantización: El estudio encontró que, aunque la cuantización reduce los requisitos de memoria, no siempre mejora el rendimiento de forma lineal debido al costo computacional de la descuantización. Por ejemplo, pasar de q8_0 a q4_k_m mejoró el rendimiento, pero pasar a q2_k lo disminuyó en clústeres de gama media.

Significado y Afirmaciones

El artículo afirma que LlamaWeb representa un paso significativo hacia hacer que la inferencia de LLM basada en navegadores sea accesible, privada y de alto rendimiento. Al integrarse con el ecosistema llama.cpp, soporta una vasta biblioteca de modelos (177,691 modelos compatibles en formato GGUF) que otros marcos de navegadores no pueden igualar.

Los autores enfatizan que su trabajo demuestra que la portabilidad de rendimiento es alcanzable en el navegador mediante un ajuste cuidadoso de kernels y una planificación de memoria estática. Si bien reconocen que los backends nativos actualmente mantienen una ventaja de rendimiento, LlamaWeb proporciona una alternativa competitiva multiplataforma que evita la necesidad de instalaciones nativas específicas de la plataforma. El trabajo también destaca la necesidad de mejoras adicionales en las especificaciones de WebGPU, como un mejor soporte para tipos de enteros más pequeños y la eliminación de comprobaciones de seguridad en tiempo de ejecución que obstaculizan el rendimiento, para realizar plenamente el potencial de la IA basada en navegadores.

El código es de código abierto y se ha integrado en el código de llama.cpp mediante más de 100 solicitudes de extracción (pull requests), soportando 72 operadores principales de ML y 23 formatos de datos.