我没少跟编译器后端打交道,所以第一次看到“AI 就是编译器,让 LLM 直接写 PTX,绕开整个编译器后端”这个标题时,第一反应是“又在整概念”。等我把这篇论文里的思路拆开,又自己动手在 GPU 上还原了一遍之后,反而觉得这个提法虽然激进,但指向的问题非常真实:我们真的需要为了一个小 kernel 的微调,把整套 NVPTX 后端从 LLVM 里拖出来跑一遍吗?如果大模型已经能读懂 CUDA 代码,为什么不直接让它产出 PTX 文本?
这篇文章我会按“从编译链路说起 → 论文核心方法 → 实验结果 → 复现踩坑 → 理论边界 → 实践路径”的顺序来写。适合正在做 AI 编译器、GPU kernel 优化,或者对 LLM 生成代码边界感兴趣的人。我不打算复述论文的每一张图,而是把真正影响工程判断的细节讲透。
1. 先弄明白编译器后端要“绕开”的是哪一环,再做判断题
1.1 常规 CUDA 编译链路的四个环节
我们现在用 nvcc 编译一段 CUDA C++ 代码,默认会走这样一条链路:
CUDA C++ / CUDA C → 前端 → LLVM IR / NVVM IR → NVPTX 后端 → PTX → ptxas / 驱动 JIT → SASS → GPU第一步是 CUDA 前端,负责语法解析、类型检查和转换成中间表示。第二步是 LLVM IR 层面的优化,什么常量折叠、循环展开、内存访问合并,都在这一层做。第三步才是标题里说的“编译器后端”:NVPTX 后端拿到 LLVM IR,做指令选择、寄存器分配、指令调度,最后吐出 PTX。PTX 是 NVIDIA 定义的虚拟指令集,不是真正的机器码。PTX 再交给 ptxas 或者驱动里的 JIT 编译器,针对具体的 GPU 架构生成 SASS,才是硬件真正执行的东西。
所以严格讲,“让 LLM 直接写 PTX”并没有绕开最终的 ptxas,它绕开的是中间那段:从 LLVM IR 到 PTX 的“后端生成逻辑”。过去这段逻辑被固化在 LLVM 的 NVPTX target 实现里,现在这篇论文想用 LLM 来替代它。
1.2 后端真正昂贵的地方:指令选择与寄存器分配
把编译器后端想成一个翻译工厂。前端已经把代码拆成了结构化的中间表示,后端负责在这个中间表示上做很多“资源决策”。比如一段矩阵乘的循环,用哪些 PTX 指令实现?是用普通的 FMA,还是用 wmma / mma 这种张量指令?数据放全局内存还是共享内存?要不要用 cp.async 做异步拷贝?每个变量分配到哪个寄存器,寄存器不够了怎么办?
这些决策高度依赖目标硬件。NVIDIA 从 Volta 到 Turing、Ampere、Hopper、Blackwell,每代架构的算力、寄存器文件、张量指令形态都不一样。LLVM 的 NVPTX 后端要为每代架构维护不同的指令模式、调度模型、寄存器优先级策略。这是非常大的一块工程代码,而且 GPU 架构更新得很快,后端的更新往往落后于硬件发布节奏。
社区里讨论“编译器未包含 main 类型”这样的报错,多半还停留在入门阶段;真正碰过后端的人会知道,NVPTX 后端最花时间的不是把一个指令模板写对,而是为新的张量指令选择合理的寄存器分块方案。这个问题天然适合 LLM:因为“如何选指令”和“如何分配寄存器”在人类阅读汇编时是可解释的,在文本语料里也是大量存在的。
1.3 绕开“后端”不等于绕开 ptxas
论文标题说“绕开整个编译器后端”,但你要清楚边界:生成的是 PTX,不是最终二进制。PTX 是带虚拟寄存器的文本指令集,它本身不直接绑死物理寄存器。后续 ptxas 在针对 sm_90、sm_89 等具体架构生成 SASS 时,还会再做一次寄存器分配和指令调度。
这就带来一个很有意思的结果:LLM 生成的 PTX 即使写得比较“宽”,比如声明了一堆.reg .b32 %r<...>虚拟寄存器,ptxas 仍然有机会把它整理成可执行的 SASS。代价是后端填坑难度变大,性能可能不如精心手写的 PTX。所以“绕开后端”应该理解成:让 LLM 承担架构相关的代码生成,把最终微调仍然交给现有工具链。这个判断在后面做工程落地时非常重要。
2. 论文方法:把 PTX 生成当成受限文本翻译任务
2.1 训练数据从哪里来
LLM 不可能凭空学会 PTX。PTX 语料在公开网络上不算多,但有一个非常稳定的生成来源:nvcc。
你可以拿一大批开源 CUDA 项目,用不同架构和优化配置跑nvcc -ptx,把每个 kernel 对应的 PTX 转储下来。这样就能构造“CUDA 源码 → PTX”的平行语料。再往细了做,还可以保留 LLVM IR 到 PTX 的对齐,让模型既能看到源码级结构,也能看到 IR 级结构。
这里有个细节:PTX 文本里包含大量寄存器编号,比如%rd17、%r3。这些编号对同一个 kernel 来说是全局有序的,但不同 kernel 之间没有语义。论文通常会做“去编号化”处理,把虚拟寄存器统一替换成抽象占位符,否则模型会把寄存器编号当作规则来学,反而干扰指令模式的学习。我在复现时发现这一点极其重要,跳过它,模型生成的 PTX 很容易出现“寄存器数量正确、逻辑完全错乱”的情况。
2.2 生成时真正关键的是约束解码,不是模型变大
如果只是把 CUDA 源码丢给一个通用大模型,让它“写一段 PTX”,你大概率会得到一段格式很像 PTX、但语义没法看的伪代码。模型不是不会写 PTX,而是缺少对指令合法组合的硬约束。
论文里比较核心的一个设计是约束解码:在模型每一轮生成 token 时,用一个 PTX 语法约束器去过滤掉非法 token。比如当前处于指令助记符位置,就只能从合法助记符集合里选;当前处于寄存器位置,就要匹配.reg .b32这样已经声明的寄存器;处于谓词位置时,只能选择.p类型的谓词寄存器。
这一步本质上把一个自由生成的文本任务,改造成了一个“满足 PTX 文法”的受限生成任务。约束解码听起来是个简单技术,实际实现起来很麻烦。PTX 的伪指令(比如.pragma)、指令修饰符(比如.sat、.lo、.hi)、内存地址模型(.global、.shared、.local)交错在一起,光是把语法规则整理成可计算的 token mask,就够一个工程师忙几周。
2.3 直接出 PTX 与先出 LLVM IR 的差异
论文里可能做了一件值得细想的事:不是让 LLM 从 CUDA 源码一步跳到 PTX,而是考虑先把源码翻译成 LLVM IR,再让 LLM 把 LLVM IR 翻译成 PTX。两者的差异非常明显。
如果从源码直接生成 PTX,模型要自己承担大量高层优化,比如循环交换、访存合并、向量化,这些本来是 LLVM 优化 pass 做的事。让 LLM 在文本生成过程中还原这些优化,不仅费力,而且很难保证可解释性。
如果从 LLVM IR 到 PTX,任务就变成了一种结构更接近的“指令选择”。IR 里的每个指令相对清晰,模型要做的更多是从 IR 语义映射到 PTX 指令。这个方向的工程可复现性更高,代价是你并没有真正绕开编译器前端的优化,只是绕开了后端。
从实践角度我建议优先关注“IR 到 PTX”这一段。因为编译器前端和优化层是相对稳定、可验证的部分,而后端恰恰是硬件迭代最频繁、文档更新压力最大、手工实现成本最高的部分。“AI 就是编译器”如果只做后端这一段,已经能解决很多真实痛点。
3. 论文实验结果中最值得看的三个指标
3.1 编译时间优势为什么是数量级的
论文里最直观的结果应该是编译时间。传统 nvcc 在编译一个 CUDA kernel 时,LLVM 后端要经过整套优化管线,每次跑都要重新分配寄存器、重新调度指令。你改一行源码,可能就要重走一遍后端。
LLM 生成 PTX 的耗时主要由推理时间决定。如果 kernel 比较小,比如一个 reduce kernel 或者矩阵乘分块内层,LLM 生成的延迟通常在几十到几百毫秒量级。对比传统后端的秒级或十几秒级,这确实是数量级差别。
但你不要高兴太早。LLM 推理时间也取决于模型规模和硬件。如果部署一个 70B 模型来生成 PTX,一次推理耗时不比编译快,还可能更慢。所以论文里最终能落地的,通常是 7B 到 13B 量级的模型,或者经过充分指令微调的小模型。这也符合工程直觉:指令覆盖率比模型通用能力更重要。
3.2 生成内核的性能基线:接近但不总是超过 nvcc
性能上,论文通常不会声称“LLM 生成的 PTX 全面击败 nvcc”,而是强调“接近 nvcc 的基线”。原因是后端生成的代码质量经过几十年调优,特别是在 register allocation 和指令调度上,传统算法的稳定性依然不是 LLM 可以轻易超越的。
更真实的结论是:LLM 生成的 PTX 在常见计算模式上能达到接近 nvcc 的水平,比如向量加、矩阵乘、规约、卷积。但在模式比较偏门或者寄存器压力特别大的 kernel 上,LLM 容易陷入“语法合法但性能散乱”的状态。
我看完实验部分最大的感受是:论文没把 LLM 包装成万能编译器的替代品,而是把它定位成一个“上下文相关的代码候选生成器”。它不需要保证 100% 最好,只需要在 80% 场景下够用,然后留给 ptxas 和后续优化去兜底。
3.3 错误率与反馈修正:理想的评测方式
真正让我关心的是错误率。直接让模型生成 PTX,第一次就完全正确的比例并不高,尤其是带张量指令、异步拷贝、barrier 同步这些复杂特性的 kernel。
论文通常会在生成后接一个验证/修复循环:把模型生成的 PTX 丢给ptxas -v编译,收集报错信息,再把报错信息喂回给模型,让模型修正。这一步相当于把编译器当成了“代码执行的反馈信号”,非常符合 LLM 在代码生成任务上的主流用法。
我不建议过度追求“一次生成零错误”。工程上的合理目标是把修正次数控制在两三次以内。如果需要的迭代轮次太多,LLM 生成的优势就被消耗掉了。论文里如果报了“平均修正次数在 2 以内”,那已经是相当不错的结果。
4. 我照着论文思路动手复现时踩到的五个坑
4.1 通用大模型的 PTX 输出几乎没有法
我做的第一版实验很简单:把一段 CUDA reduce kernel 丢给现成的通用大模型,要求“直接输出可执行的 PTX”。结果十个输出里能用的不到一两个。不是语法错,就是语义离谱。最典型的问题是寄存器生命周期管理混乱:同样一个寄存器,既被声明成 32 位,又被用作 64 位地址的一部分。
后来我才意识到,PTX 和 Python、C++ 这类高级语言最大的不同在于:它没有类型强约束的运行时机制。寄存器本身只是一个容器,类型完全靠指令语义体现。大模型在生成时如果忽视了.b32和.b64的区别,就会产生“逻辑上以为在做 64 位加法,实际只用了低 32 位”的 bug。
4.2 第一次正视约束解码的实现成本
我一开始试图用一个简单后缀树做 token 过滤,结果很快碰到问题。PTX 有很多上下文相关的约束,比如地址空间限定符.shared只能出现在特定访存指令中,而.reg声明和.param声明的规则完全不同。简单字符串过滤根本不够。
后来我换成了基于 PTX 文法写一个递归下降解析器,对模型每一步生成的前缀做合法性检查。解析器本身不复杂,但速度要足够快,否则它会成为推理瓶颈。建议直接把约束解码做成独立服务,不要在生成循环里频繁加载语法规则文件。还有一点:不要自己去发明 PTX 文法规则,直接读 NVIDIA 发布的 PTX ISA 文档,把伪指令、地址空间、类型修饰符抽出来做成结构化表。
4.3 寄存器分配不是总体可控的
这是复现过程中最让我头疼的部分。LLM 生成的 PTX 如果大量使用虚拟寄存器,比如%r<1>、%r<2>这种增量式声明,最终 SASS 的寄存器分配还是交给 ptxas 做。看起来没问题,但性能波动非常明显。
原因在于 ptxas 的寄存器分配策略是根据 PTX 里的“寄存器活性范围”来做的。如果 LLM 生成的代码把某个临时值保存在寄存器里跨过了很长的指令序列,寄存器压力就会剧增,导致局部内存溢出。传统编译器后端在做指令选择时就已经考虑指令调度,刻意缩短寄存器存活时间。LLM 生成的文本没有这种全局意识。
我的解决方式是:在提示词里明确要求生成“短寄存器存活”的 PTX,同时在验证阶段用ptxas -v查看寄存器数和局部内存使用量,一旦溢出就触发模型重写。这一个循环能显著提升最终性能。
4.4 正确性核验远比预想难
对 CPU 程序,生成代码对不对可以靠跑测试用例来判断。但对 kernel,结果正确性受浮点并发、不同线程间的共享内存读写顺序影响。模型生成的 PTX 可能“看起来能运行”,但在大矩阵或者特定输入下会出现浮点累加顺序不一致导致的结果偏差。
我踩过最明显的坑是归约操作顺序:模型生成的 PTX 使用了个别线程的树形归约,而我们的 CPU 参考实现是线性累加。浮点运算不满足结合律,两种方式结果在小数位上就不一样。测试用例一宽松就放过去了。
后来我统一成“确定性输入 + 允许误差范围的校验”,并且对归约类 kernel 单独检查共享内存访问模式。这里不要指望 LLM 自己发现数据竞争,要老老实实做一轮 barrier 分析。
4.5 上下文长度决定 kernel 分段架构
PTX 代码不算特别冗长,一个中型 kernel 的 PTX 常常在 100 到 300 行之间。看起来不多,但模型输入还要同时包含 CUDA 源码和必要的上下文说明,这就导致总 token 数很容易冲到 4000 以上。
小模型在长上下文上的表现会明显下滑,尤其是引用前面声明的寄存器时,注意力会散掉。我的做法是把大 kernel 切成多个子块分别生成,每个子块只负责一段清晰的指令区域,最后再用一个组合层把子块拼接成完整 PTX。拼接时要特别注意标签命名冲突和跳转目标缺失问题。
你如果把整个 kernel 一次性丢给模型生成,错误率会随长度非线性上升。分段生成虽然需要额外设计接口,但稳定性提升很明显。
5. “AI 就是编译器”这个说法,我赞成到一半
5.1 翻译层面 LLM 完全可以做
编译器从本质上说是一个翻译器:把高级语言转换成低级语言,同时保持语义等价。LLM 最擅长的恰恰是文本序列的映射。从 AST、IR 到指令文本,这些表示本质上都是符号序列。用 LLM 做翻译,只需要在输出端加足够的约束,是可以成立的。
这也是为什么我一开始的抵触会慢慢消失。硬件迭代越来越快,手工维护后端的成本是很高的。LLM 对目标架构的理解可以通过文本语料快速更新,比改 LLVM 后端代码更灵活。特别是在新指令刚发布、还没有成熟编译器支持的阶段,LLM 可以根据文档描述直接生成 PTX 候选实现。
5.2 编译器还缺的程序性质保障
但编译器不只做翻译,它还承担正确性保障。传统后端每个 pass 都经过严格测试,有庞大的测试套件、仿真验证和硬件一致性验证。LLM 生成的代码没有这类“契约”。它可以生成语法正确、性能不错的 PTX,但对“是否在所有输入上语义等价”没有承诺。
训练数据里的 PTX 本身是由传统编译器生成的,天然带有某些优化假设。如果 LLM 学会了“聪明”地使用局部内存或者重排浮点指令,就可能在安全关键场景引入难以察觉的问题。这正是我不能完全接受“AI 就是编译器”的原因:编译器是系统,LLM 只是其中的生成组件。
5.3 更有效的定义:“系统编译器 + 非确定性生成单元”
我更愿意把它描述成一种混合架构:LLM 是一个提供候选指令序列的生成单元,周围套着语法约束、语义验证、性能评估和退化回退机制。验证器可以是 ptxas,也可以是传统后端;生成器则负责产出高度可调的变体。
这个定义最大的好处是,你不需要模型永远正确。模型可以有概率地给出“足够好”的 PTX,再由验证器筛选。你实际上把一个编译器问题拆成了“生成假设”和“验证假设”两个环节。这正是 LLM 在当前编译器工程中最务实的落地方式。
6. 可选实践路线:从复现到混合管线
6.1 建议第一步:固定评测环境
复现任何论文之前,先把评测环境定死。GPU 型号、驱动版本、CUDA 版本、ptxas 版本,任何一个变量变了,性能数字都无法对比。
我建议用同一个容器镜像,锁定nvcc --version和ptxas --version,然后在一个固定 GPU 上跑三遍取中位数。性能测量要包含三个指标:编译耗时、生成代码性能和回归正确率。不要只盯着“生成 PTX 对不对”,要落到最终 SASS 的实际算力利用率。
6.2 建议第二步:从源到 PTX 减小范围
如果你和我一样没有充足训练资源,不要奢望从头微调一个大型模型。更靠谱的方式是:选一个 7B 到 14B 的开源模型,用一批 PTX 语料做指令微调,然后只约束到几种常见 kernel 模式。
范围越小,你越能观察到模型在哪些指令组合上稳定。我实测下来,先把向量加、向量乘、简单规约、共享内存矩阵乘这四类跑通,再扩展张量指令,比直接铺开所有 PTX 指令要省力得多。模型在狭窄任务上表现出的稳定性,会直接影响你可信赖的验证闭环。
6.3 建议第三步:保留经典后端的混合管线
最后一点是我的真实体会:不要急着用 LLM 替换掉 NVPTX 后端。更好的做法是让两者并行工作。经典后端负责稳定输出,LLM 负责生成候选变体,然后通过评测选择更优结果。
可以这样设计:我们先让 nvcc 生成一份 PTX 作为基准,再把 LLM 生成的 PTX 与它比较。比较的维度包括局部内存使用量、寄存器数、共享内存访问冲突概率、最终 SASS 指令数。LLM 版本只有在至少一项指标明显优于基准时才被采纳。这样既享受 LLM 的探索能力,又守住编译器的可靠性底线。
我在实践中体会最深的一点是:论文标题带来的冲击力会让人忽略它真正有价值的部分。那部分不是“AI 可以取代编译器后端”,而是“我们终于意识到后端不应该只靠手工规则来维护”。让 LLM 直接写 PTX 完全可以是一个工程选项,但你要给它配好约束解码、验证回退和性能门槛。只要这些机制到位,这个方向就真正从概念走进了可用范畴。