MaxKernel: 面向 TPU 的智能体内核生成

TL;DR MaxKernel 是一个开源多智能体系统，通过 HITL、自主优化和基于图的全局搜索三种范式，自动生成 TPU 高性能 kernel，在 JaxBench 50 项任务上匹配专家手写性能。

问题

问题背景

随着大模型 scaling 与专用加速器普及，TPU 等硬件上的 kernel 性能成为训练/推理吞吐上限。顶尖团队常需手工编写 Pallas/JAX kernel 来压榨算力与带宽。

现有方法局限

• 手工编写：依赖多年硬件经验，迭代慢、难规模化，且不可移植。
• 模板/编译器自动调优：如 XLA 融合、AutoTVM 类方法，通常受限于预设 schedule 空间或局部搜索，无法跳出编译器保守决策。
• 单 Agent LLM 代码生成：能写可编译 kernel，但缺乏闭环硬件 profiling 反馈；生成结果多是“能用”而非“接近 hand-tuned 性能”，遇到逻辑/数值错误时自修复不稳定。

为什么难且重要

• 设计空间维度高：tiling、vectorization、DMA、pipelining、内存布局互相耦合。
• 评估成本高：TPU 编译+运行时间长，错误反馈稀疏且延迟，agent 必须学会从 trace/metric 中反推优化方向。
• 业界关注度：开源模型需要极致 kernel 来降低单位 token 成本，kernel 质量直接决定集群 ROI。

行业类比

类似 AutoML/NAS 从手工设计网络转向自动化架构搜索，kernel 生成也在从手工调优走向以 agent + 真实硬件反馈 驱动的自动优化循环。

核心洞察

• MaxKernel 的核心设计是把 Human-in-the-Loop、Autonomous 和 Graph-Based Search 三种开发范式统一到同一个多智能体框架中，共享相同的规划、编码、调试、测试和 profiling 子智能体。以往内核生成工作要么全自动生成，要么依赖人工逐步指导，缺乏灵活切换机制；MaxKernel 允许工程师根据任务难度和信任度选择协作模式，并在不同阶段迁移，这种渐进自动化思路更贴近真实工程迭代流程。
• Auto 代理的 metric/trace 驱动闭环容易陷入局部最优，MaxKernel 进一步提出 Graph-Based Autonomous Search，将内核变体作为图中的节点，通过扩展和剪枝实现全局设计空间探索。相比简单的 beam search 或随机重启，该方法在 JaxBench 上验证了并行搜索与 beam search 的权衡，能够在保持优化效率的同时覆盖更多优化方向。对工程启示：对于编译调优类任务，显式构建搜索图并利用 LLM 作为节点生成器，可能比逐步局部优化更稳健。
• MaxKernel 显著区别于以往 LLM 生成内核的工作在于，它将 TPU 的实时编译器反馈和硬件 profiling 指标作为第一级优化信号，而不只是依赖代码静态特征或合成基准。这种反馈闭环使智能体能够在 Pallas/JAX 的特定约束下（如内存布局、tile 配置）进行自我调试，从而匹配专家手调性能。这表明，在 AI for systems 场景中，将编译器或运行时作为可交互环境，让 LLM 在真实反馈中迭代，比一次性生成代码更能应对底层硬件的复杂约束。

方法

MaxKernel 以高层 kernel 任务描述为输入（如运算类型、张量形状、性能目标），通过多智能体协作输出经过硬件剖析验证的高性能 TPU kernel 实现。核心模块包括共享子代理池、Knowledge Store 与三种工作范式。

• 共享子代理池：每个子代理专注一个环节，包括 planning（任务分解与方案设计）、implementation（代码生成）、self-debugging（基于错误反馈修复）、testing（正确性验证）、hardware profiling（采集性能指标与 trace）。
• Knowledge Store：积累历史 kernel 实现、优化策略与硬件特性，供各子代理检索复用。
• 三种范式：
  1. Human-in-the-Loop (HITL)：开发者分步确认设计决策，系统在每步生成候选并收集反馈，适合复杂或需专业判断的任务。
  2. Autonomous (Auto)：全自动循环，根据 metric/trace 反馈持续修改代码、重新测试，直至达到性能阈值或迭代上限。
  3. Graph-Based Autonomous Search：将 Auto 代理扩展为图搜索，节点表示 kernel 版本，边表示修改操作；通过并行搜索与 beam search 在全局设计空间中探索，避免局部最优。

三种范式共享子代理池，但编排方式不同：HITL 引入人类决策节点，Auto 仅依靠自动反馈，图搜索则在 Auto 基础上增加全局路径管理。输出为经过 profiling 验证的 kernel 代码及对应性能报告。

