Агент на вопросах довёл CUDA-ядро до 1,13× от скорости FlashAttention-3 против 0,50

Исследователи из Carnegie Mellon University предложили необычный рецепт для агентов, которые оптимизируют GPU-ядра: вместо того чтобы снабжать модель разобранными отчётами профайлера, поиском по документации и хранилищем «выученных уроков», ей в нужные моменты задают короткие вопросы. На одной и той же задаче, с тем же оценщиком, на том же узле H200 и в ту же неделю такой агент дошёл до 1,13 от скорости FlashAttention-3 по среднему геометрическому на 26 конфигурациях. Автономный агент, собранный по господствующей сегодня схеме, остановился на 0,50 после 133 шагов.

Между скриптом и свободным поиском

Авторы описывают два лагеря. Скриптованные агенты заранее фиксируют, что модель делает на каждом шаге: обученная политика, меню оптимизаций, жёсткие роли. Они экономно расходуют попытки там, где рецепт подходит, и буксуют там, где не подходит. Например, у Autocomp производительность рушится без меню оптимизаций, а RL-подход с фиксированным циклом, конкурентный на отдельных операторах, проваливается на fused-моделях.

Открытые агенты вроде AVO дают модели свободу мутировать популяцию ядер и платят за это числом попыток. По данным статьи, агент AVO за семь дней изучил более 500 направлений, чтобы зафиксировать 40 версий, и нуждался в супервизоре для выхода из циклов неудачных правок. Около 30% добавленных такими агентами строк байт-в-байт совпадают с ранее удалёнными, а часть эволюционно найденных CUDA-кандидатов оказывается некорректной.

Оба лагеря реагируют на сложные цели одинаково: докладывают агенту больше информации. Авторы считают, что не хватает другого — направления, то есть гипотезы о том, где остался разрыв. LLM плохо справляются с открытой оптимизацией: ищут жадно, цепляются за решение из промпта, рано останавливаются и сообщают о завершении вопреки данным среды.

Как устроен метод

Название самого агента в доступном тексте не сохранилось, но конструкция описана ясно. Берётся готовое CUDA-ядро, набор сценариев и цель. Дальше три части: минимальный обвес-оценщик, универсальный кодовый агент без сценария и источник вопросов.

Оценщик одной командой собирает кандидата, сверяет выход с эталоном и с fp32-референсом на каждом сценарии и замеряет время против исходного ядра на той же GPU. Ограничения проверяются механически: запрет на CUTLASS, CuTe, cuBLAS и cuDNN контролируется по символам в собранном бинарнике. Цель оформлена как Stop-хук: когда агент пытается завершить ход, предикат вычисляется, и при невыполненных условиях агента возвращают к работе. Предикат должен покрывать каждый сценарий, а не среднее, иначе агент подгонит среднее на удобных формах. Единственный раз, когда вердикт отдали модели-судье, она приняла преждевременное заявление о завершении, которое человек потом отменил.

Вопросы задаются в пяти типовых точках: в начале кампании, на плато, когда агент объявляет предел или завершение, после подозрительно большого выигрыша и когда он объясняет away поломку корректности. Пример из статьи: «Выигрывает ли базовое ядро по тактам или только по времени на стене?» Такой вопрос ограничен и проверяем, и чтобы ответить, модель способна переписать часть ядра, к которой сама бы не притронулась.

Из 162 сообщений человека чаще всего встречались проверки измерений (13), сравнение с оценкой границы (10), валидность бенчмарка (7) и разбор разрывов (7). Между вопросами агент работает сам: примерно 71 вызов инструментов на одно сообщение. Все крупные приросты, по словам авторов, следовали за вопросом, который переводил агента на другой уровень системы или оспаривал заявленный предел, но не за вопросом, сообщавшим новую информацию.

Что получилось и где оговорки

Итоговые ядра прямого прохода FlashAttention для Hopper и Blackwell — одиночные CUDA-файлы без CUTLASS, CuTe, cuBLAS и cuDNN. Они корректны на всех 182 сценариях для Hopper (для Blackwell — на всех запускаемых) и дают в среднем геометрическом 1,16 от FlashAttention-3 и 1,46 от FlashAttention-4, включая 1,10 от FA4 на производственных формах. Всё это заняло 16 сессий и около 11,5 тысячи вызовов инструментов.

Повторяющиеся вопросы авторы свели в банк: 101 основная запись в десяти группах, пять для случая, когда есть другая реализация для сравнения, и 18 привязанных к архитектурам. Но только 29 записей подтверждены трассами, а 72, взятые из публичных источников, пока не проверены. Первая версия банка из 48 вопросов оказалась перекошена под один кейс: 21 вопрос ссылался на эталонное ядро, 13 — на структуру внимания. Из релизной версии это убрали.

Оговорки очевидны. Это один кейс — внимание, для которого есть сильный эталон. Вопросы в экспериментах задавал человек, а самостоятельное применение банка проверено слабее. Сравнение с «оснащённым» агентом строилось авторами самостоятельно, пусть и в стиле AVO. Зато код ядер, оценщик, банк вопросов и все трассы обещано опубликовать, так что результат можно перепроверить.