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

Inside the modern computer, the graphics processing unit, or GPU, acts as a powerful engine designed to handle massive amounts of data simultaneously. To get the most out of this engine, software developers write specialized instructions that tell the hardware exactly how to move data and perform calculations. These instructions are written in a low-level language called PTX, which gives the programmer direct control over the machine's inner workings. While higher-level tools exist to make this process easier, they sometimes miss the unique features of the newest hardware generations. As these machines evolve rapidly, keeping software perfectly tuned to their specific capabilities becomes a difficult challenge. The question arises: can artificial intelligence, specifically large language models, learn to write these precise, low-level instructions on its own, or does it still rely on older, generic methods that leave performance on the table?

A team of researchers set out to answer this by creating a new testing ground called PTXBench. They wanted to see if these AI models could not only write code that works but also code that actively uses the specific, advanced features of the latest GPUs. The researchers focused on two of the most powerful graphics cards available, the H100 and the B200, and tested the models on tasks involving complex mathematical operations used in artificial intelligence, such as matrix multiplication and attention mechanisms. The goal was to measure three things: whether the code ran without errors, whether it actually triggered the specific hardware instructions the researchers asked for, and whether it ran faster than the best existing software libraries.

The results revealed a significant gap between what the models could do and what was required for top-tier performance. While the AI models were often successful at writing code that worked correctly for simpler tasks, they struggled considerably with more complex operations, particularly those involving the backward pass of attention mechanisms, which are crucial for training AI systems. Even when the models managed to write code that executed the requested hardware instructions, that code rarely matched the speed of the leading professional libraries. In many cases, the AI-generated code was slower, suggesting that simply getting the instructions to run was not enough to achieve true optimization. The study found that the models frequently fell back on generic coding patterns rather than leveraging the unique, architecture-specific features of the newer chips, even when those features were explicitly described to them.

To understand how to improve these models, the researchers conducted a controlled experiment where they taught a specific AI model using a method called supervised fine-tuning. They showed the model examples of its own mistakes, along with corrected versions of the code and explanations of why the corrections worked. This approach, which they called Fixit, helped the model improve on several tasks, particularly in generating code that was functionally correct. However, the improvement was not uniform. The model still struggled to generalize its new skills to different problem sizes or variations of the tasks it had not seen during training. The researchers discovered that the quality of the corrections and the balance of the training data mattered just as much as the sheer amount of data. A model trained on a diverse set of problems performed better than one trained on a larger but less varied set, indicating that the type of learning experience was more important than the volume.

The study also highlighted the importance of providing the AI with the right context. When the researchers gave the models detailed information about the hardware's specific capabilities, the models were much more likely to use the requested instructions. Without this specific knowledge, the model rarely uses the requested PTX. This suggests that while AI can learn to program at this low level, it requires precise guidance and high-quality examples to do so effectively. The researchers concluded that while current AI models show promise, they are not yet ready to replace expert human engineers in optimizing code for the newest hardware. The path forward involves better training data, more balanced learning experiences, and a deeper understanding of how to teach these models to reason about the specific constraints of evolving computer architectures.