ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

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

CAKE:AI Agent与编译器协同进化,生成超越专家的GPU Kernel 最近在社群看到不少人在讨论AI 能不能直接写出 GPU Kernel这个话题。我自己的测试结论很一致大模型能轻松生成一个看起来正确的 CUDA 代码但离专家手写版往往还差一截尤其是当算子稍微复杂一点、并发调优要求高一点的时候。而这次看到的论文 CAKECompiler-Assisted Kernel Evolution恰好切中这个痛点它把 Agent 和编译器放在同一条进化轨道上让两者互相反馈、互相迭代最终目标是生成超越专家水平的 GPU Kernel。这篇文章就把我对 CAKE 的整体理解和思考拆开聊聊适合在做 GPU 编程、编译器优化、AI Agent 应用的朋友作为参考。1. 背景与问题GPU Kernel 为什么这么难写1.1 AI 编程浪潮下的 GPU 内核开发最近两年用大模型写代码已经从玩具级走向工程级。Python 脚本、SQL 查询、甚至部分业务逻辑实测下来都能省不少时间。但 GPU Kernel 是个例外它不是普通的顺序代码而是要在成千上万个线程上做并行计算同时要考虑访存局部性、占用率、寄存器压力、bank conflict 这些东西。一个 CUDA Kernel 跑起来很容易跑得快很难跑得比专家手写版本更快更是难上加难。我在实际项目里写过各类算子包括矩阵乘、LayerNorm、Softmax、各种 pooling 和 scan 操作。说实话哪怕是看起来很简单的向量加优化的空间都比想象中大——是否用向量化加载shared memory 怎么放线程块大小怎么定循环展开的程度这些组合起来就是一个巨大的搜索空间。而 LLM 生成 Kernel 的问题也在于此它见过很多代码能拼出一个合理骨架但它并不清楚目标 GPU 上到底哪种组合最优。所以大家很快发现纯靠 LLM 读题和凭经验生成得到的结果大概率只能算能跑离最优还有距离。1.2 编译器在这个问题里被严重低估传统上写 GPU Kernel 的流程是写代码 - 编译 - 跑 - 看性能 - 改。这个流程的反馈信号极其稀疏编译器只告诉你能不能编译过或者性能计数器大概多少。不会有太多人把编译器内部做的优化过程当作可学习的信息但编译器本身就是最了解底层硬件的组件。LLVM 或者 NVCC 的优化流水线知道指令选择了什么、寄存器分配情况、循环展开了几轮、向量宽度够不够、哪些访存被合并了。问题是这些信息都被深深埋在中层表示IR和调试文件里普通开发者不会去看更不要说让大模型去看了。这就造成了一个断层大模型和编译器各自掌握一部分知识但没有形成联动。1.3 共同进化到底说的是什么CAKE 这篇论文的核心观点就是不要再把 Agent 当成一次性代码生成器把编译器当成只报错不解释的裁判而是让 Agent 和编译器在一个循环里共同进化。类比一下生物进化就很好懂了。Agent 负责变异它不断生成各种 Keyboard variants编译器负责选择压力它把偶发性能优先的候选从衡量中凸显出来。变异和选择交替进行经过若干代之后种群的平均水平会越来越高甚至可能在某几个算子上超过由人类专家留下的特定位实现。这正是 CAKE 这个框架里共同进化的含义Agent 提供探索的广度编译器提供反馈的精度两者缺一不可。2. CAKE 的系统设计与核心组件拆解2.1 总体架构的四个关键环节CAKE 的总体架构并不神秘拆开看就是四个部分的闭环生成器Agent、编译器插桩层、评估器、进化引擎。我用一个表格先做个整体轮廓梳理组件核心职责对应我的理解Agent生成器根据反馈生成/修改 Kernel 候选代码不是一次性写代码而是听取反馈再改编译器插桩层把编译过程的优化信息、错误信息结构化导出把编译器知道但没说出口的秘密暴露出来评估器在真实硬件上跑得出正确性和性能指标唯一能拍板是否超越专家的角色进化引擎选择、保留、变异候选维持种群多样性防止过早收敛到局部最优这四个组件缺一个都跑不起来。没有 Agent 就没有探索能力没有编译器插桩就是盲人摸象没有评估器就没法判断好坏没有进化引擎就会在同一个死胡同里反复纠结。2.2 Agent 端不止猜代码更要学会读取反馈我前面说过普通 LLM 生成 Kernel 的问题在于单轮生成、没有反馈。CAKE 里面的 Agent 被设计成一个能持续学习的角色每一轮生成后它会收到来自编译器和评估器的反馈把反馈消化之后在下一轮生成时调整策略。这里的关键点是反馈的消化方式不能简单地把编译错误信息复制粘贴给 Agent 就完事。CAKE 的设计中Agent 需要读到的反馈至少包含几类编译错误和警告的精炼摘要优化 pass 的触发情况和未触发情况IR 层面发生的变化摘要比如某个循环被自动向量化了、某个访存被合并了编译产物和运行时统计例如寄存器用量、Shared Memory 使用率、实际执行时间我把这类信息统称为结构化的编译反馈。它的好处是让 Agent 知道自己的改动在编译器眼里产生了什么效果从而在下一次变异时朝着正确的方向走。比如 Agent 发现上一轮因为一个数组没有对齐导致编译器无法生成向量化访存它下一轮就会尝试加对齐属性或者改循环结构。这里想补充一个我在本地实验里的体会把编译器的原始错误信息直接丢给大模型效果很差。原因是原始信息非常碎有的 slot 可能有几百行噪声很大。CAKE 的思路是把这些信息压缩成有语义的要点再给 Agent 读相当于帮 Agent 提炼了编译器的诊断结果。这个细节我觉得比框架本身更值得借鉴。2.3 编译器端让内部知识外化成文本特征一般开发者接触到的编译器就是一个黑盒翻译器——代码进去二进制出来错了就报个错。但在 CAKE 的设计里编译器被赋予了一个新角色反馈发生器。为了实现这个目标需要在编译器上做插桩通常是在 LLVM/Clang/NVCC 的优化流水线上增加 log 输出把 IR 的变化过程记录下来。这样做的好处是Agent 能看到编译器的推理过程。举个例子编译器在做循环展开时它可能会评估展开因子是 4 还是 8 对指令调度的好处。普通开发者只看到最终性能数字但 CAKE 中的 Agent 能看到展开因子 8 导致寄存器压力过大spill 增多这样的中间信息。这已经超出了人类手工调优时能获取的信息量等于给 Agent 配了一副显微镜。当然这个插桩层也有实现成本。不同编译器版本、不同 GPU 后端日志格式可能完全不同。论文里的做法是把这些信息统一格式化为文本摘要以便 Agent 用语言模型理解。我的理解是这个插桩层越通用越好因为 GPU 架构迭代很快如果每次换硬件都要重写反馈逻辑那这套系统的维护成本就非常高了。2.4 进化与选择如何避免在优化空间里迷路进化引擎听起来高大上实际上做的是非常决策性的事每一轮产生一批候选 Kernel到底谁留下、谁淘汰、以什么变异策略产生下一代。这里有个我不止一次踩过的坑如果每一轮都只保留性能最好的那个候选很容易陷入局部最优。比如矩乘计算可能某个候选用了很激进的 tiling 方式在某个矩阵尺寸下表现好但稍微换一下输入尺寸就崩了。CAKE 的做法是维持一个候选集合保留一定程度的多样性让最终的种群既包含当前最优解也包含一些看起来不是最优但保留了不同优化思路的个体。评估器的标准也要仔细设计。除了运行时间必须同时考虑正确性否则 Agent 很容易学歪。我见过太多案例某个 Kernel 跑得极快但算出来的结果是错的因为它在快速路径里跳过了边界处理。浮点误差也是另一个坑GPU 上运算顺序变化会带来误差判断正确性时不能直接用等值比较要用容差来判定。所以 CAKE 的进化引擎里的选择压力是多维度的正确性 性能 收敛速度不是简单地比谁耗时短。3. 效果评测怎么证明超越专家不是吹的3.1 评测基准的选择逻辑CAKE 想让别人信服超越专家就不可能只在某一个算子上表现好。评测基准得覆盖 GPU Kernel 的主要类型访存密集型比如 vector add、copy、stencil、计算密集型比如 GEMM 矩阵乘、卷积、访存-计算混合型比如 softmax、layernorm以及带 scan/前缀和这类有数据依赖的算子。我特别关注的是基准里必须包含专家手写优化版本作为 baseline。这个 baseline 通常来自 CUB、CUTLASS 这类经过工业级调优的库或者来自论文作者团队手写的优化版本。如果只是拿没优化过的朴素实现来比那超越专家的门槛就太低了参考价值不大。另外还要看评测显卡的范围。单张卡上表现好不代表所有卡都表现好因为不同 GPU 的 SM 数量、显存带宽、寄存器文件大小、L2 缓存行为完全不同。CAKE 的实验如果能覆盖两三代硬件那说服力就强很多至少说明这套框架不是针对某个硬件专门重写的。3.2 超越专家到底体现在哪个维度超越专家这个词很容易被误解成所有算子都比专家快或每个尺寸都最快。从我看到的论文实验逻辑来理解更合理的表达是在大部分测试算子/尺寸上CAKE 生成的 Kernel 达到或接近专家版本性能在若干特色算子和规模下通过 Agent 和编译器的组合变异发现了专家平时不会想到的优化组合从而明显超过专家版本在自动化和鲁棒性上系统能在无人干预情况下自动搜索和收敛这种局部超越 整体追平的结果才是现实中有价值的目标。毕竟专家 Kernel 库是几代人调优的成果一个大模型加编译器进化系统能在一个晚上搜索到同等水平这本身就是革新。这里还想提一句论文实验里展示的超越很多是百分比提升比如 10%-30%。但实际工程里不同算子的时间占比不一样如果一个算子只占整个模型的 2%那即便提升 30%对端到端影响也不大。所以在解读结果时不能只看相对提升率还得看算子本身的热度。3.3 成本与资源约束工程化最绕不开的现实问题Agent 编译 运行这个循环看起来很美但每一次迭代都是有成本的。Agent 调用需要 GPU 算力来跑大模型推理每次 Kernel 要重新编译编译完还要在真实 GPU 上跑评测。如果一个算子需要搜 100 轮每轮 20 个候选那就是 2000 次编译和运行。我自己在做一个简化版框架的时候最初的版本非常粗糙一轮生成一个候选编译一次跑一次。结果发现大量时间浪费在等待上尤其是集中在编译过程中。后来我学到的经验是批量编译充分复用增量信息把多个候选的 Kernel 放进一个翻译单元编译器就能只花一次开销处理多种变体。CAKE 应该也有类似的考虑否则进化速度跟不上大模型推理和编译的时间成本。所以CAKE 的价值不仅在于找到好 Kernel也在于它把整个搜索过程的成本控制在一个可以接受的范围。如果搜一个算子要花一周那还不如让专家手写三天。只有当搜索成本大幅降低时这套系统才有工程化的空间。4. 从 CAKE 到我们自己的实践怎么落地一个简化版循环4.1 最小闭环Agent 脚本 编译报告我没有办法直接复制 CAKE 的完整代码但里面的核心思想完全可以用一个最小闭环在本地跑通。说白了就是三步让 Agent 生成 Kernel用脚本编译并运行把结果反馈给 Agent 再生成。这里的最小闭环我用 Python 脚本就能搭起来循环里做四件事调用大模型的 API给定算子描述和上一轮反馈生成一个或多个 CUDA Kernel 候选把候选写入 .cu 文件用 nvcc 编译成可执行文件同时收集编译器的 warning 和优化日志执行可执行程序用固定输入跑多次取最小耗时并校验输出结果整理成反馈文本注入下一轮的 prompt代码框架大概是这样实际部署时可以按自己的情况替换import subprocess, json, re from openai import OpenAI client OpenAI(base_urlyour_llm_endpoint, api_keyyour_key) def build_feedback(build_log, run_result): lines [] if build_log and error in build_log.lower(): # 提取出错行和错误类型 lines.append(COMPILE_ERROR: extract_error(build_log)) elif run_result and run_result.get(passed): lines.append(fOK elapsed{run_result[time_ms]:.3f}ms) lines.append(Compiler notes: summarize_compiler_notes(build_log)) else: lines.append(RUNTIME_ERROR: str(run_result.get(error))) return \n.join(lines) prompt_template 你是一位 GPU Kernel 优化专家。请根据下面的基线代码和上一轮的编译/运行反馈改进这个 CUDA Kernel。 要求保持正确性减少运行时间。只输出 CUDA 代码不要解释。 算子说明{spec} 上一轮反馈{feedback} 请输出改进后的完整 CUDA kernel current_code baseline_cuda_code for round_idx in range(15): prompt prompt_template.format(specop_spec, feedbacklast_feedback or 无这是第一轮) resp client.chat.completions.create(messages[{role: user, content: prompt}], modelyour_model) current_code extract_code(resp.choices[0].message.content) write_file(kernel.cu, current_code) build_log run([nvcc, -archsm_80, kernel.cu, -o, kernel]) run_result run_kernel_and_validate(kernel, expected_output, threshold1e-5) last_feedback build_feedback(build_log, run_result)这个循环本身不复杂真正要想清楚的是反馈格式。我一开始直接把 nvcc 的完整 stderr 丢进下一轮 prompt结果模型开始纠结一些无意义 warning效果很差。后来学乖了用正则把错误提取成几行关键信息效果立刻提升。这和论文里讲把编译器信息提炼成特征是同一个道理。4.2 反馈的格式、粒度与信噪比如果要给这个循环排优先级我觉得反馈设计的优先级高于生成策略。因为 Agent 再有想象力如果接收到的反馈是编译失败这种没头没尾的信息它也只会随机乱改。我在实践里总结出几条可以抄的规则编译错误必须带行号和错误类型比如第 24 行变量未声明比编译失败有用得多运行结果必须区分错误类型段错误、浮点异常、结果错误要分开描述性能数据要多次运行取中位数或最小值减少噪声干扰编译器优化日志要挑关键变化说不要全文复制如果 Agent 连续三轮没有改进把种群内其他候选的信息也一并送给 Agent 作为参考这最后一条很有意思。我发现一个现象Agent 在单条路径上如果连续失败往往会陷入重复修改同一个地方。这时给它看种群里的另一个候选为什么更快等于给它新的搜索方向。这其实就是 CAKE 进化引擎里多样性保持在 prompt 层面的体现。4.3 正确性与浮点容差容易被忽略的细节在代码生成循环里性能指标很显眼正确性验证却很容易做糊。GPU 上的并行规约、矩阵乘等运算浮点数计算顺序一变结果就会有微小差异。如果验证脚本直接用相等比较会有大量误判。我现在的做法是提前做好一个参考实现用 double 精度在 CPU 上跑一遍拿到参考输出。GPU 结果和参考输出做相对误差比较容差放宽到 1e-4 或者 1e-5同时处理 NaN 和无穷大。这个阈值可以根据算子的数值特性调整但切忌用正好相等的判断。另外还有一个冷门细节随机数的种子。如果验证时使用了随机输入每次测试要用同一个种子否则即使 Kernel 本身没问题之前失败的候选也可能在下一轮突然通过导致反馈信号不可靠。4.4 工具选型Clang 的优化报告比想象中有用CAKE 类系统的关键在于从编译器拿到反馈。除了修改编译器源码自己动手时其实还有更轻量的路子。Clang/LLVM 本身就支持-Rpass系列选项可以在编译时把优化 pass 的触发信息打印出来。我用过一个命令组合效果很好clang -O3 -Rpassloop-vectorize -Rpass-missedloop-vectorize -Rpass-analysisloop-vectorize -c kernel.cu -o /dev/null 21这会把哪些循环被向量化、哪些循环没有向量化、为什么没有向量化比如无法证明无别名都打出来。这些信息结构化成文本后Agent 就明白下一轮要把指针加 restrict或者要加对齐 hint。这比我最初直接把 nvcc 的标准输出丢给模型靠谱太多了。如果你用的是 N 卡的 CUDA 工具链nvcc的--ptxas-options-v可以额外打印寄存器用量和 shared memory 使用量也是很好的反馈源。比如它的输出里显示某个 kernel 用了 255 个寄存器且 spill 了很多那 Agent 下一轮就知道要降低并行度或者减少每个线程的工作量。5. 落地过程中的常见误区和避坑指南5.1 从我踩过的坑里整理出的对照表前面讲了具体实现这里把我在本地实践 CAKE 简化版时遇到的最典型问题列成一张速查表方便大家对照自查误区典型表现后果正确做法反馈不结构化把 nvcc 完整 stderr 扔给模型Agent 被无关 warning 干扰乱改用正则提炼行号错误类型关键信息只测性能不验正确性某候选运行时间最短但算错Agent 会把错误结果当成正确的最优必须用参考实现 容差对比每轮只保留最优个体种群多样性归零容易收敛到局部最优同时保留前几名的不同思路候选不考虑编译时间每轮单候选独立编译大量时间浪费迭代太慢多候选合并到一个编译单元用增量编译使用固定输入尺寸只在一个 shape 上优化换个 shape 性能骤降多种输入尺寸轮流测试忽略浮点误差结果对比用等值比较误判正确性反馈不可靠固定随机种子使用 1e-4 级相对误差编译器日志冗余把优化 pass 日志全文输出Prompt 过长关键信号被淹没只挑向量化、寄存器、占用率等摘要信息这里面最核心的还是反馈信噪比。我意识到一件事AI 生成代码的质量很大程度上取决于你给它提供的上下文质量。编译器内部的优化日志本来是为工程师设计的直接当 prompt 的一部分用是不可行的因为干扰太多。做一层摘要和过滤本质上就是把编译器知识翻译成 Agent 能理解的语言。5.2 算力消耗与超时控制Agent Kernel 搜索循环真正跑起来后最大的现实问题是超时。某一次编译可能因为错误的模板膨胀变得极慢某个生成的 Kernel 可能因为极端 block 配置导致 GPU 驱动超时。我在实践里给编译、执行分别设置了超时时间编译超时 60 秒就杀执行超时 5 秒就杀。如果用的是公有云 GPU 实例还要注意计费和并发问题。建议本地调试时用最小参数量模型验证循环逻辑确认反馈有效之后再切到大模型做正式搜索。直接上大模型 大批量候选很容易把资源耗尽而且难以定位问题。5.3 泛化性一个算子的经验如何迁移写了一个优化循环不能只在一个算子上有效就欢呼。我发现让 Agent 在不同算子之间做迁移学习是 CAKE 这套思路里很好用的能力。比如 Agent 在优化 GEMM 时学会了对 shared memory 做 padding 避免 bank conflict这个经验在优化卷积时同样有效。实现方式也很简单把历史有效的反馈 优化动作组合成经验案例在下一个新算子的第一轮 prompt 里附上这些案例。这等于给 Agent 建了一个可复用的调优记忆库虽然比 CAKE 完整版朴素一些但已经能解决很多重复劳动的问题。6. 从更广的角度看 CAKE 的意义6.1 对 Agent 开发模式的影响CAKE 对 Agent 开发者的启发比单纯生成 GPU 代码更大。它把 Agent 从一个对话式代码生成器变成了可进化的优化系统。这种 AI Agent 的开发模式跟现在很多 agent 框架里的 ReAct 循环是一脉相承的但其中的反馈设计和工具调用深度要比大多数偏向网页搜索和 API 调用的 agent 项目扎实得多。我理解 CAKE 的中心思想不是用 Agent 替代编译器而是Agent 和编译器之间建立高频双向通信。这个思想可以迁移到所有Agent 生成代码的场景数据库查询优化、CPU 内在函数优化、甚至 Makefile/CMake 配置优化。只要编译器能够提供结构化反馈Agent 就有可能在闭环中逐步逼近最优解。6.2 编译器作为一个观察仪器这几年大家讨论 AI 编译器时方向大多集中在怎么用编译器来优化神经网络算子。但 CAKE 提醒我们编译器除了做优化还有一个重要功能是观察——观察代码如何被翻译成硬件指令观察哪些优化假设成立观察寄存器压力和访存调度。这种编译器作为观察仪器的思路我认为会是未来 AI 辅助编程的重要方向。如果编译器能把它的推理过程以更结构化的方式暴露出来那么 Agent 就能理解代码在硬件层面的真实表现而不只是停留在语义层面。6.3 对未来编程工具链的展望把 CAKE 的思路推到极限可以看到一种未来的可能开发者只描述算子意图Agent 在编译器引导下自动探索最优实现最后生成的代码不仅经过真机验证还有完整的优化轨迹记录。人类专家的角色从写码者变成设目标和审核者。这种前景很诱人但离真正成熟还有很长的路。当前的 GPU 生态还分散在 CUDA、HIP、OpenCL、WebGPU 等多种编程模型上编译器反馈格式也没有统一标准更不用说 Agent 本身的可复现性和稳定性问题。但方向已经很明确值得持续观察。7. 最后的落地经验CAKE 论文对我来说最大的价值不是让我拥有一个一键生成超优 Kernel 的工具而是让我重新理解了 Agent 和编译器的关系。在此之前我总把两者分开看Agent 负责创意编译器负责执行。CAKE 的实验结果则说明把编译器的内部反馈暴露给 Agent让 Agent 的创意在硬件知识引导下搜索系统的能力上限会远高于任何单一方。如果你想尝试我的建议是不要一上来就复刻整套 CAKE。先用一个小算子比如 Softmax 或者 LayerNorm搭一个我前面描述的 15 轮循环看看 Agent 能不能在编译器反馈的帮助下稳定超过 baseline。等这个循环稳定了再引入多候选、种群保留、编译器插桩逐步丰富系统的能力。我在实际操作中最后的体会是这个系统最大的瓶颈不是模型不够聪明而是反馈回路设计得不够好。模型的能力是固定的但只要你把反馈做得结构化、及时、信噪比高模型的输出质量就会肉眼可见地提升。哪怕不写 GPU 代码把这个原则用于日常的代码生成工作流也会觉得AI 编程这件事变得更靠谱。
返回列表