Проблема ручного оптимизации GPU
Большая часть времени выполнения современных ML-моделей тратится на узкий набор вычислительных ядер: умножение матриц, свертки и нормализацию. Традиционно для их ускорения требовалась ручная переписывание кода на CUDA опытными инженерами. Существующие инструменты на базе LLM часто работают с изолированными тензорами, генерируют код, который сложно интегрировать, и не имеют средств отладки.
Решение: Kernel Forge
Команда исследователей во главе с Joshua Brodsky представила Kernel Forge — open-source агентную систему, которая принимает любую неизмененную PyTorch-модель. В отличие от линейных цепочек генерации, Kernel Forge использует Monte Carlo Tree Search (MCTS) для исследования множества путей оптимизации. Система также оснащена графическим интерфейсом для мониторинга прогресса и отладки.
Результаты бенчмарков
Эксперименты проводились на NVIDIA DGX Spark с GPU GB10. Всего за 50 итераций оптимизации Kernel Forge удалось ускорить ключевые операции в четырех различных моделях (Vision, Diffusion, LLM) по сравнению с режимом PyTorch eager mode:
| Модель / Операция | Ускорение | Тип задачи |
|---|---|---|
| Gemma 4 E2B (softmax) | 2.83× | LLM |
| Qwen 3.5 35B-A3B (softmax) | 1.54× | LLM |
| Stable Diffusion 3.5 Medium (group_norm) | 1.70× | Diffusion |
| ResNet-50 (adaptive_avgpool2d) | 1.52× | Vision |
Почему это важно
Kernel Forge демонстрирует, что агенты на базе LLM могут не просто генерировать код, но и достигать производительности, сопоставимой с ручной оптимизацией экспертов, но с гораздо меньшими затратами времени. Это открывает путь к автоматизации низкоуровневой оптимизации для широкого круга разработчиков, работающих с GPU.
Источник: arXiv cs.AI ↗