Разработчики представили CAKE — систему совместного проектирования компилятора и ИИ-агента для GPU-ядер
Создание высокопроизводительных GPU-ядер обычно требует ручной настройки и экспертных знаний. До сих пор инструменты ИИ и языки программирования развивались отдельно, из-за чего воспроизводить такие ядра было сложно. ИИ-агенты обычно работают с компилятором как с черным ящиком, получая только ошибки или замеры времени, а предметно-ориентированные языки скрывают ключевые решения по планированию вычислений.
Новая система CAKE предлагает совместный дизайн компилятора и агента. Вместо традиционного подхода агент пишет на промежуточном представлении CAKE IR — типизированном и явно описывающем аппаратные особенности. Это представление раскрывает роли варп-нитей, перемещение данных, синхронизацию и конвейеры, при этом поддерживает верификацию, оценку стоимости и локализованную диагностику.
Важная особенность CAKE — самообучающаяся инфраструктура. Повторяющиеся ошибки превращаются в правила верификатора, новые примитивы IR, калибровки моделей и переиспользуемые тактики оптимизации. Это позволяет системе эволюционировать от одной задачи к другой, накапливая опыт.
В тестах на задаче Flash-KMeans с нуля на GPU B200 лучший кандидат CAKE IR при бюджете в 80 миллионов токенов показал результат в 1,144 раза быстрее оптимизированного базового варианта FlashML. Для сравнения, прямое написание на CUDA/PTX дало лишь 0,928 от базового уровня.
На другом бенчмарке агентное ядро Kimi Delta Attention достигло среднего геометрического ускорения в 2,05 раза относительно официальной реализации FlashKDA и прошло сквозную проверку инференса. Это подтверждает, что CAKE применим не только к отдельным операциям, но и к более сложным сценариям.
Для диспетчерских версий KNN и KMeans производительность выросла в 1,42–2,12 раза на более чем 400 формах входных данных. Четыре изменения ядер уже доступны в виде upstream-запросов, что говорит о практической готовности разработки.
CAKE ориентирован на NVIDIA GPU от архитектуры Ampere до Blackwell и разделяет задачу эволюции одиночной формы от обобщения библиотеки и диспетчеризации. Такой гибридный подход может стать шагом к более автономной разработке высокопроизводительного кода.
