PTXBench：基于架构特定PTX的GPU内核优化基准测试与LLM适配

# PTXBench：基于架构特定PTX的LLM GPU内核优化基准测试与适配
来源：https://arxiv.org/html/2608.17379
耿汉张¹同等贡献。通讯作者：耿汉张 [email protected]。代码获取地址（https://github.com/zhang677/PTXBench）。
董怡昕¹¹脚注标记：1。
范承泽 隶属机构：独立研究者。
曾志辰 隶属机构：RadixArk。
袁月明 隶属机构：RadixArk。
朱邵威 隶属机构：独立研究者。
昆勒·奥卢科图恩[0.5em] 斯坦福大学

###### 摘要

我们推出了 PTXBench，一个用于评估和调整大型语言模型（LLMs）以使用架构特定PTX进行GPU内核优化的基准测试框架。PTXBench 在H100和B200 GPU上，针对GEMM和注意力工作负载，评估功能正确性、所选目标指令在运行时是否执行，以及相对于前沿库的加速比。我们的评估表明，架构特定的PTX能力仍不均衡：在复杂的注意力反向工作负载上成功率显著下降，执行目标指令并不总能转化为有竞争力的性能。没有任何被评估模型在全套测试中始终匹配前沿库的性能。我们进一步使用监督微调适配了Qwen3.6-27B。基于修复条件的训练改进了若干任务，但泛化能力仍不均衡；除了数据集规模，数据覆盖范围、平衡性以及推理教师的质量也很重要。PTXBench 为度量和提升LLM利用演进中的GPU架构的能力提供了一个可审计的测试平台。

## 1 引言

