让AI自己写TPU芯片的代码，谷歌这次玩明白了

你有没有想过，写代码这件事，最难的部分往往不是写"逻辑"，而是写"给硬件听得懂的逻辑"？

普通程序员写Python、写Java，脑子里想的是业务逻辑，编译器帮你把这些逻辑翻译成机器能懂的东西。但如果你要给TPU这种专门为AI计算设计的芯片写"内核"（也就是最底层、最贴近硬件的计算代码），事情就完全不一样了。你得亲自去管理内存怎么在不同的存储层级之间搬运，得手动安排数据传输的流水线,得算清楚一块一块的数据要怎么切分才能让芯片的每一个计算单元都不闲着。这活儿，全世界能干好的人不多。

谷歌这次拿出的答案是让AI自己干这个活。他们做了一个叫MaxKernel的系统,用大语言模型配合真实的编译器反馈,自动生成、调试、优化TPU上的高性能内核代码。听起来简单,但这背后藏着不少值得琢磨的设计取舍。

**为什么写内核这么难，难到需要专门造一套AI系统**

先说清楚问题有多棘手。

TPU*：谷歌研发的张量处理器，一种专门为深度学习计算优化的芯片，和GPU类似但架构不同

在TPU上写内核，程序员通常要用JAX配套的Pallas语言。这个过程里，你得手动区分HBM（芯片外的大容量但慢速内存）和VMEM（芯片内的小容量但极快的内存），得设计DMA（直接内存访问，一种硬件级别的数据搬运机制）怎么排队搬数据，还得琢磨多维度的"tiling"策略，也就是把一大块数据切成小块分别处理，切法不对，性能可能差好几倍。

这已经够难了，更麻烦的是，就算你写出来的代码能跑，也不代表它跑得快。要真正压榨出硬件的性能，往往需要反复做实测、分析瓶颈、调整参数，这是个需要大量经验积累的迭代过程。

那如果直接让AI来写呢？论文里给出了一个挺打脸的数据：用最朴素的方式，也就是让大语言模型直接根据参考代码生成100份候选答案，选出跑得最快且正确的一份，结果50个测试任务里只有10个能编译成功，10个是正确的，整体速度提升几乎为零，只有1.08倍。

这个结果说明了一件事：大语言模型光靠自己脑子里的知识瞎猜，是写不出能跑的硬件级代码的。它缺的不是"聪明"，而是缺一个能告诉它"你刚才这步错在哪"的反馈通道。这就像让一个从没进过厨房的人凑着菜谱做满汉全席,菜谱写得再详细,他也不知道火候不对的时候锅里发生了什么,因为他闻不到烧焦的味,看不到油温冒烟的样子。如果没有这些实时反馈,他只能一遍遍从头猜,而且猜错了都不知道错在哪。

**三种工作模式：给人类留多少控制权**

MaxKernel没有选择一条路走到底，而是设计了三种不同的运作模式，分别对应不同的信任程度和使用场景。

第一种叫Human-in-the-Loop，简称HITL

HITL*：人在回路中，指系统在关键节点会暂停，等待人类审核或干预后再继续执行

这个模式遵循"一个agent做完，然后等待"的原则。系统把整个内核开发流程拆成规划、实现、编译验证、测试生成、测试执行、自动调参、性能分析这几个阶段，每完成一个阶段就停下来，把结果扔给用户看，等用户点头或者修改之后才继续往下走。

这个设计的意义在哪？在于复杂内核的开发里，专家的直觉往往能在关键节点省下大量弯路。如果完全交给AI自己跑，它可能在某个方向上一头扎进死胡同，跑了半天才发现思路根本错了。而HITL模式相当于给AI装了一个"随时刹车"的按钮，人可以在每个关键路口看一眼路况再决定往哪拐。

第二种是Autonomous Loop，简称Auto agent，也就是全自动模式

Auto agent*：一种闭环优化流程，系统自动循环执行规划、生成代码、编译验证、测试、调参、性能分析这一整套流程，不需要人类介入

它的运作逻辑是这样的：先根据参考代码生成一套测试用例，并且把这套测试冻结起来，不允许后续的实现阶段偷偷修改测试标准来"作弊"过关。然后系统进入循环，每一轮都是先做优化规划，再翻译成代码，再验证能不能编译、数值是否正确，最后做自动调参和硬件层面的性能剖析。这一轮跑完拿到的性能数据，会直接喂给下一轮的规划阶段，让下一次的方案设计更有针对性。

