JAXBench: Benchmarking Autonomous TPU Kernel Optimization

论文解决了什么问题

严格的基准测试通过提供共同的优化目标，已显著推动了 GPU 内核自动性能优化的进展。然而 TPU 作为另一类主流 AI 加速器，至今缺乏等价的评测框架来衡量 AI 生成内核的质量。这导致研究者无法系统评估大模型在 TPU 内核编写与优化上的真实能力，也难以比较不同自动化方法的优劣。

核心方法

JAXBench 构建了一套 TPU 原生的基准测试套件，包含多种代表性算子和优化场景。它要求 AI 系统直接生成可在 TPU 上编译执行的内核代码，并从功能正确性和运行性能两个维度进行自动化评测。整个流程提供标准化的输入输出规范和可复现的执行环境，使不同模型和策略可以在同一框架下公平比较。

实验结论

• 当前大模型在 TPU 内核生成上的正确率和性能均显著低于 GPU 同类基准
• TPU 特有的向量寄存器与内存层级对模型理解硬件抽象提出更高要求

局限性与证据边界

适合谁阅读