Материал подготовлен автоматически по первоисточникам: ссылки на них — в конце статьи.
Маленькую LLM научили выдавать целиком сценарий, который превращает исходное GPU-ядро в более быстрое, а также переписывать уже найденную оптимизацию.
В работе UCSD и Amazon Annapurna Labs модель с 8 млрд параметров обошла GPT-6 Astra, Opus 5, Kimi-K3 и GLM-5.3 по правильности и ускорению ядер; это нерецензированный препринт, в котором все числа получили сами авторы.
Результат предлагает другой путь для команд, которые оптимизируют повторяющиеся вычисления: не просить универсальную модель написать CUDA-код, а обучить небольшую модель управлять проверяемыми преобразованиями компилятора.
Модель меняет сценарий оптимизации, а не код ядра
GPU-ядро редко ускоряет одно изолированное преобразование. Разбиение циклов, размещение данных в памяти и привязка вычислений к потокам зависят друг от друга: удачный по отдельности шаг может сломать программу или замедлить её после следующего шага.
Little Dragon поэтому не генерирует новое ядро и не выбирает лучший вариант из готового списка. Модель получает исходное представление программы в TVM TensorIR и выдаёт полный сценарий преобразований. При улучшении существующего варианта она тоже строит новый сценарий от исходной программы, а не дописывает команды в конец старого.
Для вывода создали компактный типизированный язык преобразований. Он сохраняет порядок команд, параметры и зависимости между временными объектами, а детерминированный декодер восстанавливает точную трассу TVM. Это позволяет компилятору воспроизвести предложение модели, отклонить недопустимую команду и проверить получившееся ядро.
Отдельный граф связывает факты о программе, объяснение выбора, исполняемые преобразования и ожидаемый эффект. Объяснение помогает обучению, но не управляет компиляцией: точный результат задаёт только формальная часть. Поэтому неточная формулировка от модели-наставника не может незаметно изменить целевой сценарий.
Обучающие данные получили через поиск TVM MetaSchedule на 54 116 разных нагрузках. Из найденных и проверенных сценариев собрали 108 594 примера: часть учила строить оптимизацию с нуля, остальные показывали, как заменить приемлемый сценарий более быстрым.
Сначала модель обучали с учителем на этих примерах. Затем применили обучение с подкреплением: сценарий восстанавливали, выполняли, компилировали и сравнивали с исходным вычислением. Награду за низкую задержку модель получала только после проверки численного результата, а промежуточные баллы помогали ей научиться хотя бы выдавать разбираемый и компилируемый сценарий.
Обратная связь от выполнения улучшила скорость, не сломав правильность
После обоих этапов доля правильных первых ответов достигла 91,27% при построении сценария и 92,01% при улучшении. Правильным считали вариант, который удалось воспроизвести и скомпилировать, а его результат совпал с результатом исходной программы.
При построении оптимизации преимущество над сильнейшей внешней моделью, Opus 5, составило 27,26 процентного пункта. Little Dragon также чаще выдавала ядро с требуемым ускорением, чем сравниваемые универсальные модели. Существенная часть разрыва появилась уже после обучения с учителем, а обратная связь от выполнения прежде всего улучшила скорость получаемых ядер.
В сравнении с отдельным запуском TVM MetaSchedule медианное время получения оптимизации оказалось в 39 раз меньше. Поиск здесь не исчез: его стоимость перенесли в подготовку корпуса, после чего модель стала повторно использовать найденные закономерности для новых задач.
Работа меняет устройство контура оптимизации, но не отменяет проверку
Практический вывод состоит не в замене компилятора моделью. Рабочий контур выглядит иначе: LLM предлагает полный сценарий, TVM воспроизводит его, компилятор собирает ядро, проверяющий модуль сравнивает результат и только затем измеряет задержку. Обычный поиск можно оставить запасным путём для неудачных предложений.
Такой подход подходит продуктам с устойчивым набором операторов и понятной целевой платформой. Команда может заранее оплатить поиск, измерения и обучение, а затем сокращать время оптимизации новых вариантов похожих вычислений. Особенно полезно то, что модель умеет не только начинать с нуля, но и пересматривать все решения в существующем сценарии.
Проверку провели на NVIDIA H100: тестовый набор включал 310 заданий и запросов, охватывавших 39 задач. Численную правильность каждого варианта проверяли на одном случайно сгенерированном входе для каждой нагрузки. Формат сценариев также привязан к зафиксированным версиям TVM и кодека, поэтому перенос на другой компилятор или ускоритель потребует собственного корпуса и повторных замеров.
Для планирования это аргумент в пользу небольшой специализированной модели рядом с компилятором, а не более крупной универсальной LLM. Однако менять производственный контур стоит после пилота на собственных ядрах: выигрыш зависит от того, насколько рабочие нагрузки и оборудование совпадают с теми, на которых модель получила опыт поиска.
Источники
Похоже на вашу задачу?
Расскажите, что собираете. За полчаса разложим на этапы и назовём сроки — это бесплатно и ни к чему не обязывает.
