Google Cloudの研究チームは、TPU（Tensor Processing Unit）向けの自動カーネル最適化を評価するための包括的なベンチマークスイート「JAXBench」を発表しました。これまでGPUカーネルの性能最適化分野では厳密なベンチマークが進歩を牽引してきたのに対し、TPU領域ではそうした共通の評価基準が存在していませんでした。JAXBenchはこの空白を埋めるため、50個のJAXワークロードで構成されており、これらは実用的でありながら最適化の余地を十分に備えています。 ベンチマークスイートには、公開されているMaxTextライブラリから抽出した17個の本番向けMLオペレータが含まれており、Llama-3.1、DeepSeek-V3、Mixtral、Mamba-2、AlphaFold2といった主要なアーキテクチャに対応しています。さらにKernelBenchから33個のオペレータを翻訳・検証し、TPU v6eのMXU（Matrix Multiplication Unit）利用率を最大化するよう新たな問題規模が設定されました。うち8つのオペレータには、公開されているTokamax library由来の手動最適化済みPallasカーネルが付属し、専門家による上限値のベースラインを確立しています。 評価ではGemini 3 Flashを用いて4つのフィードバック駆動型手法をテストしました。結果として、Pallasのような文書が不十分なDSL（ドメイン固有言語）では、モデルスケールよりもターゲット固有のコンテキストが重要であることが判明しました。キュレーション済みのTPUドキュメントを条件として追加すると、サンプルごとの正確性が5.8%から37.3%に向上し、48個中50個のベンチマークを解決できるようになります。その結果、AutocompのビームサーチパイプラインはXLA比で1.36倍の高速化を達成しました。手動チューニングされた8つのカーネルではAutocompが1.60倍の高速化に到達し、Tokamaxの上限値2.08倍の大部分を回復していますが、ページング化およびラグドアテンションオペレータでは若干の遅れが見られます。研究チームはJAXBench、評価ハーネス、基準結果をオープンソースコミュニティに公開し、さらなる貢献を促進する方針です。