Добавить в корзинуПозвонить
Найти в Дзене
NeuroTechnus

AutoKernel: Open-Source фреймворк для оптимизации GPU PyTorch

Написание быстрого кода для GPU одна из самых трудоемких и узкоспециализированных задач в области машинного обучения, требующая от
инженеров многолетнего опыта и глубокого понимания архитектуры аппаратного обеспечения. Оптимизация даже одного ядра может превратиться
в недели кропотливой ручной настройки. Компания RightNow AI решила кардинально изменить эту ситуацию, представив AutoKernel open-source фреймворк, который автоматизирует этот сложнейший процесс. Проект использует LLM агент для оптимизации GPU, который автономно настраивает ядра в любых моделях PyTorch, делая высокопроизводительные вычисления доступными для широкого круга разработчиков без необходимости в глубоких экспертных знаниях. Ключевое обещание AutoKernel звучит почти фантастически: «загрузите модель перед сном, а утром проснетесь с заметно более быстрыми ядрами». В основе лежит простой, но мощный автономный цикл, в котором ИИ-агент итеративно вносит изменения в код, тестирует
производительность и сохраняет только
Оглавление

Написание быстрого кода для GPU одна из самых трудоемких и узкоспециализированных задач в области машинного обучения, требующая от
инженеров многолетнего опыта и глубокого понимания архитектуры аппаратного обеспечения. Оптимизация даже одного ядра может превратиться
в недели кропотливой ручной настройки. Компания RightNow AI решила кардинально изменить эту ситуацию, представив AutoKernel
open-source фреймворк, который автоматизирует этот сложнейший процесс. Проект использует LLM агент для оптимизации GPU, который автономно настраивает ядра в любых моделях PyTorch, делая высокопроизводительные вычисления доступными для широкого круга разработчиков без необходимости в глубоких экспертных знаниях. Ключевое обещание AutoKernel звучит почти фантастически: «загрузите модель перед сном, а утром проснетесь с заметно более быстрыми ядрами». В основе лежит простой, но мощный автономный цикл, в котором ИИ-агент итеративно вносит изменения в код, тестирует
производительность и сохраняет только лучшие результаты, эффективно
выполняя работу эксперта без его участия.

Почему оптимизация GPU-ядер - это искусство для избранных

Чтобы понять, почему создание AutoKernel стало важным событием,
необходимо сначала разобраться в самой проблеме. Написание быстрого
GPU-кода — одна из самых сложных и трудоемких специализаций в области
машинного обучения. В основе этого процесса лежит работа с GPU-ядрами.
GPU-ядро это функция или программа, которая выполняется параллельно на тысячах вычислительных блоков графического процессора (GPU). Оно лежит в основе большинства высокопроизводительных вычислений в машинном обучении, таких как умножение матриц или операции нормализации. Когда вы
запускаете трансформерную модель, например LLaMA, основное время
вычислений тратится именно на выполнение ядер для таких операций, как
умножение матриц (matmul), softmax и attention.

Проблема в том, что выжать максимальную производительность из этих
ядер — задача, требующая от инженера одновременного учета множества
взаимозависимых факторов. Это настоящее искусство. Необходимо
балансировать арифметическую интенсивность, обеспечивать коалесценцию
памяти для эффективного доступа к данным, управлять регистровым
давлением, подбирать оптимальные размеры тайлов (блоков данных) и
синхронизировать потоки на уровне варпов. Чтобы представить масштаб,
отметим, что одно
высокопроизводительное ядро для matmul может содержать более 200 строк сложного кода на CUDA или Triton.

На освоение этого набора навыков уходят годы, а экспертов, способных
писать действительно быстрый GPU-код, на рынке единицы. Такой ручной
процесс настройки плохо масштабируется по мере того, как архитектуры
моделей постоянно усложняются. Насколько эта задача сложна, наглядно
демонстрирует бенчмарк KernelBench, который оценивает способность
передовых LLM справляться с оптимизацией GPU-ядер. Результаты
показывают, что даже лучшие модели смогли сравниться с базовой
производительностью PyTorch менее чем в 20% случаев. Именно этот
зияющий пробел в производительности и призван закрыть AutoKernel, предлагая автоматизированный подход к этому элитарному искусству.

Архитектура AutoKernel: Как фреймворк оптимизирует GPU ядра

