ThunderKittens, da Together AI: uma nova linguagem para kernels de GPU eficientes

A Together AI lançou o ThunderKittens, uma linguagem de domínio específico (DSL) para escrever código otimizado para GPU. O projeto se posiciona como uma tentativa de simplificar o desenvolvimento de redes neurais, que hoje exige conhecimento profundo da arquitetura do hardware.

Por que isso era necessário

Escrever código eficiente para GPU é magia negra para a maioria dos engenheiros de ML. Chips como a NVIDIA H100 têm tensor cores especiais, que entregam 94% de todo o poder computacional. Mas, para usá-los, é preciso escrever em CUDA, uma linguagem complexa e de baixo nível que só poucos na indústria dominam.

Em 2023, o popular FlashAttention2 mostrou que era possível simplificar bastante o trabalho com atenção em transformers. Mas isso é só um caso particular: um algoritmo para um operador específico. Os desenvolvedores precisam de uma forma universal de escrever código rápido sem esse CUDA infernal, que funcione para diferentes kernels.

Como funciona o ThunderKittens

O ThunderKittens fica entre dois extremos. De um lado, o CUDA puro: muito rápido, mas muito complexo, com barreira de entrada alta. Do outro, o Triton, que esconde os detalhes da GPU e facilita a vida, mas às vezes não consegue extrair o máximo de desempenho do hardware.

O ThunderKittens propõe um terceiro caminho. Sua API é parecida com a do PyTorch, familiar a todo desenvolvedor de ML acostumado a operações com tensores. Ao mesmo tempo, é transparente o bastante para o desenvolvedor entender o que acontece no nível do hardware. Os autores dizem que, se você conhece CUDA, consegue "compilar" o ThunderKittens de cabeça.

O objeto fundamental do ThunderKittens é o tile (ladrilho): uma matriz de tamanho que cabe em um tensor core. É a unidade de computação que permite carregar totalmente o núcleo especializado e extrair todo o potencial do hardware moderno.

Desempenho e resultados

Em chips A100 e RTX 4090, o ThunderKittens iguala o FlashAttention2 em velocidade, com um código um pouco mais curto e claro. Na H100, ele leva vantagem: é mais rápido que o FlashAttention2 tanto na passagem direta quanto na retropropagação. Ou seja, não há nenhum compromisso entre código limpo e velocidade.

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

• Based, uma versão otimizada da atenção linear
• Hedgehog e outros kernels especializados para transformers
• Várias soluções que se destacam em velocidade em relação às versões em Triton

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

Projeto aberto e ensino

Os autores da Together AI dizem com franqueza que isto é um art project. Não esperem atualizações regulares nem atendimento a todas as suas reclamações no bug tracker. O projeto foi liberado em acesso aberto porque os desenvolvedores acham interessante compartilhar ideias e ferramentas com a comunidade.

Junto com o ThunderKittens, foi lançado o NanoGPT-TK, uma versão do icônico projeto NanoGPT, de Andrej Karpathy, em que os principais kernels de computação foram reescritos em ThunderKittens. Isso foi feito especificamente para ensino e para demonstrar as capacidades da linguagem. O NanoGPT é considerado há tempos um dos melhores projetos de IA para entender como funciona o treinamento de transformers do zero.

O que isso significa

O ThunderKittens mostra que existe, na IA, uma lacuna real entre a conveniência das abstrações (PyTorch, Triton) e o controle sobre o hardware de fato (CUDA). Os desenvolvedores, ao que parece, aceitam escrever um 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 migrar um modelo treinado de um chip para outro: bastaria reescrever alguns kernels de GPU, em vez de refazer metade da infraestrutura. Para pesquisadores, é uma ferramenta para experimentar rapidamente com algoritmos especializados, sem exigir um curso de CUDA de meses.