Source-linked AI summary
PTXBench: Benchmark and Adapt LLMs for GPU Kernel Optimization with Architecture-specific PTX
Genghan Zhang, Yixin Dong, Chengze Fan, Zhichen Zeng, Yueming Yuan, Shaowei Zhu, Kunle Olukotun
TL;DR
现有证据表明，LLM 尚未稳定展现出利用架构特定 PTX 进行 GPU kernel 优化的能力，尤其是在复杂 attention 工作负载上。PTXBench 同时评估正确性、运行时指令执行和性能，并研究修复条件适配，发现能力差距持续存在，训练收益也不均衡。
问题
关于 LLM 能否利用高性能 GPU kernel 通常所需的不断演进的架构特定 PTX 机制，现有证据仍然有限。
方法
PTXBench 将多轮 CUDA–PTX 生成 benchmark 与结构化执行反馈相结合，并分别衡量正确性、指令执行和性能，同时开展针对性后训练。
结果
当前 LLM 有时能够执行所要求的指令并解决 forward 工作负载，但难以处理 attention backward；修复条件训练提升了若干任务，却未能一致优于直接解法监督。
结论与局限
不同工作负载、模型和训练方案之间，架构特定 PTX 能力与适配收益仍然不均衡，而非始终达到 frontier 性能。
结论与局限
当前 PTXBench 聚焦于 H100 和 B200 上的 BF16 GEMM 与 attention kernel，无法代表 GPU 算子和硬件的完整多样性。
Abstract
from arXiv · showhide
We introduce PTXBench, a benchmark for evaluating and adapting large language models (LLMs) to use architecture-specific PTX for GPU kernel optimization. PTXBench measures functional correctness, whether selected target instructions execute at runtime, and speedup over frontier libraries across GEMM and attention workloads on H100 and B200 GPUs. Our evaluation shows that architecture-specific PTX capability remains uneven: success rates fall substantially on complex attention backward workloads, and executing the target instructions does not necessarily translate into competitive performance. No evaluated model consistently matches frontier libraries across the suite. We further adapt Qwen3.6-27B using supervised fine-tuning. Repair-conditioned training improves several tasks, but generalization remains uneven; data coverage, balance, and the quality of the reasoning teacher matter in addition to dataset size. PTXBench provides an auditable testbed for measuring and improving LLMs' ability to exploit evolving GPU architectures.
1 引言
PTXBench评估LLM能否利用架构特定的PTX，生成正确、执行目标指令且性能良好的GPU内核。它将受控的架构知识与MiniPTXAgent基于反馈的CUDA–PTX生成相结合，支持跨模型和GPU的能力刻画，以及修复条件后训练。
- 动机: PTXBench针对架构特定优化的需求，因为可移植CUDA虽然能够保持正确，却可能无法利用面向张量核心、内存单元和同步机制的新PTX机制。PTX是CUDA开发者能够显式控制的最低层可编程接口，高性能内核通常需要PTX级优化 。
- 基准设计: PTXBench分别测量功能正确性、所需目标指令的运行时执行情况，以及相对于前沿库的性能。该基准评估LLM生成的内核能否跨架构和工作负载利用指定的底层GPU机制。
- 基准设计: MiniPTXAgent提供架构特定知识和结构化执行反馈，使模型能够迭代生成并修改带有内联PTX的CUDA内核。最终的CUDA–PTX内核通过多轮agent循环生成。
- 适配研究: 该框架刻画当前模型和GPU上的架构特定PTX能力，随后测试有针对性的修复条件后训练。PTXBench既是受控的测量框架，也是收集适配数据的环境。
2 PTXBENCH
PTXBench 通过受控的架构知识、目标指令的运行时验证、功能正确性以及相对于 frontier libraries 的效率，评估 architecture-specific PTX kernel 生成。其流程从零生成 kernels，迭代编译并分析候选方案，并在固定 workloads 下测量正确性与 speedup。
- Benchmark 设计: PTXBench 将 frontier-library reference operators [17] [18] 与固定 workloads、目标 GPU architectures 以及所需的 architecture-specific PTX instruction families 配对，同时从零生成带 inline PTX 的 CUDA kernels。FlashInfer-Trace schema 表示整个 benchmark 流程中的 workloads、solutions 和 correctness checks。
- Benchmark 设计: 每条 trajectory 都接收内容完全相同的 architecture-specific knowledge packs，其中包含 hardware parameters、PTX wrappers，以及关于 layouts、synchronization 和 memory consistency 的 contracts。这些 packs 占用 20k–30k tokens；部分研究充分的 operators 还会接收经专家验证的 scheduling principles。
- Evaluation 流程: MiniPTXAgent 利用 execution feedback 迭代生成、编译和修改 kernels，仅将编译成功的 kernels 发送至 profiling，以进行安全检查、功能评估、运行时错误检测和 latency 测量 [19] [20] [21]。每条 trajectory 允许固定次数的 model calls，并保留此前的 kernels 和 feedback。
- Evaluation 标准: Target-instruction correctness 要求功能正确性，并要求选定 tensor paths 在运行时执行：Hopper 上的 GMMA 或 UTMA，以及 Blackwell 上的 TCGEN05。在 Blackwell 上，仅使用 TMA 不满足要求；执行情况通过 SASS analysis 验证。
- Evaluation 标准: 正确的 candidates 使用 torch.allclose 与 references 进行比较，其中 atol=rtol=1e-2；speedup 定义为 reference latency 除以 candidate latency，测量基于 10 次 warmups 以及 50 次计时迭代的中位数。Latency 测量使用 CUPTI ；候选 CUDA code 中包含 cuDNN 或 cuBLAS headers 会被视为不正确。
3 面向架构特定 PTX 编程适配 LLM
本节介绍 Fixit：一种面向架构特定 PTX 编程的监督策略，利用模型失败、执行反馈以及教师生成的修复结果和理由。该策略旨在缓解专家知识稀缺和 kernel 验证成本高造成的数据瓶颈。
- 动机: 该方法针对一个持久的优化目标，其主要障碍是高质量 kernel 所需的专家知识稀缺，以及反复验证成本高。已部署的 GPU 多年来保持固定的底层能力，因此有必要针对架构特定 PTX 进行后训练。
- Fixit: Fixit 利用待适配模型产生的失败构建监督数据，将问题上下文、失败的 kernel 以及编译或执行反馈结合起来。失败样本取自适配前的模型，并根据功能正确性进行评估。
- Fixit: 修复教师根据每个失败样本生成修正后的 kernel，仅保留通过功能正确性检查的修复结果。修复过程以与失败 kernel 相同的问题、失败信息和反馈为条件。
- Fixit: 随后，推理教师综合生成一份理由，将观察到的失败与保留的修复结果联系起来。每个训练样本都用教师生成的理由和修正后的 kernel 对学生模型进行监督。
4 基准测试结果
基准测试结果表明，模型对架构特定 PTX 的能力会随模型、工作负载和 GPU 代际而显著变化：Blackwell attention 尤其困难，而较高层的 Triton 往往仍比直接使用 CUDA-PTX 更可靠。显式 PTX 知识能提升目标指令的执行率和正确性，但无法消除性能差距。
- 模型比较: Claude Opus 4.8 在 Blackwell GEMM 上达到 cuBLAS 的 1.012×，而 Gemini 3.1 Pro 达到 0.892×，尽管其截止时间与 Blackwell PTX 发布时间重合。Claude 在 Blackwell 上的目标指令正确率也显著更高。
- 模型比较: GLM-5.2 在 H100 和 B200 上的 GEMM，以及未掩码 H100 attention 上，都能与 Gemini 3.1 Pro 竞争，但在 B200 attention 上落后。GLM-5.2 和 Claude Opus 4.8 在 Blackwell 上经常退回使用通用 CUDA，而 Gemini 更成功地执行 Blackwell 指令。
- 实现比较: 在 Blackwell 上，Triton 在两个 attention backward 工作负载上达到 0.484× 和 0.436×，而 CUDA-PTX 仅达到 0.133× 和 0.015×。在 Hopper 上，CUDA-PTX 在 GEMM 上几乎追平 Triton；在因果 MHA forward 上则达到 0.768×，高于 Triton 的 0.759×。
- 架构影响: Blackwell 的专用机制和有限的 Blackwell 专项训练数据，或许可以解释为何直接生成 CUDA-PTX 会落后于较高层语言，而 CuTeDSL 的成功率仍极低。高层语言有助于在新架构上生成正确的 kernel，但直接的低层实现一旦正确，有时也能达到或超过它们的性能。
- 知识消融: 仅使用架构参数时，模型的正确轮次达到 26.0%，但没有任何目标指令执行成功；模板函数能够启用执行，而架构契约则能提升正确性。消融实验表明，显式 PTX 知识是必需的，不能自动从架构参数中推断出来。
5 适配结果
使用 Qwen3.6-27B 进行 Fixit SFT 可改善部分 PTX kernel 结果，但收益取决于训练格式、问题均衡性、推理教师能力，以及 workload 或语言迁移。适配后的模型扩大了可执行 kernel 的搜索范围，并提升了部分峰值 speedup，但正确性和泛化能力仍存在显著不足。
- 训练格式: Fixit 条件训练提升了 GEMM、MHA-Fwd-Causal 和 MHA-Bwd 上的正确性，但在 MHA-Fwd 和 MHA-Bwd-Causal 上落后于直接解监督。比较使用相同的八类问题、GLM-5.2 推理和相近的记录数；因此，按失败条件构造的监督只能选择性地发挥作用，而非普遍有效。
- 覆盖范围与均衡性: 均衡的问题覆盖范围，而非单纯的记录数，决定了适配广度：s1 和 s5 能解决全部五个问题，而更大的不均衡配方仍会在 MHA-Bwd-Causal 上失败。将 s4 增加 1.5× 并调整为均衡的 s5 后，四个问题的八轮正确性得到提升，MHA-Fwd 持平，且四个问题的峰值 speedup 得到提升；但最佳配方仍因问题而异。
- 推理合成器: 更强的推理教师十分重要：在 Fixit 示例完全相同的情况下，s5 能解决全部五个问题，而使用 Qwen3.6-27B 作为合成器的 s6 只能解决 GEMM。目标模型的失败样本是有用的训练输入，但生成修复理由仍受益于更强的教师。
- 迁移与泛化: S1 能迁移到部分相关 workload，但并不一致：在四个 d128 MHA 任务上训练后，它在两个 d96 backward 任务和 GQA 上都无法生成正确 kernel。d96 backward 任务尤其困难，因为 d96 与 H100 WGMMA 的 m64 tile 不匹配。
- 迁移与泛化: 在 Triton 迁移设置下，s1 在每个 workload 上都降低了轮次级正确性，但在 MHA-Fwd-Causal 上将最佳正确 speedup 从 0.238× 提升至 0.632×，在 MHA-Bwd-Causal 上则从 0.043× 提升至 0.331×。使用 CUDA-PTX 训练的配方可能在损害跨语言正确性的同时提升峰值性能。
- 搜索过程与提示时监督: SFT 延长了推理过程，并使失败从编译阶段转向运行时和数值错误；专家指导也无法让基础模型编写出正确的 MHA kernel。S1 在第 0 轮以及之后七轮中的至少一轮生成了正确的 GEMM kernel，而基础模型一例也没有；只有在检索提供固定 kernel 时，检索才会带来显著帮助。
6 相关工作
已有 GPU kernel 基准主要强调功能正确性和整体效率，但不会单独考察指定架构机制的执行情况 [34] [35] [36] [38] [39] [40]。现有适配方法利用执行反馈优化正确性和速度，而本文面向更具挑战性的受限场景，从头生成 kernels [41] [42] [43] [44] [45] [46] [47] [48]。
- 6 相关工作: GPU kernel 基准涵盖翻译、Triton 生成、设备可移植性、服务轨迹、部署和硬件限制，但主要评估正确性和整体效率，而非指定架构机制的执行情况 [34] [35] [36] [38] [39] [40]。因此，这些基准无法单独判断生成的 kernel 是否执行了目标机制，而不是依赖通用方法。
- 6 相关工作: SFT 和 RL 适配方法利用 Triton 和 CUDA 的执行反馈，通常通过 DSL、编译器、常规 CUDA 或 libraries 实现；本文则面向更具挑战性的受限场景，从头生成 kernels [41] [42] [43] [44] [45] [46] [47] [48]。CUDA Agent 允许使用现有的 cuDNN functions（参见 [50] 中的 Figure 15），而 CUDA-L2 允许使用 CUTLASS 和 CuTe [51]。
7 局限性
本研究受限于使用单个 27B 模型开展的规模有限的适配实验，以及 PTXBench 当前仅聚焦于 H100 和 B200 GPU 上的 BF16 GEMM 与 attention kernel。
- 适配范围: 由于实验采用规模有限的 LoRA 数据集和单个 27B 基础模型，适配结果可能无法原样迁移到业界规模的后训练或其他模型系列。这限制了关于 repair conditioning、数据平衡和教师质量的结论在所评估设置之外的适用性。
- 基准范围: 由于当前仅评估 H100 和 B200 GPU 上的 BF16 GEMM 与 attention kernel，PTXBench 无法代表 GPU 算子和硬件的完整多样性。该基准重点考察近期的架构特定 tensor core 和异步内存；同时，其将工作负载、架构上下文与评测流程分离的设计，可以扩展到更广泛的工作负载、模型和未来 GPU。
8 结论
PTXBench 是一个用于架构特定 PTX 编程的可审计基准与适配环境。它表明，当前 LLM 有时能够执行指定指令并解决前向工作负载，但在 H100 和 B200 上处理反向 attention 及实现有竞争力的性能方面仍存在困难。
- 8 结论: PTXBench 提供了一个用于架构特定 PTX 编程的可审计基准与适配环境。
- 8 结论: 该基准分别衡量功能正确性、目标指令的运行时执行情况，以及相对于 frontier libraries 的加速比。
- 8 结论: 当前 LLM 有时能够执行指定指令并解决前向工作负载，但在处理反向 attention 方面仍存在困难。
- 8 结论: 在 H100 和 B200 上，当前 LLM 无法始终实现有竞争力的性能。
A 附录 · A.1 实验设置
实验使用面向架构的编译和固定工作负载分析，在 H100 和 B200 GPU 上评估生成的 CUDA kernel。每条优化轨迹保留此前的 kernel 和反馈，允许进行八次模型调用，并可在达到 1.2× 加速比或耗尽上下文时提前终止。
- A.1 实验设置: Kernel 评估覆盖 NVIDIA H100 80 GB（Hopper）和 B200（Blackwell）GPU，目标架构分别为 sm 90a 和 sm 100a。MiniPTXAgent 使用 nvcc -O3 编译每个候选 kernel，并在固定的 FlashInfer-Trace 工作负载上对其进行评估。
- A.1 实验设置: 每条轨迹允许进行八次模型调用，同时保留此前的 kernel 和执行反馈。保留的历史记录使后续调用能够使用早期的候选 kernel 和分析反馈。
- A.1 实验设置: 当加速比达到 1.2× 或 LLM 耗尽上下文时，轨迹会提前终止。
A.2 目标指令执行测量 … A.6 补充结果
附录详细说明可审计的目标指令测量、profiling 服务优化、架构特定 prompt 构造，以及经过校正的 FlashAttention-3 上下文。补充图表从 speedup、推理长度、SFT、训练配方和专家指导等方面扩展评估。
- A.2 目标指令执行测量: Runtime qualification 使用 Nsight Compute 计数器重新执行静态判定为阳性的候选，同时排除无关指令并保留原始计数以供审计。论文采用布尔型 runtime-execution 指标，并排除与 tensor-computation 路径无关的 Hopper TMA cache-control 和 prefetch 指令。
- A.2 目标指令执行测量: Dynamic profiling 会拒绝所选指令仅出现在未启动 kernel 中的候选，并将未定义行为案例保守地标记为 unknown。在 stream-ordered execution 期间发生的 use-after-free 解释了 MHA-Bwd-Causal 八轮评估中 Qwen3.6-27B-s1 的 4.2% 目标指令执行率与 5.2% 不受限正确率之间的差异。
- A.3 Profiling 服务: 将生成与 profiling 流水线化可使 rollout 吞吐提升 2.78×，而复用驻留 GPU 的 workload state 可使 /evaluate 吞吐提升 2.24×。随着并发度升高，吞吐趋于饱和；四张 profiling GPU 提供的端到端吞吐约等于八张。
- A.4 架构特定 Prompt Token 数量: 架构特定 prompt pack 将参数、PTX wrapper 以及 layout 和内存一致性契约组合在一起；B200 prompt 还增加了针对 shared memory、peer CTA 和 tensor memory 的描述符。这些 pack 会根据文档、报告和执行结果进行验证，因为 PTX 文档可能存在不一致、规范不充分或错误。
- A.5 已发表与校正后的 FlashAttention-3 伪代码: 受控上下文修正了已发表 FlashAttention-3 Algorithm 2 [62] 中的两处内部不一致：不完整的循环边界和缺失的概率状态更新。校正版迭代至 j = Tc −1，向前传递 pipeline state，使用旧的行最大值重新缩放，并在 epilogue 中进行归一化。
- A.6 补充结果: 补充结果报告了 H100 和 B200 上的 Fastp/FastInst、执行目标指令时的最佳 speedup，以及五个问题上的 SFT 正确率和 speedup。所提供材料通过 Figures 13–14 和 Tables 7–9 标识了这些评估，但未提供具体单元格数值。
- A.6 补充结果: 其他补充分析涵盖推理长度、训练数据配方，以及 MHA 评估中使用的专家指导。Figure 15 比较 base model、KernelGen s0、Fixit s3 和 Gemini 3.1 Pro；Figure 16 记录专家指导。