Чтобы понять, как AutoKernel оптимизирует GPU ядра, нужно взглянуть
на его архитектуру. Центральная идея фреймворка AutoKernel для PyTorch
моделей заключается в полной механизации рутинного и трудоемкого
рабочего процесса инженера по оптимизации графических процессоров. В
основе архитектуры лежит простой, но невероятно эффективный автономный
цикл
«создать-протестировать-сохранить/откатить». На
практике это выглядит следующим образом: LLM-агент, концепция которого
перекликается с темами нашего материала «Что такое NAI: Мультимодальный
фреймворк Google AI на базе Gemini»
[1], вносит изменения в единственный целевой файл — kernel.py. Затем специализированный бенчмарк-стенд строго проверяет предложенный код на математическую корректность и измеряет его пропускную способность.
Результат этого тестирования определяет дальнейшую судьбу кода. Успешное
улучшение производительности навсегда фиксируется в истории проекта
через стандартную команду git commit, продвигая ветку вперед. Любая
регрессия или ошибка компиляции мгновенно и бесследно отменяется с
помощью git reset. Вся история экспериментов сохраняется в простом
текстовом файле формата TSV, который легко читается как человеком, так и
машиной.

Подобная архитектура была напрямую вдохновлена известным проектом
autoresearch Андрея Карпатого, однако здесь она филигранно адаптирована
под специфику низкоуровневого программирования GPU. Каждая итерация
этого непрерывного цикла занимает всего около 90 секунд. Это время
делится поровну: 30 секунд уходит на жесткую проверку корректности, 30
секунд — на профилирование производительности, и еще 30 секунд
выделяется на аналитическую работу агента по генерации нового фрагмента
кода. Благодаря такой высокой скорости фреймворк использует цикл
«создать-протестировать-сохранить/откатить» и проводит сотни
экспериментов за ночь на одном GPU — в среднем от 300 до 400 итераций за
сессию. Это подчеркивает
колоссальную масштабируемость
подхода, недоступную при ручном труде. Подобные стремительные прорывы в
сфере ИИ автоматизации, как мы уже подробно обсуждали в статье
«Aletheia: ИИ агент Google DeepMind для научных открытий»
[2], кардинально меняют ландшафт современной разработки.

Однако было бы серьезной ошибкой считать, что система работает в
абсолютном информационном вакууме. Вся логика агента опирается на
фундаментальную
экспертную базу знаний
детализированный документ program.md, состоящий из 909 строк. Этот файл
представляет собой исчерпывающее шестиуровневое руководство по
оптимизации. Оно охватывает весь спектр задач: от базовой настройки
размеров блоков и количества варпов до сложных паттернов доступа к
памяти, таких как программная предвыборка. Далее следуют вычислительные
оптимизации, продвинутые техники вроде разделения K-измерения и
специфичные для конкретных аппаратных архитектур стратегии. На вершине
этой пирамиды находятся узкоспециализированные алгоритмы для конкретных
ядер, такие как онлайн-вычисление softmax для механизмов внимания или
алгоритм Уэлфорда для нормализации.

Наличие такого объемного и детального свода правил обнажает важный
нюанс современных нейросетевых инструментов. Жесткое требование к
909-строчному документу с инструкциями (program.md) для агента указывает
на то, что «автономность» все еще требует значительного человеческого
экспертного ввода и поддержки. Фреймворк не изобретает принципиально
новые методы оптимизации из ничего. Вместо этого он методично, с
машинной скоростью и абсолютной неутомимостью, применяет заложенный в
него концентрат человеческого опыта, перебирая сотни гипотез до тех пор,
пока не найдет
идеальную конфигурацию для конкретной модели нейросети.

Стратегия и надежность: От закона Амдала до пятиступенчатой проверки

В отличие от подходов, рассматривающих оптимизацию ядер изолированно,
AutoKernel применяет более интеллектуальную и целостную стратегию
«сначала профилируем, потом оптимизируем».
Это не просто перебор вариантов, а продуманная система, нацеленная на
максимальный эффект на уровне всей модели. В основе этого подхода лежит
использование `torch.profiler` для полного анализа модели и выявления
наиболее ресурсоемких операций. Затем система ранжирует цели для
оптимизации, руководствуясь законом Амдала. Этот фундаментальный для
оптимизации GPU принцип описывает максимальное теоретическое ускорение
выполнения задачи, когда только часть этой задачи оптимизирована. Он
показывает, что общее ускорение системы ограничено долей времени,
которую оптимизируемый компонент занимает в общем времени выполнения,
подчеркивая важность фокусировки на наиболее ресурсоемких частях. Такой
подход гарантирует, что усилия агента направлены туда, где они принесут
наибольшую пользу.

