ThunderKittens da Together AI: nova linguagem para kernels GPU eficientes

A Together AI lançou o ThunderKittens — uma linguagem de programação (DSL) para escrever código otimizado para GPU. O projeto é posicionado como uma tentativa de simplificar o desenvolvimento de redes neurais, que atualmente requer conhecimento profundo da arquitetura de hardware.

Por que isso era necessário

Escrever código eficiente para GPU é uma arte obscura para a maioria dos engenheiros de ML. Em chips como o NVIDIA H100, existem tensor cores especializados que fornecem 94% de toda a capacidade computacional. Mas para usá-los, é necessário escrever em CUDA — uma linguagem complexa de baixo nível que apenas poucos na indústria dominam.

Em 2023, a popular solução FlashAttention2 mostrou que é possível simplificar significativamente o trabalho com atenção em transformers. Mas este é apenas um caso especial — um algoritmo para um único operador específico. Os desenvolvedores precisam de uma maneira universal de escrever código rápido sem este infernal CUDA, que funcione para diferentes kernels.

Como funciona o ThunderKittens

O ThunderKittens fica entre dois extremos. De um lado há CUDA puro — muito rápido, mas muito complexo, alta barreira de entrada. Do outro lado está Triton, que oculta os detalhes da GPU e simplifica a vida, mas às vezes não consegue extrair o máximo de desempenho do hardware. O ThunderKittens oferece um terceiro caminho.

Sua API se parece com PyTorch, o que é familiar para todos os desenvolvedores de ML acostumados com operações de tensor. Ao mesmo tempo, é transparente o suficiente para que o desenvolvedor entenda o que está acontecendo em nível de hardware. Os autores dizem: se você conhece CUDA, pode 'compilar' ThunderKittens em sua mente.

O objeto fundamental no ThunderKittens é o tile (bloco), uma matriz de tamanho que cabe em um tensor core. Esta é uma unidade de computação que permite carregar completamente um kernel especializado e extrair todo o potencial do hardware moderno.

Desempenho e resultados

Em chips modernos A100 e RTX 4090, o ThunderKittens corresponde em velocidade ao FlashAttention2 — enquanto o código é um pouco mais curto e claro. No H100, esta solução vence: mais rápido que FlashAttention2 tanto na passagem direta quanto na retropropagação. Ou seja, não há compromisso entre limpeza de código e velocidade.

Os autores já escreveram vários kernels em ThunderKittens para outros algoritmos:

• Based — uma versão otimizada de atenção linear
• Hedgehog e outros kernels especializados para transformers
• Várias soluções que se destacam favoravelmente das versões Triton em velocidade

Um detalhe interessante: engenheiros que completaram apenas um curso de CUDA de duas horas conseguiram escrever seu próprio código em ThunderKittens. Isso sugere que a linguagem realmente simplifica o desenvolvimento, embora não oculte os detalhes do hardware.

Projeto aberto e educação

Os autores da Together AI são honestos: é um projeto artístico. Não espere atualizações regulares e suporte para todas as suas reclamações no rastreador de bugs. O projeto foi lançado em acesso aberto porque os desenvolvedores acham interessante compartilhar ideias e ferramentas com a comunidade.

Junto com ThunderKittens, lançaram o NanoGPT-TK — uma versão do icônico projeto NanoGPT de Andrej Karpathy, onde os kernels computacionais principais foram reescritos em ThunderKittens. Isso foi feito especificamente para fins educacionais e para demonstrar capacidades.

NanoGPT há muito é reconhecido como um dos melhores projetos em IA para entender como funciona o treinamento de transformers do zero.

O que isso significa

O ThunderKittens mostra que existe um hiato real em IA entre a conveniência das abstrações (PyTorch, Triton) e o controle sobre o hardware real (CUDA). Acontece que os desenvolvedores estão dispostos a escrever código um pouco mais complexo se isso lhes der controle e velocidade real na prática.

Para engenheiros de ML, isso pode significar que no futuro será mais fácil portar um modelo treinado de um chip para outro — é necessário apenas reescrever alguns kernels GPU, em vez de refazer metade da infraestrutura.

Para pesquisadores, é uma ferramenta para rápida experimentação com algoritmos especializados que não requer um curso de mês em CUDA.