JAXBench：自主 TPU 内核优化基准测试

arXiv:2607.20466v1 公告类型：新 摘要：通过建立爬山共享目标，严格的基准测试推动了自主 GPU 内核性能优化的进步，但 TPU 不存在类似的目标。我们推出了 JAXBench，这是一个 TPU 原生基准测试套件，用于在 Google Cloud TPU 上优化 AI 生成的内核。 JAXBench 包含 50 个 JAX 工作负载，这些工作负载既相关又提供优化空间。我们从公共 MaxText 库（例如 Llama-3.1、DeepSeek-V3、Mixtral、Mamba-2 和 AlphaFold2）中的架构中提取 17 个生产 ML 算子，并从 KernelBench 翻译 33 个算子，这些算子经过正确性验证，并设置了新的问题大小，以实现高 TPU v6e MXU 利用率。 17 个生产运营商中的 8 个配备了来自公共 Tokamax 库的手动优化的 Pallas 内核，并调整了块大小以建立专家上限基线。我们评估了四种反馈驱动的方法来为 JAXBench 生成候选 Pallas 内核。在 Gemini 3 Flash 的全套套件中，我们发现特定于目标的上下文比 Pallas 等文档稀疏的 DSL 上的模型规模更重要。以精选的 TPU 文档为条件，将每个样本的正确性从 5.8% 提高到 37.3%，并以 1.28 倍几何平均加速解决了 50 个基准测试中的 48 个。一旦实现正确性，搜索结构就会产生显着的收益，Autocomp 的波束搜索管道的几何平均加速比 XLA 快 1.36 倍。在 8 个手动调整的内核上，Autocomp 达到了 XLA 的 1.60 倍 geomean，恢复了 2.08 倍 Tokamax 上限的大部分，但落后于专门的分页和不规则注意力算子。高质量 TPU 内核优化仍然是一项具有挑战性的任务，我们发布了 JAXBench 基准、评估工具和基线结果来支持开源贡献。