Для управления этим процессом в AutoKernel предусмотрен специальный
оркестратор. Его задача — не позволить агенту зацикливаться на одном
ядре с убывающей отдачей. Оркестратор переключает задачи, как только
достигается один из критериев: пять последовательных неудачных попыток,
достижение 90% пиковой производительности GPU или истечение временного
лимита. Это обеспечивает
сбалансированное распределение ресурсов и планомерное улучшение всей модели.

Однако производительность не имеет ценности без надежности. Понимая это, разработчики встроили в AutoKernel строгую пятиступенчатую систему проверки корректности, которую должен пройти каждый кандидат на оптимизацию. Первый этап — быстрый дымовой тест для отсеивания базовых ошибок компиляции. Второй — проверка на различных конфигурациях входных данных и типах (FP16, BF16, FP32) для выявления багов, зависящих от размера.
Третий этап — тестирование на аdvеrsаriаl inputs (например, значения с
экстремальным динамическим диапазоном) для проверки численной
стабильности. Четвертый — верификация детерминизма путем многократного
запуска с одними и теми же входными данными для выявления состояний
гонки. Наконец, пятый этап — проверка на размерах, не являющихся
степенями двойки, что позволяет обнаружить скрытые ошибки маскирования и
обработки остатков. Этот многоуровневый контроль гарантирует, что
ускорение кода не достигается ценой его корректности.

Производительность AutoKernel на NVIDIA H100: Впечатляющие победы и сложные вызовы

Чтобы оценить реальную производительность AutoKernel на NVIDIA H100,
команда RightNow AI провела всестороннее тестирование на одном из самых
мощных ускорителей современности — NVIDIA H100 80GB, сравнивая
результаты с базовыми реализациями PyTorch (eager mode) и его
компилятором torch.compile в режиме максимальной автонастройки.
Результаты для ядер, производительность которых ограничена пропускной
способностью памяти, оказались по-настоястоящему впечатляющими.
AutoKernel демонстрирует
значительные приросты производительности:
для RMSNorm ускорение достигло 5.29×, для Softmax — 2.82×, а для
Cross-entropy — 2.21× по сравнению с eager-версией. Ключ к такому успеху
лежит в способности фреймворка объединять несколько стандартных
операций в однопроходные ядра, написанные с использованием Triton для
GPU ядер. Triton — это предметно-ориентированный язык программирования,
разработанный для написания высокоэффективных GPU-ядер. Он позволяет
разработчикам создавать оптимизированный код для GPU, который часто
превосходит стандартные библиотеки, при этом оставаясь более простым в
использовании, чем низкоуровневые языки вроде CUDA C++. Такой подход
радикально снижает количество обращений к HBM (High Bandwidth Memory) —
типу высокопроизводительной оперативной памяти, используемой в
современных GPU, что критически важно для задач, интенсивно использующих
данные.

Однако картина меняется, когда речь заходит о задачах, интенсивно
использующих вычислительные ресурсы, таких как матричные умножения
(matmul). Здесь разработчики AutoKernel честно признают существующие
вызовы: эффективность для сложных, вычислительно-интенсивных ядер пока
уступает высокооптимизированным библиотекам вроде cuBLAS, которые NVIDIA
годами оттачивала под каждую конкретную архитектуру GPU. Это
накладывает определенные ограничения на универсальность фреймворка в его
текущем виде. Тем не менее, даже в этой сложной дисциплине AutoKernel
демонстрирует свое преимущество над torch.compile, превосходя его в 1.55
раза для матриц размером 2048³. Этот факт показателен: он подчеркивает,
что даже
передовые компиляторы не всегда находят
оптимальные стратегии, и пространство для улучшений, которое может
использовать автономный агент, все еще огромно.

Пожалуй, самым убедительным доказательством потенциала AutoKernel
служат его успехи за пределами лабораторных тестов, в реальных задачах
сообщества. Оптимизированное с помощью AutoKernel ядро заняло первое
место в таблице лидеров vectorsum_v2 для грядущего ускорителя B200.
Более того, один из пользователей сообщил, что с помощью всего одного
запроса к агенту, что заняло около трех минут, ему удалось сгенерировать
ядро для матричного умножения в формате FP4, которое превосходит по
производительности знаменитую библиотеку CUTLASS в 1.63 — 2.15 раза на
различных размерах матриц. Этот пример делает сравнение AutoKernel,
cuBLAS и CUTLASS особенно примечательным, поскольку CUTLASS представляет
собой вершину ручной оптимизации кода на C++ для тензорных ядер NVIDIA,
и превзойти ее — задача высочайшей сложности, с которой справился
автономный агент.

