Преобразуйте AI с помощью команды Together AI

Прорыв произошел в праздничные выходные на День памяти в 2022 году. Пока большинство людей в Силиконовой долине наслаждались барбекю, Дан Фу, Три Дау и их коллеги собирались доказать, что мнение экспертов в области искусственного интеллекта ошибочно. Существовало мнение, что внимание трансформеров уже оптимизировано, и специалисты по GPU выжали максимум производительности из оборудования. Однако команда Фу и Дау представила FlashAttention, что изменило ситуацию.

Андрей Карпаты, тогдашний старший директор по AI в Tesla, упомянул об этом в своем твите, и вскоре информация о FlashAttention распространилась по исследовательским каналам AI. Дан вспоминает, что они не ожидали такой реакции, когда выпустили свою работу, но твит Карпаты заставил их понять, что это привлекло внимание.

Ранее исследования в области разреженности и низкоранговых методов показывали теоретическое ускорение, но лишь 10% реального прироста производительности. Команда FlashAttention выбрала другой подход, сосредоточив внимание на реальном движении памяти GPU и вычислительных паттернах. Применяя принципы классических систем баз данных, они достигли 2-3-кратного ускорения.

Для исследователей стало очевидно, что существует огромный неиспользованный потенциал в оптимизации GPU. Эта работа стала основой для одной из самых влиятельных исследовательских команд в области AI, а также критически важным элементом AI Native Cloud.

Многие не понимают, что для успешного AI недостаточно иметь лучшие модели и оборудование. Узкое место заключается в программном обеспечении, которое переводит математические операции в инструкции для кремния. Кернелы играют ключевую роль в этом процессе, позволяя максимально эффективно использовать аппаратное обеспечение. Если их реализовать неправильно, то оборудование остается неиспользуемым.

В марте 2025 года команда кернелов Together AI насчитывала около 15 человек, включая исследователей машинного обучения и ветеранов GPU. Мы получили доступ к новым GPU Blackwell от NVIDIA, и задача была ясна: создать оптимизированные кернелы за одну неделю, в то время как NVIDIA потратила на это год.

Для решения этой задачи мы разработали библиотеку ThunderKittens, которая значительно упростила процесс создания кернелов для нового поколения оборудования. ThunderKittens использует тензорные ядра NVIDIA, что позволяет сократить объем кода с более чем 1000 строк до 100-200. В результате, всего через неделю мы создали одни из самых быстрых кернелов для Blackwell с ускорением до 2 раз по сравнению с cuBLAS.