LLM直接生成PTX汇编：跳过编译器后端的AI编译新范式

1. 这篇论文到底想干什么：把编译器后端整个拿掉

第一次看到“AI 就是编译器”这个说法，我的反应是：又是一个标题党。但把论文翻完，我发现它讲的事情其实非常具体——让大语言模型直接输出 PTX（Parallel Thread Execution）汇编，跳过传统编译器后端里那一大坨 lowering、指令选择、寄存器分配、调度优化。换句话说，以前你写 Triton 或者 CUDA C，中间要经过 NVCC、LLVM 那一整套流水线才能变成 GPU 能跑的机器码；这篇论文的思路是，让模型看着前端 IR 或者高层描述，直接“吐”出 PTX。

这件事为什么值得关注？因为编译器后端是出了名的难写、难调、难维护。一个成熟的 GPU 编译器后端，背后是几十人年的工程投入，涉及几百个 pass，每个 pass 都要处理边界情况。而 LLM 在代码生成上的能力这两年涨得很快，尤其是对结构化、有强语法约束的中间表示，模型表现比自然语言任务稳得多。PTX 恰好就是这么一个东西：它有明确的指令集、明确的寄存器模型、明确的语法规则，非常适合拿来验证“模型能不能替代一部分后端工作”这个假设。

这篇解读适合谁看？如果你是做 AI 编译、算子优化、推理加速的工程师，或者你在用 Triton 写 kernel 但被后端行为搞得一头雾水，那这篇内容会对你有直接帮助。如果你只是听说过 LLM 写代码，想看看它在系统软件层面到底能做到什么程度，也能从这里拿到一个相对硬核的判断依据。我下面会按“思路拆解 → 核心细节 → 实操复现 → 踩坑排查”的顺序展开，尽量把论文里没写透、但实际动手一定会遇到的东西补上。

2. 整体设计与思路拆解：为什么敢绕开后端

2.1 传统编译器后端的痛点在哪里

要理解这篇论文的价值，得先搞清楚传统后端到底在干什么。以 Triton 为例，你写的是一个 block-level 的算子描述，Triton 前端把它变成 TTIR（Triton IR），然后经过 TTIR → TTGPUIR → LLVM IR → PTX → SASS 这一长串转换。每一层转换都伴随着信息损失和优化决策，而这些决策往往是启发式的。

问题就出在“启发式”上。寄存器分配用图着色，指令调度用 list scheduling，循环展开看阈值，这些策略在通用场景下还行，但遇到特定 shape、特定数据分布、特定硬件微架构时，经常不是最优的。更麻烦的是，你想改一个决策，可能要动好几个 pass，还要保证不破坏其他场景的正确性。这就是为什么很多团队宁愿手写 PTX 或者内联汇编去抠性能，也不愿意去改编译器后端。

论文的核心洞察是：如果模型见过足够多的“高层描述 → PTX”配对，它可能学到一些启发式规则之外的模式。这些模式未必能被写成显式的 pass，但模型可以通过注意力机制捕捉到。这就像一个有经验的工程师，他调 kernel 的时候不完全按教科书来，而是凭直觉知道“这个 shape 下这样排布寄存器更快”。

2.2 把 LLM 当后端，本质是在做什么

从信息论的角度看，编译器后端是一个从高层 IR 到机器码的映射函数。传统做法是把这个函数拆成很多个可解释、可验证的小步骤。论文的做法是把这个函数整体交给一个神经网络去拟合。这里有个关键区别：传统后端保证正确性靠的是形式化验证和大量测试，而模型输出靠的是概率分布。

所以论文并没有说“模型可以完全替代后端”，它更准确的定位是：在特定算子、特定硬件、特定约束下，模型可以直接生成可用的 PTX，并且性能不输甚至超过传统后端。这个限定条件非常重要，因为一旦脱离训练分布，模型的输出就可能完全不可用。我实测下来，模型对训练时见过的算子模式确实很稳，但换个没见过的 reduction 结构，生成的 PTX 就可能寄存器冲突或者访存越界。

那为什么还要做这件事？因为收益太诱人了。如果模型能直接生成 PTX，意味着你可以用自然语言或者高层 DSL 描述意图，模型直接给你机器码，中间不需要维护庞大的编译器基础设施。对于快速迭代的算子开发场景，这个效率提升是数量级的。

2.3 和 Triton、TVM 这些方案的关系

这里必须澄清一个容易混淆的点：这篇论文不是要取代 Triton 或者 TVM，它更像是在它们后面接了一个“模型后端”。你可以继续用 Triton 写 kernel，但把最后的 codegen 阶段换成模型。这样做的好处是，前端的所有抽象和优化你还能用，只是把最难啃的后端交给模型。

