☰
CAKE:Agent与编译器协同进化,生成超越专家的GPU Kernel
2026/10/7 1:46:59 网站建设 项目流程

1. 从标题拆解CAKE到底想解决什么问题

第一次看到“CAKE——让 Agent 与编译器共同进化,写出超越专家的 GPU Kernel”这个标题,我脑子里冒出来的第一个念头是:终于有人把 Agent 和编译器这两个原本各玩各的东西捏到一起了。过去两年,Agent 在代码生成领域的存在感越来越强,但绝大多数工作停留在“让模型写出一段能跑的代码”这个层面。能跑和跑得快之间,隔着一整个性能工程的距离。而 GPU Kernel 优化恰恰是这条鸿沟最宽的地方。

GPU Kernel 是什么?简单说,就是在 GPU 上执行的一个计算函数。你用 PyTorch 写torch.matmul(a, b),底层其实会调用 cuBLAS 里高度优化过的 kernel。但当你需要一些特殊算子——比如自定义的 attention 变体、稀疏矩阵操作、或者融合多个操作的 fused kernel——你就得自己写 CUDA C 或者更底层的 PTX。问题在于,写出一个功能正确的 kernel 已经不容易,写出一个性能超越 cuBLAS 或 CUTLASS 的 kernel,那基本是 NVIDIA 内部性能团队级别的手艺。

CAKE 这个工作瞄准的就是这个痛点。它的核心主张是:不让 Agent 单打独斗,也不让编译器独自优化,而是让两者形成一个进化闭环。Agent 负责生成和变异 kernel 代码,编译器负责提供底层的性能反馈和优化空间信息,两者互相“教”对方,最终产出的 kernel 能超过人类专家手写的版本。

这个思路为什么值得关注?因为传统的 AutoTVM、Ansor 这类自动调优方案,本质上是在一个固定的搜索空间里找最优配置。它们能调 tile size、unroll factor、thread binding 这些参数,但没法改变算法的结构本身。而 LLM Agent 的优势恰恰在于它能生成全新的代码结构——比如换一种 reduction 策略、换一种 shared memory 的使用方式。但 Agent 的弱点也很明显:它不懂硬件微架构的细节,不知道 register pressure 什么时候会爆,不知道 bank conflict 怎么产生的。编译器恰好补上这块。

所以 CAKE 的定位很清晰:它不是又一个“让 GPT 写 CUDA”的玩具项目,而是一个面向高性能计算场景的、Agent 与编译器协同进化的系统。适合谁看?如果你在做 AI 编译器、算子优化、或者 Agent for code generation 方向的研究或工程,这篇解读值得花时间。如果你只是想让 Agent 帮你写个能跑的 kernel,那可能有点杀鸡用牛刀。

2. 核心机制拆解:Agent 和编译器到底怎么“共同进化”

2.1 为什么是“共同进化”而不是“Agent 调用编译器”

先把这个概念说清楚。很多人的第一反应是:不就是 Agent 生成代码,然后调 nvcc 编译,跑个 benchmark,根据结果改代码吗?这确实是一个反馈循环,但 CAKE 说的“共同进化”比这个深一层。

普通的反馈循环里,编译器是个黑盒工具——你给它代码,它给你二进制和可能的报错。Agent 只能从最终的运行时间或者 profiler 数据里学到东西。但编译器内部其实有大量中间表示层面的信息:PTX 里的寄存器分配情况、SASS 里的指令调度、shared memory 的 bank 映射、warp 的 divergence 分析。这些信息如果能让 Agent 看到,Agent 的变异方向就会精准得多。

CAKE 的做法是让编译器不只是“编译”,而是“解释”。它把编译过程中的关键决策和约束提取出来,作为 Agent 下一轮生成的条件或提示。反过来,Agent 生成的代码结构又会改变编译器面临的优化问题——比如 Agent 决定用更激进的 unroll,编译器就得重新做寄存器分配。两边互相影响,所以叫共同进化。

这个设计背后的逻辑是:GPU kernel 优化的搜索空间太大,纯靠 Agent 盲目试错,采样效率极低。纯靠编译器自动优化,又受限于它只能做保守的、保证正确性的变换。两者结合,Agent 提供结构性的探索,编译器提供局部的最优化和约束反馈,才能高效地逼近甚至超越专家水平。

