PTXBench: アーキテクチャ固有のPTXを用いたGPUカーネル最適化のためのLLMのベンチマークと適応

要旨

本稿では、GPUカーネル最適化のためにアーキテクチャ固有のPTXを利用する大規模言語モデル（LLM）を評価し、適応させるためのベンチマークであるPTXBenchを紹介する。PTXBenchは、機能的正確性、選択された対象命令が実行時に実行されるかどうか、およびH100およびB200 GPU上のGEMMおよびアテンションのワークロードにおける最先端ライブラリに対する高速化を測定する。我々の評価は、アーキテクチャ固有のPTXの能力が依然として不均一であることを示している。すなわち、複雑なアテンションの逆伝播ワークロードでは成功率が大幅に低下し、対象命令を実行しても必ずしも競争力のある性能にはつながらない。評価されたモデルのいずれも、スイート全体にわたって最先端ライブラリに一貫して匹敵するものはなかった。さらに、教師ありファインチューニングを用いてQwen3.6-27Bを適応させる。修復条件付きトレーニングはいくつかのタスクを改善するが、汎化は依然として不均一である。データセットサイズに加えて、データのカバレッジ、バランス、推論教師モデルの質が重要である。PTXBenchは、進化するGPUアーキテクチャを活用するLLMの能力を測定し向上させるための、監査可能なテストベッドを提供する。

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.