PTXBench:利用架构特定PTX对GPU内核优化进行基准测试并适配大语言模型
PTXBench: Benchmark and Adapt LLMs for GPU Kernel Optimization with Architecture-specific PTX
August 18, 2026
作者: Genghan Zhang, Yixin Dong, Chengze Fan, Zhichen Zeng, Yueming Yuan, Shaowei Zhu, Kunle Olukotun
cs.AI
摘要
我们推出PTXBench,这是一个用于评估和适配大语言模型(LLMs)以利用架构特定的PTX指令进行GPU内核优化的基准测试。PTXBench在H100和B200 GPU上,针对GEMM和注意力负载,衡量功能正确性、所选目标指令是否在运行时实际执行,以及相较于前沿库的加速比。我们的评估表明,架构特定的PTX能力仍不均衡:在复杂的注意力反向负载上,成功率显著下降,且执行目标指令并不一定能转化为具有竞争力的性能。所有被评估的模型均未能在整个测试套件中持续匹敌前沿库。我们进一步使用监督微调适配Qwen3.6-27B。修复条件训练改善了若干任务,但泛化能力仍不均衡;除了数据集规模外,数据覆盖度、平衡性以及推理教师的质量同样重要。PTXBench为衡量和提升LLMs利用不断演进的GPU架构的能力,提供了一个可审计的测试平台。
English
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.