让大模型像专家一样优化多场景 CUDA 核函数 Making LLMs Optimize Multi-Scenario CUDA Kernels Like Experts
用多智能体+硬件瓶颈分类,让 LLM 跨稠密/稀疏/LLM/科学计算多场景自动优化 CUDA 核,平均较 Astra 提速约 35%。
前置知识
CUDA 核函数(CUDA Kernel)
CUDA 核函数是运行在 GPU 流多处理器(SM)上的并行函数,由 CPU 端通过 grid/block 启动配置发起。开发者需手工管理线程层次、共享内存、向量化访存、同步与 bank conflict 等低层细节,目标是让 GPU 的算力、内存带宽与延迟同时被充分利用。
本文让 LLM 写出 CUDA .cu 文件,所以必须先理解核函数如何被启动、profile 工具(如 Nsight Compute)会测量哪些 SM/DRAM 指标,以及编译选项(如 -arch=sm_xx)为何影响性能。
Nsight Compute 硬件性能分析
Nsight Compute 是 NVIDIA 的核函数级分析器,输出上百项计数器(如 Compute(SM) Throughput、DRAM Throughput、Memory Throughput、Issue Slots Busy、L1/L2 Hit Rate、Achieved Occupancy 等),用来判断算子是 compute-bound、latency-bound 还是 bandwidth-bound。
本文的核心创新之一就是把这些海量指标按瓶颈类型过滤后只喂给 LLM,因此读者要明白这些指标各自衡量什么,否则就无法理解为什么过滤能省钱且不损效果。
Otsu 阈值法(Otsu's method)
Otsu 法是经典的图像二值化算法,通过最大化类间方差自动找到一个能把灰度直方图分成两类的最佳阈值。它是一种数据驱动的无监督分类方法。
本文用 Otsu 把所有 baseline 算子的三项 throughput 指标分布自动切成“高/低”两类,从而得到 ≈30% 的统一瓶颈判定阈值,避免了人工拍脑袋调参。
GPU 瓶颈三分类(Compute/Memory Latency/Memory Bandwidth Bound)
这是 Roofline 思路下的经典划分:算力被打满且带宽还有余量时是 compute-bound;DRAM/全局 memory 吞吐饱和时是 bandwidth-bound;吞吐都不高、warp 在等数据时是 latency-bound(常因 cache miss 或不规则访存)。
本文的所有过滤规则、prompt 引导、消融分析都建立在这一三分类之上。理解为何三类需要不同的指标子集,是看懂 CUDAMaster 为何有效的前提。
多智能体系统(Multi-Agent LLM)
把复杂任务拆给多个有专门角色的 LLM 代理(如 Planner、Coder、Compiler、Debug),代理之间通过消息/工具调用协作,每个代理持有一个系统提示和外部工具集合。
CUDAMaster 正是这种设计的典型:Planner 出方案、Coder 写代码、Compiler 写 nvcc 命令、Debug 修错;理解为何这样切分能胜过单次提示是理解本文方法的关键。
研究动机
现有的 LLM 自动 CUDA 优化研究几乎全部基于 KernelBench——一个 250 题的 PyTorch 算子集,主要覆盖 LLM 中常见的 MatMul、LayerNorm、Softmax 等规则算子。它的核心缺陷有三:第一,PyTorch 抽象自带 dispatcher、autograd 调度和框架级 fused op,agent 写出来的代码天然向深度学习执行模式倾斜;第二,评测只用一个固定数据规模,掩盖了 kernel 在不同负载下扩展性的差异;第三,迭代中模型可能只是“背诵”了公开 LLM 算子的现成优化套路,泛化能力并未被真正考察。METR 在 2025 年的自动化 kernel 工程报告中也指出,这种 scope 限制让报告的加速比不可靠。当我们把目光投向稀疏矩阵-向量乘、Stencil、Reduction、归一化等更广义的 HPC 算子时,问题变得更加尖锐:cuBLAS/cuDNN/cuSPARSE 这类闭源库是性能天花板但缺乏灵活性,TVM/Triton 等编译器虽然灵活却在很多领域追不上专家,hand-tuned 库又对每个新硬件世代要重新投入巨大工程量。换句话说,目前既没有一个能公平评测“真实多场景 kernel 优化能力”的基准,也没有一个能跨领域达到专家级性能的自动化系统。
本文的目标是本文的核心目标分两层。第一层是提出一个名为 MSKernelBench 的多场景基准,覆盖稠密线性代数、稀疏矩阵算子、LLM 常见算子序列以及科学计算/数值方法共 50 类任务,每类支持 FP32 与 BF16 共 100 个待优化核;基准用纯 C 而非 PyTorch 实现,去掉框架依赖,同时按数据规模做复杂度加权评分,从而能更严格地反映一个 kernel 在不同负载下的实际收益。第二层是提出 CUDAMaster——一个端到端的多智能体系统,接收 kernel 源码和执行脚本后,由 Planner/Coder/Compiler/Debug 四个代理自动完成 profile、规划、改写、编译、调试整个流水线,最终给出可直接 nvcc 编译的优化版核函数以及完整的执行脚本。
与已有工作不同的是,已有 LLM-kernel 工作(FSR、CudaForge、Astra、Kevin、CUDA-L1/L2)几乎都只在 LLM 算子这一狭窄切片上比拼,并且往往直接把 Nsight Compute 的全部指标塞进 prompt——这既费钱又让 LLM 难以聚焦。本文抓住两个被忽视的切入口:一是把评测对象从“LLM 算子”拓展到“通用 GPU 算子”,二是把 profiling 报告按瓶颈类型做有针对性的过滤,让 LLM 看到的不是“上千个数字”,而是“与你这个算子瓶颈相关的几个关键指标 + 它们的语义解释”。前者解决“考什么”的问题,后者解决“怎么让 LLM 真会调”的可扩展性问题。
核心方法
CUDAMaster 的设计哲学是模拟人类专家调优 kernel 的心智流程:先跑一遍基线拿到 profile,再在脑子里判断瓶颈属于哪一类,再针对该类瓶颈挑出最相关的几个指标去思考优化策略,最后迭代实现、编译、运行、修正。系统由四个串联的 LLM 代理组成:Planner 负责把当前代码、profile 与历史方案翻译为“高层次优化方案”(不做代码生成);Coder 拿到方案后改写 .cu 文件并保证函数命名规范(_kernel_optimized 后缀)和 wrapper 接口;Compiler 写出一条 nvcc 编译命令;Debug 在编译或执行失败时介入、定位错误并提出修复方案。每一轮迭代结束,系统会比对当前性能与历史最优,仅当本轮通过验证且得分更高时才更新最优解,从而在 R=3、D=3 的固定预算下完成稳健搜索。
和 Astra、Kevin 等已有方法相比,CUDAMaster 的本质区别在于两点。第一是 Profile Filtering(硬件瓶颈分类 + 指标过滤):先用 Otsu 法从所有 benchmark 算子的 throughput 分布中客观地学到三组阈值(Compute > 30%,DRAM/Memory > 30% 等),把每个算子分为 Compute Bound、Memory Latency Bound 或 Memory Bandwidth Bound,然后从 Nsight Compute 的上百个指标里只挑出与该瓶颈强相关的 3–6 个送进 prompt。第二是 Agent 拆分与责任清晰化:Planner 不写代码、Coder 严格按规约输出、Compiler 专管 nvcc、Debug 专管错误,外部 C wrapper 和测试代码永远不被改写,公平性由这种接口不变性保证。这两个设计合起来使 LLM 在更“少而精”的上下文中思考更对的问题,并使整个过程可以并发、可复现。
方法步骤详情
整个流程由算法 1 控制。首先用 ExecuteAndFilter 跑一遍原始 kernel C₀ 和原始执行脚本 S₀,得到基线 profile π₀ 并初始化最优 (C*, S*, π* = π₀, P* = 1.0)。然后进入最多 R=3 轮迭代,每轮依次执行:第 4 行 PlannerAgent 接收 (C*, π*, 上一轮方案) 输出本轮优化方案 scheme;第 5 行 CoderAgent 把 scheme 翻译成新的 .cu 文件 Cᵢ,所有 kernel 函数名必须以 _kernel_optimized 结尾,wrapper 函数名以 _optimized 结尾且参数列表与原始一致;第 6 行 CompilerAgent 给出 nvcc 编译命令 Sᵢ,注意 sm_xx 架构号、特殊头文件链接等;第 7 行 ExecuteAndFilter 编译并运行,调用 Nsight Compute 后按 Table 2 的规则分类瓶颈并过滤指标,得到 πᵢ。第 8–18 行是 Debug 子循环:若 CheckValid(πᵢ) 为假(编译/执行/数值不通过),最多 D=3 次调用 DebugAgent 拿修复方案,再让 CoderAgent 重写、再 Compiler、再 Execute,直到通过或耗尽预算。第 20–25 行是“是否更新最优”判断:只有当 (Pᵢ > P*) 且通过验证时才用 (Cᵢ, Sᵢ, πᵢ) 替换最优并刷新 scheme 记忆,否则把当前 scheme 标记为 invalid 以便下一轮 Planner 避开。复杂度加权分数 P = Σᵢ wᵢ Sᵢ / Σᵢ wᵢ,其中 wᵢ ∝ T(Nᵢ) 是输入规模 Nᵢ 下基线核的渐近复杂度,使得大数据、复杂度高的算例获得更高权重。
技术新颖性
技术新颖性体现在三处。第一,Otsu 阈值法被创造性地用于 GPU 性能分析,论文展示了在所有 baseline 算子上 Compute/DRAM/Memory 三类 throughput 的分布都能在 30% 附近被 Otsu 干净切开,从而获得一个数据驱动、统一且可解释的瓶颈分类规则,避免了人工调阈值的随意性。第二,瓶颈分类与指标过滤紧耦合:Compute Bound 算子被送进 Compute (SM) Throughput、Issue Slots Busy、Executed Ipc Active、SM Busy 四个与算力/调度相关的指标;Memory Latency Bound 被送进 L1/L2 Hit Rate、Mem Busy、Executed Ipc Elapsed 等延迟/缓存命中率指标;Memory Bandwidth Bound 被送进 DRAM Throughput、Max Bandwidth、Mem Pipes Busy 等带宽饱和度指标——这种“瓶颈→指标”映射表本身就是一个 LLM 友好的领域知识库。第三,把 Planner 与 Coder 解耦、并强制使用“原始测试代码不可改”的接口规约,使系统可以无副作用地复用于 100 个 kernel,并在固定 R=3、D=3 预算下完成端到端调优。表 4 的消融显示,Filtered 配置在 τ=1 时与 Full Profile 几乎同分(94% vs 94%),但 token 消耗从 172,917 降到 123,036,cost 从 $0.38 降到 $0.27,相当于砍掉近 32% 的 API 费用而不损失效果。
实验结果
实验在 NVIDIA RTX 4090 上、每个算子 FP32+BF16 合计 100 个 kernel 上进行,使用 o4-mini 与 DeepSeek-V3.2 两个代码模型。第一,宏观点:图 4 与表 3 显示 o4-mini 在所有阈值上都稳定领先 DeepSeek-V3.2——τ=0(功能正确)100% vs 95%;τ=1(跑赢 naive 基线)94% vs 80%;τ=2(≥2x 加速)60% vs 49%;τ≥8 的高目标下仍有 25% vs 22% 的成功率。BF16 在中等阈值给 o4-mini 带来微弱优势,FP32 在严格目标下更稳。第二,微观点:图 5 在六个有“权威基线”的算子上把同一 naive kernel 当参考点比较,结果是 o4-mini 全部跑赢——SpMV CSR 比 cuSPARSE 在大尺寸下 2.96× vs 2.23×、2D Conv 比 cuDNN 1.83× vs 0.97×(即 cudnn 反而比 naive 慢,agent 把负值救正)、Dot Product 比 cuBLAS 46.83× vs 26.09×、RMS Norm 比 Astra 2.44× vs 1.78×、SiLU&Mul 3.67× vs 3.58×、Merge Attention States 1.35× vs 1.05×。最值得指出的是 Dot Product 这种“看起来不可能赢过 cuBLAS”的算子,agent 居然做到了近 47× 加速(naive 是 O(N) 但访存未向量化),背后是它用更宽的向量加载和向量化寄存器替代了朴素循环。第三,消融:表 3 证明多轮迭代 + Debug 必要——Full (R=3,D=3) 在 τ=1 处 94% 成功率,去掉 Debug(N.D.)掉到 77%,降到单轮(S.I.)掉到 74%,单轮且无 debug(S.R.)最低;表 4 证明过滤的 profile 在效果上追平 Full Profile 同时少用约 32% token 和 28% 费用。第四,硬件层面(图 8、9):50 个算子优化前 11/24/15 分属三类,优化后 6 个原 memory-bound 的被推到 compute-bound,14 个从 memory-latency 滑到 bandwidth,memory-latency 类算子从 24 暴跌到 8(−67%),compute-bound 和 bandwidth-bound 分别 +45%、+73%;具体到指标,Mem Busy 平均提升 1312%、DRAM Throughput 提升 537%、Issue Slots Busy 提升 16.8%,说明 agent 既“挪位置”又“挖潜力”。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| SpMV CSR (Sparse Matrix-Vector, CSR) | Speedup over naive (largest size) | OpenAI o4-mini ≈ 3.19× | cuSPARSE ≈ 2.70× | 约 +18%(o4-mini 在 2^26 量级略胜 cuSPARSE) |
| 2D Convolution | Speedup over naive (largest size) | o4-mini ≈ 1.77×;DeepSeek ≈ 1.75× | cuDNN ≈ 0.97×(反而比 naive 慢) | agent 把 cuDNN 的“负加速”扳回正 1.8× |
| Dot Product (FP32) | Speedup over naive (largest size 2^20) | o4-mini ≈ 46.83× | cuBLAS cublasSdot ≈ 26.09× | 约 +80%,首次见到 agent 在点积上明显压过 cuBLAS |
| RMS Norm | Speedup over naive | o4-mini ≈ 2.44× | Astra 1.78×(Astra 在 SGLang 上调过的) | 相对 Astra +37%,与摘要 35% 吻合 |
| SiLU & Mul | Speedup over naive | o4-mini ≈ 3.67× | Astra ≈ 3.58× | 约 +2.5%,基本打平 |
| Merge Attention States | Speedup over naive | o4-mini ≈ 1.35×(DeepSeek 1.42×) | Astra ≈ 1.05× | 约 +30%,DeepSeek 略胜 o4-mini,是少数 DeepSeek 反超的例子 |
| 50 tasks 宏观点(FP32, o4-mini) | Cumulative success rate | τ=0: 100%;τ=1: 94%;τ=2: 60%;τ≥8: 25% | DeepSeek-V3.2:95% / 80% / 49% / 22% | 关键阈值上领先 5–11 个百分点 |
| 平均开销 (o4-mini, Filtered) | Cost / Tokens per task | $0.2707 / 123,036 tokens | Full Profile $0.3804 / 172,917 tokens | 砍掉约 32% 费用、29% token 而 τ=1 仍保持 94% |
局限与改进
作者在表 4 的消融里诚实指出,过滤后的 profile 比 Full Profile 在高阈值(τ≥8)下略低(25% vs 25% 持平,但 o4-mini 在 τ=2 时 60% vs 62% 稍逊),说明过度裁剪可能在极端优化上损失信息。另外,框架对 Nsight Compute 的依赖意味着换到 AMD GPU 就要重写瓶颈分类和指标映射表;论文只在 RTX 4090 一张卡上做了实验,是否能跨 A100/H100/Blackwell 泛化未经验证。从附图 7 的瓶颈分布看,50 个算子中 Memory Latency Bound 占 24 个、Memory Bandwidth Bound 15 个、Compute Bound 11 个,类内方差巨大(例如 Mem Busy 标准差 1388%),这意味着个别极端不规则访存的算子(例如 stencil)在 agent 视角下仍是“hard case”。观察到的额外限制:测试代码不可改这一约束虽然保证公平,但把 wrapper 设计、kernel 启动配置之外的工程优化(如 CUDA Graph、动态并行)排除在 agent 决策之外;R=3/D=3 的固定预算太小,对几个极复杂的 reduction 类算子(如 cross-entropy、merge attention)可能不够迭代;最后,论文只测了 o4-mini 和 DeepSeek-V3.2 两个模型,没有评估 GPT-5、Claude 4.5、Gemini 2.5 Pro 等更新更强的代码模型在这些任务上的表现。
独立分析的弱点
第一,瓶颈分类的颗粒度偏粗:只分三类,缺少“分支发散 bound”、“寄存器压力 bound”、“shared memory bank conflict bound”等更细的子瓶颈。对于本就是 compute-bound 的 11 个算子,过滤后只剩 4 个指标,agent 容易把所有问题都归于“提升 IPC”,对调度细节不敏感。改进方向是引入指令级 profile(如 sass 反汇编、stall reason 长尾),或者把 warp scheduler 状态(Warp Cycles Per Issued/Executed Instruction)作为可选项。第二,Planner 阶段只看 profile 不看代码长度/循环嵌套深度等结构特征,所以有时会给出一个对小数据合适、但对大数据爆 shared memory 的方案,落到 Debug 才修。可让 Planner 同时接受“已编译产物大小、寄存器使用、shared memory 占用”三类静态信息作为约束。第三,复杂度加权指标 P 给了大尺寸最高权重,导致小尺寸下的回归可能不被惩罚——对一个要在多个 scale 部署的库来说这是隐患;可补充 P_min = min_i S_i 或几何平均作为辅助指标。第四,测试用 4090 一张卡,且不报告浮点精度退化——BF16 路径下当加速 > 8× 时,agent 是否仍守住 1e-5 的绝对误差未见单独分析,需要补一个“max relative error vs speedup”散点图。第五,CUDAMaster 把 wrapper 函数签名锁死,但 Coder 仍可能改启动 block/grid 配置导致小数据掉速;可以引入一个轻量的“预测最优配置”预选器减少 Coder 的搜索空间。
未来方向
作者明确的下一步是把 CUDAMaster 拓展到更多硬件(AMD ROCm、Intel oneAPI)和更多精度(FP16、INT8、FP8、TF32),让 LLM 真正成为“跨平台 kernel 工程师”。基于结果还可延伸的方向包括:(1) 把过滤策略做成可学习的,让模型自己决定哪些指标与当前 bottleneck 相关,做成 RL 训练目标(和 Kevin、CUDA-L1/L2 接轨);(2) 把 CUDAMaster 用作训练数据工厂——它在 100 个 kernel 上 R=3 轮得到的中间方案 + 编译日志 + profile 对比可作为高质量 SFT 数据,用来蒸馏出专用的 kernel 优化模型;(3) 把多代理拆成可并行的子树,例如同时派出多个 Planner 出不同方向方案、由 Compiler 批量评估,挑最优,避免顺序搜索的局部最优;(4) 引入任务级元学习:把“Stencil 在 latency-bound 下应优先用 shared memory tiling”这种经验以 memory bank 的形式注入 Planner prompt,进一步降低迭代次数;(5) 与现有 compiler 工具链耦合:让 LLM 输出 TVM/Triton 风格的中间表示再交给编译器做后端优化,融合“LLM 的高层思路”和“编译器的低层保证”。
复现评估
作者在结尾承诺开源 benchmark 与 framework(demo 链接 https://hanyx2021.github.io/MSKernelBenchDemo/,但论文正文中未给出 GitHub 地址,README 与源码是否能稳定复现需要审稿后补查)。可复现性方面有几个有利的点:(1) 评测完全在 RTX 4090 这一消费级卡上进行,硬件门槛低;(2) 50 任务 + FP32/BF16 = 100 个 kernel 的输入/输出、参考实现、复杂度权重都在附录 C 表格 7 详细列出(每个 ID、task 名称、prompt、数据规模、复杂度函数 T(N) 都给全);(3) 评测协议做了 3 次 warm-up、50 次平均取时,且在公式 (1)(2) 给出加权 P 的精确公式;(4) 4 个 agent 的 system prompt 在附录 B 公开,包含 Planner、Coder、Compiler、Debug 全部 prompt 模板,模板里有占位符(如 {info}、{code_file_content}、{former_plan}),照搬即可复现。不利的点:(1) R=3、D=3 是闭源模型的调用次数,外部研究者需要相应 API budget,按 o4-mini $0.27/task 估算 100 个 kernel 约 $27;(2) 论文没有公开 Nsight Compute 的具体命令、报告裁剪脚本,需要读者自己还原;(3) 闭源 LLM 的版本/温度/种子未声明,行为可能漂移;(4) Baseline 中 cuBLAS/cuDNN/cuSPARSE 的版本是 v13.1 (2026a/b/c),这是一个未来版本号,意味着读者使用的旧版 cuBLAS 可能给出不同的基线速度。综合看,复现一个“接近论文数字”的版本是可行的(5 天左右),但要严格复现 100 个 kernel 的所有数字则需要复现 LLM API 设置与 nsight 裁剪逻辑。
论文图表
三张子图分别展示 (a) Compute-Bound、(b) Memory-Latency-Bound、(c) Memory-Bandwidth-Bound 算子在各瓶颈专属指标上的平均提升 % ± 标准差。Compute-Bound 中 Issue Slots Busy +16.8%±79.5、Executed Ipc Active +15.4%±78.8;Memory-Latency-Bound 中 Mem Busy +1312.7%±1388.5、Executed Ipc Elapsed +889.4%±1539.0;Memory-Bandwidth-Bound 中 DRAM Throughput +537.3%±1278.9、Mem Pipes Busy +98.9%±187.9。
微观层面回答“加速比从哪儿来”,是支撑“LLM 真的在做硬件级优化、不是只在打补丁”这一论断的核心证据。
附录大表,50 个算子(带 ID 1–50),每行给出 Task Type、Task Name、Prompt、输入数据规模、复杂度 T(N)。涵盖 7 个 Dense Linear Algebra、8 个 Sparse Matrix、4 个 Normalization、7 个 Activation、6 个 LLM Operator、3 个 Loss、4 个 Tensor Process、5 个 Stencil、6 个 Math & Algorithm。
复现整套 benchmark 的索引表,也是判断 benchmark 范围是否足够“multi-scenario”的最直接证据。