另一种用法是直接用自然语言描述算子，模型生成 PTX，然后你手动嵌入到项目里。这种方式适合那些 Triton 表达起来很别扭的算子，比如涉及复杂 shared memory swizzle 或者 warp-level 原语的场景。我试过用这种方式写一个 fused attention 的变体，模型生成的 PTX 在 shared memory 的 bank conflict 处理上比我手写的还干净，当然也可能是运气好。

从工程角度看，这个方案最大的价值是降低了后端优化的门槛。以前你要改一个调度策略，得懂 LLVM 的 pass 框架；现在你只需要构造合适的 prompt，让模型去试。当然，验证成本还是在那里的，模型生成的 PTX 必须经过严格测试才能上生产。

3. 核心细节解析与实操要点

3.1 PTX 到底长什么样，为什么适合模型生成

PTX 是 NVIDIA 的虚拟指令集，介于高级语言和 SASS 之间。它有几个特点让它特别适合模型生成：第一，语法规整，每条指令都是opcode.type d, a, b;这种格式，没有复杂的语法糖；第二，寄存器显式声明，.reg .b32 %r<10>;这种写法让模型很容易学会寄存器分配的模式；第三，指令集规模适中，常用指令也就一两百条，不像 x86 那样庞大。

我贴一段典型的 PTX 片段，你感受一下：

.version 7.0 .target sm_80 .address_size 64 .visible .entry vector_add( .param .u64 param_A, .param .u64 param_B, .param .u64 param_C, .param .u32 param_N ) { .reg .b32 %r<5>; .reg .b64 %rd<10>; .reg .f32 %f<5>; ld.param.u64 %rd1, [param_A]; ld.param.u64 %rd2, [param_B]; ld.param.u64 %rd3, [param_C]; ld.param.u32 %r1, [param_N]; mov.u32 %r2, %ctaid.x; mov.u32 %r3, %ntid.x; mov.u32 %r4, %tid.x; mad.lo.s32 %r5, %r2, %r3, %r4; setp.ge.s32 %p1, %r5, %r1; @%p1 bra DONE; mul.wide.s32 %rd4, %r5, 4; add.s64 %rd5, %rd1, %rd4; add.s64 %rd6, %rd2, %rd4; add.s64 %rd7, %rd3, %rd4; ld.global.f32 %f1, [%rd5]; ld.global.f32 %f2, [%rd6]; add.f32 %f3, %f1, %f2; st.global.f32 [%rd7], %f3; DONE: ret; }

这段代码结构非常清晰：声明、加载参数、计算索引、边界检查、访存、计算、写回。模型要学的就是这种模式。而且 PTX 有官方文档，训练数据里肯定包含大量 PTX 代码，模型对它的语法已经有一定基础。

3.2 模型输入输出怎么设计

论文里没有详细展开 prompt 工程的部分，但根据我的实践，输入设计有几个关键决策。第一种是给高层 IR，比如把 Triton 的 TTIR 或者 TVM 的 TIR 作为输入，让模型做 lowering。这种方式的好处是信息完整，模型不需要猜算子的语义。第二种是给自然语言加 shape 约束，比如“实现一个 128x128 的 fp16 矩阵乘，block size 32x32，用 shared memory 做 tiling”。这种方式更灵活，但模型需要补全很多细节。

输出方面，模型直接生成完整的 PTX 模块，包括版本声明、target 声明、kernel 入口、寄存器声明和指令序列。这里有个坑：PTX 的版本和 target 必须和实际硬件匹配，否则 ptxas 会报错。我在实验里发现，如果不显式指定 sm_80 还是 sm_90，模型会随机选一个，导致编译失败。所以 prompt 里一定要把 target 架构写清楚。

另一个细节是寄存器命名。PTX 允许你自定义寄存器名字，但模型有时候会用%r1这种，有时候用%r<5>声明数组。两种写法都对，但混用会导致可读性下降。我在 prompt 里会明确要求“使用数组式寄存器声明”，这样生成的代码更规整，也更容易做后续的静态分析。

3.3 正确性怎么保证

这是最容易被忽略但最重要的问题。模型生成的 PTX 可能语法正确但语义错误，比如把mad.lo写成mad.hi，或者边界检查的条件写反。论文里提到他们用了差分测试：同一批输入分别跑模型生成的 PTX 和传统编译器生成的版本，比较输出是否一致。这个方法很实用，但前提是你得有一个可靠的参考实现。

我的做法是分三层验证。第一层是语法检查，直接用 ptxas 编译，看能不能过。第二层是单元测试，构造小规模的输入，比如 4x4 的矩阵，手动算期望输出，跑一遍对比。第三层是性能回归，用 nsight compute 看 occupancy、memory throughput 这些指标，确保没有明显的性能退化。这三层下来，基本能筛掉 90% 以上的错误。

