我过去几年一直在做 GPU 内核优化日常工作就是跟 CUDA C、PTX、SASS 这些东西打交道。说实话最磨人的环节不是写出第一版能跑的 kernel而是后面那一轮又一轮的调优——循环展开几级、共享内存怎么 padding、寄存器占用压到多少、要不要拆成多个 kernel这些决策往往要靠 profiler 反馈来回试。后来大模型火起来身边不少人开始琢磨能不能让 LLM 直接改内核代码但试过几轮之后大家普遍碰壁模型生成的 CUDA C 要么编译不过要么编译过了但性能还不如原来手写的版本更别提直接优化 PTX 这种汇编级代码了。PTXBench 这个项目我关注了很久它的核心思路很明确把 LLM 优化 GPU kernel 这件事从运气式生成变成一个有 benchmark、有适配方法、有评估指标的闭环。它解决的是真实痛点——PTX 是 NVIDIA 虚拟指令集架构夹在 CUDA C 和硬件微码 SASS 之间既保留了程序员对寄存器和指令的控制权又不像 SASS 那样绑死特定 GPU 架构。让 LLM 直接在这个层级做优化理论上能拿到比 CUDA C 更细的调优空间同时也比反编译 SASS 靠谱得多。这篇文章我想拆开聊聊 PTXBench 这类系统背后的设计逻辑、评估陷阱以及真正落地时你会遇到的坑。1. 为什么偏偏是 PTXLLM 优化内核的最优战场先说清楚 PTX 到底是什么、它凭什么值得让大模型去学。CUDA 的编译链路是三层你写的.cu文件先进过 nvcc 前端变成 PTXParallel Thread Execution然后由 ptxas 把 PTX 编译成 SASSSASS 才是 GPU 硬件真正执行的微码。绝大多数开发者只接触最上层写 CUDA C 然后交给编译器但如果你追求极致性能就得往下走。1.1 一段代码的三层编译旅程我举一个非常简单的例子向量加法 kernelCUDA C 写出来就几行__global__ void vec_add(float *a, float *b, float *c, int n) { int i blockIdx.x * blockDim.x threadIdx.x; if (i n) c[i] a[i] b[i]; }经过 nvcc 编译成 PTX 之后会变成类似这样的形式.reg .b32 %r4; .reg .f32 %f3; ld.global.f32 %f1, [%rd1]; ld.global.f32 %f2, [%rd2]; add.f32 %f3, %f1, %f2; st.global.f32 [%rd3], %f3;你会看到每个线程的执行过程被拆成了明确的加载、计算、存储指令寄存器分配也是显式的。到这一步你还能通过调.reg的数量、改指令顺序来影响性能。再往下到 SASS指令编码变成二进制不同代架构Ampere、Ada、Hopper的指令格式还都不一样人类几乎没法直接写。1.2 LLM 面对 CUDA C 的无力感为什么不直接让 LLM 优化 CUDA C我试过实在太多次了总结下来三个问题第一CUDA C 离硬件太远。编译器帮你做了大量指令调度、寄存器分配、内存访问优化模型改了 C 代码实际生成的 SASS 可能完全没变性能原地踏步。这就像你改了一个函数的 Python 源码但解释器编译出的字节码跟你改之前一模一样。第二LLM 生成的 CUDA C 往往看起来对、跑起来错。潜台词是它会用很自然的 C 语法写 kernel想当然地加一些并行结构结果要么边界条件算错要么 race condition 满天飞。编译能过跑出来数值不对排查半天才发现是共享内存不同步的问题。第三CUDA C 层面的优化空间其实已经不大。像循环展开、向量化加载这些手段编译器基本都会自动做真正还没被编译器解决的部分恰恰在更底层的选择上。PTX 恰好是那个甜点位。它比 CUDA C 更贴近硬件寄存器、指令类型、内存访问模式都是显式的又比 SASS 更可读、可移植、可生成。ptxas 会再做一轮优化但你对寄存器和指令的显式控制仍然能显著影响最终 SASS 的质量。这个特性对 LLM 格外友好——生成目标是文本指令规则明确反馈及时ptxas 能直接告诉你哪里编译不过这是天然的 LLM 应用场景。2. 基准测试的坐标系正确性、编译率与性能增益如何权衡PTXBench 这类项目最重要的贡献不是让 LLM 生成了 PTX而是建立了一套能区分模型真的会优化内核和模型只是碰巧生成了一段能跑的汇编的评测体系。做 benchmark 本身就很难尤其是针对汇编级代码生成这里有三个层层递进的问题。2.1 三层评估从能编译到跑得比基线快第一层是编译通过率。给定一个输入的 CUDA kernel 或 PTX让模型输出优化后的 PTX然后调用ptxas -c尝试编译看看能过多少。这一层看似基础实际淘汰率极高。LLM 很容易生成不存在的指令修饰符、写错寄存器类别、漏掉类型后缀。我印象里通用模型直接生成 PTX 的编译通过率经常只有 10%-30%这还是在有参考示例的情况下。第二层是运行时正确性。编译过了不算数还得真的拿到 GPU 上跑跟基准输出做对比。浮点运算要考虑容差不能要求 bitwise 相等整数和地址计算则必须完全一致。这一层会再淘汰一批能编译但语义错的样本常见的错误包括把有符号当无符号处理、错误处理内存对齐、把某些特殊值NaN、Inf路径改错。第三层才是性能增益。把通过前两层的 PTX 拿去和原始 kernel 的 SASS 性能做对比算加速比。这一步要小心因为 GPU 性能测量充满了噪声——时钟频率的动态变化、L2 缓存的冷热、同卡上其他任务的干扰都会让结果抖动。规范做法是多次运行取中位数并且用同一块 GPU、同一种散热环境做 A/B 对比。下面是我认为一套合理评测体系必须包含的维度评估维度具体指标说明编译可行性passk 编译成功率采样 k 个候选至少一个能通过 ptxas语义正确性数值一致性、边界样本覆盖需要设计包含特殊数值和边界尺寸的测试集性能表现相对基线的加速比多次运行取中位数必须控制时钟、温度、缓存状态等变量泛化能力训练集之外的 kernel 表现防止模型死记硬背测试集硬件可迁移性在不同代 GPU 架构上的表现PTX 可移植但优化效果可能随架构变化2.2 评测集怎么构建从真实工作负载里挖测试用例评测集的质量决定 benchmark 的可信度。我在搭建这类数据集时遵循一个原则测试用例必须来自真实世界而不是人工拼凑的玩具 kernel。人工写的简单 kernel比如 vector add、matrix multiply 的基础版本太规整了LLM 轻松就能优化真实场景里的 kernel 往往有复杂的数据布局、奇怪的边界条件、多层循环嵌套这些才是检验优化能力的试金石。具体来源可以是几个方向开源项目里的 CUDA 代码比如 cuDNN 的某些参考实现、Torch 扩展里的自定义 op、NVIDIA 官方文档和示例里的 kernel、以及从编译器测试套件里抽出的一部分中间表示。每个测试用例都要标注清楚原始性能基线、硬件平台、输入数据特征、优化目标延迟还是吞吐。没有这些元信息评测结果就是一笔糊涂账。2.3 为什么性能指标容易骗人这里必须多说两句。很多人在评估 LLM 优化效果时只看最终加速比这个指标极其容易被误导。一个常见陷阱是基线太弱。如果拿来对比的原始 kernel 本身写得很差比如完全没用向量化加载、共享内存随便用那 LLM 随便做点改进就能有 2 倍以上的加速看起来成果斐然实际上并没有什么含金量。更合理的做法是拿 cuBLAS、cuDNN 这类经过手工极致调优的库做对比基线或者至少拿nvcc -O3下编译器能产出的最优版本做基线。另一个陷阱是挑选性的成功。一个大模型每次生成的结果都不完全一样如果从 100 个采样结果里挑性能最好的那个去宣传这个数字没有统计学意义。正确做法是报告 passk 下的性能分布而不是单点最优值。3. 让 LLM 真正读懂硬件手册领域适配的关键操作Benchmark 解决的是怎么测的问题Adapting 解决的是怎么让模型变强的问题。直接拿通用大模型去生成 PTX效果通常很差原因不复杂通用模型的训练数据里 PTX 的占比极低它见过 Python、Java、C但没怎么见过.reg .b32这种指令风格的代码。要让它在 PTX 优化领域发挥作用必须做领域适配。这就好比你不能让一个文学编辑直接上手审阅 Verilog 代码得先让他学硬件描述语言。3.1 让模型先学会PTX 语法再谈优化领域适配的第一个层级是语法学习。PTX 的语法和高级语言差别很大它有自己的一套指令集、寄存器命名规则、类型系统、状态空间修饰符。模型如果没有见过足够的 PTX 文本生成出来的东西基本是胡编乱造。一个关键操作是扩展模型的词表tokenizer。通用模型的 tokenizer 里通常有add、load这样的词但没有ld.global.f32、st.shared.u32这类紧凑的指令 token。有人做过实验直接拿标准 tokenizer 去处理 PTX 代码指令名会被切成碎片模型很难学到整条指令是一个语义单元。正确的做法是把这些高频 PTX 指令序列加入词表或者用 BPE 在 PTX 语料上重新训练一个词表。然后是构建高质量的 PTX 语料。来源包括 NVIDIA PTX ISA 文档里的示例、开源项目比如一些直接用 PTX 内联汇编的库、以及从大量.cu文件编译出的 PTX。语料规模不用特别大——毕竟 PTX 的语法空间远小于自然语言——但质量要极高每一条都得是能通过ptxas编译的合法代码。3.2 结构感知的领域知识注入像带研究生一样教模型语法学会之后更深入的适配是让模型理解 PTX 的程序结构。PTX 和高级语言不同它的基本块、控制流、数据流都是显式的而且充斥着大量在高级语言里感知不到的约束——比如 warp 内所有线程必须执行同一条指令SIMT 模型、共享内存的 bank 冲突会影响访问效率、寄存器数量直接限制占用率。我在实践中最有效的做法是结构感知注入在提示词或训练样本里不只给模型看一段孤零零的 PTX 代码而是把代码的上下文结构信息一并塞给它。比如输入 kernel 对应的 CUDA C 源码让模型知道人眼里这段代码原本长什么样PTX 代码的寄存器分配报告和基本块划分可以通过工具自动生成profiler 给出的性能热点标注哪条指令是瓶颈、哪块内存访问延迟最高。这个概念可以类比为带研究生的过程。你不会直接甩给新生一段代码说优化它而是先给他看需求文档、架构设计、profiling 报告告诉他瓶颈在哪儿、约束是什么再让他动手改。结构感知注入做的事情就是给 LLM 补上这些上下文。一个具体的 prompt 设计模式是这样的下面是一段在 A100 上性能不达标的 PTX 内核代码瓶颈为共享内存 bank 冲突 profiler 显示平均 8-way conflict [PTX 代码] 请分析该代码的共享内存访问模式识别可能引发 bank 冲突的指令 并重写为优化版本。优化约束 1. 不得改变算法复杂度 2. 共享内存大小维持在当前数值 3. 保持寄存器数量不超过 64 个这样模型就有足够信息做有依据的修改而不是瞎猜。对比一下只给它一段代码说优化它生成质量差距非常大。3.3 微调策略LoRA 还是全量微调语法和上下文知识都准备好了接下来是选择微调策略。我在不同项目里试过全量微调和 LoRA说说我的经验。全量微调在 PTX 场景下有个天然劣势高质量 PTX 优化数据太少了可能只有几千到几万条样本全量微调很容易过拟合模型会死死记住训练集里的模式遇到新的 kernel 结构就不知道怎么处理。而且全量微调成本高PTX 这种特殊领域数据还没有多少公开的预训练权重可以直接继承。LoRA低秩适配是更合适的选择。它只训练一小部分参数在保持模型原有通用能力的同时注入 PTX 领域知识。我在实际对比中观察到LoRA 在编译通过率上比全量微调高出不少回忆一下大概从 40% 提升到 70% 左右在性能增益上也更稳。原因是 LoRA 本质上是在原有能力旁边加了一个专家通道模型调用通用知识比如对算法复杂度的理解和领域知识PTX 指令规则时互不干扰。当然LoRA 也有它的边界。如果目标任务和原模型能力差异太大比如一个以自然语言为主的模型要全靠 LoRA 学会汇编优化那也顶不住。所以选基座模型时尽量选代码能力本身就强的比如在 CodeLlama、DeepSeek-Coder 这类模型基础上做适配而不是从通用聊天模型硬来。3.4 强化学习反馈让模型自己对比汇编结果监督微调只能让模型学会生成合法的 PTX但它不知道什么 PTX 是性能更好的。要教会模型优化而不是改写得引入性能反馈。这也是 Adapting 阶段最有价值的一步。具体做法是类似的 RLHF 流程但奖励信号不是人类打分而是机器实测数据。流程如下给模型一个输入 kernel让它生成多个 PTX 候选每个候选都经过ptxas编译编译不过的直接给低奖励编译通过的运行 benchmark 测出实际性能和基线对比算出加速比把加速比映射成奖励信号通过强化学习算法常见的是 PPO 或 GRPO更新模型。这个过程有个很妙的点奖励信号完全来自客观测量不掺入任何主观判断。模型会被逼着去探索那些真正能提升性能的改动——比如重新排列指令以减少 bank 冲突、调整寄存器分配以提升占用率——而不是只顾着把代码写得更整洁。有个我在实操中发现的细节对奖励做归一化很重要。不同 kernel 的优化空间差异巨大有的 kernel 基准版本已经接近理论极限再怎么改也就提升 5%有的 kernel 写得稀烂随便优化就能快 3 倍。如果不做归一化模型会疯狂优化那些容易捡分的简单 kernel忽视难啃的硬骨头。正确的做法是按 kernel 分别归一化——每个 kernel 的加速比除以该 kernel 在验证集上见过的最大加速比让所有样本的奖励分布在一个可比的范围里。4. 实测效果与翻车现场哪些优化真的有效哪些是伪优化理论讲了一堆最终还得看实际效果。我在这类项目上跑了不少实验有成功也有翻车下面把这些经验原原本本写出来。需要提前说明的是不同硬件、不同 baseline、不同 seed 下结果会有波动我只谈趋势和规律。4.1 有效案例寄存器重排与指令合并先说我看下来最有效的几类优化指令合并。模型学到的最有价值的模式是把多条指令合并成一条功能更强大的指令。比如mul.f32加add.f32合并成fma.rn.f32融合乘加一步操作同时完成乘法和加法减少指令发射数量和中间寄存器占用。这类优化对计算密集型的 kernel 效果极其显著有时候能把指令数降低 30% 以上。寄存器压力优化。LLM 通过观察 PTX 中.reg声明的数量可以推断出这个 kernel 用了多少寄存器然后相应调整代码结构来压寄存器用量。在 A100 上用满 255 个寄存器会让 occupancy 降到极低每个 SM 只有 2 个线程块能并行但如果压到 64 个以内就能塞下更多 block。模型学到的做法包括复用寄存器而不是频繁声明新的、把大循环改成小循环以缩小存活变量集合、优先使用fma而不是muladd后者需要多一个临时寄存器。内存访问向量化。把 4 次 32 位加载合并成 1 次ld.global.v4.f32128 位访问本质上就是 CUDA 里float4的做法但在 PTX 层面做更直接。这种优化对带宽受限型 kernel 的加速效果特别明显。4.2 翻车案例看似优雅实则更慢的重写说完成功的说说翻车的。而且我必须强调翻车比例一点都不低。第一个典型的翻车案例是模型把精确计算改成近似计算。它可能学到了一招——用乘倒数替代除法确实能加速但它不会判断什么情况下精度损失是不可接受的。比如div.rn.f32在某些 kernel 里替换成mul.approx.f32会引入肉眼可见的误差最终导致数值验证失败。模型不懂业务场景的误差容忍度只看到指令便宜了就用这在 benchmark 阶段就会被正确性检查拦下来——但在真实生产环境里这种错误往往藏得更深可能只在特定数据集上才暴露。第二个翻车案例是过度展开循环。LLM 学到循环展开能减少分支开销于是把 4 次迭代的循环全展开生成了巨长的 PTX 代码。这会导致指令缓存I-cache压力剧增代码膨胀反而拖慢执行。实测中有些 kernel 展开 2 次是加速展开 4 次就开始回退展开 8 次直接比原版慢。模型很难理解这个甜蜜点的存在因为它取决于 GPU 微架构的指令缓存大小而这种信息不会出现在 PTX ISA 文档里。第三个翻车案例是移动瓶颈而非消除瓶颈。模型把计算密集的指令从主循环里挪出去放进一个看似只执行一次的初始化区但那个初始化区实际上会被每个线程都执行一遍等于什么都没优化。更隐蔽的情况是把共享内存加载移到分支里面结果导致某些线程浪费等待时间。这类问题只有在完整的运行时 profiling 下才能暴露单看 PTX 静态结构很难发现。4.3 性能复现的方差问题最后提一个非常容易被低估的问题同一段 PTX 的性能在不同实机环境下的方差可能比你的优化幅度还大。我记得有一次测试同一个优化后的 kernel 在冷启动状态下跑出 1.8 倍加速等 GPU 温度上来之后重测只剩 1.1 倍再隔几分钟测又变回 1.5 倍。原因很复杂——GPU 的 boost clock 依赖温度L2 缓存状态影响不可控CPU 侧的数据传输也会干扰。如果 benchmark 过程不严格控温、不做多次重复、不取中位数你得到的加速比就是一个随机数。这也是为什么我在评估 LLM 优化效果时强烈建议引入基准线多次测量 优化多次测量 排序后取中位数的流程。宁可测试时间翻倍也不要被单次测量的噪声误导。另外跨架构测试也很重要——在 A100 上优化的结果迁移到 H100 或消费级卡如 RTX 4090上可能完全不适用因为 SASS 指令调度策略不同寄存器文件大小也不同。5. 后续怎么用 PTXBench 思路扩展自己的优化流水线聊了 benchmark 设计和模型适配最后这部分我想讲讲更落地的事就算你不用 PTXBench 这个现成框架它的核心思路也可以迁移到自己的 GPU 优化工作流里。我自己的做法是这样的供你参考。5.1 接入现有编译工具链自动验证与回归测试第一件事是把 LLM 生成的 PTX 嵌入到现有编译流程里做一个预过滤器。我在本地搭了一个简单流水线.cu源码 → nvcc 编译出 PTX → 交给 LLM 做优化 →ptxas编译优化后的 PTX → 在 GPU 上跑自检包括正确性断言和性能计时→ 通过则替换原来的 kernel失败则自动回退到原始版本。这里的核心是自动回退机制。AI 生成的代码不可能每次都是对的生产环境的铁律是不通过验证的代码绝不进入发布分支。这个流水线可以放在 CI 里每个 PR 触发一次让 LLM 优化作为一个可选的性能优化阶段而不是替代人工 review。这么做的好处是能持续收集哪些优化有效、哪些无效的数据为后续优化积累训练材料。5.2 把失败样本做成增量训练集让模型越用越聪明第二件事是建立失败样本的回流机制。每次模型生成的 PTX 编译失败或者性能回退我都会把这个样本存档包括输入 kernel、模型输出、编译报错信息或 profiling 数据。这批负面样本是宝贵的训练数据——它们能让模型学习避免同类错误。具体做法上可以把这些样本攒到一定量我一般攒 500 条左右就做一次小规模增量微调。增量微调用 LoRA 就行不必重新全量训练。我观察到一个很有意思的趋势经过两三轮这样的失败驱动迭代后模型生成的 PTX 编译通过率有明显提升而且它开始学会主动避免一些常见错误——比如不再给标量指令加.v4后缀、不再混淆.b32和.s32、不会在不需要的地方插入bar.sync。这个循环本质上是 OpenAI 那套通过反馈不断改进生成质量的做法只是把反馈信号从人类打分换成了编译器报错和硬件计时这在工程上是完全可闭环的。5.3 通用经验别指望一次推理就给出最终内核最后分享一个心态层面的建议别指望 LLM 一次推理就能给出最终的生产级内核。我现在工作流里对 LLM 的定位是生成候选的优化思路而不是直接交付可上线的代码。具体操作是给定输入 kernel 后让模型一次性生成 8-16 个不同的优化候选通过调高 temperature 保持多样性然后批量跑编译和 benchmark。最后把结果做排序挑表现最好的候选和人工写的版本做对比。这样一个流程下来选到可接受结果的概率比单次生成高一个量级。从统计角度讲这是 passk 思想的实际应用与其追求单次生成的成功率不如接受一次不行就多试几次反正 GPU 编译和跑 benchmark 的成本远比人工调优低。还有一个小技巧让模型多轮自我反思。第一轮生成后把ptxas的报错和 profiler 的瓶颈信息反馈给模型让它基于这些反馈修自己上一轮的代码。这种生成—反馈—修正的循环结构往往比让它一次性直接生成更成熟的代码。本质上这就是把像人类学生一样学习的思路用到了模型身上先有个初稿再看错误反馈然后逐步修订最终收敛到合格方案。这样一来LLM 就不再是碰运气生成代码的投币机而是真正参与到了优化迭代的环节里成为整个 GPU 内核调优流水线中一个靠谱的协作者。