ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

AI写PTX:让LLM绕开编译器后端,挑战LLVM NVPTX

AI写PTX:让LLM绕开编译器后端,挑战LLVM NVPTX 1. 这个命题的第一印象掀掉编译器后端凭什么昨晚睡前刷到个标题一下子就清醒了AI 就是编译器——让 LLM 直接写 PTX绕开整个编译器后端。 乍一看像标题党再琢磨一下又觉得这可能是接下来一两年编译器方向最值得盯的议题。大家都知道 AI 编译器最近有多热但绝大多数工作要么让 LLM 帮编译器做算子匹配要么在 IR 层做优化真正敢把 LLVM 后端的指令选择、寄存器分配、指令调度这些脏活直接交给生成模型这个命题算是把话说到了底。这篇论文具体投在哪、作者是谁并不是我想在这里复述的重点我下面要拆的是题目背后的核心主张顺带夹带点我写编译器、调 CUDA kernel 的个人看法。想真正读懂它得先知道一件事传统编译器后端不是锦上添花的部分而是把一个人能读的中间表示变成机器能跑的指令的关键路径。绕开它等于把编译器里几十年积累的启发式优化全掀了让 LLM 用统计方式重新回答一段计算应该变成什么指令。1.1 传统编译器后端到底在干什么先用最通俗的话把这套流程说清楚。假设你有一段 CUDA C 代码比如向量加法。编译的第一步叫前端负责词法分析、语法分析、类型检查最后产出一个中间表示也就是 LLVM IR第二步叫中端在 IR 上跑循环展开、常量传播、公共子表达式消除这些优化第三步才轮到后端它要把优化完的 IR 映射到目标指令集里面包含指令选择、寄存器分配、指令调度、窥孔优化、发射代码等等。对 NVIDIA GPU 来说传统后端的输出就是 PTX。你平时用nvcc -archsm_80编译一个 CUDA kernel它内部实际是先把 C/C 转成 LLVM IR再用 NVPTX 后端把 IR 变成 PTX 文本最后把 PTX 丢给驱动里的 ptxas编译成当前 GPU 真正执行的 SASS。所以让 LLM 直接写 PTX绕开整个编译器后端这句话严格来说绕开的是 LLVM 的 NVPTX 后端那个环节ptxas 这一跳你绕不开也不应该绕因为它是把 PTX 变成 sm_80、sm_90 各自可执行代码的最终适配器。1.2 为什么不直接让 LLM 写 SASS看到这里你可能会问既然都要过 ptxas干嘛不干脆让 LLM 直接生成 SASS这个问题我在这个方向的工作里反复想过答案是现阶段不现实。SASS 是真正绑定到具体微架构的指令集合sm_80 和 sm_90 的编码方式、可用指令、寄存器压力特征都不一样而且 SASS 通常以二进制 cubin 形式存在你拿来做训练也好、做验证也好都得先反汇编。PTX 则是 NVIDIA 提供的一个虚拟 ISA它足够低层保留了并行语义和访存意图但又不至于低到每个 GPU 世代都要重新学。更关键的是PTX 允许你使用虚拟寄存器寄存器分配这种最麻烦的活可以交给 ptxas 再做一遍。这对 LLM 来说是极其友好的模型不需要像真正的后端那样精确管理几十个物理寄存器它只要生成逻辑上正确的文本剩下的物理资源分配由确定性工具兜底。这个性质我认为正是这个方向能够成立的真正理由。2. PTX 到底是个什么东西看不懂它就不知道这个方向有多大胆如果你平时只写 CUDA C可能对 PTX 的印象就是编译中间产物看都没看过几行。真要评估AI 写 PTX这个命题你得先知道 PTX 的语法形态和它的能力边界。2.1 一段 PTX 长什么样先看一段示意性的 PTX说明我不是在空谈。真实的 PTX 通常长得比较啰嗦函数名带完整的参数签名我下面做了简化.visible .entry reduce_kernel( .param .u64 in_ptr, .param .u64 out_ptr, .param .u32 n ) { .reg .u32 %idx; .reg .u64 %addr; .reg .f32 %val; .reg .f32 %acc; mov.u32 %idx, %tid.x; cvta.to.global.u64 %addr, %in_ptr; // 把参数地址转成全局地址 ld.global.f32 %val, [%addr %idx*4]; add.f32 %acc, %acc, %val; // 示意真实要有 init 和 sync st.global.f32 [%addr %idx*4], %acc; }这段代码单独拿出去肯定不能通过 ptxas 的严格检查它只是用来展示语法形态的示意片段。真正想看的可以拿一段简单的 CUDA kernel 跑nvcc -archsm_80 -ptx你会在文件头看到.version 8.0、.target sm_80、一堆.reg .b32、.reg .f64之类的东西。熟悉这些之后你会发现PTX 本质上是文本化、带虚拟寄存器、以 SIMT 线程模型为中心的汇编。2.2 PTX 的半虚拟属性正好降低了 LLM 的下场门槛我盯着 PTX 看了很久最终得出一个判断PTX 是这个方向唯一合理的中间目标。原因有四个第一它是文本格式天然适配语言模型第二它允许虚拟寄存器模型不需要在一开始就解决物理寄存器分配的 NP 难题第三它有明确、公开、相对稳定的指令手册指令数量和编码规律比 SASS 有迹可循第四它的最终正确性由 ptxas 和执行结果做双重把关模型可以在这个闭环里迭代。当一个语言模型拿到 LLVM IR要它输出一段 PTX本质上是在做一个跨语言翻译任务只不过目标语言换成了带并发语义的虚拟汇编。这和当年 LLM 做代码生成、做反编译面临的困难有相似之处也有更麻烦的地方IR 里的每一条指令语义和 PTX 的指令语义并不是一一对应的中间隔着大量优化决策。2.3 一张表看懂 PTX 与 LLVM IR、SASS 的差别很多读者会把这三种东西混在一起我用一个表把它们的关系摆清楚维度PTXLLVM IRSASS本质NVIDIA 虚拟指令集编译器中间表示真实 GPU 微架构指令形态文本.ptx文本/bitcode.ll/.bc二进制.cubin或汇编目标由 ptxas 编译为 SASS由后端转换成目标指令直接可执行寄存器虚拟寄存器为主SSA 虚拟寄存器物理寄存器/固定编码稳定性跨代较稳定与编译器版本绑定每个微架构都可能变是否适合 LLM 直接生成非常适合目前已经有大量尝试难度最高还没人敢当主路从这个表格可以看出PTX 恰好卡在足够低层和足够文本化的交界处。它既不像 LLVM IR 那样抽象需要模型理解大量类型系统和控制流语义也不像 SASS 那样琐碎需要精确到每个机器周期。这让LLM 直接写 PTX变成一个虽然难、但边界清楚的生成任务。3. 如果真让 LLM 去写 PTX系统该怎么设计顺着题目的思路下面我要谈的是我眼里这套系统最可能、也最合理的工程形态。注意这部分是按这个标题反推出来的设计不是论文原样的复述但我觉得任何一个认真做这个方向的人最后都会走到这些模块跟前。3.1 输入我更看好 LLVM IR而不是源代码标题说的是绕开编译器后端那说明前端大概率还是保留的。如果是直接拿 C/C 源代码让 LLM 生成 PTX那绕开的就不只是后端还包括整个前端语义跨度大太多模型要同时解决解析、类型推导、内联决策、并行语义分析这基本是把一个不可能任务塞给生成模型。把输入限定在 LLVM IR意味着你用 Clang 或 NVCC 先把源代码变成结构化的 IRLLM 只需要负责后端映射混进问题里的变量从全世界所有 C 语法缩小到规则化、带类型信息、SSA 形式的 IR。这个选择还有一个实际好处语料容易造。你可以用同一份 CUDA 源文件分别让工具链产出 LLVM IR 和 PTX它们天然形成一对平行语料用来做指令微调。如果你坚持源代码到 PTX 的映射反而要处理大量模板代码、宏、头文件展开数据对齐成本高得吓人。3.2 生成不是重点验证才是现在很多 LLM 编译器工作最大的误区是把注意力全放在怎么生成一段好看的代码上。真做这个方向的人都知道生成只是起点闭环验证才是核心。PTX 有一个非常好的优势它有一个快速、严格的确定性检查器就是 ptxas。只要你的 PTX 语法错了、寄存器宽度对不上、指令不存在ptxas 会直接报错就算 ptxas 过了还可以把生成的 cubin 加载到 GPU 上跑真实数据跟 host 端的参考结果做对比。所以系统设计应该长这样第一遍让 LLM 根据输入的 LLVM IR 和少量示例生成 PTX第二遍调用 ptxas 尝试编译第三遍如果编译通过就加载进 GPU 跑若干组随机数据检查数值正确性第四遍任何一步失败都把错误信息拼回 prompt让 LLM 再做一轮修复。代码长得像下面这样import subprocess from pathlib import Path ir_text Path(kernel.ll).read_text() feedback for attempt in range(20): ptx_text llm_complete(promptir_text feedback \n// output PTX:\n) Path(kernel.ptx).write_text(ptx_text) ret subprocess.run( [ptxas, -archsm_80, -o, kernel.cubin, kernel.ptx], capture_outputTrue, textTrue, ) if ret.returncode 0: runtime_err run_kernel_on_gpu(kernel.cubin) if not runtime_err: break feedback \n// runtime error:\n runtime_err else: feedback \n// ptxas error:\n ret.stderr这个闭环里的关键不是 prompt 写得有多花哨而是把 ptxas 和 GPU 当成一个带反馈的习题集模型每猜一次环境就告诉它错在哪。把这个反馈过程规模化以后哪怕 LLM 初始准确率只有 10%也能迭代到能用的程度。我个人觉得这个方向真正的工作量就是把这个闭环的自动化程度做扎实。3.3 语料从哪来老编译器是最好的老师没有语料上面这套设计全是空话。可能的语料来源有这些从 GitHub 收集开源 CUDA/OpenCL 项目用nvcc -ptx批量编译得到源代码 PTX的配对对同一份源码开不同优化参数比如-O2、-O3、-maxrregcount32、不同-arch得到同一逻辑的多种 PTX 表达如果把输入定位为 LLVM IR则先用 Clang 生成 IR再用 NVCC 生成 PTX得到IR - PTX的映射语料对已有 PTX 做有限的程序变换比如交换指令顺序、替换访存指令、调整寄存器分配再把运行时间和正确性记录下来形成改写 性能标注的数据。我特别想强调最后一点。只靠源代码到 PTX的平行语料LLM 学到的只是翻译要让它学会优化还需要大量带性能标签的样本。传统编译器几十年里产出的不同优化决策本身就是最好的监督信号。换句话说老编译器不需要被淘汰它可以先当老师的老师。4. 真按这条路走哪些坑是我看到就会踩的这个方向听起来很性感但如果你真的打算动手复现下面这些坑几乎躲不掉。我一个个说顺序就是它们出现的顺序。4.1 正确性概率式生成和确定性编译器之间的裂缝第一个绕不开的问题是正确性。编译器被信任的前提是确定性和可证明性同一个输入永远产生同一个输出而且有大量测试套件保证语义一致。LLM 不是这样它是从概率分布里采样的可能前一次生成完全正确下一次只是把add.f32写成了mul.f32整个 kernel 就算错了。更麻烦的是ptxas 只检查 PTX 的合法性不检查它和原始 IR 的语义是否一致。所以光有 ptxas 还不够必须做运行时验证。我的建议是测试数据不能只跑一组要跑大量随机向量包含零值、负值、极大极小值、NaN 和 Inf因为这些边界情况最容易暴露指令选错的问题。对涉及共享内存、同步的 kernel还要额外检查 bank conflict 和 race因为这些不一定会导致数值错误但会严重影响性能。4.2 性能能编译通过不等于能打赢 NVCC就算 LLM 生成的 PTX 全部通过验证它可能比 NVCC 默认生成的慢三倍甚至十倍。原因很简单传统编译器后端花了大量时间做指令选择和调度而 LLM 在没有训练到位的情况下生成的往往是结构上正确但优化得很差的指令序列。比如它可能频繁使用局部内存来搬数据导致寄存器溢出可能不使用向量化访存指令比如ld.global.v4.f32可能把一些循环不变的计算放到 block 内部反复执行。所以评估时必须区分两个指标正确率之外还要有 kernel 耗时、寄存器数、spill 数、指令数。NVCC 的-O3输出就是最自然的 baseline。如果 LLM 生成的 PTX 在同样的硬件和输入规模下既正确又只慢 10% 以内那已经是个很了不起的结果。目前我看过的同类型实验大多还没稳定达到这个水平。4.3 幻觉、上下文长度和微调成本第三批坑来自 LLM 本身的工程特性。幻觉是其中最大的一个模型可能发明一个不存在的 PTX 指令比如mad.f64或者奇怪的memory.barrier可能把.reg .u32写成.reg .f32之后手滑用整型加法。好在 ptxas 的报错信息通常非常明确用迭代反馈能压住大部分幻觉但代价是推理次数变多一个 kernel 可能要生成几十次才通过。上下文长度同样头疼。一个中等复杂度 kernel 的 LLVM IR 可能几千行对应 PTX 可能也有几百行到上千行塞进 prompt 之后7B 级别的模型很难在长上下文里保持注意力稳定。解决思路无非是分块生成但分块又会影响寄存器分配和跨块优化这是一个结构性矛盾。至于微调成本一个 7B 模型用几万条 IR-PTX 语料做指令微调至少需要单卡 A100 级别的机器跑一天到几天不等如果你还想要更高性能需要引入 RL 类训练那整体成本又是一个量级。5. 抛开论文我的几条实验建议和判断最后这部分是我真正想写给准备动手的人。别一上来就冲从 IR 生成完整 PTX的大目标我建议你按下面的路径一点点做。5.1 最低成本复现路线先改写 PTX再尝试从 IR 生成第一步先让 LLM 做 PTX 的改写任务。你拿 NVCC 自己编出来的正确 PTX故意引入一些次优写法比如取消向量化访存、增加不必要的寄存器占用然后要求 LLM 把它改回接近 NVCC 的风格。这一步不需要额外造语料验证工具也现成能让团队快速确认 LLM 到底有没有能力理解 PTX 的语义。第二步再做单 kernel 的 IR 到 PTX 生成。选一个很简单的 kernel比如向量加法、归约、矩阵转置这类控制流简单的算子用 Clang 生成 IR 作为 prompt让 LLM 生成 PTX然后走 ptxas GPU 验证闭环。等准确率打满了再逐步增加 kernel 复杂度。第三步才是铺数据和调优。这一步用到的命令非常朴素我直接给你# 用 NVCC 生成 PTX nvcc -archsm_80 -ptx vector_add.cu -o vector_add.ptx # 用 ptxas 把 LLM 生成的 PTX 编成 cubin /usr/local/cuda/bin/ptxas -archsm_80 -o vector_add.cubin vector_add.ptx # 反汇编 cubin看到底生成的是不是合理 SASS cuobjdump -sass vector_add.cubin一个小提醒ptxas一般不在系统默认 PATH 里直接用/usr/local/cuda/bin/ptxas最省事。另外加载 cubin 跑数值对比可以用 CUDA Driver API 写个几十行的小程序也可以用 PyCUDA 快速搭关键是要把 ptxas 的错误信息和运行结果都回灌给模型。5.2 你们真正该比的基线是什么做实验的时候基线设置直接影响结论说服力。我见过不少文章拿NVCC 默认编译当作唯一对比对象这其实不够。除了 NVCC你至少还要对比这三条线LLM 直接生成 CUDA C再用 NVCC 编译这条线可以说明 PTX 介入的价值如果 LLM 写 CUDA 再由编译器优化效果可能比直接写 PTX 更好也可能更差这个对比很关键LLVM 自带的 NVPTX 后端在不同优化级别下的输出-O0和-O3都要当基线因为 LLM 生成的代码如果只比-O0好那并没有证明什么人工编写的 PTX 模板对固定算子类型人工模板往往能把性能和功耗调到很极致LLM 至少要接近这个水平才能谈通用性。比较指标建议用 kernel 耗时、寄存器溢出数、指令数、显存占用、以及失败率。不要只报一个平均加速比它是编译类工作最容易骗到自己的数字。5.3 我对AI 是编译器这个论断的最终看法回到标题这个有点挑衅的说法我的看法是与其把AI 就是编译器理解成LLM 要取代编译器不如理解成编译器这个传统的确定性系统第一次有了一个可以随时对话、可以接受错误反馈、可以按用户意图改写的智能协作者。传统编译器擅长的是在形式化约束下搜索最优解LLM 擅长的是从海量经验中联想和生成候选解。PTX 这个目标语言的出现刚好让后者的联想能力有了一个可以安全试错的沙盒错误被 ptxas 拦下来正确性被 GPU 验证性能被基准测试量化。我在实际动手写这些实验的时候最深的感受是这个方向最大的价值不是让模型写出一个能跑的内核而是把编译器里几十年靠人肉积累的优化直觉放到一个可反馈、可迭代、可自动化的系统里重新学一遍。这篇论文如果真的只让人记住一句话我更愿意记住它证明了编译器后端的搜索空间可以被生成模型正面挑战——剩下的就看后面的人怎么把正确率和性能打磨到可以交付的水平了。
返回列表