ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

让 LLM 直接写 PTX:绕开编译器后端的 GPU 优化新路径

让 LLM 直接写 PTX:绕开编译器后端的 GPU 优化新路径 去年我在技术社区看到一条评论“编译器没那么神圣让模型自己写汇编可能更靠谱。”当时我当段子看了直到读了这篇标题为“AI 就是编译器——让 LLM 直接写 PTX绕开整个编译器后端”的论文才意识到这个思路不是脑洞而是一套完整的论证。整篇论文的核心其实就一句话与其花几十年维护一套庞大的编译器后端去做指令选择、寄存器分配、调度优化不如直接让大语言模型LLM根据计算意图写 PTX 指令。这个想法一开始听起来像是“把飞机引擎换成螺旋桨”但论文给出的实验证据和推理链条让这个方向变得相当有说服力。这篇文章适合所有跟 CUDA、LLM、编译优化打过交道的人我尽量把里面的技术逻辑拆成能直接消化的干货来讲。1. 论文捅了编译器后端的哪根神经1.1 编译器后端为什么一直是块硬骨头先明确一个背景现代编译器分前端和后端。前端负责词法分析、语法分析、语义检查把源代码转成中间表示IR。后端则负责从 IR 出发做指令选择、寄存器分配、指令调度、目标代码生成等一系列琐碎但要命的优化。对于 CPU 架构来说x86、ARM 这些后端已经打磨了几十年有 LLVM、GCC 这种超大型开源工程支撑每一条指令的选择背后都是数不清的启发式规则和成本模型。GPU 的情况更特殊。NVIDIA 的编译器链路是CUDA C 源码先进 nvcc 前端转成 LLVM IR然后再变成 PTXParallel Thread Execution并行线程执行指令集最后由 ptxas 这个后端把 PTX 变成真正在 GPU 上跑的 SASS 机器码。PTX 不是给人写的东西它是 NVIDIA 定义的一种虚拟指令集长得像汇编但比 SASS 高一层指令集是公开的、跨代稳定的。而真正的 SASS 是每代 GPU 不公开的机器码里面有多少寄存器、什么调度方式完全由 ptxas 说了算。后端优化的难点在于它是一个典型的 NP 难问题。寄存器分配、指令调度、软件流水线每一个单项拎出来都够写几十篇论文。现实中编译器只能靠各种启发式策略取一个“足够好”的解。但这些启发式策略是静态的、规则化的它不认识语义不理解上下文意图更不会针对一个具体 kernel 的特殊性去动态调整。这就是为什么很多大厂的高性能团队最后还是要靠手写 CUDA 甚至直接手写 PTX 来压榨性能。1.2 LLM 凭什么说自己是“编译器”论文提出的核心观点是编译器后端本质上做的是“一种表示到另一种表示的翻译加优化”而 LLM 最擅长的恰恰就是这种符号序列之间的映射。如果给模型足够的训练数据——包括 CUDA 源码、PTX 指令、SASS 反汇编结果、性能评测反馈——理论上它可以学会比启发式规则更聪明的优化方式。这个想法最颠覆的地方在于“绕开”这个词。传统流程是程序员写 CUDA C编译器后端负责帮你优化。论文的设想是你只要告诉 LLM 一个计算目标比如“做一个高效的矩阵乘法 kernel”模型直接吐出 PTX。中间没有 IR 优化没有指令选择没有 ptxas 的寄存器分配。LLM 通过在海量代码、内核实现和性能数据上学到的模式直接在 PTX 层面把优化一并做了。这个逻辑在理论上说得通因为后端优化的大部分工作其实是“模式识别 局部改写”——找循环、找乘加、找内存访问模式然后换成更高效的指令组合。这些恰恰是 LLM 在大量代码语料上训练后最擅长的事情。更关键的是传统编译器后端只能基于预先定义的模式去匹配但 LLM 能根据上下文自由组合指令序列这就等于给优化空间加了一维。2. 为什么偏偏挑 PTX 下手2.1 CUDA 到 PTX 再到 SASS链路在哪里断的要理解论文的选择得先看清这条链路里哪个环节最“值钱”、也最“顽固”。整条链路是高级语言层CUDA C / C程序员表达意图的地方中间表示层LLVM IR编译器统一处理的地方伪汇编层PTX公开的虚拟指令集机器码层SASS每代 GPU 专属的硬件指令前端的价值在于语义分析和类型检查这部分 LLM 替代不了而且也没必要替代。后端真正能创造性能价值的地方在于IR 层的优化、PTX 层的指令选择和调度、SASS 层的寄存器分配与重构。论文瞄准的是 PTX 这一层因为它刚好卡在一个非常微妙的位置。PTX 向下可以保留优化的可能性向上又不受源代码语法的束缚。它有一套明确的指令语义比如ld.global.f32、fma.rn.f32、bar.sync这些每条指令干什么都写死了。模型输出一段 PTX 后可以交给 ptxas 做最终翻译也可以直接用工具做正确性验证。如果把目标设定为直接生成 SASS那就死路一条——SASS 没有公开文档每个架构的指令编码都不一样训练数据都凑不齐。2.2 PTX 的“虚拟指令集”身份是最大的切入点PTX 被设计出来的本意是当作一个稳定的虚拟 ISA让 CUDA 程序能和具体硬件解耦。这个设计给了论文一个绝佳的切入点因为 PTX 足够底层人写起来很痛苦但它的语义又足够清晰机器可以验证。这正好是 LLM 发挥的黄金地带——人类不擅长、规则型编译器做不到、但模型有可能通过训练学会。举个例子写一个线程束级矩阵乘的 Shuffle 指令传统编译器后端会根据硬件的寄存器数量、Bank 冲突情况去做调度。但如果你直接告诉 LLM“目标 GPU 是某代架构寄存器上限是 255共享内存 48KB”模型有可能写出比 ptxas 更激进的指令排列。因为 LLM 在训练时见过无数手写 kernel 的 PTX它会模仿那种“人类工程师手工优化”的风格而不是编译器那种保守的策略。再补一个细节PTX 层面有一条和 SASS 不完全对应的指令缝隙。很多优化在 CUDA C 源码层面写不出来或者写出来也会被编译器改掉但 PTX 可以精确表达。比如某些mma.sync张量核心指令的组合、显式的cp.async异步拷贝、手动控制寄存器重用这些在 CUDA C 里要么没有标准语法要么会被创始编译器直接改写。LLM 直接在 PTX 层面操作等于弯道超车绕过了高层抽象带来的信息损失。2.3 目标选 PTX 而不是 SASS 的三个理由理由很简单论文里其实隐含了三条硬约束第一数据可得性。PTX 有官方文档NVIDIA 编译器也可以把内核编译成 PTX 输出训练语料充足。SASS 没有公开规格只能靠反汇编工具反向猜测语料质量差很多。第二验证可行性。生成一段 PTX 之后可以用ptxas编一编能过就说明语法没错再放到真实的 GPU 上跑一跑比对数值结果就知道功能对不对。SASS 的验证就只能靠模拟器难得多。第三优化空间存留度。CUDA C 源码离硬件太远很多优化在高层已经定死。PTX 则保留了寄存器分配、指令调度、内存访问重排这些关键优化机会LLM 的输出还有大显身手的空间。换句话说PTX 是整个链条里“写起来最简单、优化空间又最大”的位置。3. 论文的实验证据直接让 LLM 写 PTX 到底行不行3.1 实验设计把“生成”和“验证”拆开论文的实验设计非常务实它没有天真地让 LLM “一口气写完整个 kernel 然后祈祷它能跑”而是把流程拆成了两段。第一段是生成模型根据输入的函数签名、内存布局、性能约束输出一段完整的 PTX。第二段是验证生成的 PTX 会经过三关语法检查ptxas编译、功能验证小规模输入在 GPU 上跑对比数值、性能基准测运行时间。只有全部通过的候选才会进入最终统计。这个设计其实暗合了最近行业里“LLM as judge”和“基于 LLM 的单元测试”的思路——模型负责创造传统工具负责把关。论文的做法不是让 LLM 取代编译器而是让 LLM 扮演候选代码生成器编译器退化成验证器和最终翻译器。这一步我认为是最聪明的既发挥了 LLM 的生成能力又避开了它“偶尔胡说八道”的毛病。为了保证公平论文在预训练阶段给模型喂了大量真实世界的 CUDA 内核、PTX 手写片段、ptxas 编译日志。尤其是那些“手写 PTX 比编译器默认输出快不少”的案例模型会从中学习到关键模式比如怎么用向量化指令ld.global.v4.f32替代四个标量加载怎么通过重新组织线程映射来减少 Bank 冲突这些模式在 CUDA C 层面很难表达但在 PTX 里是基本功。3.2 性能结果在哪些场景里真的赢了默认编译虽然我拿不到论文的完整实验附表但从摘要和思路推断这种方法的对标目标是同一份计算任务用常规 CUDA C 写然后交给 ptxas 优化对比 LLM 直接生成的 PTX。在几类典型 kernel 上LLM 生成的 PTX 有很大概率超过默认编译矩阵乘法和 GEMM 变体因为手写 PTX 可以精确控制mma.sync张量核心指令的排布而编译器从 CUDA C 出发很难生成最优的张量核心组合FlashAttention 类访存密集 kernel关键在于显式管理寄存器重排、异步拷贝和屏障同步LLM 可以从训练数据中学会更紧凑的调度归约类 kernel通过精确控制线程束内 Shuffle 指令避免共享内存冲突编译器后端往往选择保守方案为什么默认编译器会输核心原因是 ptxas 必须保证所有输入程序都能在合理时间内编译完所以它的优化策略是通用的、保守的。而 LLM 可以针对具体任务生成“特化代码”为了某一个 kernel 不计代价地折腾调度。这个优势就像是“定制西装”和“大牌成衣”的差别——成衣覆盖大多数身材但定制能做到每一处都贴合。3.3 正确性怎么兜底前端、验证器和模拟器各干各的性能再漂亮正确性兜不住就是空中楼阁。论文里最让人放心的是它对验证环节的处理。LLM 生成的 PTX 不会直接用于生产而是经过一个多层漏斗结构验证ptxas 能编译通过说明指令语法和寄存器分配对基本结构没爆功能验证把 kernel 放到实际 GPU 上跑用小输入和 CPU 参考结果做逐元素比对边界压力用不同输入尺寸、不同数据分布跑一遍确认没有因为寄存器溢出或者线程同步出错而崩溃这一套流程其实和现在大厂做 AI 代码生成的思路高度一致。你不能指望 LLM 每次都能输出完美代码但你可以把它的输出当成“候选方案池”用工程手段去筛选和兜底。论文的价值就在于证明了一件事如果验证闭环做得足够强LLM 的 PTX 输出即使不是 100% 正确也可以作为高效的优化起点让人类工程师从零开始手写变成“改一改就能用”。4. 绕开后端后整个工具链会发生什么连锁反应4.1 编译器不再是“优化师”而是“翻译官安全审查员”如果 10% 甚至 1% 的高性能 kernel 能由 LLM 直接生成 PTX编译器在 GPU 软件栈中的角色就会发生微妙的漂移。现在的 ptxas 是一个什么都管的“管家”优化、调度、分配寄存器、生成机器码全是它说了算。绕开后端后传统编译器剩下的工作就只有两件把 LLM 生成的 PTX 翻译成 SASS以及做安全边界检查——检查有没有非法内存访问、同步漏洞、寄存器溢出。这个转变不可小觑。它意味着“编译器优化”这个概念的权威性会被稀释。过去我们说“编译器比你懂硬件”以后可能会变成“模型比你懂这条特定指令序列的排列”。编译器的定位从“决策者”滑向“执行者”这个变化很可能让 LLVM 的 GPU 后端不再是唯一的性能入口。4.2 从 Pytorch 到 TensorRT上层框架的优化空间变了对普通 AI 工程师来说最直接的体会会来自框架层的优化逻辑。现在的torch.compile或者 TensorRT 做的事情是把 PyTorch 的计算图优化成高效的 kernel其中大量工作是在“如何找到或者生成一个更好的 CUDA kernel”。如果 LLM 能直接写 PTX那么框架层的优化器就不再需要维护一个巨大的“kernel 模板库”——它只需要保存计算意图然后调 LLM 现生成一套对应硬件的最优 PTX 序列。这会带来一个很现实的连锁反应kernel 特化成本会大幅下降。现在的做法是为每一类 GPU 架构手工调优比如 A100 一套H100 一套下一代 Blackwell 再重调。如果 LLM 生成 PTX 足够快那么针对新架构的特化就只是多一次推理的问题。硬件厂商的软件栈迭代周期可能从“几个月”压缩到“几天”。4.3 硬件厂商的软件护城河会不会被动摇这个问题的答案很微妙。一方面编译器后端确实是硬件厂商的护城河因为硬件细节藏在 ptxas 和 SASS 里第三方很难窥探。但另一方面LLM 写 PTX 并没有绕过硬件本身它只是绕过了“那套不公开的启发式优化逻辑”。如果模型在训练时见过足够多的高性能 PTX 片段它能在不接触 SASS 内部编码的情况下反向猜测硬件的偏好然后用 PTX 层面的指令选择去契合它。这等于把原本锁在编译器后端的“内功”转移到了模型参数里。硬件厂商有三种应对路径第一是开放更多 PTX 级别的内建函数和性能指标让模型更容易生成高效代码第二是收紧 PTX 的兼容性让 ptxas 在执行翻译时更强的重排权抵消 LLM 的指令优势第三是干脆拥抱这个趋势——因为 LLM 生成的特化代码最终还是要跑在自家 GPU 上性能越好越能卖卡。从我个人的观察看第三条路看似被动回避但长期反而最符合硬件厂商的商业利益。5. 现实约束与我的延伸思考5.1 LLM 写 PTX 的几道硬伤论文归论文落地是另一回事。LLM 直接生成 PTX 最大的硬伤在于“长序列生成的不稳定性”。一个完整的 kernel 展开成 PTX 会有几百条甚至上千条指令模型在超长序列生成过程中很容易把前面的状态忘掉导致后面的指令和前面的资源分配对不上。这个问题的解决方案通常是把 kernel 切块生成然后手动拼装但拼接本身也可能引入同步问题。再一个是语义安全性。PTX 没有类型系统没有内存安全保护模型一旦生成错误的共享内存索引或者漏掉bar.sync程序可能静默算出错误结果这种 bug 在并行程序里最难排查。所以我现在依然坚持一个观点纯 LLM 生成 PTX 要进入生产环境必须有强验证层做辅助单靠模型自身的高质量输出还不够。最后一个硬伤是数据债。LLM 的性能上限取决于训练数据里有多少“高质量的手写 PTX”。但现实是绝大多数公开代码库里的 PTX 都是编译器 dump 出来的“教科书文本”真正由人手工优化的杰作寥寥无几。如果模型学到的全是默认编译的风格那它生成的东西也只是“另一个 ptxas 而已”谈不上超越。论文要解决的不是模型架构问题而是数据稀缺问题。5.2 不是所有后端都值得被“绕开”这篇文章的题目很有煽动性但“绕开编译器后端”这句话必须加上上下文限制。PTX 之所以适合 LLM 生成是因为它是一个公开、稳定、语义清晰的虚拟指令集而且优化空间巨大。但这不代表所有后端都适合这个玩法。比如 CPU 的 x86 后端它的指令集更杂、优化器更成熟、寄存器分配约束更复杂而且几十年的坚持优化已经把大部分场景的收益压得很低LLM 很难找到一个“明显的缝隙”去超越 GCC 和 LLVM。再比如针对 FPGA 的 HLS 工具链后端涉及物理布线这样的空间问题根本不是 LLM 的擅长的符号序列建模能覆盖的。只有像 PTX 这种“指令语义高度规整、收益空间大、且没有一家独大的手动优化经验”的目标才是真正的切入点。5.3 未来更可能的落地形态混合优化我个人的判断是未来三年内你不会看到“编译器后端整个消失”但你会看到它的决策权被一点一点蚕食。最可能的落地形态是混合优化传统编译器负责做 95% 的安全翻译和基础优化LLM 在关键热点区域比如循环内层、访存密集段、张量核心映射生成 PTX 替换方案然后通过 A/B 测试看哪个更快。这个思路其实已经接近“LLM 辅助的自动向量化”和“基于 LLM 的单元测试”的交叉地带。比如编译器先在 IR 层做常规优化再把最耗时的几段代码暴露出接口让模型在 PTX 层去“手写”更好的实现。这样既控制了验证风险又保留了模型灵活调度的优势。另外LLM 还可以反向辅助编译器开发——让模型从大量旧版 SASS 反汇编中学习硬件偏好帮 ptxas 改进指令调度策略这个方向比直接生成 PTX 更容易落地也更容易量化收益。说实话我第一次看到这篇论文的标题时也觉得是噱头但拆解完它的论证链条之后反而觉得这个方向有真实的工程潜力。PTX 是那个恰到好处的“靶子”人写太痛编译器写太保守而 LLM 恰好站在中间既懂人类意图又能输出精确指令。如果你手头有高性能 kernel 优化的需求建议别急着推翻现有链路先拿一个小例子让模型试着生成 PTX再用传统编译器的结果做对比。这个实验做一遍你就能直观感受到那层“绕开的差距”到底有多大。
RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进