CUDA-агент обучает LLM написанию GPU-ядер, превосходя по скорости torch.compile

Введение

Команда Seed из ByteDance совместно с Институтом искусственного интеллекта Университета Цинхуа (AIR) представила CUDA Agent — масштабную систему обучения агентов с подкреплением, предназначенную для обучения языковых моделей генерации высокопроизводительных CUDA-ядер.

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

Именно вокруг этой второй проблемы и построен CUDA Agent: он не просто генерирует корректный CUDA-код, но и учится итеративно проводить профилирование производительности, верификацию и оптимизацию ядер до достижения значимых улучшений времени выполнения.

Почему генерации только корректного CUDA-кода недостаточно

До внедрения обучения агентов с подкреплением базовая модель Seed1.6 уже демонстрировала достаточно сильные способности к написанию CUDA-кода на бенчмарке KernelBench.

На 250 задачах KernelBench уровней 1–3, использованных исследователями, Seed1.6 показала 74,0% прохождения тестов. Иными словами, модель могла генерировать функционально корректные решения для большинства задач базового бенчмарка.

Однако производительность — это совсем другая история.

Лишь 27,2% ядер, сгенерированных базовой моделью, оказались быстрее torch.compile, а среднее геометрическое ускорение относительно компиляторного базового уровня составило всего 0,69×.

Этот разрыв подчёркивает главную проблему автоматической генерации GPU-ядер.

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

Профилирование производительности

• Рекомендации по оптимизации CUDA
• Ограниченная песочница
• Сигнал вознаграждения, привязанный к фактическим результатам выполнения

Файл SKILL.md также ограничивает обходные пути, которые могли бы исказить результаты бенчмарка. Например, агент не может просто использовать произвольные операции PyTorch внутри своей реализации пользовательского ядра.

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

Обучение PPO расширяет возможности агента для длинных траекторий оптимизации

Исследователи обучали CUDA Agent с помощью Proximal Policy Optimization (PPO).

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

Поэтому в работе используются различные длины контекста и этапы обучения для стабилизации процесса.

Базовой моделью является Seed1.6 — модель на основе смеси экспертов с 230 миллиардами общих параметров и 23 миллиардами активных параметров.

Для обучения агента с подкреплением:

• Контекстное окно составляет 131 072 токена.
• Воспроизведение в обучении допускает до 150 раундов взаимодействия агента.
• Оценка допускает до 200 раундов взаимодействия агента.
• Модель обучается 150 шагов RL.
• Модели actor и critic используют многоэтапную стратегию разогрева перед полной длинной траекторией оптимизации.

Это гораздо сложнее, чем просто тонкая настройка модели на паре CUDA-подсказок и ответов.

Цель — научить модель улучшать собственные ядра путём многократного взаимодействия с обратной связью от фактического выполнения.

Песочница разделяет компиляцию и профилирование GPU

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

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

Песочница на основе Docker выполняет задачи, связанные с CPU (например, компиляцию), тогда как верификация и профилирование производительности направляются в выделенный пул GPU-песочниц, содержащий 128 графических процессоров NVIDIA H20.

Такое разделение служит двум целям.

Во-первых, оно изолирует компиляцию и взаимодействие агента от бенчмаркинга на GPU. Во-вторых, выделенное распределение GPU снижает помехи между параллельными рабочими нагрузками, делая измерения задержек более стабильными.

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

CUDA Agent достигает 98,8% прохождения тестов и превосходит компилятор на 96,8% задач

Финальная система была оценена на 250 задачах KernelBench уровней с 1-го по 3-й.

CUDA Agent достиг:

• 98,8% общего прохождения тестов
• 98,4% ускорения относительно PyTorch eager execution
• 96,8% ускорения относительно torch.compile
• Среднее геометрическое ускорение в 2,60× относительно eager execution
• Среднее геометрическое ускорение в 2,11× относительно torch.compile

Эти результаты представляют собой значительный скачок по сравнению с исходной точкой Seed1.6.

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

Результаты по уровням сложности KernelBench

На всех трёх уровнях сложности KernelBench улучшение производительности остаётся значительным.

Особенно впечатляющими являются результаты на уровне 2.

Эти задачи включают последовательности операторов, где возможности оптимизации за счёт фьюжена и перемещения данных не всегда эффективно используются статическими эвристиками компилятора.

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

Данные для обучения скоро будут опубликованы

В настоящее время полные веса обученной модели не включены в публичный выпуск проекта.

Тем не менее исследователи опубликовали несколько важных компонентов обучающего стека.

CUDA-Agent-Ops-6K

Датасет CUDA-Agent-Ops-6K содержит 6000 синтетических обучающих задач уровня операторов.

Процесс его построения включает:

1. Сбор исходных операторов из библиотек torch, transformers и других.
2. Использование LLM для объединения нескольких операторов в более сложные задачи фьюжена.
3. Фильтрация сгенерированных задач на основе проверок по результатам выполнения.

Этап фильтрации удаляет задачи со случайным поведением, слишком простые, некорректные или слишком похожие на задачи оценки.

Датасет опубликован на Hugging Face под лицензией CC BY 4.0.

SKILL.md и среда агента

В репозитории GitHub также содержатся инструкционный файл SKILL.md для CUDA и рабочая среда агента.

Эти материалы документируют рабочий процесс оптимизации, доступные инструменты, ограничения и настройку среды выполнения, используемую для инициализации агента.

Механизм вознаграждения и схема обучения

В статье и репозитории описаны дизайн вознаграждений, этапы разогрева, инициализация критика и схема обучения на основе PPO для стабилизации длительного процесса оптимизации.

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

Почему это важно не только для KernelBench

