尧图网站设计 尧图网站设计YAOTU DESIGN
ARTICLE DETAIL

资讯详情

深耕网站设计与一线实操的经验洞察。

CAKE:让Agent与编译器协同进化,自动生成超越专家的GPU Kernel

CAKE:让Agent与编译器协同进化,自动生成超越专家的GPU Kernel 做 GPU Kernel 的同行应该都有一种体验你辛辛苦苦写了十几版 CUDA 代码能碰到一版跟 CUTLASS 默认配置差不多快的已经算运气很好。更打击人的是很多你觉得“写得更聪明”的版本实际测出来反而更慢慢的原因往往不是算法思路而是寄存器分配、共享内存 bank conflict、编译器某个保守优化决策。最近我读了一篇名为 CAKE 的论文标题直白得让人没法忽视让 Agent 与编译器共同进化写出超越专家的 GPU Kernel。我先把关键词翻译成大家熟悉的场景Agent 指的是 LLM 驱动的代码生成智能体编译器就是我们每天都在用的 nvcc 这类工具GPU Kernel 就是你在 CUDA/HIP 里写的那一小段跑在显卡上的函数。CAKE 想解决的问题非常具体我们能不能让 Agent 不仅仅是“生成一段能编译的 Kernel”而是像专家那样在一个能感知硬件行为的闭环里持续优化它最终达到甚至超过专家手写 Kernel 的性能。如果你正在做 AI Agent 开发、编译器优化或者头疼如何让大模型写出高性能 GPU 代码这篇文章值得你停下来细看。我会把这套设计的核心机制拆开讲再补充一些我在类似思路落地时踩过的坑。关于论文里没有细说的部分我会基于常见实践做一些补充尽量让大家看完能直接借鉴到自己的项目里。1. 为什么 GPU Kernel 是 LLM 编程最难攻克的堡垒1.1 “能跑”和“跑得快”之间隔着一座山先说一个基本事实写 GPU Kernel 的难点从来不在“语法正确”而在“性能可预测”。CPU 上的代码你大致能根据算法复杂度估算运行时间GPU 不一样同样的 GEMM 算法线程块划分不同、共享内存 tile 大小不同、循环展开因子不同性能可能差出 5 到 10 倍。我把这里面最折磨人的几个维度列一下线程组织线程块的维度怎么设决定了 SM 的占用率也决定了访存合并的粒度共享内存tile 数据要搬到 shared memory 里做复用但 bank conflict 一出现带宽直接打折寄存器压力每个线程能用多少寄存器是有限的一旦溢出局部变量会被丢到全局内存慢上一个数量级指令级优化是否生成 FMA 指令、是否向量化 load/store、是否避免了多余的地址计算全靠编译器心情。一个专家 Kernel 并不是“写出来”的而是“试出来”的。我记得之前调一个 Flash Attention 风格的 Kernel光是调整 double buffering 的流水线节奏就试了二十多组配置。这就是 CAKE 要面对的真实战场性能不是静态属性而是代码、编译器、硬件三者交互后的结果。1.2 现有 LLM 生成 Kernel 的流程缺了什么现在大多数 LLM 生成代码的流程是这样的给模型一个函数签名和功能描述模型生成代码拿编译器编译如果有语法错误就再问一遍模型修复。这个循环停留在“能不能编译”的层面对性能几乎没有反馈。问题在哪儿我举个例子。你让模型生成一个 SGEMM Kernel模型给你一版看起来很标准的实现编译通过结果跑出来只有 cuBLAS 的 20% 性能。这时候你问模型“为什么这么慢”它根本答不上来因为它看不到自己的 Kernel 在 GPU 上到底发生了什么寄存器溢出了几十个、共享内存访问有严重的 bank conflict、SM 占用率只有 30%。这不是模型笨而是流程里压根没有性能反馈回路。编译器在传统流程里只扮演一个安检员的角色检查通过就放行检查不通过就打回重写。可对于高性能代码来说安检员说“通过”只是起点真正的挑战是如何在几十项连续的性能指标上逼近硬件极限。1.3 CAKE 的核心命题从“生成代码”转向“优化代码”CAKE 的第一步就是把“编译加运行”变成一个 Agent 可以反复交互的环境。编译器不再是一次性工具而是一个持续输出信号的传感器。Agent 生成一版 Kernel编译器编译、运行、打分把结果反馈给 AgentAgent 根据反馈调整策略再生成下一版。这个转变意味着两件事发生了变化。第一目标变了不再追求“生成一版完美代码”而是追求“在有限预算内搜索到性能足够好的代码”前者是创作后者是优化。第二反馈变多了编译器给的不是一个简单的对错标签而是一堆可操作的中间信号比如寄存器溢出次数、bank conflict 警告、实际运行时间、occupancy 数据。我最初读这篇论文时的感受是它没有发明什么全新的算法而是把强化学习里“环境”这个概念正确地安放到了编译器身上。就这一个视角转换整个问题的性质都变了。2. CAKE 的核心转变编译器从工具变成环境2.1 “工具视角”和“环境视角”差在哪里很多人会把“编译器作为工具”和“编译器作为环境”混为一谈但这两者差别非常大。我用一个表说清楚对比维度工具视角传统流程环境视角CAKE交互方式Agent 生成代码编译器返回通过/不通过Agent 生成代码环境返回多维度观测值反馈粒度稀疏的二进制信号密集的连续指标状态记忆无状态每次独立有状态上一轮的结果影响下一轮决策可控对象只有代码文本代码文本、编译选项、优化策略都可调目标通过编译在预算内逼近性能上限这个视角转换特别像下棋。传统方法是让 Agent 凭空想一步“最好的棋”然后裁判告诉它这步棋合不合法完了。CAKE 是让 Agent 真正坐到棋盘前每走一步都能看到棋盘的变化根据变化调整下一步直到棋局结束。你说哪种方式更容易下出好棋2.2 环境反馈了哪几类信号CAKE 里的编译器环境会返回哪些信号根据我看到的论文描述和这类系统的常见设计大致有四层编译期反馈编译时间、告警数量、IR 优化前后的指令数变化、是否发生寄存器溢出静态分析反馈共享内存占用、寄存器使用量、线程块大小、局部变量是否被 spill运行期反馈Kernel 执行时间、SM occupancy、访存吞吐、计算吞吐甚至包括功耗数据结构反馈生成的汇编里是否出现了目标指令组合比如 FMA 链、向量化 load/store、没有分支发散。这些信号单独看都有噪声但合在一起就能构成一个“性能画像”。Agent 不需要盯着某一个数字看而是通过一组特征判断当前代码离硬件极限还有多远、瓶颈在哪个方向。2.3 密集反馈为什么比 RLHF 稀疏奖励更实用现在很多人训练代码 Agent 还在用 RLHF 的思路生成一批代码人工或者规则打分然后做偏好优化。这套做法的问题在于奖励信号太稀疏。你生成一百个 Kernel可能只有五六个能跑出“还不错”的性能剩下的都是噪声模型很难从这么稀疏的反馈里学到东西。CAKE 的思路是绕开“人为定义好坏”直接让编译器环境提供密集的引导信号。每一个 Kernel 无论好坏都会产出一堆中间指标这些指标本身就携带了“往哪个方向改”的信息。比如说某个版本寄存器溢出严重下一轮 Agent 自然会去减少每个线程的私有数组大小或者调整线程块尺寸。这是传统的“对错标签”给不了的指导。3. 双环路进化到底在进化什么3.1 内环路Agent 的 Kernel 生成策略更新CAKE 框架里最核心的机制是双环路进化。第一个环路围绕着 Agent 本身转Agent 生成 Kernel 代码编译器环境返回性能反馈Agent 根据反馈优化下一版代码。这个内环路里的 Agent 并不只是机械地改参数。它要做几件事诊断瓶颈分析反馈信号判断当前瓶颈是访存、计算、还是并行度不足选择操作决定下一步是调整线程块大小、改 tile 尺寸、增加展开因子还是换成不同的数据布局保留有效信息上一轮已经表现好的部分不要破坏只修改有问题的部分避免来回震荡。我在自己的场景里跑过类似的多轮优化循环一个深刻体会是Agent 能不能记住“什么不该动”和“什么该动”同样重要。如果每轮都大改搜索过程就退化成了随机采样。3.2 外环路编译策略的搜索与调整第二个环路是我觉得 CAKE 最有意思的部分编译器本身也可以被改变。同一个 Kernel 源码在-O2和-O3 -use_fast_math下性能可能差出两三倍换一组编译选项甚至可能让某个版本的 Kernel 从不可用变成可用。CAKE 把编译器的参数空间也纳入了搜索范围。Agent 或者一个独立的优化模块会根据 Kernel 的表现去调整编译选项、pass 顺序、优化开关。这意味着“写 Kernel”和“配编译选项”不再是被割裂的两件事而是同一个搜索问题里的两个维度。为什么这个外环路很关键因为手写高性能库的专家从来不会只用默认编译配置。他们会在多个编译选项组合里做 A/B 测试找出最适合某个 Kernel 的配置。而 LLM 生成代码的传统做法几乎完全忽略了这一层等于把一大块性能空间白白放弃了。3.3 两个环路如何协同而不振荡双环路同时进化的一个天然风险是振荡。Agent 可能为了迎合某个编译器配置而生出结构畸形的代码编译器配置又可能反过来适应这种畸形代码。两个优化者互相追逐最后收敛到一个谁都不满意的怪圈。CAKE 处理这个问题的方式在我理解里很像交替优化alternating optimization每轮迭代只修改其中一个环路保持另一个环路固定。要么这轮只让 Agent 根据固定编译配置调整代码要么只调整编译策略而固定代码版本。这样就把双变量联合优化拆成了两个单变量优化稳定性会好很多。我在做类似的搜索系统时也验证过这个思路同步更新两个参数几乎必然震荡交替更新则收敛得又快又稳。如果你也想复现这套设计我建议从交替优化开始不要一开始就上联合更新。4. 超越专家 Kernel 的关键模板空间、奖励与预算4.1 搜索空间不是自由文本而是结构化模板GAKE 能高效搜索的一个关键前提是它没有让 Agent 在无限大的文本空间里乱逛而是把 Kernel 结构模板化。比如写 GEMM Kernel 时可变的维度被限定在几个明确的旋钮上可调参数典型取值范围影响方面BLOCK_SIZE16 / 32 / 64 / 128SM 占用率、访存并发度TILE_M / TILE_N64 / 128 / 256数据复用率、寄存器压力TILE_K8 / 16 / 32流水线深度、共享内存用量循环展开次数1 / 2 / 4 / 8指令级并行、指令缓存命中向量化宽度1 / 2 / 4访存带宽利用、bank conflict 概率double buffering开 / 关访存计算重叠程度Agent 的任务从这个巨大的文本生成问题变成了一个参数推荐和结构选择问题。这种设计有两个好处一是生成的代码结构完整不太会出现语法残缺二是搜索空间虽然也不小但每一维都有明确的物理意义Agent 可以根据反馈定位到具体是哪个旋钮出了问题。4.2 性能奖励和评测预算怎么平衡性能评估是整个框架里最贵的操作。一次评估要编译、加载、跑多个 batch、取中位数在消费级显卡上怎么也要几十秒。如果搜索空间有几千个组合全部跑完是不现实的。所以 CAKE 这类框架一般会做两件事。第一用相对性能而不是绝对时间做奖励把当前版本和 baseline往往是一个专家库的 Kernel对比得到加速比再映射成奖励值。这样能抵消不同硬件间的性能差异。第二严格控制预算设定最大评估次数或者设定一个时间上限。在预算内跑完搜索就停止哪怕还有潜力没挖完。这个“预算意识”特别重要。我自己一开始就没注意放任搜索跑了一整夜第二天起来看结果发现优化的收益还不如我在 baseline 上手动调两个参数来得快。没有预算约束的搜索很容易变成没有意义的烧机器行为。4.3 结果怎么看超越专家是“特定条件下的超越”论文标题说的是“超越专家的 GPU Kernel”但我不建议大家把这个理解成全面碾压。更准确的说法是在特定算子的特定 shape 下在有限的搜索预算内CAKE 能找到一个 Kernel性能接近甚至略微超过专家手写库的默认实现。为什么是“特定条件”因为专家手写库比如 CUTLASS、cuBLAS通常面向通用 shape 做优化而实际业务里的 shape 往往很“歪”比如某个矩阵维度是 1027、某个特征大小是 333这些 case 恰好是通用库优化不够好的盲区。自动搜索的优势恰恰在于它可以针对特定 shape 和特定硬件定制一个方案。我在实际工作中也遇到过这种情况。某个内部算子的 shape 很特殊cuBLAS 里没有对应的高性能路径我手写了个 Kernel 也只到库性能的 60%。如果把搜索跑起来多试几组 tile 配置反而能找到性能翻倍的方案。CAKE 的价值在于把这个过程自动化了。5. 我在落地类似方案时踩过的三个坑5.1 全长评估太慢搜索预算被低估我最早尝试类似思路时犯的第一个错误就是低估了性能评估成本。你以为一次评估是“跑一下就看结果”实际上每次评估包括编译、加载、profiling、取多次运行中位数一套流程下来小一分钟很正常。搜索 500 次就是 8 小时等结果出来已经是第二天了。后来我的做法是分阶段控制预算先用小规模输入快速筛掉明显差的配置只对排名靠前的候选做完整评估。这种两段式评估能省掉至少一半时间。再一个技巧是在基准测试里减少重复次数先跑一轮拿粗粒度估计最后对前三名多跑几轮精测而不是每次评估都跑满十次。5.2 编译器的 warning 到底算信号还是噪声这是我踩过最深的一个坑。一开始我贪多把所有编译输出都塞给 Agent结果模型被大量无关 warning 淹没性能反而没有提升。比如说 nvcc 对某些未初始化变量会报 warning但这跟你 Kernel 的性能几乎没有关系而寄存器溢出 warning 可能很关键。后来我学到的处理方式是“信号分级”。把反馈分成三档强相关信号寄存器溢出次数、bank conflict 计数、SM occupancy、实际运行时间中相关信号是否生成向量化指令、共享内存使用率、block 调度模式弱相关信号编译时间、非关键 warning、IR 指令数变化。只把强相关和中相关信号放进 Agent 的上下文里弱相关信号直接丢掉。这一步优化之后Agent 的决策质量明显提升搜索过程的收敛速度也快了很多。5.3 硬件特异性和过拟合问题在 RTX 4090 上搜出来的一组最优配置拿到 A100 上可能性能掉 30% 甚至更多。这不是说结论错了而是一个很现实的问题Kernel 优化高度依赖具体硬件的 SM 数量、共享内存大小、寄存器文件和访存带宽。所以如果你要做跨硬件部署一定不能只在单一显卡上跑搜索。我的建议是至少在两到三块有代表性的 GPU 上分别做搜索然后比较配置的差异。如果某个配置在这几块卡上都很强那它大概率是一个鲁棒的方案如果每个卡的强配置差异巨大那就老老实实按硬件分别保存配置部署时按 GPU 型号加载对应的 Kernel 版本。6. 从 CAKE 里能直接抄走的三个工程思路6.1 给 Agent 加一个“可反馈执行环境”即使你完全不做 GPU KernelCAKE 最重要的工程启示也值得借鉴代码生成 Agent 应该跑在一个能给出密集反馈的执行环境里而不是生成完就结束。我见过不少团队做 SQL Agent模型生成 SQL、验证能不能跑、能跑就交付。本质上跟“能编译就通过”是一样的问题SQL 能跑和跑得快之间同样隔着一条巨大的性能鸿沟。如果你能把执行计划解读、耗时分析、索引命中情况这些信号喂回给模型让它基于反馈修正 SQL效果会比单纯生成强很多。6.2 把编译告警量化成信号而不是文本噪音很多做代码 Agent 的人把编译告警当成需要“消除”的异常。CAKE 给我的另一个启示是告警本身是信号关键是你怎么处理它。与其把几十条 warning 原样塞给模型不如先做一层量化统计每类警告出现的次数、严重级别筛选出与性能相关的警告类别如寄存器溢出、bank conflict、未合并访存把筛选后的结果格式化成简洁的结构化文本再放进模型上下文。我试过之后模型对问题的定位准确率提升很明显。原因很简单信息不是越多越好可操作的信息才是。6.3 生成器和工具链一起迭代CAKE 最“反常规”的一点是它不只优化生成结果还优化生成工具本身。这个思路可以迁移到很多场景。举个例子。你给 Agent 配了一套代码评审规则规则刚写出来的时候大概率不够完善。传统做法是拿着规则去评审 Agent 的输出然后不断调整 AgentCAKE 式的做法是在调整 Agent 的同时也根据 Agent 的表现反向修改评审规则。比如评审规则里有一条“禁止使用动态内存分配”但 Agent 在某个性能敏感场景里用动态分配反而更快这时候规则本身就应该被修订。生成器、评价器、环境三者共同进化而不是只优化其中一个环节。这是我读 CAKE 后最想推荐给大家的方法论。7. 遗留问题和我的个人判断7.1 正确性验证还没有真正解决自动搜索出来的 Kernel 跑得快是好事但正确性验证是个绕不开的问题。目前这类框架对正确性的检查大多停留在样例级别的对比验证跑一组输入跟 baseline 的 Kernel 对比输出误差在一定范围内就算通过。这在大多数场景够用但严格来说并不是形式化的正确性证明浮点误差累积到一定程度后某些极端输入可能出问题。如果你的业务对数值精度要求非常高比如科学计算、金融计算用自动搜索出来的 Kernel 之前一定要额外做全套的差分测试和误差边界分析不能只跑几个样例就放心。7.2 部署成本不低更适合离线生成CAKE 整个框架跑起来一边是 LLM Agent一边是编译器加 profiling 工具都是重量级组件。把这样一个系统部署到在线服务里、每来一个请求就现场搜索一遍 Kernel成本显然太高也不现实。更合理的落地方式是把 CAKE 做成一个离线工具链新模型发布、新算子出现的时候离线跑一次搜索把找到的高性能 Kernel 固化下来作为静态库的一部分随软件发布。搜索是在发布前做的运行时只加载结果这样成本可控收益实实在在。7.3 我对这个方向能走多远的看法如果把 CAKE 放在一个大趋势里看它其实是非常务实的一步它没有去挑战“大模型能不能独立思考”这种宏大命题而是老老实实地把编译器、硬件、LLM 三个组件拼成了一个有反馈的闭环。这个方向短期之内最有可能落地的地方就是各厂商的 Kernel 自动调优引擎。长期来看把编译器当成 Agent 可以干预的对象这个思路一定会扩展到更多场景。未来可能出现的不只是“自动写 Kernel 的 Agent”而是“自动改编译器优化策略的 Agent”、“自动配置分布式训练方案的 Agent”。它们的共同点都是生成器 可反馈环境 交替进化。读 CAKE 这篇论文之前我一直觉得代码生成领域的突破应该来自更大的模型、更聪明的采样策略。读过之后我调整了自己的看法模型的能力只是起点环境反馈的密度和质量才是决定优化天花板的关键。让编译器学会说话比让模型多背几天代码更有用——这句话算是这篇论文带给我的最大收获。最后分享一个小技巧如果你也想在自己项目里做类似的事情不要一上来就复刻全套框架。先做一个最小闭环——模型生成环境评估拿到密集指标再生成。哪怕没有双环路进化没有复杂的搜索策略单单是把这个闭环跑起来效果就已经能超过“生成即结束”的朴素流程。
返回列表