PTX(Parallel Thread Execution)是 CUDA 开发者能显式控制的最低级可编程接口。最近几代 NVIDIA GPU 每次迭代都在这层加入新的架构专属指令——Hopper 的 GMMA 计算与 UTMA 数据搬运、Blackwell 的 TCGEN05 tensor 路径。一个追求可移植的 CUDA kernel 可以在新硬件上跑得完全正确,却把新一代硬件的定义性能力全部闲置。高性能 kernel 要吃到新卡的算力,绕不开 PTX 级的手工优化,而这恰恰是稀缺专家的领地。Stanford 的 Kunle Olukotun 组联合 CMU、RadixArk 在 8 月 18 日提交的 PTXBench(arXiv:2608.17379)就把问题摆上了台面:LLM 能不能直接干这活?
三层指标:先问对不对,再问指令跑没跑,最后问快不快
现有 GPU kernel 基准(如 KernelBench)大多让模型把 PyTorch 算子替换成更快的 kernel,只看端到端结果。但这测不出模型是否真的会写架构专属代码——性能可能来自通用 CUDA,或者干脆是调了厂商库。PTXBench 把探针扎得更深,拆成三层:
- 功能正确性:输出与参考实现对齐(torch.allclose,atol=rtol=1e-2);
- 目标指令是否真的执行:先静态查 SASS 汇编里有没有目标指令族,再用 Nsight Compute 取 predicate-enabled 线程数做动态验证,排除死代码里的假命中;
- 性能:CUPTI 计时,50 次迭代取中位数(10 次 warmup),基线是 cuBLAS v13.1.0(GEMM)、cuDNN v9.20.0 与 FlashInfer v0.6.14(attention)这些前沿库。
测试在 H100 和 B200 上进行,覆盖 GEMM 与 MHA 前向/反向及 causal 变体。模型在 MiniPTXAgent 多轮循环里从零生成带内联 PTX 的 kernel,禁止 include cuBLAS/cuDNN 头文件,靠结构化执行反馈迭代修改。每个任务附 20k–30k token 的架构知识包;消融实验显示,没有这份上下文,模型基本不会去用被要求的 PTX。
结果:forward 尚可,backward 跳水,Blackwell 更难
论文的核心结论很直白:架构专属 PTX 能力严重不均衡,且没有任何一个被测模型能在全套件上稳定追平前沿库。几个具体数字(H100,目标指令口径,首轮 turn 正确率):
- Claude Opus 4.8 在 GEMM 上首轮 91.7%,8 轮内到 94.8%,全场最稳;
- 同一个 Claude,在 MHA backward 首轮只有 8.3%,要 8 轮才爬到 79.2%;
- Gemini 3.1 Pro 与 GLM-5.2 的 GEMM 首轮都只有 33.3%,后者的 MHA backward 8 轮内仅 5.2%;
- 换到 B200,全员进一步下滑:Gemini 3.1 Pro 的 GEMM 首轮掉到 8.3%,attention 前向 8 轮也只有 15.6%。
更微妙的一点:目标指令执行了,不代表性能有竞争力。即便验证了 kernel 在运行时真的跑到了指定指令,大部分实现仍然慢于前沿库。会用新指令,和用出新指令的价值,中间隔着一整层调度工程。
Fixit:让模型把自己踩过的坑变成训练数据
Benchmark 之外,团队还做了(据其所知)首个针对特定架构 CUDA/PTX 生成的修复条件 SFT 受控研究,起名 Fixit:从被适配模型(Qwen3.6-27B)自己产生的失败 kernel 出发,repair teacher 生成通过正确性检查的修复版本,reasoning teacher 再合成从失败到修复的推理链,拼成监督数据。
效果方向明确:基础版 Qwen3.6-27B 在 GEMM 上任何一轮都没有正确 kernel,Fixit 训练后第 0 轮就有 3 个正确,后续 7 轮里 6 轮至少 1 个正确。更有意思的是 SFT 与上下文监督的对比:基础模型即便配上专家指导也写不对 MHA kernel,而 Fixit 版不带指导也能写对——说明 SFT 提升的是底层 PTX 能力本身,不只是读懂指导的能力。反过来看,给基础模型检索修复笔记毫无用处,只有连修复后的 kernel 一起给才大幅起效——但那基本等于把答案放进 prompt。
局限作者也摆明:适配实验只用了一个 27B 模型加适度的 LoRA 数据集;工作负载限于 H100/B200 上的 BF16 GEMM 与 attention。跨语言迁移到 Triton 时,Fixit 版正确率反而下降,只是峰值加速比明显改善(causal 前向从 0.238× 提到 0.632×)。
所以呢
PTXBench 的价值不在排行榜,在于把 LLM 能否利用快速演进的 GPU 架构变成一个可审计的测试床。对模型厂商,它指出了 coding benchmark 之下的真空地带:会写能跑的 kernel 和会写吃满硬件的 kernel 是两种能力。对推理基础设施团队,当 H100 换 B200、B200 再换下一代,这种架构专属能力直接决定迁移成本。代码在 GitHub 开源(github.com/zhang677/PTXBench)。下一个值得盯的问题:当修复条件 SFT 从 27B 规模走向工业级后训练,这条断层会不会被填平——还是说,PTX 终究是编译器和少数专家的自留地。