用PTX指令优化GPU核函数,评测大模型表现并改进其适配能力。
PTXBench: Benchmark and Adapt LLMs for GPU Kernel Optimization with Architecture-specific PTX

- 构建PTXBench基准,评估大模型生成特定架构PTX代码的能力。
- 复杂注意力反向传播任务成功率低,指令执行不保证性能提升。
- 微调后模型泛化仍不均衡,数据质量与教师推理至关重要。
我们提出PTXBench,一个用于评估和适配大语言模型(LLMs)使用架构特定PTX进行GPU核函数优化的基准。该基准在H100和B200 GPU上,针对GEMM和注意力计算任务,衡量功能正确性(目标指令是否在运行时执行)以及相对于前沿库的速度提升。评估显示,架构特定PTX能力仍不均衡:复杂注意力反向工作负载的成功率显著下降,且目标指令执行并不必然带来竞争力性能。所有评估模型均未在全套任务中一致超越前沿库。我们进一步对Qwen3.6-27B进行监督微调。修复条件训练提升了多个任务表现,但泛化能力仍不均衡;数据覆盖度、分布平衡性以及推理教师的质量,除数据规模外,同样关键。PTXBench为测量和提升大模型利用演进式GPU架构的能力提供可审计测试平台。
原文摘要 · Abstract (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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。