CUDA Agent от ByteDance: ИИ-агент пишет CUDA-ядра быстрее torch.compile

ByteDance Seed и Tsinghua AIR представили CUDA Agent — систему агентного обучения с подкреплением, которая тренирует большую языковую модель писать GPU-ядра, превосходящие по скорости результат работы компилятора. Разрыв, который система призвана закрыть, узкий, но упрямый: современные модели уже генерируют корректный CUDA, но медленный. На бенчмарке KernelBench базовая модель Seed1.6 проходит 74.0% задач, но обгоняет torch.compile лишь на 27.2% из них, показывая среднее геометрическое ускорение 0.69× — то есть её ядра в среднем медленнее того, что компилятор генерирует сам.

CUDA Agent устраняет этот разрыв, помещая модель в реальную среду разработки CUDA с профилированием, проверками корректности и sandbox с ограниченными правами, а затем обучая её с помощью PPO на протяжении 150 шагов с контекстом в 131 072 токена. Результат — 98.8% пройденных задач и 96.8% задач, где ядра быстрее torch.compile, при среднем геометрическом ускорении 2.11×. На самом сложном уровне Level-3 это примерно на 40 пунктов опережает Claude Opus 4.5 и Gemini 3 Pro.

Можно ли использовать CUDA Agent прямо сейчас

Частично. Обученный агент не выпущен. Он построен на Seed1.6 — проприетарной модели MoE с 23B активных и 230B общих параметров, и веса не опубликованы. Публично доступны: датасет CUDA-Agent-Ops-6K, спецификация SKILL.md, а также рецепты вознаграждения и warm-up.

Полное воспроизведение системы требует серьёзных ресурсов: только песочница профилирования использовала 128 GPU NVIDIA H20, что по силам фронт-лабораториям, GPU-облакам и крупным инфраструктурным командам. Средние команды могут внедрить отдельные компоненты — датасет, milestone-вознаграждение, ограничения против reward hacking, спецификацию навыков — поверх открытой базовой модели. Такой подход к переносу навыков между агентами напоминает недавний Microsoft SkillOpt, где навыки ИИ-агентов тоже переносятся между разными средами.

Как устроена генерация данных и обучение

Исследователи собирают эталонные операторы из библиотек torch и transformers. LLM выбирает до пяти классов операторов и объединяет их в один fused-слой. Фильтр оставляет только операторы, которые выполняются и в eager-, и в compile-режимах, детерминированы, дают непостоянный вывод и работают от 1 мс до 100 мс в eager-режиме. Сэмплы со сходством AST выше 0.9 к любой задаче KernelBench удаляются. Итог — датасет CUDA-Agent-Ops-6K: 6000 сэмплов, 83.77% из которых — композиции из двух операторов.

Цикл агента повторяет инструменты OpenHands: Bash, Read/Write, Edit/MultiEdit, Glob, Grep, NotebookEdit, BashOutput, KillBash — в паттерне ReAct. Инструкции по CUDA поставляются в формате Agent Skills. SKILL.md предписывает модели профилировать модель PyTorch, переписать model_new.py с пользовательскими ядрами, скомпилировать в GPU-песочнице и итерировать, пока ядро не станет минимум на 5% быстрее torch.compile при atol=1e-2, rtol=1e-2.

Вознаграждение дискретное, а не сырое отношение ускорений: r ∈ {−1, 1, 2, 3}. −1 за провал проверки корректности, 3 если ядро обходит и eager, и torch.compile более чем на 5%, 2 если обходит только eager, 1 во всех остальных случаях.

Результаты на KernelBench

Сводные результаты по таблице 1: 98.8% пройденных задач, 98.4% быстрее eager, 96.8% быстрее torch.compile, при среднем геометрическом ускорении 2.60× и 2.11× соответственно. Уровень 2 (последовательности операторов) — самый сильный: 100% прохождения, 100% ускорений, 2.80× относительно torch.compile. Уровень 3 — 94.0% прохождения, 90.0% ускорений и 1.52×, что примерно на 40 пунктов выше Claude Opus 4.5 (50.0%) и Gemini 3 Pro (52.0%) по доле ускорений относительно компилятора.

