Пробел в экосистеме TPU закрыт
Несмотря на доминирование GPU в сфере автоматической оптимизации нейросетевых ядер, для TPU отсутствовал аналогичный стандартный набор тестов. Команда исследователей из Google и Университета Корнелла представила JAXBench — специализированный бенчмарк для оценки AI-генерации оптимизированных ядер Pallas на архитектуре Google Cloud TPU v6e. Цель набора — создать общую точку отсчета для развития автономных систем оптимизации, аналогично тому, как это произошло с GPU.
Структура и метрики JAXBench
Набор данных включает 50 рабочих нагрузок JAX, разделенных на две категории: 17 операторов из реальных production-моделей (Llama-3.1, DeepSeek-V3, Mixtral, Mamba-2, AlphaFold2) и 33 оператора из KernelBench. Особое внимание уделено достижению высокой утилизации MXU (Matrix Multiplication Unit) на TPU v6e. Для 8 ключевых операторов созданы вручную оптимизированные ядра (Tokamax), служащие «золотым стандартом» (upper-bound baseline).
| Метрика / Метод | Результат (Geomean Speedup) | Примечание |
|---|---|---|
| Базовый XLA | 1.0x (референс) | Стандартная компиляция без ручной оптимизации |
| Autocomp (Beam Search) | 1.36x | Лучший результат на полном наборе JAXBench |
| Autocomp (на 8 hand-tuned) | 1.60x | Восстановление большей части потенциала Tokamax |
| Tokamax (Upper Bound) | 2.08x | Ручная оптимизация экспертами |
| AI-генерация (Gemini 3 Flash) | 1.28x | Без донастройки структуры поиска |
Ключевой инсайт: Контекст важнее масштаба модели
Эксперименты с различными LLM показали, что для работы с узкоспециализированным DSL Pallas критически важна не мощность модели, а качество контекста. Простое условие на курируемой документации TPU повысило долю корректно сгенерированных кодов с 5.8% до 37.3%. Это позволило решить 48 из 50 задач бенчмарка.
Дальнейшее улучшение достиглось за счет структуры поиска: пайплайн Autocomp с beam-search превзошел базовый XLA на 36% (1.36x). На задачах с ручными оптимизациями AI-подход достиг 1.60x, что составляет около 77% от потенциала ручной работы экспертов (2.08x). Отставание сохранилось только в специфичных операторах внимания (paged и ragged attention).
Открытый доступ для сообщества
Авторы опубликовали не только результаты, но и сам бенчмарк, инструменты оценки и базовые результаты, чтобы стимулировать развитие open-source решений для оптимизации TPU. Это шаг к демократизации доступа к высокопроизводительным вычислениям на специализированном железе Google.
Источник: arXiv cs.AI ↗