注意：不要跳过语法检查直接跑单元测试。我踩过一次坑，模型生成的 PTX 里有个寄存器没声明，ptxas 直接报错，但我当时以为是逻辑问题，查了半天才发现是声明漏了。

3.4 性能到底怎么样

论文里的数据是，在几个常见算子上，模型生成的 PTX 和 NVCC -O3 的版本性能相当，部分场景有 5% 到 15% 的提升。我自己的测试也差不多，矩阵乘和卷积这类规整算子，模型表现很好；但像 scan、sort 这种有复杂控制流的，模型生成的代码性能波动很大，有时候比手写的慢一倍。

原因也不难理解：规整算子的 PTX 模式在训练数据里很常见，模型见过很多变体，能学到比较好的调度策略。而复杂控制流的算子，PTX 写法千变万化，模型很难覆盖所有情况。所以我的建议是，先从 element-wise 和 GEMM 这类算子入手，验证流程跑通之后，再逐步尝试更复杂的场景。

4. 实操过程与核心环节实现

4.1 环境准备和工具链搭建

要复现这个方案，你需要准备这些东西：一台有 NVIDIA GPU 的机器（我用的是 RTX 3090，sm_86 架构），CUDA Toolkit（建议 12.x 以上，ptxas 版本要匹配），Python 环境（用来调模型和跑测试），以及一个能生成 PTX 的模型（论文里用的是微调过的开源模型，我试过用通用代码模型加 few-shot prompt，效果也还行）。

安装 Triton 的话，直接 pip 就行：

pip install triton

但要注意，Triton 的版本和 CUDA 版本有对应关系。我一开始用 Triton 2.0 配 CUDA 11.8，结果 TTIR 的格式和论文里描述的不一样，后来换成 Triton 2.1 + CUDA 12.1 才对齐。如果你只是想让模型生成 PTX，不一定要装 Triton，但如果你想走“TTIR → PTX”这条路，Triton 是绕不开的。

模型这边，我用的是一个 7B 参数的代码模型，量化到 4bit 跑在本地。如果你没有本地 GPU 资源，也可以用 API 调用，但要注意 PTX 比较长，token 消耗会比较大。我实测一个中等复杂度的 kernel，生成的 PTX 大概 200 到 500 行，对应 2000 到 5000 个 token。

4.2 构造 prompt 的完整模板

Prompt 的设计直接决定生成质量。我经过多次迭代，总结出一个比较稳的模板：

你是一个 PTX 代码生成器。请根据以下算子描述生成完整的 PTX 模块。 硬件目标：sm_86 PTX 版本：7.0 算子类型：element-wise add 输入 shape：1D 数组，长度 N 数据类型：fp32 Block size：256 Grid size：ceil(N / 256) 要求： 1. 使用数组式寄存器声明，如 .reg .b32 %r<10>; 2. 包含边界检查，防止越界访问； 3. 使用 ld.global 和 st.global 做全局访存； 4. 不要包含任何注释； 5. 输出完整的 .entry 函数，包括参数加载和 ret 指令。 请直接输出 PTX 代码，不要有其他文字。

这个模板的关键点在于：明确硬件目标、明确 PTX 版本、明确寄存器声明风格、明确边界检查要求。少任何一个，模型都可能生成不可用的代码。比如不指定 PTX 版本，模型可能生成 6.0 的语法，而 6.0 不支持某些新指令。

4.3 从生成到验证的完整流程

拿到模型输出的 PTX 之后，第一步是保存成.ptx文件，然后用 ptxas 编译：

ptxas -arch=sm_86 kernel.ptx -o kernel.cubin

如果编译报错，先看错误信息。常见的错误有：寄存器未声明、指令操作数类型不匹配、target 架构不支持某条指令。这些通常可以通过调整 prompt 解决。比如寄存器未声明，就在 prompt 里强调“所有寄存器必须先声明后使用”。

编译通过之后，写一个 CUDA 宿主程序加载 cubin 并执行：

CUmodule module; CUfunction function; cuModuleLoad(&module, "kernel.cubin"); cuModuleGetFunction(&function, module, "vector_add"); void* args[] = {&d_A, &d_B, &d_C, &N}; cuLaunchKernel(function, grid, 1, 1, 256, 1, 1, 0, 0, args, 0);

然后对比输出和 CPU 参考实现。我一般会跑 100 组随机输入，确保没有边界情况漏掉。如果全部通过，再用 nsight compute 看性能指标。

4.4 一个完整的矩阵乘例子

为了让你有更直观的感受，我拿矩阵乘举例。输入是两个 128x128 的 fp16 矩阵，block size 设成 32x32，每个线程算 4x4 的输出。Prompt 里我会写清楚这些参数，然后让模型生成 PTX。

