Проблема ручного оптимизации
Разработка кастомных ядер для ускорителей (TPU) традиционно требует глубоких знаний архитектуры «железа» и занимает много времени. MaxKernel решает эту проблему, используя Large Language Models (LLM) в связке с обратной связью от компилятора в реальном времени. Система позволяет автоматизировать процесс создания кода, снижая порог входа и ускоряя итерации.
Три парадигмы работы агентов
MaxKernel — это многоагентная система, которая использует пул специализированных под-агентов для планирования, реализации, отладки, тестирования и профилирования. Архитектура предлагает три режима работы:
- Human-in-the-Loop (HITL): Агент работает в режиме совместной работы, помогая разработчику шаг за шагом проектировать ядро.
- Autonomous (Auto): Полностью автоматизированный цикл оптимизации, управляемый метриками и трассировками (traces).
- Graph-Based Autonomous Search: Масштабирование Auto-агента для глобального поиска по пространству дизайнов (design space exploration).
Результаты на JaxBench
Система была протестирована на JaxBench — наборе из 50 разнообразных задач для ядер TPU, а также на сложных реальных рабочих нагрузках современных open-source моделей. MaxKernel демонстрирует стабильное генерирование высокооптимизированных реализаций, которые сопоставимы с базовыми значениями, настроенными экспертами вручную.
Детали производительности
Хотя в аннотации акцент сделан на сопоставимости с экспертными решениями, ключевым достижением является способность системы достигать этих показателей без участия человека в цикле оптимизации (в режиме Auto). Это открывает путь к массовому созданию оптимизированного кода для специфичных аппаратных ускорителей.
| Параметр | Значение / Описание |
|---|---|
| Платформа | TPU (Tensor Processing Units) |
| Бенчмарк | JaxBench (50 задач) |
| Архитектура | Multi-agent system (LLM + Compiler Feedback) |
| Режимы | HITL, Autonomous, Graph-Based Search |
| Статус | Open-sourced (код доступен) |
Значение для индустрии
Публикация MaxKernel знаменует сдвиг в сторону «agentic kernel generation». Вместо того чтобы писать низкоуровневый код вручную, инженеры могут делегировать рутинную оптимизацию ИИ-агентам, которые используют компилятор как источник истины для проверки производительности. Это критически важно для эффективного использования ресурсов TPU в эпоху больших моделей.
Источник: arXiv cs.AI ↗