摘要
严谨的基准测试推动了自主GPU内核性能优化的发展，通过建立共同目标进行攀登，但TPU领域尚无等效工具。我们推出了JAXBench，这是一套针对Google Cloud TPU的原生基准测试套件，旨在进行AI生成的内核优化。JAXBench包含50个相关的JAX工作负载，并为优化提供了空间。
我们从公共MaxText库中提取了17个生产级机器学习操作符，涵盖Llama-3.1、DeepSeek-V3、Mixtral、Mamba-2和AlphaFold2等架构，并将KernelBench中的33个操作符进行翻译，确保正确性并设置新的问题规模，从而实现TPU v6e MXU的高利用率。8个生产操作符附带来自公共Tokamax库的手动优化Pallas内核，经过块大小调优，以建立专家级上限基线。
我们评估了四种基于反馈的方法，以生成JAXBench的候选Pallas内核。在使用Gemini 3 Flash的完整套件中，我们发现目标特定上下文的影响大于稀疏文档DSL（如Pallas）的模型规模。通过基于策划的TPU文档进行条件处理，单样本正确率从5.8%提升到37.3%，并在1.28倍几何平均加速下解决了50个基准中的48个。一旦实现正确性，搜索结构带来了显著的收益，Autocomp的束搜索管道在XLA上达到了1.36倍的几何平均加速。在8个手动调优的内核上，Autocomp达到了1.60倍的几何平均加速，恢复了大部分2.08倍的Tokamax上限，但在专业的分页和不规则注意力操作符上落后。
高质量的TPU内核优化仍然是一项具有挑战性的任务，我们发布了JAXBench基准测试、评估工具和基线结果，以支持开源贡献。
博主点评： JAXBench的发布填补了TPU内核优化领域的空白，通过提供丰富的基准测试和反馈机制，显著提升了内核性能的优化效率。这一举措将激励更多开发者参与到TPU优化的开源生态中，推动技术的进步和创新。