差异点：相比单一 LLM 直接生成或仅依赖编译反馈的方法，MaxKernel 通过多代理分工、显式知识存储和图搜索策略，实现了更系统的设计空间探索与更高自动化程度，同时保留 HITL 接口以兼顾人工经验。

实验

实验设计

论文在 JaxBench 上评估 MaxKernel，该基准包含 50 个多样化的 TPU 内核任务，覆盖从常见算子到复杂工作负载。实验同时纳入来自 SOTA 开源模型 的真实内核场景。系统采用三种代理范式：Human-in-the-Loop (HITL)、Autonomous (Auto) 以及 Graph-Based Autonomous Search。所有范式共享一组专用子代理，负责规划、实现、自调试、测试与硬件剖析。评估重点关注端到端内核性能、并行搜索与 beam search 的对比，以及对开源模型内核的改进。

关键发现

• MaxKernel 持续生成高度优化的内核实现，能够匹配专家手工调优的基线，并在基准测试上带来显著性能提升。
• Graph-Based Autonomous Search 通过全局探索设计空间，扩展了 Auto 代理的能力，避免局部最优。
• 在真实开源模型的工作负载中，MaxKernel 同样展现出优化潜力，生成的内核可直接用于生产环境。

与基线对比解读

传统内核开发依赖工程师深入的硬件知识，而 MaxKernel 通过多代理协同与实时编译器反馈，实现了与专家手写内核相当的性能。自动化范式的优势在于无需人工干预即可迭代优化，且可并行探索多个设计路径，显著降低开发成本。与标准编译器流水线相比，MaxKernel 能够针对 TPU 特定特性进行细粒度调优，这是通用编译器难以覆盖的。

行业影响

落地场景

MaxKernel 可直接嵌入云端大模型推理/训练服务、AI 编译器基础设施及硬件厂商的算子库建设。在电商推荐系统中，特征交叉算子（如 embedding 查找、attention 变体）普遍依赖手写内核，MaxKernel 可自动化生成并调优；在自动驾驶感知模型部署中，点云卷积、voxel 操作等非标准算子同样受益。

商业价值

传统 TPU 内核开发需专家数天到数周，MaxKernel 的自主优化循环压缩到小时级，显著降低研发人力成本。生成的实现匹配专家手写基线，意味着在不损失性能的前提下缩短模型上线周期。对按 TPU 小时付费的云用户，内核效率提升 10% 即可带来明显成本节约；对硬件厂商，加速算子库完善有助于生态构建。

跟现有产品/工作流的接口

MaxKernel 以 Python 包形式开放，可嵌入基于 JAX/Pallas 的编译流程：开发者定义高层计算图后，由 MaxKernel 生成优化内核并回调 XLA 编译器验证。HITL 模式适合逐步调试新算子，Auto 模式可集成 CI 做批量优化，Graph-Based 搜索用于离线探索大设计空间。内置知识存储能沉淀团队优化经验，与现有代码仓库及 JAX profiler 等剖析工具协同工作。

局限

• - **评估硬件范围窄**：MaxKernel 的实验与优化流程完全围绕 TPU 与 JAX/Pallas 构建，未涉及 GPU 上的 CUDA/Triton 或其他 AI 加速器。虽然论文标题已明确限定 TPU，但这也意味着其在不同硬件后端上的泛化能力未经验证。多智能体架构本身可能与硬件无关，但 `profiling` 子代理、编译器反馈解析等模块深度依赖 TPU 工具链与 XLA 编译特性；若要迁移到 GPU 或其他加速器，需要重新实现大量适配层，预期工作量不小。
• - **缺少成本与效率分析**：论文没有报告生成一个 kernel 所需的 LLM token 数、端到端时延或计算成本。特别是 **Graph-Based Autonomous Search** 会并行探索大量设计候选，结合实时编译与硬件 profiling，可能产生高昂的 LLM 调用和 TPU 编译开销。对于实际工程落地，成本与耗时往往与性能同等重要；缺乏与手工调优或简单模板方法的系统成本对比，使得读者难以判断该方法的综合收益。
• - **对 LLM 能力鲁棒性依赖过强**：MaxKernel 的表现依赖于底层 LLM 的代码生成与推理能力，但论文未在不同 LLM 版本或基座模型上做消融，也未分析当 LLM 输出错误或幻觉时各个子代理的自修复效果。真实场景中硬件特性复杂多样，LLM 可能缺乏相应知识；如果缺少失败案例的定性归因与容错策略评估，系统的鲁棒性和可维护性仍然存疑。