Перспективы, риски и сценарии будущего

Несмотря на впечатляющие результаты, для всесторонней оценки
AutoKernel необходимо рассмотреть потенциальные ограничения и риски,
которые могут повлиять на его широкое внедрение. Зависимость от
LLM-агента может приводить к непредсказуемым результатам или
«застреванию» в локальных оптимумах, а заявленные ускорения сильно
привязаны к конкретным конфигурациям оборудования. Кроме того, хотя
фреймворк является open-source, его интеграция в существующие
производственные пайплайны может потребовать
значительных инженерных усилий и специфических навыков.

Помимо этих практических сложностей, можно выделить несколько ключевых категорий рисков. Технологический риск
заключается в неспособности агента находить оптимальные решения для
всех типов ядер, особенно для новых архитектур. Экономический риск
связан с высокими затратами на вычислительные ресурсы для длительных
прогонов оптимизации. Существует и риск зависимости, при котором
разработчики могут утратить глубокое понимание низкоуровневой
оптимизации GPU, чрезмерно полагаясь на автоматизацию. Не стоит забывать
о риске безопасности, так как генерация низкоуровневого кода в
open-source проекте открывает потенциал для внедрения уязвимостей.
Наконец, по мере усложнения моделей и архитектур GPU, время итерации
агента может увеличиваться, снижая общую эффективность фреймворка.

Учитывая эти факторы, можно представить три основных сценария будущего развития AutoKernel. В позитивном сценарии,
AutoKernel становится де-факто стандартом для оптимизации GPU-ядер в
PyTorch, значительно ускоряя разработку и демократизируя доступ к
высокопроизводительным вычислениям. Нейтральный сценарий предполагает,
что фреймворк находит свою нишу как полезный инструмент для оптимизации
определенных типов ядер, особенно memory-bound, но не заменяет полностью
ручную экспертную настройку для самых критичных задач. В негативном
сценарии, AutoKernel сталкивается с трудностями в масштабировании и
стабильности, не оправдывая ожиданий по производительности, и его
использование остается ограниченным узкоспециализированными
применениями.

Автоматизация экспертных знаний и будущее высокопроизводительных вычислений

AutoKernel представляет собой не просто очередной инструмент, а фундаментальный сдвиг
в подходе к оптимизации GPU. Его главная ценность — трансформация
недель кропотливой ручной настройки, требующей глубокой и редкой
экспертизы, в полностью автономный ночной процесс. Эта автоматизация
экспертных знаний открывает доступ к высокопроизводительным вычислениям
для широкого круга разработчиков, ранее ограниченных чрезвычайно высоким
порогом входа в мир низкоуровневого программирования.

Успех фреймворка зиждется на нескольких ключевых принципах.
Во-первых, это стратегический подход к оптимизации, направляемый законом
Амдала, который гарантирует, что усилия агента сосредоточены на
наиболее значимых участках кода. Во-вторых, это впечатляющие результаты
для ядер, ограниченных пропускной способностью памяти, где прирост
производительности по сравнению с PyTorch eager и torch.compile наиболее
ощутим. И, что самое важное, —
бескомпромиссное отношение
к надежности. Корректность здесь не обсуждается: прежде чем будет
зафиксировано любое ускорение, каждое ядро-кандидат должно пройти
строгую пятиступенчатую проверку. При этом важно понимать текущие
границы: в задачах, где доминируют вычисления, узкоспециализированные
библиотеки вроде cuBLAS все еще задают планку производительности.

Таким образом, AutoKernel следует рассматривать не как полную замену инженера-эксперта, а как мощнейший мультипликатор
его усилий и инструмент демократизации. Он берет на себя рутинный, но
критически важный процесс перебора и тестирования конфигураций,
освобождая человеческий интеллект для решения более сложных
архитектурных задач. Этот проект — яркий пример того, как ИИ-агенты
становятся полноценными партнерами в решении сложнейших инженерных
проблем, знаменуя собой важный шаг к будущему, где создание
высокоэффективного кода станет доступным не только узкому кругу
специалистов, но и любому исследователю в области машинного обучения.