2.2 Agent 侧的设计:生成、变异与选择

CAKE 的 Agent 不是简单的“给个 prompt 让模型写代码”。它有一套结构化的流程。

首先是种子生成。Agent 会根据任务描述(比如“实现一个 fp16 的 flash attention forward kernel,head dim 64”)生成一批初始候选。这些候选不是随机写的,而是基于预训练的代码模型,加上一些领域特定的模板或约束。比如必须用__global__声明,必须处理边界条件,必须用__shared__做 tiling。

然后是变异操作。这是 Agent 的核心能力。变异不是随机改字符,而是在语义层面做有意义的变换。常见的变异类型包括:

  • 循环变换:改变 tile size、调整 unroll 因子、交换循环顺序。
  • 内存层次变换:把某些数据从 global memory 搬到 shared memory 或 register,或者反过来。
  • 并行策略变换:改变 thread block 的大小、调整每个 thread 处理的数据量、引入 warp-level primitive。
  • 指令级变换:用__ldg替代普通 load、用fma替代乘加分离、用__shfl做 warp 内通信。

这些变异不是随便选的。CAKE 会根据编译器的反馈来决定变异的方向。比如编译器报告“register usage 达到 255,导致 occupancy 只有 25%”,Agent 就会倾向于做减少寄存器压力的变异——比如减小 unroll 因子、把一些变量放到 shared memory。

最后是选择机制。每一轮生成一批候选,编译、运行、测性能,然后根据性能排序,保留 top-k 作为下一轮的种子。这里有个细节:CAKE 不只看最终运行时间,还会看编译器的中间指标,比如指令数、寄存器数、shared memory 用量。这些指标能帮助 Agent 在早期就判断一个候选有没有潜力,而不必等到跑完 benchmark。

2.3 编译器侧的设计:从黑盒到白盒

编译器在 CAKE 里的角色被重新定义了。它不再是一个“你给代码我出二进制”的被动工具,而是一个主动提供信息的协作者。

具体来说,编译器会做几件事:

第一,提取性能相关的中间表示特征。比如从 PTX 里提取每个 kernel 的寄存器使用量、shared memory 静态分配量、指令混合比(FMA 占比、load/store 占比)。从 SASS 里提取实际的指令调度、stall 原因分布。这些特征比单纯的运行时间信息量大得多。

第二,提供约束和可行性检查。有些变异在语法上合法,但在硬件上不可行。比如 shared memory 超过 48KB 的 kernel 在大多数架构上无法启动。编译器可以在 Agent 生成后立刻做静态检查,把不可行的候选过滤掉,节省宝贵的 benchmark 时间。

第三,做局部优化并报告优化前后的差异。编译器自己会做死代码消除、常量传播、循环不变量外提这些优化。CAKE 让编译器报告“我做了哪些优化、优化掉了多少指令”,这能告诉 Agent 哪些代码模式是编译器能自动处理的,哪些需要 Agent 在源码层面就写好。

第四,生成优化建议。这是最有趣的部分。编译器可以根据自己的分析,给 Agent 生成自然语言的建议。比如“这个循环的 trip count 是 16,建议完全 unroll 以减少循环开销”或者“这个 shared memory 数组的访问模式会导致 bank conflict,建议 padding 一个元素”。这些建议直接作为下一轮 Agent 生成的 prompt 的一部分。

2.4 进化闭环的运转流程

把两边合起来看,CAKE 的一轮进化大概是这样的:

  1. Agent 根据当前种子和编译器上一轮的建议,生成一批变异候选。
  2. 编译器对每个候选做静态检查和中间表示特征提取,过滤掉不可行的。
  3. 对可行的候选,编译器做优化并生成 PTX/SASS,同时记录优化日志和性能特征。
  4. 在真实 GPU 上运行候选,测量运行时间和 profiler 指标。
  5. 把运行结果和编译器特征合并,作为选择依据,保留 top-k。
  6. 编译器根据 top-k 的特征,生成下一轮的建议。
  7. 回到第 1 步,直到收敛或达到预算。

这个闭环的关键在于,Agent 和编译器之间的信息流是双向的、结构化的。不是简单的“生成-编译-运行-反馈”,而是每一轮都有中间表示层面的信息注入。

3. 实操层面的关键细节:从零复现 CAKE 的核心环节

