MaxKernel：面向TPU的Agent内核生成

为加速器设计和编写高性能定制核是一项需要深入硬件层面专业知识的复杂任务。可以将大语言模型（LLM）与实时编译器反馈相结合，构建用于核生成的智能体系统。本文提出MaxKernel，这是一个多智能体系统，实现了三种不同的TPU核开发范式：（1）人机协同（HITL）智能体，支持分步协作设计；（2）自主（Auto）智能体，执行完全自动化的、基于指标/轨迹驱动的优化循环；（3）基于图的自主搜索，扩展Auto智能体以实现设计空间的全局探索。这三种范式均利用共享的专用子智能体来处理规划、实现、自我调试、测试和硬件性能分析。我们在JaxBench基准上评估了MaxKernel，该基准包含50个多样化的TPU核任务，以及来自前沿开源模型的复杂真实世界工作负载。实验表明，MaxKernel能持续生成高度优化的实现，性能与专家手工调优的基线相当，并在基准中实现显著提升。该智能体已开源，相关代码可访问此 https URL。