Google JAXBench 给 TPU 内核优化补上 50 项实战基准
omarsar0 · x · 2026-07-27
Google 推出 JAXBench：给 TPU 自动内核优化立基准
Google、Harvard 和 UC Berkeley 提出了 JAXBench，这是一个面向 TPU 自动内核优化的基准套件，目标是让 agent 生成的 TPU kernel 有一个可以公平对比的“爬坡靶子”。
论文的主要内容包括：
- 50 个真实 JAX 工作负载，来自 Llama-3.1、DeepSeek-V3、Mixtral、Mamba-2、AlphaFold2 等生产级架构。
- 17 个算子由 KernelBench 迁移并验证正确性。
- 另有 8 个算子直接提供 Tokamax 的手工调优 Pallas kernel 作为专家 baseline。
- 在 Gemini 3 Flash 下，补充整理过的 TPU 文档后，单样本正确率从 5.8% 提升到 37.3%。
- 模型最终在 50 个 benchmark 中解决了 48 个，达到 1.28x 的几何平均加速；再加 beam search 后提升到 1.36x。
作者的结论很直接：正确性更像是文档问题，速度更像是搜索问题。论文同时开放了 benchmark、评测 harness 和 baseline 结果，鼓励社区继续贡献。
「Infra」频道最新
- Rat Stack：把应用和云写成一个类型化程序，让 AI agent 可靠地搭建部署 — samgoodwin89 · 2026-09-23
- optimAIzr:本地 CLI 找出 AI 用量浪费,支持 Claude Code 与 Codex 省钱 — stichstichstich · 2026-09-23
- 256GB 内存版 Mac Studio 二手溢价 6000 美元仍一机难求 — GabGarrett · 2026-09-23
- Ben Bajarin 解析 AI 数据中心:CPU 与内存如何限制 GPU 推理服务能力 — BenBajarin · 2026-09-23
- Epoch AI 报告:「思考」的价格正在暴跌,智能成本持续陡降 — Proper_Actuary2907 · 2026-09-23
- 从燃气到吉瓦:Diligence Stack 专家访谈拆解 AI 数据中心供电难题 — BenBajarin · 2026-09-23