ThunderKittens от Together AI: новый язык для эффективных GPU kernels

Together AI выпустила ThunderKittens — язык программирования (DSL) для написания оптимизированного кода на GPU. Проект позиционируется как попытка упростить разработку нейронных сетей, которая сейчас требует глубоких знаний архитектуры железа.

Почему это было нужно

Написание эффективного кода для GPU — это чёрная магия для большинства ML-инженеров. На чипах вроде NVIDIA H100 есть специальные tensor cores, которые дают 94% от всей вычислительной мощности. Но чтобы их использовать, нужно писать на CUDA — сложный низкоуровневый язык, который понимают только единицы в индустрии. В 2023 году популярное решение FlashAttention2 показало, что можно значительно упростить работу с вниманием в трансформерах. Но это лишь частный случай — алгоритм для одного конкретного оператора. Разработчикам нужен универсальный способ писать быстрый код без этого адского CUDA, который бы работал для разных ядер.

Как работает ThunderKittens ThunderKittens расположен между двумя крайностями.

С одной стороны есть чистый CUDA — очень быстро, но очень сложно, высокий входной барьер. С другой — Triton, который скрывает детали GPU и упрощает жизнь, но иногда не может выжать максимум производительности из железа. ThunderKittens предлагает третий путь. Его API похож на PyTorch, что знакомо всем ML-разработчикам, привыкшим к тензорным операциям. Одновременно он достаточно прозрачный, чтобы разработчик понимал, что происходит на аппаратном уровне. Авторы говорят: если вы знаете CUDA, вы можете "скомпилировать" ThunderKittens в голове. Фундаментальный объект в ThunderKittens — это tile (плитка), матрица размером, который влезает в tensor core. Это единица вычисления, которая позволяет полностью загрузить специализированное ядро и извлечь весь потенциал из современного железа.

Производительность и результаты

На современных чипах A100 и RTX 4090 ThunderKittens совпадает по скорости с FlashAttention2 — в то время как код несколько короче и понятнее. На H100 это решение выигрывает: быстрее FlashAttention2 как на прямом проходе, так и на обратном пропагейшене. То есть нет никакого компромисса между чистотой кода и скоростью. Авторы уже написали на ThunderKittens несколько ядер для других алгоритмов: Based — оптимизированная версия линейного внимания Hedgehog и другие специализированные ядра для трансформеров * Несколько решений, которые выгодно отличаются от Triton-версий по скорости Интересный момент: инженеры, которые прошли всего двухчасовой курс по CUDA, смогли написать собственный код на ThunderKittens. Это говорит о том, что язык действительно упрощает разработку, хотя и не скрывает детали железа.

Открытый проект и обучение Авторы из Together AI честно говорят: это art project.

Не ждите регулярных обновлений и поддержки всех ваших жалоб в багтрекере. Проект выпущен в открытом доступе, потому что разработчикам кажется интересным делиться идеями и инструментами с коммьюнити. Вместе с ThunderKittens выпустили NanoGPT-TK — версию культового проекта NanoGPT от Андрея Карпатьи, где основные вычислительные ядра переписаны на ThunderKittens. Это сделано специально для обучения и демонстрации возможностей. NanoGPT давно признан одним из лучших проектов в AI для понимания того, как работает обучение трансформеров с нуля.

Что это значит

ThunderKittens показывает, что в AI есть реальный зазор между удобством абстракций (PyTorch, Triton) и контролем над реальным железом (CUDA). Оказывается, разработчики готовы писать немного более сложный код, если это даст им контроль и реальную скорость на практике. Для ML-инженеров это может значить, что в будущем переносить обученную модель с одного чипа на другой будет проще — достаточно переписать несколько GPU ядер, а не переделывать половину инфраструктуры. Для исследователей это инструмент для быстрого экспериментирования со специализированными алгоритмами, который не требует месячного курса CUDA.