PTXBench: Benchmark and Adapt LLMs for GPU Kernel Optimization with Architecture-specific PTX

📄 arXiv: 2608.17379v1 📥 PDF

作者: Genghan Zhang, Yixin Dong, Chengze Fan, Zhichen Zeng, Yueming Yuan, Shaowei Zhu, Kunle Olukotun

分类: cs.CL, cs.AI

发布日期: 2026-08-18

💡 一句话要点

提出PTXBench以优化GPU内核的LLM适配问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: GPU内核优化 大型语言模型 架构特定PTX 基准测试 性能评估 监督微调 深度学习

📋 核心要点

  1. 现有方法在复杂的GPU内核优化中存在性能不均衡和适配困难的问题，尤其是在注意力反向工作负载上表现不佳。
  2. PTXBench通过引入架构特定的PTX，提供了一种新的基准测试方法，旨在评估和适配LLMs以优化GPU内核性能。
  3. 实验结果显示，尽管某些任务的性能有所提升，但在复杂任务上的泛化能力仍然不足，且没有模型在所有测试中超越前沿库。

📝 摘要（中文）

我们介绍了PTXBench，这是一个用于评估和适配大型语言模型（LLMs）以利用特定架构PTX进行GPU内核优化的基准测试工具。PTXBench测量功能正确性、所选目标指令在运行时的执行情况，以及在H100和B200 GPU上针对GEMM和注意力工作负载的速度提升。评估结果表明，架构特定PTX的能力仍然不均衡：在复杂的注意力反向工作负载上成功率显著下降，执行目标指令并不一定能转化为竞争力的性能。经过监督微调后，我们进一步适配了Qwen3.6-27B，修复条件训练在多个任务上有所改善，但泛化能力仍然不均衡；数据覆盖、平衡以及推理教师的质量在数据集大小之外也起着重要作用。PTXBench为测量和提升LLMs利用不断发展的GPU架构的能力提供了可审计的测试平台。

🔬 方法详解

问题定义：论文要解决的问题是如何有效评估和适配大型语言模型（LLMs）以利用特定架构的PTX进行GPU内核优化。现有方法在复杂的注意力反向工作负载上表现不佳，导致性能不均衡。

核心思路：论文的核心思路是通过PTXBench基准测试工具，系统地评估LLMs在不同GPU架构上的性能表现，并通过监督微调来提升模型的适配能力。这样的设计旨在确保模型能够更好地利用GPU的架构特性。

技术框架：PTXBench的整体架构包括功能正确性测量、目标指令执行情况评估以及与前沿库的速度比较。主要模块涵盖了GEMM和注意力工作负载的测试，使用H100和B200 GPU进行评估。

关键创新：最重要的技术创新点在于引入了架构特定的PTX，使得LLMs能够针对特定GPU架构进行优化。这与现有方法的本质区别在于，PTXBench提供了一个系统化的评估平台，而不仅仅是单一模型的性能测试。

关键设计：在关键设计方面，论文采用了监督微调的方法来适配Qwen3.6-27B模型，修复条件训练被用于提升模型在多个任务上的表现。同时，数据覆盖、平衡和推理教师的质量被认为是影响模型泛化能力的重要因素。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，尽管在某些任务上通过修复条件训练实现了性能提升，但在复杂的注意力反向工作负载上，成功率显著下降，且没有任何模型在所有测试中超越前沿库，显示出当前方法的局限性。

🎯 应用场景

该研究的潜在应用领域包括高性能计算、深度学习模型优化以及GPU架构设计等。通过提升LLMs在GPU内核优化中的适配能力，PTXBench能够为开发更高效的计算框架和算法提供支持，推动相关领域的技术进步。

📄 摘要（原文）

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.