ThunderKittens de Together AI: un nuevo lenguaje para kernels de GPU eficientes

Together AI ha lanzado ThunderKittens, un lenguaje específico de dominio (DSL) para escribir código optimizado para GPU. El proyecto se presenta como un intento de simplificar el desarrollo de redes neuronales, que hoy exige un conocimiento profundo de la arquitectura del hardware.

Por qué hacía falta

Escribir código eficiente para GPU es magia negra para la mayoría de los ingenieros de ML. Chips como la NVIDIA H100 incluyen tensor cores especiales que aportan el 94% de toda la capacidad de cómputo. Pero para usarlos hay que escribir en CUDA, un lenguaje complejo y de bajo nivel que solo entienden unos pocos en la industria.

En 2023, la popular FlashAttention2 demostró que se podía simplificar bastante el trabajo con la atención en los transformers. Pero es solo un caso particular: un algoritmo para un operador concreto. Los desarrolladores necesitan una forma universal de escribir código rápido sin ese infernal CUDA, que funcione para distintos kernels.

Cómo funciona ThunderKittens

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

ThunderKittens propone un tercer camino. Su API se parece a PyTorch, algo familiar para cualquier desarrollador de ML acostumbrado a las operaciones con tensores. Al mismo tiempo, es lo bastante transparente para que el desarrollador entienda qué ocurre a nivel de hardware. Los autores dicen que, si conoces CUDA, puedes "compilar" ThunderKittens en tu cabeza.

El objeto fundamental de ThunderKittens es el tile (baldosa): una matriz de un tamaño que cabe en un tensor core. Es la unidad de cálculo que permite cargar por completo el núcleo especializado y aprovechar todo el potencial del hardware moderno.

Rendimiento y resultados

En chips A100 y RTX 4090, ThunderKittens iguala en velocidad a FlashAttention2, con un código algo más corto y claro. En la H100 gana: es más rápido que FlashAttention2 tanto en el paso hacia delante 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 en ThunderKittens varios kernels para otros algoritmos:

• Based, una versión optimizada de la atención lineal
• Hedgehog y otros kernels especializados para transformers
• Varias soluciones que salen bien paradas frente a las versiones en Triton en velocidad

Un detalle interesante: ingenieros que solo habían hecho un curso de CUDA de dos horas lograron escribir su propio código en ThunderKittens. Eso indica que el lenguaje realmente simplifica el desarrollo, aunque no oculta los detalles del hardware.

Proyecto abierto y formación

Los autores de Together AI dicen con franqueza que es un art project. No esperen actualizaciones periódicas ni atención a todas sus quejas en el gestor de incidencias. El proyecto se publica en 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 emblemático proyecto NanoGPT de Andrej Karpathy en la que los principales kernels de cálculo se han reescrito con ThunderKittens. Se hizo específicamente para la formación y para demostrar las capacidades del lenguaje. NanoGPT se considera desde hace tiempo uno de los mejores proyectos de IA para entender cómo funciona el entrenamiento de transformers desde cero.

Qué significa

ThunderKittens muestra que en la IA existe 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 algo más complejo si eso les da control y velocidad real en la práctica.

Para los ingenieros de ML, esto puede significar que en el futuro será más fácil trasladar un modelo entrenado de un chip a otro: bastaría con reescribir unos pocos kernels de GPU, en lugar de rehacer la mitad de la infraestructura. Para los investigadores, es una herramienta para experimentar rápido con algoritmos especializados que no exige un curso de CUDA de meses.