开源|Google发布Accelerator Agents：用Gemini将PyTorch工作负载迁至JAX并优化TPU内核
talkingdev • 2026-09-08
3398 views
Google在GitHub开源了Accelerator Agents工具包，面向Google Cloud TPU开发场景，利用Gemini模型帮助开发者将PyTorch工作负载迁移到JAX，并优化自定义内核。该工具包包含MaxCode和MaxKernel两个核心模块：MaxCode负责模型转换，MaxKernel则用于编写、移植、剖析和调试Pallas内核。随着AI芯片竞争加剧，TPU的软件生态长期面临从GPU/PyTorch迁移成本高、自定义内核开发门槛高等问题。Accelerator Agents将大模型能力引入底层编译器与内核开发流程，有望自动处理框架差异、生成并调优内核代码，缩短适配周期。对寻求替代GPU方案或需要大规模训练推理的团队而言，这一开源项目具备较高的实践价值。
核心要点
- Google开源Accelerator Agents，利用Gemini降低TPU开发与迁移门槛
- MaxCode负责PyTorch模型向JAX转换，MaxKernel聚焦Pallas内核编写与调试
- 该工具面向Google Cloud TPU生态，助力AI工作负载迁移和自定义内核优化