В статье есть несостыковка: в аннотации и введении заявлены показатели 100% / 100% / 92% ускорений для уровней 1–3, тогда как таблица 1 сообщает 97.0% / 100.0% / 90.0%. Таблица 1 — основная таблица результатов.

Абляции показательны. Удаление агентного цикла роняет долю ускорений относительно compile с 96.8% до 14.1%. Сырое вознаграждение за ускорение даёт 60.4%, отсутствие RFT — 49.8% и коллапс вознаграждения, отсутствие value pretraining — 50.9% и неуправляемые траектории.

Примеры из кейсов показывают, чему обучается политика. Диагональное умножение матриц, переписанное как построчное масштабирование: 73.31× относительно torch.compile. Цепочка matmul-divide-sum-scale, переупорядоченная и слитая: 24.04×. ResNet BasicBlock со свёрнутым в свёртку BatchNorm и cudnnConvolutionBiasActivationForward: 3.59×.

Где пригодится

Область применения — ИИ-инфраструктура и инференс-сервинг, GPU-облака, автономное вождение, количественная торговля, медицинская визуализация и рекомендательные системы — везде, где fused-ядра находятся на критическом пути задержки. Конкретные сценарии: слияние последовательностей операторов, с которыми torch.compile справляется плохо, снижение стоимости токена и перенастройка ядер под новые поколения GPU. Если вам интересно, как другие агентные модели справляются с похожими задачами, стоит посмотреть на OpenWorker от Эндрю Нга — открытого агента, который выполняет работу за пользователя.

Частые вопросы

Что такое CUDA Agent простыми словами?

Это система, которая обучает языковую модель писать CUDA-ядра в реальной среде разработки: модель сама профилирует код, проверяет корректность и итерирует, пока не обгонит компилятор. Обучение идёт методом подкрепления с дискретным вознаграждением за преодоление порога в 5% ускорения.

Какие результаты показывает система?

На KernelBench — 98.8% пройденных задач и 96.8% задач, где ядра быстрее torch.compile, при среднем ускорении 2.11×. На самом сложном уровне Level-3 — 94.0% прохождения и 90.0% ускорений.

Можно ли скачать и использовать обученную модель?

Нет. CUDA Agent построен на проприетарной модели Seed1.6, веса не опубликованы. Открыты только датасет CUDA-Agent-Ops-6K, спецификация SKILL.md и рецепты обучения.

Что открыто для сообщества?

Датасет CUDA-Agent-Ops-6K на 6000 сэмплов, спецификация SKILL.md, а также рецепты вознаграждения и warm-up. Их можно применять поверх открытой базовой модели.

Какие ресурсы нужны для полного воспроизведения?

Только песочница профилирования использовала 128 GPU NVIDIA H20. Полное воспроизведение по силам фронт-лабораториям, GPU-облакам и крупным инфраструктурным командам.

Почему дискретное вознаграждение лучше сырого ускорения?

Дискретное milestone-вознаграждение даёт на 36.4 пункта больше по доле ускорений, чем сырое отношение скоростей. Сырое вознаграждение даёт 60.4% против 96.8% у финальной системы.

Что такое reward hacking и как с ним борются?

Reward hacking — это когда модель находит способ получить высокое вознаграждение, не выполнив реальную задачу. Здесь применяется пять контрмер: проверочные скрипты с ограниченными правами, запрет fallback на torch.nn.functional, проверки на пяти случайных входах, профилирование с синхронизацией устройств и warm-up, отсутствие веб-поиска.

В чём разница между CUDA Agent и обычной генерацией кода?

Обычная модель выдаёт код один раз и не видит результата. CUDA Agent работает в цикле: пишет ядро, компилирует, профилирует, сравнивает с torch.compile и переписывает, пока не достигнет цели. Это и есть агентное обучение.

Есть ли противоречия в результатах самой статьи?

Да. В аннотации и введении заявлены показатели 100% / 100% / 92% ускорений для уровней 1–3, а таблица 1 сообщает 97.0% / 100.0% / 90.0%. Авторы называют таблицу 1 основной.

Кому это может быть полезно на практике?

Командам, которые разрабатывают инференс-сервинг, GPU-облака, системы автономного вождения, торговые алгоритмы, медицинскую визуализацию и рекомендательные системы — везде, где fused-ядра влияют на задержку и стоимость токена.

Источник: marktechpost.com