JAXBench: Новый стандарт для автооптимизации ядер TPU
В мире вычислительной графики и искусственного интеллекта долго существовал разрыв: алгоритмы могли автономно оптимизировать работу видеокарт (GPU), но для специализированных ускорителей Google — процессоров TPU — аналогичных стандартов не было. Новая работа на arXiv представляет JAXBench, первую в своем роде наборы тестов, созданный нативно для архитектуры TPU. Этот инструмент позволит исследователям и инженерам проводить честные замеры производительности для ИИ-генерированного кода, который управляет вычислениями на GPU-подобных кристаллах Google.
# JAXBench: Заполнение пустоты в тестировании TPU
Долгое время сообщество машинного обучения имело четкие метрики для оценки видеокарт и алгоритмов их оптимизации. Для процессоров_tensor_-core (TPU) ситуация была иной: отсутствовала единая цель для "поднятия" эффективности кода. Группа исследователей из Google и партнеров опубликовала новый репозиторий, который меняет этот статус-кво.
Архитектура тестовых нагрузок
JAXBench (Java/JSON Architecture Benchmark) — это не просто набор тестов, а комплексная среда, включающая 50 рабочих нагрузок, написанных на библиотеке JAX. Выбор этих задач был продиктован необходимостью репрезентативности: они покрывают архитектуру моделей Llama-3.1, DeepSeek-V3, Mixtral, Mamba-2 и AlphaFold2. Дополнительно в набор включено 33 оператора из библиотеки KernelBench, которые прошли валидацию на корректность. Эти тесты настроены так, чтобы максимально загружать MXU (матричные вычислительные блоки) процессоров TPU v6e.
Для создания верхней границы производительности (expert upper-bound baseline), команда использовала библиотеку Tokamax. Восьмь из 17 основных операторов в JAXBench сопровождаются вручную оптимизированными ядрами Pallas с настроенным размером блоков. Это позволяет понять, насколько далеко ИИ-системы могут приблизиться к усилиям человеческих инженеров.
Искусственный интеллект против ручной настройки
Центральным вопросом исследования стало: может ли ИИ написать код так же эффективно, как человек? Исследователи оценивали четыре метода, использующие обратную связь для генерации кандидатов на ядра Pallas. Ключевым открытием стало то, что для специфических языков программирования (DSL) с низкой документацией, таких как Pallas, контекст конкретной задачи важнее масштаба модели ИИ.
Включение в контекст ИИ-модели (в данном случае Gemini 3 Flash) отсортированной документации TPU кардинально улучшило результат: корректность генерации кода выросла с 5,8% до 37,3%. При этом скорость выполнения (geomean speedup) увеличилась в 1,28 раза.
Особый интерес представляет метод поиска структур. Когда ИИ достигал необходимого уровня корректности, структура поиска (beam-search pipeline в Autocomp) давала дальнейший рост, достигая ускорения в 1,36 раза по сравнению с стандартным XLA. На оптимизированных вручную ядрах Autocomp показал ускорение в 1,60 раза, однако отстаивал на сложных задачах с paginated и ragged attention. Исследователи подчеркивают, что качественная оптимизация ядер TPU остается сложной задачей, но JAXBench предоставляет открытый инструмент для поддержки сообщества.
Технический глоссарий для читателей
Читателю может быть интересно, что такое ядро (kernel) в данном контексте. Это базовый блок кода, который выполняется параллельно на графическом процессоре. В мире TPU этот код пишется на DSL (Domain-Specific Language) под названием Pallas. Оптимизация ядра — это процесс перестройки этого кода, чтобы он занимал минимум ресурсов и выполнял максимум вычислений за единицу времени. Стандарт XLA (Accelerated Linear Algebra) часто выступает базовой линией сравнения.
Использование JAXBench важно для экосистемы ИИ, так как позволяет создавать более быстрые и эффективные модели без необходимости постоянно привлекать дорогостоящих инженеров для ручной наводки кода.
Исследовательский проект открывает новые горизонты в области аппаратного обеспечения для ИИ, позволяя обществу совместно работать над улучшением производительности Google Cloud TPUs.