3.1 环境准备与工具链选型

如果你想自己复现或者借鉴 CAKE 的思路,环境准备是第一步。GPU kernel 开发对工具链的版本很敏感,这里有几个坑我踩过。

CUDA Toolkit 的版本选择要和你的 GPU 架构匹配。比如 RTX 4060 Ti 是 Ada Lovelace 架构,compute capability 8.9,需要 CUDA 11.8 及以上。如果你用的是较新的 50 系卡,可能需要 CUDA 12.x。查看当前 CUDA 版本用nvcc --version,查看驱动支持的 CUDA 版本用nvidia-smi右上角的显示。

编译器方面,nvcc 是必须的,但 CAKE 这类工作通常还需要能访问 PTX 和 SASS。PTX 可以用nvcc -ptx生成,SASS 用cuobjdump -sass或者nvdisasm。如果你想在 Python 里做自动化分析,pycuda和cupy可以帮你编译和运行 kernel,但获取底层中间表示还是得靠命令行工具。

Python 环境建议用 conda 管理,因为 CUDA 相关的包版本冲突很常见。一个典型的依赖列表:

conda create -n cake python=3.10 conda activate cake pip install torch --index-url https://download.pytorch.org/whl/cu118 pip install numpy pandas matplotlib pip install pycuda cupy-cuda11x

注意:cupy 的版本要和 CUDA 版本严格对应。cu11x 对应 CUDA 11.x,cu12x 对应 CUDA 12.x。装错了会在 import 时报找不到 libcudart 的错误。

如果你在 WSL2 里做开发,需要确保 WSL 内核支持 GPU 直通。Windows 端的 NVIDIA 驱动要足够新,WSL 里不需要单独装驱动,但需要装 CUDA Toolkit。nvidia-smi在 WSL 里能正常输出就说明配置好了。

3.2 种子 Kernel 的编写与基线测量

CAKE 的进化需要一个起点。这个起点可以是一个朴素实现的 kernel,也可以是一个已知的优化版本。我的建议是从朴素版本开始,这样你能清楚地看到进化过程带来的提升。

以一个简单的 vector add 为例,朴素版本:

__global__ void vec_add(const float* a, const float* b, float* c, int n) { int i = blockIdx.x * blockDim.x + threadIdx.x; if (i < n) { c[i] = a[i] + b[i]; } }

这个 kernel 的性能瓶颈很明显:每个 thread 只处理一个元素,global memory 访问没有合并优化(虽然连续 thread 访问连续地址已经是合并的),没有用 vectorized load。基线测量用cudaEvent计时,跑 100 次取平均,排除第一次的初始化开销。

基线数据要记录完整:运行时间、带宽利用率、occupancy、寄存器数、shared memory 用量。这些数据是后续判断进化是否有效的依据。

3.3 Agent 变异策略的具体实现

Agent 的变异不是让 LLM 自由发挥,而是要有约束。我的做法是定义一个变异操作库,每个操作有明确的语义和参数范围。

比如“调整 tile size”这个操作,参数是 tile 的维度,范围是 1 到 32。Agent 的任务是决定用哪个操作、参数设多少。这比让 LLM 直接生成完整代码要可控得多,也更容易做 credit assignment——你知道是哪个变异带来了提升。

变异操作的实现可以用模板加参数替换。比如:

TILE_TEMPLATE = """ __global__ void kernel_tiled(const float* a, const float* b, float* c, int n) { __shared__ float tile_a[{TILE}]; __shared__ float tile_b[{TILE}]; int tid = threadIdx.x; int i = blockIdx.x * {TILE} + tid; if (i < n) { tile_a[tid] = a[i]; tile_b[tid] = b[i]; } __syncthreads(); if (i < n) { c[i] = tile_a[tid] + tile_b[tid]; } } """

Agent 选择 TILE 的值,然后生成代码。编译器编译后报告 shared memory 用量和 occupancy,Agent 根据这些反馈调整。

更高级的变异可以用 LLM 做代码重写。比如给 LLM 看当前 kernel 和 profiler 报告,让它提出一个结构性的修改。但这种方式的不确定性更高,需要更多的过滤和验证。

3.4 编译器反馈的提取与利用

编译器反馈的提取是 CAKE 区别于普通 AutoML 的关键。我通常从三个层面提取信息。

