Команда Стэнфордского университета показала, что большие языковые модели способны автоматически создавать очень эффективные GPU-ядра, которые в ряде случаев работают быстрее стандартных функций популярного фреймворка машинного обучения PyTorch.
Эти CUDA-C-ядра — маленькие программы, выполняющиеся напрямую на видеокартах Nvidia. Они обрабатывают задачи вроде перемножения матриц или обработки изображений и являются основой многих ИИ-процессов.
В ходе экспериментов языковые модели создавали CUDA-ядра, а затем их сравнивали по скорости с встроенными функциями PyTorch. Фреймворк PyTorch, среди разработчиков которого Meta, широко используется в машинном обучении и содержит библиотеку заранее подготовленных GPU-операций, нужных для ИИ-задач.
В нескольких тестах автоматически сгенерированные ядра работали заметно быстрее стандартных функций PyTorch. Например, одно из них ускорило операцию layer normalization (стандартизация значений в слоях нейросети) в 4,8 раза.
Улучшения также отмечены при обработке изображений (Conv2D), функции softmax (преобразует выходные значения в вероятности) и комбинированной операции — сверка (convolution), активация ReLU (обнуляет отрицательные числа) и max-pooling (выбор максимального значения в области изображения). Во всех этих случаях автоматически созданный код оказался быстрее.
Исследование опиралось на бенчмарк KernelBench: языковая модель пыталась заменить отдельные операторы PyTorch своими CUDA-ядрами для ускорения расчетов на видеокарте.
В работе задействовали две большие языковые модели: OpenAI o3 и Gemini 2.5 Pro. Они используют параллельные стратегии оптимизации. Каждый созданный ядро проверяли на корректность и скорость.
Стэнфордский подход отличается от привычных ручных: идеи для оптимизации писали простым языком, генерировали сразу несколько вариантов кода, запускали их параллельно и выбирали только самые быстрые решения для следущего этапа.
Такой разветвленный поиск дал большой разброс решений. Лучшие ядра основывались на известных техниках: более эффективный доступ к памяти, одновременное выполнение арифметики и обращений к памяти, снижение точности данных (например, FP32 на FP16), эффективное использование вычислительных блоков GPU и упрощение циклов.
Исследователи подчеркивают: их подход не заменяет существующие способы обучения, а работает в дополнение. Синтетические данные, которые появляются в ходе поиска, могут использоваться для обучения новых моделей. Итог — ускорение готового кода и ценные обучающие данные.
В одном из примеров AI-сгенерированное ядро для Conv2D стало с 20% до 180% от скорости PyTorch после 13 циклов улучшения. Conv2D — ключевая операция для обработки изображений, когда исходные данные складываются с фильтрами.
Показаны важные приемы: превращение свертки в матричное умножение для использования ядер GPU, двойная буферизация для одновременной загрузки и вычислений, и заблаговременное вычисление индексов в общей памяти.
Финальная версия включала сложные методики CUDA-программирования, до этого доступные только опытным разработчикам.
Есть и ограничения. Автоматически созданные ядра хуже справляются с задачами на новых типах данных, таких как FP16: например, одна из матричных операций FP16 выдала скорость только 52% от PyTorch, а для Flash Attention — сложной техники оптимизации памяти — получилось всего 9%.
Авторы отмечают: раньше никто не мог надежно генерировать такие ядра автоматически. Даже при ограниченных ресурсах новый подход уже показывает результат.
Этот подход согласуется с другими исследованиями, где параллельный поиск и мощные языковые модели позволяют создавать эффективные компоненты систем. Аналогичные итоги показали проекты DeepMind AlphaEvolve и Deep Think (Gemini 2.5 Pro).
*Meta — запрещенная в РФ организация.
