Naukowcy zaprezentowali JAXBench, pierwszą całościową paczkę testów przeznaczoną do automatycznej optymalizacji kerneli na TPU Google Cloud. Benchmark zawiera 50 workloadów JAX, które są zarazem praktycznie znaczące i oferują miejsce na dalsze ulepszenia wydajności.
Zestaw obejmuje 17 kerneli produkcyjnych wyekstrahowanych z architektury MaxText library, zawierających operatory z Llamy-3.1, DeepSeek-V3, Mixtral, Mamba-2 i AlphaFold2. Dodatkowo 33 operatory pochodzą z KernelBench i zostały dostosowane do wysokiego wykorzystania MXU na TPU v6e. Osiem z siedemnastu produkcyjnych operatorów zawiera ręcznie zoptymalizowane kernele Pallas z biblioteki Tokamax, które służą jako wzorzec ekspertów. Wszystkie kernele zostały zweryfikowane pod kątem poprawności.
Badania czterech metod opartych na feedback'u pokazują, że dla słabo udokumentowanego DSL-a jak Pallas, dostęp do kontekstu specyficznego dla TPU jest ważniejszy niż sama wielkość modelu. Gdy model Gemini 3 Flash otrzymał dostęp do wyselekcjonowanej dokumentacji TPU, poprawność kodu na próbie wzrosła z 5,8% do 37,3%, a system rozwiązał 48 z 50 benchmarków z przyspieszeniem 1,28x w geometrycznej średniej. Metoda Autocomp z beam-search pipeline'em osiągnęła przyspieszenie 1,36x w stosunku do XLA. Wyniki sugerują, że struktura wyszukiwania daje znaczące zyski wydajności po osiągnięciu poprawności kodu.