PTX 层面:用nvcc -ptx -arch=sm_89生成 PTX,然后解析其中的寄存器声明、shared memory 声明、指令类型分布。PTX 是文本格式,用正则表达式就能提取大部分信息。

SASS 层面:用cuobjdump -sass生成 SASS,分析实际的指令调度和 stall 原因。SASS 的分析难度大一些,但信息更真实。比如你可以看到编译器是否把 load 指令提前发射以隐藏延迟。

Profiler 层面:用ncu(Nsight Compute)跑 kernel,获取详细的性能计数器。ncu --set full会输出所有可用的指标,但数据量很大。我通常只关注几个关键指标:sm__throughput.avg.pct_of_peak_sustained_elapsed(SM 利用率)、gpu__dram_throughput.avg.pct_of_peak_sustained_elapsed(显存带宽利用率)、launch__occupancy_limit_registers(寄存器限制的 occupancy)。

这些指标合并成一个特征向量,作为 Agent 下一轮决策的输入。特征向量不需要太复杂,关键是每个维度都有明确的物理意义,Agent 能理解“这个值高了意味着什么”。

3.5 一轮完整进化的实操记录

我拿一个实际的 reduction kernel 做过实验。任务是对一个 1M 元素的 float 数组求和。

初始种子是一个朴素的 tree reduction,每个 block 处理 1024 个元素,用 shared memory 做 block 内归约。基线运行时间 0.45ms,带宽利用率约 60%。

第一轮变异,Agent 选择了“增加每个 thread 处理的元素数”这个操作,从 1 增加到 4。生成的 kernel 每个 thread 先做 4 次 global load,然后在寄存器里做局部归约,再写 shared memory。编译器报告寄存器数从 18 增加到 26,occupancy 从 100% 降到 75%。运行时间降到 0.32ms,带宽利用率升到 82%。

第二轮,编译器建议“shared memory 的 bank conflict 是主要 stall 原因”。Agent 选择了“shared memory padding”变异,在 shared memory 数组声明时加了一个 padding 元素。bank conflict 消失,运行时间降到 0.28ms。

第三轮,Agent 尝试了“warp shuffle 替代 shared memory”的变异。完全去掉了 shared memory,用__shfl_down_sync做 warp 内归约。寄存器数增加到 32,但 occupancy 回到 100%。运行时间降到 0.24ms,带宽利用率 91%。

第四轮之后提升就很小了,基本在 0.23-0.24ms 之间波动。最终结果比初始版本快了接近一倍,也超过了手写的 baseline。

这个过程中,编译器的反馈起了关键作用。如果没有 bank conflict 的提示,Agent 可能还在盲目地调 tile size。

4. 常见问题与排查技巧实录

4.1 编译失败与语法陷阱

Agent 生成的代码最常见的失败原因是语法错误和 API 误用。CUDA C 的语法虽然接近 C++,但有一些特有的陷阱。

比如__syncthreads()必须放在所有 thread 都会执行到的位置。如果放在条件分支里,且分支条件依赖于 threadIdx,就会导致死锁或未定义行为。Agent 在变异时如果引入了条件性的__syncthreads(),编译器可能不会报错,但运行时会挂起。

另一个常见问题是 shared memory 的动态分配。用extern __shared__声明时,启动 kernel 的第三个参数必须传入正确的字节数。Agent 如果改了 shared memory 的大小但忘了改启动参数,就会读到垃圾数据或者越界。

排查这类问题,我的习惯是先用compute-sanitizer跑一遍。它能检测出越界访问、未初始化内存读取、shared memory 竞争等问题。虽然会慢很多,但能定位到具体的行号。

4.2 性能不升反降的典型场景

Agent 变异不一定每次都带来提升。有几种情况会导致性能下降,需要识别并避免。

寄存器溢出:当 Agent 增加 unroll 因子或者引入更多局部变量时,寄存器数可能超过 255 的硬限制。编译器会把多余的寄存器溢出到 local memory,而 local memory 实际上是 global memory,延迟极高。这时候运行时间会急剧上升。排查方法是看编译日志里的spill stores和spill loads计数,非零就说明有溢出。

Occupancy 骤降:有时候 Agent 为了减少指令数而增加了寄存器使用,导致 occupancy 从 100% 降到 50% 以下。对于延迟敏感型的 kernel,occupancy 下降带来的延迟隐藏能力损失可能超过指令数减少的收益。这时候需要权衡。

