MaxKernel: мультиагентная система создаёт ядра для TPU не хуже экспертов
Написание высокопроизводительных кастомных ядер для ускорителей, сложная задача, требующая глубокой экспертизы на уровне железа. Авторы работы показывают, что большие языковые модели (LLM) в связке с обратной связью от компилятора в реальном времени можно превратить в агентную систему для генерации таких ядер, и представляют MaxKernel, мультиагентную систему для разработки ядер под TPU.
MaxKernel реализует три разных парадигмы работы. Первая, Human-in-the-Loop (HITL): агент проектирует ядро пошагово вместе с человеком. Вторая, полностью автономный агент (Auto), который выполняет цикл оптимизации без участия человека, ориентируясь на метрики и трассы выполнения. Третья, автономный поиск на графе (Graph-Based Autonomous Search), который масштабирует автономный агент для более широкого, глобального исследования пространства возможных решений. Все три режима опираются на общий пул специализированных суб-агентов, отвечающих за планирование, реализацию, самостоятельную отладку, тестирование и профилирование на реальном железе.
Систему проверили на JaxBench, наборе из 50 разноплановых задач по написанию ядер для TPU, а также на сложных реальных нагрузках из современных открытых моделей (конкретные модели в тексте не названы). По утверждению авторов, MaxKernel стабильно генерирует хорошо оптимизированные реализации, сопоставимые по качеству с вручную настроенными экспертами базовыми решениями, и показывает значительный прирост производительности в рамках бенчмарка, точные числовые показатели этого прироста в доступном тексте не приведены. Код агента выложен в открытый доступ на GitHub, в репозитории AI-Hypercomputer/accelerator-agents (папка MaxKernel).
Ключевые факты
- MaxKernel, мультиагентная система для написания и оптимизации кастомных ядер под TPU на основе LLM и обратной связи от компилятора в реальном времени.
- Реализует три парадигмы: пошаговое проектирование вместе с человеком (HITL), полностью автономный агент с циклом оптимизации по метрикам и трассам, и автономный поиск на графе, масштабирующий этот цикл для более широкого исследования пространства решений.
- Все три режима используют общий пул специализированных суб-агентов: планирование, реализация, самостоятельная отладка, тестирование, профилирование на железе.
- Проверена на JaxBench, 50 разноплановых задач по написанию ядер для TPU, а также на реальных нагрузках из современных открытых моделей.
- По заявлению авторов, сгенерированные ядра стабильно сопоставимы по качеству с вручную настроенными экспертами реализациями; код открыт на GitHub.
Почему это важно
Написание кастомных ядер для ускорителей, трудоёмкая задача, которой обычно занимается узкий круг инженеров с глубокой экспертизой на уровне железа. MaxKernel показывает, что связка LLM-агентов с обратной связью от компилятора способна автоматизировать эту работу и приближаться по качеству к результатам ручной настройки экспертами. Это потенциально снижает порог входа в оптимизацию ядер для TPU и ускоряет разработку под специализированное железо.
Кому это важно
Инженерам и исследователям, работающим с TPU и экосистемой JAX; командам, которые пишут кастомные ядра под ускорители и ищут способ автоматизировать эту трудоёмкую задачу; разработчикам мультиагентных систем, которым интересны конкретные паттерны организации агентов, совместная работа с человеком, полностью автономный цикл и поиск на графе.
Как это применить
Код MaxKernel открыт на GitHub, в репозитории AI-Hypercomputer/accelerator-agents (папка MaxKernel), архитектуру можно изучить и опробовать на собственных задачах по написанию ядер для TPU. Три режима закрывают разные сценарии: HITL, для пошаговой совместной работы с контролем со стороны человека, полностью автономный агент, для метрик-ориентированной оптимизации без вмешательства, поиск на графе, для более широкого перебора вариантов дизайна на большем масштабе.
Можно ли доверять
Работа представлена как исследовательская публикация с оценкой на бенчмарке JaxBench (50 задач) и на реальных нагрузках из открытых моделей, а код выложен в открытый доступ, что позволяет независимо проверить результаты. При этом в доступном тексте не приведены точные числовые показатели прироста производительности по сравнению с экспертными реализациями, только качественное утверждение о том, что результаты «сопоставимы» и «значительны»; имена авторов и организация-разработчик в тексте абстракта не указаны.
Риски и подводные камни
Заявление о том, что сгенерированные ядра «не хуже» экспертных, пока не подкреплено точными цифрами в доступном тексте, сравнение стоит проверять по полной статье или коду. Конкретные открытые модели, на реальных нагрузках которых тестировалась система, в тексте не названы, как и детали того, что именно даёт поиск на графе сверх автономного агента, кроме общего расширения пространства поиска.