这里有个挺聪明的补丁叫"Best-of-N状态回滚"。系统会记住每一轮跑出来的结果，包括代码、编译状态、测试结果、延迟数据，跑完所有轮次之后，自动把最终交付的版本回滚成历史上表现最好的那一版，而不是最后一轮跑出来的那一版。

这个设计背后藏着一个很实际的教训：优化过程不是单调向上的。你可能在第三轮找到了一个特别好的方案，但第四轮尝试了一个新想法结果反而更差。如果系统傻乎乎地只认"最后一轮"，那前面辛苦找到的最优解就白费了。这就跟你在网上反复修改一份简历一样，改到第五版突然手滑删错了一段重要经历，如果没有历史版本可以恢复，你辛苦攒的东西说没就没了。有备份和回滚机制，代价是要多存储这些历史记录,但换来的是不会因为一次失误把之前的好成果全部丢掉。

第三种是Graph-Based Autonomous Search，图搜索模式

这个模式是Auto agent的升级版，专门用来解决单条路径容易陷入局部最优的问题。

局部最优*：指某个方案在附近的小范围调整里已经是最好的了，但换一个完全不同的思路可能会有更好的方案，只是当前路径走不到那里

它把整个优化过程建模成一张搜索图，图上每个节点代表某一个具体的内核版本，包括它的代码、优化思路和实测性能。基于这张图，系统可以用不同的搜索策略去探索。目前实现了两种：并行搜索和束搜索（Beam Search）。

并行搜索的思路很直接：同时跑好几条完全独立的优化路线，每条路线给足够长的迭代预算，让它慢慢把复杂的bug调通、把内存布局打磨到位。因为各条路线互不干扰，所以每一条都可以走得比较深。

束搜索则相反，它更看重探索的宽度。每一轮只保留表现最好的前几个候选方案（比如beam width设为3），每个候选给较少的迭代次数，然后快速淘汰表现差的，把资源集中投给最有潜力的分支。

这两种策略的差异，其实很像两种不同的找工作策略。并行搜索像是同时深入接触三家公司，每家都认真聊上好几轮，充分了解利弊再做决定；束搜索像是先海投三十家公司拿到offer，快速筛掉明显不合适的，只对留下的少数几家深入接触。前者更稳，能挖得深，但覆盖面窄；后者覆盖面广，能快速找到"看起来还不错"的选项，但对需要长期磨合才能显现价值的机会容易误判提前放弃。论文里的实验也验证了这个直觉：束搜索因为每轮迭代预算只有2次（相比并行搜索的5次），在搜索到第三层深度时明显出现了性能平台期，这正是因为那些需要更多步骤才能调通的复杂优化方案，还没来得及被磨熟就被剪掉了。

**知识从哪里来：RAG检索增强**

单靠大语言模型自己的记忆是不够的，尤其是TPU、Pallas这类专业性极强又不断更新的领域知识。MaxKernel引入了RAG机制来解决这个问题。

RAG*：检索增强生成，指系统在生成代码或做决策之前，先从一个外部知识库里检索出相关资料，再把这些资料交给大语言模型参考

这里有个挺值得说一说的设计选择：MaxKernel的知识库里刻意排除了人类专家写的内核代码，只保留框架文档、内存布局指南、性能优化手册这些静态资料。

为什么要这么做？因为这篇研究想验证的核心问题之一是，AI能不能在不抄"标准答案"的情况下，靠自己摸索出高性能的优化思路。这就像考试的时候不让你看别人的满分答卷，只给你发教科书，看你能不能自己把知识用出来考出好成绩。如果直接把专家的内核代码喂给AI去参考模仿，那测出来的可能只是AI的"抄袭和微调"能力，而不是它真正理解硬件、独立推导优化方案的能力。

**实测结果：AI写的代码到底有多快**

说了这么多设计，最终还是要看数字。

论文在JaxBench上做了完整测试。

JaxBench*：一个专门用来评测TPU内核自动优化能力的基准测试集，包含50个多样化的任务，涵盖注意力机制、矩阵运算、损失函数等多种AI计算场景

