ThunderKittens de Together AI: nuevo lenguaje para GPU kernels eficientes

Together AI lanzó ThunderKittens, un lenguaje de programación (DSL) para escribir código optimizado en GPU. El proyecto se posiciona como un intento de simplificar el desarrollo de redes neuronales, que actualmente requiere un profundo conocimiento de la arquitectura del hardware.

Por qué era necesario

Escribir código eficiente para GPU es magia negra para la mayoría de los ingenieros de ML. Los chips como NVIDIA H100 tienen tensor cores especializados que proporcionan el 94% de toda la potencia computacional. Pero para usarlos, necesitas escribir en CUDA, un lenguaje complejo de bajo nivel que solo entienden unos pocos en la industria.

En 2023, la popular solución FlashAttention2 demostró que era posible simplificar significativamente el trabajo con atención en transformers. Pero esto es solo un caso especial: un algoritmo para un operador específico. Los desarrolladores necesitan una forma universal de escribir código rápido sin este infernal CUDA, que funcione para diferentes kernels.

Cómo funciona ThunderKittens

ThunderKittens se sitúa entre dos extremos. Por un lado está CUDA puro: muy rápido, pero muy complejo, con una barrera de entrada alta. Por el otro, Triton, que oculta los detalles de GPU y simplifica la vida, pero a veces no puede extraer el máximo rendimiento del hardware.

ThunderKittens ofrece un tercer camino. Su API es similar a PyTorch, lo que es familiar para todos los desarrolladores de ML acostumbrados a operaciones tensoriales. Al mismo tiempo, es lo suficientemente transparente para que el desarrollador entienda lo que sucede a nivel de hardware. Los autores dicen: si conoces CUDA, puedes "compilar" ThunderKittens en tu cabeza.

El objeto fundamental en ThunderKittens es tile (mosaico), una matriz de un tamaño que cabe en un tensor core. Esta es la unidad de cálculo que permite cargar completamente un núcleo especializado y extraer todo el potencial del hardware moderno.

Rendimiento y resultados

En los chips modernos A100 y RTX 4090, ThunderKittens coincide en velocidad con FlashAttention2, aunque el código es un poco más corto y más claro. En H100, esta solución gana: es más rápida que FlashAttention2 tanto en la pasada hacia adelante como en la retropropagación. Es decir, no hay ningún compromiso entre la limpieza del código y la velocidad.

Los autores ya han escrito varios kernels en ThunderKittens para otros algoritmos:

• Based: versión optimizada de atención lineal
• Hedgehog y otros kernels especializados para transformers
• Varias soluciones que se destacan favorablemente de las versiones Triton en velocidad

Un punto interesante: los ingenieros que tomaron solo un curso de dos horas sobre CUDA pudieron escribir su propio código en ThunderKittens. Esto sugiere que el lenguaje realmente simplifica el desarrollo, aunque no oculta los detalles del hardware.

Proyecto abierto y capacitación

Los autores de Together AI dicen honestamente: este es un art project. No esperes actualizaciones regulares y apoyo para todas tus quejas en el rastreador de errores. El proyecto se lanzó en acceso abierto porque a los desarrolladores les parece interesante compartir ideas y herramientas con la comunidad.

Junto con ThunderKittens, lanzaron NanoGPT-TK, una versión del icónico proyecto NanoGPT de Andrej Karpathy, donde los kernels computacionales principales se han reescrito en ThunderKittens. Esto se hizo específicamente para capacitación y demostración de capacidades.

NanoGPT hace mucho tiempo que es reconocido como uno de los mejores proyectos en AI para entender cómo funciona el entrenamiento de transformers desde cero.

Qué significa esto

ThunderKittens muestra que en AI hay una brecha real entre la comodidad de las abstracciones (PyTorch, Triton) y el control sobre el hardware real (CUDA). Resulta que los desarrolladores están dispuestos a escribir código un poco más complejo si esto les da control y velocidad real en la práctica.

Para los ingenieros de ML, esto podría significar que en el futuro será más fácil trasladar un modelo entrenado de un chip a otro, simplemente reescribiendo algunos kernels de GPU en lugar de rehacer la mitad de la infraestructura.

Para los investigadores, es una herramienta para experimentar rápidamente con algoritmos especializados, que no requiere un curso de un mes sobre CUDA.