Оптимизация CUDA-ядер лежит в основе значительной части современной инфраструктуры ИИ.

Более быстрые ядра могут улучшить:

• Пропускную способность обучения моделей
• Задержку инференса LLM
• Коэффициент использования GPU
• Экономическую эффективность обслуживания
• Пайплайны ИИ в реальном времени
• Высокопроизводительные численные вычисления

В оригинальной статье AIBase отмечается её значение для

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

Результаты CUDA Agent не означают, что традиционные компиляторы устарели.

torch.compile по-прежнему остаётся мощным автоматическим базовым решением, а оценка CUDA Agent сама по себе опирается на контролируемую тестовую среду и конкретное GPU-окружение.

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

Это означает, что низкоуровневую инженерию производительности можно превратить в полноценную область агентного обучения, а не просто в одноразовую генерацию кода.

Часто задаваемые вопросы

Что такое CUDA Agent?

CUDA Agent — это масштабная система обучения с подкреплением для агентов, разработанная ByteDance Seed, Институтом интеллектуальной промышленности Университета Цинхуа (AIR) и их совместной лабораторией SIA-Lab. Она обучает языковые модели итеративно писать, проверять, анализировать и оптимизировать CUDA-ядра.

На какой модели основан CUDA Agent?

В исследовании используется Seed1.6 в качестве базовой модели. В статье она описана как модель на основе смеси экспертов с общим количеством параметров 230 миллиардов и активируемыми параметрами 23 миллиарда.

Что такое KernelBench?

KernelBench — это открытый бенчмарк для оценки того, способны ли языковые модели генерировать корректные и эффективные GPU-ядра для программ PyTorch. CUDA Agent оценивался на 250 задачах, охватывающих уровни с 1-го по 3-й KernelBench.

Насколько CUDA Agent быстрее torch.compile?

На всём наборе бенчмарков CUDA Agent достигает среднего геометрического ускорения в 2,11 раза по сравнению с torch.compile. Сгенерированные им ядра оказались быстрее компиляторного базового решения в 96,8% оцениваемых задач.

Использует ли CUDA Agent обучение с подкреплением?

Да. Система использует PPO в сочетании с многоэтапным прогревом, предобучением модели ценности, устойчивым дизайном вознаграждений и длинными многораундовыми траекториями агента.

Сколько времени может длиться одна траектория оптимизации CUDA Agent?

Обучающие роллауты допускают до 150 раундов агента, а оценка — до 200 раундов. Обучающее окно контекста агента составляет 131 072 токена.

Является ли CUDA Agent открытым проектом?

Проект опубликовал свой репозиторий GitHub, среду агента, SKILL.md, рецепт обучения и набор данных CUDA-Agent-Ops-6K. Полные веса обученной модели не входят в текущий публичный релиз.

Что такое CUDA-Agent-Ops-6K?

CUDA-Agent-Ops-6K — это набор данных, содержащий 6 000 синтетических задач обучения на уровне операторов CUDA. Он специально разработан для масштабного агентного обучения с подкреплением и включает этапы фильтрации, обеспечивающие выполнимость, детерминированность, нетривиальность задач и их отделение от оценочного набора KernelBench.

Связанные инструменты

• CUDA Agent: официальная страница проекта системы генерации CUDA-ядер от ByteDance Seed и Института интеллектуальной промышленности Университета Цинхуа.
• CUDA Agent GitHub: официальный репозиторий, содержащий среду агента, SKILL.md и опубликованные материалы проекта.
• CUDA-Agent-Ops-6K: официальный набор обучающих данных из 6 000 образцов, выпущенный вместе с проектом.
• KernelBench: открытый бенчмарк для оценки GPU-ядер, сгенерированных большими языковыми моделями.
• PyTorch: фреймворк глубокого обучения, выполнение в режиме eager и через torch.compile которого является ключевым базовым решением в исследовании.
• NVIDIA CUDA: официальная документация NVIDIA по программированию CUDA и разработке GPU-ядер.

Связанные ссылки

• Страница проекта CUDA Agent: официальный обзор, результаты бенчмарков, диаграммы проекта, ссылки на статью, код и наборы данных.
• Статья о CUDA Agent (arXiv): полное исследование, «Масштабное агентное обучение с подкреплением для генерации высокопроизводительных CUDA-ядер».
• Репозиторий CUDA Agent на GitHub: репозиторий исходного кода с опубликованной средой и файлами проекта.
• Набор данных CUDA-Agent-Ops-6K: официальный набор данных Hugging Face, содержащий 6 000 задач на комплексные операторы.
• Репозиторий KernelBench на GitHub: среда бенчмарков и оценки для эффективной генерации GPU-ядер.
• PyTorchtorch.compile: официальная документация PyTorch для компиляторного базового решения, использованного в исследовании.
• Руководство по программированию на CUDA C++ от NVIDIA: официальное руководство NVIDIA по программированию и оптимизации CUDA.

Краткое изложение

CUDA Agent решает конкретную слабость системного кода, генерируемого большими языковыми моделями: сгенерированный CUDA-код должен быть не просто корректным, а действительно быстрее альтернатив, созданных компилятором.

Начиная с Seed1.6, исследователи использовали специализированную агентную среду CUDA, обратную связь по результатам выполнения, устойчивый дизайн вознаграждений и длинногоризонтное обучение PPO, повысив долю успешного прохождения с 74,0% до 98,8%, а долю случаев, когда результат быстрее torch.compile, — с 27,2% до 96,8%.

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

Главный вклад CUDA Agent заключается в демонстрации того, что инженерия производительности сама по себе может быть освоена через итеративный агентный цикл, а не только приближена за счёт одноразовой генерации кода.