结果显示，最朴素的Best-of-N方式几乎全军覆没，编译和正确率都只有10/50，速度提升近乎为零。而引入了迭代验证和优化闭环的Auto agent，编译率和正确率跳到了48-49/50，速度提升的中位数是1.39倍，但因为单次运行容易卡在局部最优,存在比较大的波动区间,从1.19倍到1.42倍都有可能。

再往上，MaxKernel Parallel（并行搜索）把编译率和正确率都做到了满分50/50，几何平均速度提升达到1.58倍，其中fast1指标（既保证正确又实现1倍以上速度提升的任务比例）达到34/50。Beam搜索则是1.49倍的速度提升，fast1是31/50。

这里最有意思的对比，是MaxKernel跟人类专家写的手工内核直接打擂台。

在8个有人类手工优化版本存在的生产级内核任务上，人类专家平均能做到2.02倍的速度提升，而MaxKernel的并行搜索版本做到了2.32倍，束搜索版本是1.78倍。这意味着在这8个任务里，AI生成的代码在7个任务上都超过了人类专家写的版本。

其中一个特别典型的例子是MLA Attention（一种注意力机制的变体）。人类专家在这个任务上反而写出了比不优化还慢的代码，速度只有0.69倍，也就是说人工调优的结果还不如原始的XLA编译器基线。但MaxKernel的两个搜索方法都找到了超过基线的方案，分别是1.21倍和1.23倍。

还有Paged Attention这个任务，束搜索找到的方案速度提升达到6.74倍，是人类专家版本2.41倍的将近三倍。

不过也不是所有战场AI都赢了。在Ragged Paged Attention这个任务上，人类专家做到了4.65倍的速度提升，而AI只做到了1.42倍，输得比较明显。这说明有些特别复杂或者特别依赖领域直觉的场景，人类专家的经验积累仍然有它不可替代的价值。这种此起彼伏的结果反而让人更相信这些数字是真实的，如果每一项AI都全面碾压，那多半是测试设计出了问题。

**在真实模型上的表现**

除了标准化的基准测试，研究团队还把MaxKernel用到了几个当下比较火的开源模型架构上，包括Multi-Head Latent Attention（MLA）、Qwen3-Next的Gated DeltaNet、DeepSeek-V4的Sparse Attention等。

在MLA v1架构上，跟人类写的Pallas基线相比，MaxKernel把延迟降低了8.68%，吞吐量提升了9.50%。

在Qwen3-Next的Gated DeltaNet上，跟原始JAX实现相比，前向传播延迟降低到了1.63倍的速度，整个训练步骤（包括前向和反向传播）的速度提升达到了4.70倍。

在DeepSeek-V4的Sparse Attention上，从小规模的解码场景到大规模的预填充操作，速度提升范围从2.36倍一直到7.85倍。

还有个挺实用的案例是关于代码健壮性的。MaxKernel被用来修复Ragged Page Attention v3在预填充模式下的崩溃和死锁问题，它自动在代码里加入了防护性的数值截断指令，来正确处理左侧填充的输入数据，避免出现负数的切片大小，从而保证内存搬运操作能顺利执行完毕。这几乎没有带来额外的性能开销,同时把原本会崩溃的场景救回来了。

这个细节值得单独说一说，因为它展示了这类系统的价值不只体现在"跑得更快"，也体现在"修得更稳"。在真实的生产环境里，一个每天崩溃几次的快内核，实际价值可能还不如一个稳定但稍慢的内核。

Q&A

Q1：MaxKernel是什么？

A：MaxKernel是谷歌开发的一套多智能体系统，利用大语言模型配合编译器实时反馈，自动为TPU芯片生成、调试和优化高性能内核代码，支持人机协作、全自动闭环、图搜索三种运作模式。

Q2：MaxKernel生成的代码性能能超过人类专家吗？

A：在测试的8个生产级内核任务里，MaxKernel在7个任务上超过了人类专家手工调优的版本，整体几何平均速度提升达到2.32倍，高于人类专家的2.02倍，但在个别复杂任务上仍不及人类经验。

Q3：MaxKernel和普通的AI代码生成有什么区别？

A：普通的零样本代码生成方式在TPU内核任务上几乎失败，编译和正确率仅20%左右；MaxKernel通过引入真实编译反馈、自动测试、性能剖析的闭环优化，把正确率提升到接近100%，并显著提高了运行速度。