MaxKernel：让多智能体协作编写和优化 TPU 内核

高性能加速器内核一直是 AI 工程中的“硬骨头”。它不仅要求开发者理解算子逻辑，还要熟悉内存访问、并行布局、编译器约束以及具体芯片的执行特征。MaxKernel 的思路是把大语言模型放进一个能够持续获得编译器和硬件反馈的开发闭环，让模型不只是生成一段代码，而是反复提出方案、编译、测试、分析并改进。

三种开发模式

MaxKernel 采用多智能体架构，并围绕同一组专业子智能体组织工作。不同子智能体分别承担任务规划、代码实现、自我调试、正确性测试和硬件性能分析等职责。系统提供三种使用方式：

• 人在回路（HITL）：由人类与智能体逐步协作。开发者可以参与设计过程，在关键决策处提供方向，适合需要控制方案和验证思路的场景。
• 自主优化（Auto）：智能体独立执行迭代循环，根据指标和硬件 trace 反馈调整实现，减少人工试错。
• 基于图的自主搜索：把不同实现和优化路径组织成搜索图，在更大的设计空间中探索，试图避免过早收敛到局部较优方案。

这种划分体现了内核生成的两个现实要求：一方面，工程师有时需要保留对实现的控制；另一方面，性能调优往往包含大量重复实验，适合交给自动化系统完成。图搜索模式则进一步将单一路径的优化扩展为多路径探索。

从“写代码”转向“做实验”

传统代码生成系统通常以一次性产出为主，而 TPU 内核优化很难仅凭静态知识完成。相同的计算逻辑，可能因为数据布局、分块方式或内存访问策略不同而产生明显的执行差异。MaxKernel 将实时编译反馈、测试结果和硬件 profiling 纳入流程，使智能体可以根据实际证据判断某个方案是否正确、是否值得继续优化。

论文在 JaxBench 上评估系统，该套件包含 50 个面向 TPU 的多样化内核任务；同时还测试了来自先进开源模型的复杂真实工作负载。根据论文描述，MaxKernel 生成的实现能够稳定达到较高优化水平，并在多类任务中接近专家手工调优基线。素材没有给出具体加速倍数，因此更适合将其理解为一种系统层面的验证，而不是单项性能纪录。

意义与限制

MaxKernel 的价值不只在于让模型生成 TPU 代码，更在于展示了“智能体负责实验组织、编译器负责即时反馈、硬件负责最终裁判”的协作范式。它可能降低专用加速器内核开发的门槛，并为面向 GPU、TPU 或其他 AI 芯片的自动调优系统提供参考。

不过，内核生成仍然受制于编译器能力、硬件可观测性、搜索成本和验证覆盖范围。自动系统能否在新的芯片架构、未见过的算子组合和更复杂的生产环境中稳定工作，还需要更多公开实验。MaxKernel 已开源，研究者和工程团队可以进一步检查其智能体分工、搜索策略与评测流程。

来源：Hugging Face Daily Papers