MaxKernel: ИИ генерирует высокопроизводительные TPU-ядра на уровне экспертов

Разработчики из проекта AI-Hypercomputer представили MaxKernel — мультиагентную систему автоматической генерации высокопроизводительных ядер для тензорных процессоров (TPU). Работа опубликована в архиве препринтов arXiv.

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

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

Код агента опубликован в открытом доступе на GitHub в репозитории accelerator-agents/MaxKernel. Решение может упростить разработку для TPU и сократить потребность в узкоспециализированных инженерах, однако требует дальнейших проверок на более широком спектре задач.