PTX（并行线程执行）是CUDA开发者可以显式控制的最低级可编程接口。高性能内核通常需要PTX级优化以高效利用架构特定的机制[59 (https://arxiv.org/html/2608.17379#bib.bib40)]。近年来的NVIDIA GPU系列引入了新的架构特定PTX指令，用于张量核心、内存单元和同步机制[24 (https://arxiv.org/html/2608.17379#bib.bib35), 25 (https://arxiv.org/html/2608.17379#bib.bib36), 27 (https://arxiv.org/html/2608.17379#bib.bib37), 28 (https://arxiv.org/html/2608.17379#bib.bib38), 33 (https://arxiv.org/html/2608.17379#bib.bib39)]。一个可移植的CUDA内核可能在功能上保持正确，但未能利用新硬件的定义性能力。更高级别的内核语言旨在恢复性能可移植性[49 (https://arxiv.org/html/2608.17379#bib.bib33), 48 (https://arxiv.org/html/2608.17379#bib.bib34)]，它们最终都会转换到相同的PTX接口。然而，让这些内核语言与快速演进的硬件保持一致需要持续的编译器工程[6 (https://arxiv.org/html/2608.17379#bib.bib57)]，而验证一个优化编译器需要处理远大于验证单个内核的输入和配置空间，因此难度大得多[10 (https://arxiv.org/html/2608.17379#bib.bib56)]。这为内核开发者提出了一个实际问题：LLM能否直接利用架构特定特性？有针对性的后训练能否提升这种能力？

为回答此问题，我们推出了PTXBench，一个架构特定的GPU内核基准测试框架，旨在评估LLM能否利用指定的底层GPU机制。PTXBench整合了MiniPTXAgent，一个多轮代理循环，其中模型被提供准确的架构特定知识包，并生成带有内联PTX的CUDA内核（我们称之为CUDA–PTX），并利用结构化的执行反馈进行修改。在每次轨迹中，PTXBench分别测量功能正确性、所选指令系列在评估的工作负载下运行时是否执行，以及峰值和典型性能。PTXBench基于FlashInfer-Trace[54 (https://arxiv.org/html/2608.17379#bib.bib19)]构建，这是一种用于内核工作负载、解决方案和评估的统一模式。将能力探测与问题集合分离使得PTXBench可以扩展到具有兼容接口的测试套件，例如SOL-ExecBench[20 (https://arxiv.org/html/2608.17379#bib.bib18)]。现有的GPU内核基准测试（从KernelBench[38 (https://arxiv.org/html/2608.17379#bib.bib1)]开始）要求模型用功能正确、更快的GPU内核替换参考PyTorch算子[41 (https://arxiv.org/html/2608.17379#bib.bib31), 17 (https://arxiv.org/html/2608.17379#bib.bib32)]。这些端到端结果至关重要，但它们无法隔离模型是否能直接编程特定的架构机制：性能可能来自通用CUDA代码或对现有供应商库的调用。PTXBench通过受控的能力探测补充了这些套件：模型必须直接使用指定的架构特定PTX指令系列生成高效的低级内核，并且相应的目标指令必须在运行时执行。

PTXBench提供了一个受控的测量框架和用于收集适应性数据的环境。我们首先在当前模型和GPU上表征架构特定的PTX能力，然后测试有针对性的、基于修复条件的后训练。

本工作有三项贡献：

- •**架构特定的PTX基准测试**。我们引入了PTXBench，一个多轮基准测试，用于评估LLM能否生成功能正确、并在运行时执行所需目标指令的GPU内核。它提供了受控的架构知识，并分别测量正确性、目标指令执行情况以及相对于前沿库的性能。
- •**跨模型、架构和工作负载的能力研究**。我们在H100和B200上，针对GEMM和注意力工作负载，评估了闭源和开源权重模型。模型在正向工作负载上经常成功，但在反向注意力上表现不佳，即使验证了目标指令执行的内核通常仍慢于前沿库。这些结果揭示了执行架构特定指令、生成正确内核和实现有竞争力的性能之间的显著差距。
- •**受控的适应性研究**。据我们所知，我们首次进行了针对特定架构的CUDA和PTX生成的基于修复条件的SFT受控研究，适配了Qwen3.6-27B，并消融了训练格式、问题覆盖范围与平衡性以及推理监督。基于修复条件的监督微调在若干任务上优于直接生成，但向保留形状和注意力变体的迁移仍不均衡。问题覆盖范围、数据平衡性以及推理教师都很重要。

参见图注图1：PTXBench基准测试和适应性工作流概览。
## 2 PTXBench

图1 (https://arxiv.org/html/2608.17379#S1.F1)总结了PTXBench的基准测试和适应性工作流。PTXBench围绕评估架构特定PTX编程的三个需求设计。首先，模型接收受控的架构知识，因为否则它们很少使用请求的PTX。其次，我们验证所需的指令系列在固定工作负载下运行时是否执行。第三，我们评估每个内核相对于前沿库实现的正确性和效率。

### 2.1 基准测试任务与受控上下文

一个PTXBench实例将一个由cuBLAS、cuDNN和FlashInfer[30 (https://arxiv.org/html/2608.17379#bib.bib49), 7 (https://arxiv.org/html/2608.17379#bib.bib47), 55 (https://arxiv.org/html/2608.17379#bib.bib48)]等前沿库构建的参考算子与固定的工作负载、目标GPU架构以及所需的架构特定PTX指令系列配对。模型从头生成带有内联PTX的CUDA内核，而不是从初始实现开始。PTXBench使用FlashInfer-Trace模式来表达工作负载、解决方案和正确性检查，允许相同的基准测试工作流支持其他兼容的任务集合。

为了度量架构特定的推理而非文档检索，每次轨迹在其基础提示中接收相同的架构特定知识包：架构参数、用于PTX指令的CUDA封装器，以及管理布局、同步和内存一致性的契约。对于研究充分的算子，我们还提供固定的、经专家验证的调度原则（附录A.5 (https://arxiv.org/html/2608.17379#A1.SS5)展示了一个FlashAttention的示例）。每个知识包占用20k-30k个令牌（附录A.4 (https://arxiv.org/html/2608.17379#A1.SS4)）。表3 (https://arxiv.org/html/2608.17379#S4.T3)的消融研究显示，没有此上下文，模型很少使用请求的PTX。

### 2.2 多轮评估协议

MiniPTXAgent采用多轮协议，因为迭代的内核生成、执行反馈和修正构成了内核代理的基本操作循环[4 (https://arxiv.org/html/2608.17379#bib.bib15), 58 (https://arxiv.org/html/2608.17379#bib.bib14), 14 (https://arxiv.org/html/2608.17379#bib.bib16)]。每次轨迹允许固定次数的模型调用，并保留所有先前的内核和执行反馈。MiniPTXAgent首先在CPU容器中使用`nvcc`编译每个候选内核，仅将成功编译的结果发送到分析服务进行内存安全检查、运行时错误消息、功能评估和延迟测量。

我们定义**目标指令正确性**为功能正确性加上在评估的工作负载下运行时执行选定目标指令。我们选择张量执行路径：Hopper架构上的GMMA计算或UTMA负载移动，Blackwell架构上的TCGEN05张量路径。仅TMA不足以使Blackwell内核合格，因为Blackwell从Hopper继承了它。为了测量执行情况，我们分析SASS，这是从PTX生成并由GPU执行的NVIDIA原生GPU汇编。静态SASS检查可以揭示是否存选定的指令，但不能说明它在评估的工作负载下是否运行。因此，我们使用NVIDIA Nsight Compute（NCU）获取匹配指令的谓词启用线程数。我们的检查对每个功能正确的候选内核分两个阶段进行。我们首先检查表5 (https://arxiv.org/html/2608.17379#A1.T5)中的架构特定指令系列的最终SASS；如果不存在，则不运行NCU，将目标指令指示符设置为零。否则，仅当NCU报告匹配指令的谓词启用线程数为正时，才将指示符设置为1。此动态检查排除了死亡代码或未启动代码中的匹配指令。目标指令正确性确定了目标指令是否被执行，而不是它们执行了多少有用工作或是否导致了测量到的加速。附录A.2 (https://arxiv.org/html/2608.17379#A1.SS2)提供了实现细节和边缘情况。

功能正确性检查器验证输出形状和数据类型，然后使用`torch.allclose`（atol=rtol=1e-2）比较值。候选内核通过正确性检查后，使用CUPTI[32 (https://arxiv.org/html/2608.17379#bib.bib58)]测量候选内核和参考内核的延迟。对于每次试验，我们取10次预热迭代后50次计时迭代的中位数，并计算加速比为参考延迟除以候选延迟。在CUDA源代码中包含来自cuDNN或cuBLAS[30 (https://arxiv.org/html/2608.17379#bib.bib49)]的头文件，无论执行结果如何都被视为错误。生成的内核可能崩溃、挂起或破坏后续测量，因此分析服务隔离内核执行，并将正确性、净化、调试和延迟测量作为独立操作暴露出来。有关分析服务的更多详细信息，请参见附录A.3 (https://arxiv.org/html/2608.17379#A1.SS3)。

## 3 将LLM适配到架构特定的PTX编程

部署的GPU在多年内保留固定的底层能力，使得架构特定的PTX成为后训练的持久目标。障碍在于数据：高质量内核需要稀缺的专家知识和迭代验证。因此，我们研究模型故障和执行反馈，结合教师生成的修复和解释，是否能为该领域提供有效的监督。

#### Fixit 方法。

Fixit 从被适配模型产生的故障中构建监督。设$x$表示带有其架构特定上下文的问题提示，$\pi_0$为适配前的模型，$H$为MiniPTXAgent反馈函数，$C$为功能正确性指示符。我们首先采样一个失败的内核$k^{-}$并收集其编译或执行反馈$e$：
$k^{-} \sim \pi_0(\cdot \mid x), \quad e = H(x, k^{-}), \quad C(x, k^{-}) = 0.$
修复教师$\pi_R$然后基于相同的问题、故障和反馈生成修正后的内核。我们只保留通过正确性检查的修复$k^{+}$：
$k^{+} \sim \pi_R(\cdot \mid x, k^{-}, e), \quad C(x, k^{+}) = 1.$
最后，推理教师$\pi_T$综合一个解释$r$，说明从观察到的故障到保留的修复的推理过程：
$r \sim \pi_T(\cdot \mid x, k^{-}, e, k^{+}).$
因此，每个Fixit示例将学生条件化在$(x, k^{-}, e)$上，并用$(r, k^{+})$对其进行监督。故障仅从适配前的模型中收集一次，因此监督针对该固定检查点产生的错误状态；修复和解释是由教师生成的。

## 4 基准测试结果

### 4.1 评估指标与基线

我们从四个互补维度评估模型。轮次正确率（正确内核数/轮次数）捕捉生成正确内核的能力。目标指令轮次正确率（执行选定目标指令的正确内核数/轮次数）评估模型能否生成功能正确且能运行请求的架构机制的内核。最后，在正确轮次中，最佳加速比衡量峰值优化能力，而目标指令最佳加速比则将该峰值限制在合格内核上。受KernelBench[38 (https://arxiv.org/html/2608.17379#bib.bib1)]启发，对于$N$次评估轮次（具有正确性$C_i$、运行时目标指令指示符$I_i$和加速比$s_i$），我们总结完整的加速比分布为
$\mathrm{Fast}_p = \frac{1}{N} \sum_{i=1}^{N} \mathbf{1}[C_i = 1 \land s_i > p], \qquad \mathrm{Fast}^{\mathrm{Inst.}}_p = \frac{1}{N} \sum_{i=1}^{N} \mathbf{1}[C_i = 1 \land I_i = 1 \land s_i > p].$
在阈值$p$下，$\mathrm{Fast}^{\mathrm{Inst.}}_p$是所有轮次中产生正确内核、在运行时执行选定目标指令且超过加速比$p$的比例。NCU失败仍未知且不计入目标指令分子，因此报告的值是保守的。在空间受限的表格标题和坐标轴中，我们将目标指令简写为“Target inst.”。在所有结果表格中，第一行是目标指令指标，括号内的行是无限制的，三元组报告$\leq 1 / \leq 4 / \leq 8$轮次。曲线图显示相应的轮次比例，垂直标签标记最佳合格加速比。实线显示三个提示变体的平均值，每个变体在四次八轮轨迹上评估（$N=32$），而阴影区域覆盖提示级别的最小值和最大值。在整篇论文中，Fwd表示多头注意力。