Агент с вопросами обогнал агента с инструментами: 113% против 50% скорости FlashAttention-3

Быстрые GPU-ядра для внимания (attention) — основа скорости обучения и запуска больших моделей, но писать их вручную долго: каждое поколение видеокарт требует нового переписывания, а экспертным командам на это нужны месяцы. Исследователи из Carnegie Mellon University проверили, можно ли поручить такую работу кодинг-агенту, и пришли к неочевидному выводу. Агенту, который застрял, помогают не новые инструменты и не больше свободы, а короткий вопрос в нужный момент.

На одной и той же задаче агент с полным набором инструментов остановился на 50% скорости FlashAttention-3 (эталонного ядра внимания для видеокарт Hopper) после 133 шагов. Агент, которому задавали вопросы, дошёл до 113% на тех же 26 конфигурациях. Для разработчика это значит следующее: если агент буксует, скорее всего, ему не хватает гипотезы о том, где искать проблему, а не ещё одной порции данных.

Почему инструменты не спасают

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

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

Как вопрос меняет задачу

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

Из специального кода у метода только проверяющая обвязка. Одна команда собирает кандидата, сверяет его выходы с исходным ядром и с эталоном в fp32, замеряет скорость по каждому сценарию относительно исходного ядра и считает среднее геометрическое ускорение. Цель задана как проверяемое условие остановки, а не просьба к модели. Вопросы приходят в моменты, когда прогресс стоит, когда агент объявляет «это предел» или «готово» и после подозрительно большого выигрыша. Сначала их задаёт человек, а повторяющиеся вопросы собираются в банк, которым агент пользуется сам. Метод также встроили в эволюционный поиск, где указания для мутаций сформулированы как вопросы.

Что получилось

Полные кампании дали ядра прямого прохода внимания для Hopper и Blackwell: один CUDA-файл без CUTLASS, CuTe, cuBLAS и cuDNN. Они верны на всех 182 сценариях для каждой архитектуры (для Blackwell — на всех запускаемых). В среднем геометрическом они быстрее FlashAttention-3 в 1,16 раза, то есть на 16%, и быстрее FlashAttention-4 в 1,46 раза, на 46%. На производственных формах запросов выигрыш у FlashAttention-4 составил 1,10 раза. Авторы называют их самыми быстрыми открытыми ядрами этого класса на проверенных сценариях.

Любопытна статистика по вмешательствам: 162 сообщения человека и 11,5 тыс. вызовов инструментов в 16 сессиях. Вопросы не означали, что агент пользовался инструментами меньше, он просто применял их иначе. Каждый крупный прирост следовал за вопросом, который переводил агента на другой слой системы или отвергал заявленный им предел, и ни один — за сообщением, которое просто давало информацию.

Оговорки очевидны. Сравнительного автономного агента авторы собрали сами по образцу прежних работ. Эксперимент один: на нём есть сильный эталон для измерения. Часть направляющей работы выполняет человек. Ядра, обвязку, банк вопросов и все трассы исследователи опубликовали.

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