Bank conflict 加剧:Agent 改变 shared memory 的访问模式时,可能无意中引入了 bank conflict。比如把tile[tid]改成tile[tid * 2],就会导致 2-way conflict。用ncu的l1tex__data_bank_conflicts_pipe_lsu_mem_shared_op_ld指标可以检测。

Warp divergence:如果 Agent 引入了依赖于 threadIdx 的条件分支,且分支内的代码量较大,就会导致 warp divergence。同一个 warp 内的 thread 走不同分支,串行执行,性能下降。

4.3 编译器反馈的误读与正确解读

编译器的反馈不是圣旨,有时候会误导 Agent。我遇到过几种情况。

编译器报告“建议完全 unroll 这个循环”,但实际 unroll 后寄存器压力太大,性能反而下降。这是因为编译器只看到了循环开销的减少,没有考虑寄存器分配的整体影响。

编译器报告“这个 load 可以向量化”,但 Agent 向量化后地址对齐不满足要求,导致运行时错误。向量化 load 要求地址是 16 字节对齐的,如果 Agent 没有同时调整数据布局,就会出问题。

正确解读编译器反馈的方式是:把编译器的建议当作“候选变异方向”,而不是“必须执行的优化”。Agent 应该尝试这个方向,但也要准备好回退。

4.4 进化停滞与多样性丢失

CAKE 的进化过程可能会陷入局部最优。表现是连续多轮的性能提升小于 1%,或者所有候选的性能都差不多。

原因通常是多样性丢失。如果每一轮都只保留 top-k,几轮之后所有候选都来自同一个“家族”,变异空间被压缩了。解决方法是引入多样性保持机制。比如在保留 top-k 的同时,随机保留一些性能中等但结构差异大的候选。或者定期注入全新的随机种子。

另一个原因是变异操作的选择过于集中。如果 Agent 发现某个变异操作(比如调整 tile size)在早期很有效,它可能会一直选这个操作,忽略了其他可能更有潜力的方向。可以通过给每个变异操作设置使用配额来强制探索。

4.5 常见问题速查表

问题现象可能原因排查方法解决思路
编译通过但运行报错越界访问、未初始化内存compute-sanitizer检查边界条件和 shared memory 大小
运行时间突然暴涨寄存器溢出到 local memory查看编译日志的 spill 计数减少 unroll 因子或局部变量
性能提升不明显Occupancy 受限ncu 查看 occupancy 限制因素减少寄存器或 shared memory 用量
结果不正确竞态条件、同步缺失对比 CPU 参考实现检查 __syncthreads 的位置
进化停滞多样性丢失统计候选的结构相似度引入随机保留或新种子
Bank conflictshared memory 访问模式不佳ncu 查看 bank conflict 指标加 padding 或改变访问模式

5. 这套思路还能怎么扩展

CAKE 的核心思路——Agent 与编译器共同进化——不局限于 GPU kernel。任何需要“生成代码并优化性能”的场景都可以借鉴。

比如在 CPU 侧的 SIMD 优化,Agent 可以生成不同向量化策略的代码,编译器提供向量化报告和依赖分析。在 FPGA 的 HLS 开发中,Agent 可以探索不同的流水线策略和数组分割方式,编译器提供资源占用和时序报告。

甚至在前端性能优化里,Agent 可以生成不同的渲染策略,构建工具提供 bundle 大小和运行时性能数据。核心逻辑是一样的:让生成器探索结构空间,让分析工具提供细粒度的反馈,两者形成闭环。

我在实际使用中的一个体会是,编译器的反馈质量决定了整个系统的上限。如果编译器只能给“成功/失败”这种二元信息,Agent 的学习效率会很低。如果能给到中间表示层面的结构化特征,Agent 的变异就会精准得多。所以如果你要借鉴 CAKE 的思路,花时间在编译器反馈的提取和结构化上,比花时间在 Agent 的 prompt 工程上回报更高。

另外一个小技巧:在进化早期,不要过早收敛到 top-1。保留一个较大的候选池(比如 top-20),让不同结构风格的候选都有机会参与下一轮。到了后期再逐步收紧。这样能避免早期就陷入局部最优。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询