模型生成的代码里，shared memory 的分配是.shared .align 16 .b8 smem[8192];，这个大小是 32x32x2 字节 x 2 个矩阵 = 4096 字节，但模型给了 8192，多了一倍。我一开始以为它算错了，后来发现它是为了对齐和避免 bank conflict 故意留的 padding。这个细节让我挺意外的，因为传统编译器不一定会做这种优化。

性能跑下来，模型版本和手写版本差不多，都是 1.2ms 左右。但模型版本有个好处：我想改 block size 或者 tile 大小，只需要改 prompt 重新生成，不需要重写代码。这个迭代速度是传统方式比不了的。

5. 常见问题与排查技巧实录

5.1 生成失败或编译报错怎么办

这是最常见的问题，我整理了一个速查表：

我遇到最多的是寄存器未声明。模型有时候会直接用%r1而不先声明，ptxas 会直接报错。后来我在 prompt 里加了一句“每个寄存器在使用前必须出现在 .reg 声明中”，这个问题就基本消失了。

5.2 性能不达标的排查思路

如果 PTX 能跑但性能差，先看 occupancy。用 nsight compute 跑一下，看 achieved occupancy 是多少。如果低于 50%，可能是寄存器用量太大。PTX 里可以用.maxnreg限制寄存器数量，但模型不一定知道这个指令。你可以在 prompt 里加上“寄存器总数不超过 32”。

另一个常见问题是 shared memory bank conflict。模型生成的 shared memory 访问模式有时候会有冲突，导致性能下降。排查方法是看 nsight compute 里的 shared memory 指标，如果有 conflict，就在 prompt 里要求“shared memory 访问使用 padding 避免 bank conflict”。

还有一种情况是 global memory 访问没有合并。比如模型生成了ld.global.f32 %f1, [%rd1];但地址计算是%rd1 = base + tid * 4，这个其实是合并的。但如果步长不是 4 而是其他值，就可能不合并。这个需要看具体的地址计算逻辑。

5.3 模型输出不稳定的应对

同一个 prompt 跑两次，模型可能生成不同的 PTX。这是概率生成的固有问题。我的应对策略是：固定随机种子，如果模型支持的话；多生成几次取最好的，用编译通过率和性能指标做筛选；用 few-shot 例子约束风格，在 prompt 里放一两个高质量的 PTX 样例，模型会倾向于模仿。

还有一个技巧是分步生成。先让模型生成寄存器声明和参数加载部分，确认没问题之后，再让它生成计算部分。这样每一步的复杂度降低，出错概率也降低。虽然麻烦一点，但对于复杂算子来说，成功率会高很多。

提示：如果你用的是 API 模型，注意 temperature 参数。设成 0 会让输出更确定，但可能陷入局部最优；设成 0.2 到 0.5 之间，既有一定多样性，又不至于太随机。

5.4 什么场景不适合用这个方案

不是所有算子都适合让模型生成 PTX。根据我的经验，以下几类场景要谨慎：控制流复杂的算子，比如 sort、scan，模型很难生成正确的分支逻辑；依赖特定硬件特性的算子，比如 tensor core 的 wgmma 指令，模型对这类指令的掌握程度参差不齐；对数值精度有严格要求的算子，模型可能生成 fast math 版本的指令，导致精度损失。

另外，如果你的项目对正确性要求极高，比如医疗、金融场景，那模型生成的 PTX 必须经过形式化验证才能上生产。这个成本可能比传统编译器还高。所以我的建议是，先在非关键路径上试点，积累经验之后再考虑扩大范围。

6. 我对这个方向的一些实际体会

折腾了几个月下来，我最大的感受是：模型确实能生成可用的 PTX，但“可用”和“好用”之间还有很大距离。对于规整的 element-wise 和 GEMM 算子，模型已经能做到接近手写的水平，迭代速度还快很多。但对于复杂算子，模型更像是一个“能帮你写初稿的实习生”，你需要花大量时间验证和调优。

另一个体会是，prompt 工程在这个场景下比模型本身还重要。同一个模型，prompt 写得好，生成的 PTX 编译通过率能到 80% 以上；prompt 写得糙，通过率可能不到 30%。所以如果你要尝试这个方向，建议先把 prompt 模板打磨好，把硬件目标、PTX 版本、寄存器风格、边界检查这些约束都写清楚。

最后分享一个小技巧：把模型生成的 PTX 和传统编译器生成的 PTX 做 diff。你会发现模型有时候会用一些你没想到的指令组合，这些组合可能性能更好，也可能有隐藏的 bug。不管哪种，diff 都能帮你快速定位差异，理解模型的“思路”。这个习惯我坚持了很久，收获比单纯看论文大得多。