Специалисты ByteDance Seed и Tsinghua AIR представили CUDA Agent — агентную систему обучения с подкреплением (RL), которая обучает большую языковую модель писать CUDA-ядра, превосходящие по эффективности стандартные компиляторы. Проблема, которую решают авторы, специфична и глубока: современные фронтирные модели уже способны генерировать синтаксически корректный CUDA-код, однако он зачастую оказывается медленнее, чем решения, полученные автоматическими средствами оптимизации. Статистика бенчмарка KernelBench наглядно иллюстрирует этот разрыв: базовая модель Seed1.6 успешно проходит 74,0% задач, но при этом лишь в 27,2% случаев демонстрирует преимущество над torch.compile, имея геометрическое среднее ускорение 0,69×. Это прямо указывает на то, что в большинстве сценариев сгенерированный нейросетью код уступает результатам, выдаваемым компилятором в автономном режиме.

Механика работы и обучение CUDA Agent

CUDA Agent устраняет этот разрыв, помещая модель в полноценную среду разработки с доступом к инструментам профилирования, верификации корректности и выполнения кода в изолированной песочнице с жестко ограниченными правами доступа. Обучение системы проводится с использованием алгоритма PPO в течение 150 шагов при внушительном контекстном окне 131 072 токена. В результате система достигает показателя успешности прохождения в 98,8% и превосходит torch.compile в 96,8% случаев на выборке из 250 задач. Среднее геометрическое ускорение составляет 2,11×, что примерно на 40 процентных пунктов выше результатов моделей Claude Opus 4.5 и Gemini 3 Pro на наиболее сложном уровне (Level-3) бенчмарка.

Инфраструктура и применение

На текущий момент вопрос развертывания остается открытым. Обученный агент не был представлен публично: решение базируется на проприетарной архитектуре Seed1.6 — MoE-модели (Mixture of Experts) с 23 млрд активных и 230 млрд общих параметров, веса которой закрыты. Однако разработчики открыли доступ к датасету CUDA-Agent-Ops-6K, спецификациям SKILL.md, а также рецептам вознаграждения и предварительного прогрева. Масштабы проекта значительны: одна лишь среда профилирования потребовала использования 128 графических ускорителей NVIDIA H20, что ограничивает возможности полной репликации инфраструктурами крупных исследовательских лабораторий и специализированных облачных команд.

Тем не менее, средние команды могут интегрировать отдельные элементы системы — датасет, алгоритмы начисления наград, ограничения для защиты от «взлома» вознаграждения и спецификации навыков — поверх открытых базовых моделей. Технология найдет применение там, где fused-ядра критичны для latency: в AI-инфраструктуре, системах вывода (inference serving), облачных GPU-сервисах, автономном вождении, количественном трейдинге, медицинских визуализациях и рекомендательных системах.

Синтез данных и среда исполнения

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

Алгоритм и система наград

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

Для предотвращения «взлома» награды (reward hacking) предусмотрено пять уровней контроля:

• Использование верификационных и профилировочных скриптов с ограниченными правами.
• Применение контекстных менеджеров, запрещающих использование фоллбэков torch.nn.functional.
• Проверка на пяти случайных наборах входных данных.
• Профилирование с обязательной синхронизацией устройства и предварительным прогревом (warm-up).
• Полное отсутствие инструментов поиска в сети.

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

Анализ результатов и выводы

Согласно основной таблице (Table 1), общая статистика выглядит следующим образом: 98,8% успешных прохождений, 98,4% превосходства над eager-режимом и 96,8% превосходства над torch.compile. Особенно сильны показатели на втором уровне (Level-2), где агент продемонстрировал 100% прохождения и 2,80-кратное ускорение относительно компилятора. На третьем уровне (Level-3) показатели составили 94,0% прохождения и 1,52-кратное ускорение. Важно отметить, что авторы фиксируют вклад каждого компонента: отключение агентского цикла снижает долю превосходства над компилятором с 96,8% до 14,1%, а использование raw-метрики ускорения вместо дискретной награды дает лишь 60,4% успеха.

Кейсы показывают реальные возможности системы: так, диагональное умножение матриц было ускорено в 73,31 раза, цепочка операций matmul-divide-sum-scale — в 24,04 раза, а пересборка ResNet BasicBlock позволила получить ускорение в 3,59 раза. CUDA Agent подтверждает потенциал агентного подхода в задачах, где стандартные компиляторы работают неоптимально, открывая путь к более эффективному использованию вычислительных ресурсов на всех уровнях инфраструктуры.

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