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

The Problem: LLMs Can't Reliably Write Fast GPU Code

Large language models struggle to write optimized GPU kernels using architecture-specific instructions (PTX) for modern hardware like NVIDIA's H100 and B200 chips. Even when an LLM generates correct code that runs, it often doesn't actually use the specialized instructions it claims to, and the resulting kernels run slower than hand-optimized libraries. This gap matters because GPU kernel performance directly impacts AI workload speed, yet we lack a systematic way to measure or improve LLM performance on this task.

The authors introduce PTXBench, a benchmark that exposes these failures and provides a testing ground for fixing them. Their core finding: existing LLMs are inconsistent at GPU optimization, and current fine-tuning approaches only partially solve the problem.

How the Benchmark Works

PTXBench evaluates LLMs on two GPU-intensive workloads—GEMM (matrix multiplication) and attention mechanisms—targeting H100 and B200 architectures. For each task, it measures three things:

1. Functional correctness: Does the generated code produce the right numerical result?
2. Instruction execution: Does the code actually run the target PTX instructions it claims to use, not just fall back to generic alternatives?
3. Speedup: How does the kernel's runtime compare to state-of-the-art libraries like cuBLAS and FlashAttention?

The benchmark includes both forward and backward passes for attention, with backward passes being particularly challenging. This three-tier evaluation catches a critical failure mode: code can be correct and claim to use fast instructions while still executing generic fallback code that provides no speedup.

What the Authors Found

Across all tested models (including GPT-4, Claude, and open-source LLMs), success rates drop dramatically on complex attention backward kernels. Most striking: even when target instructions execute, the resulting code rarely matches frontier library performance. No single model was consistently competitive across the full suite.

The authors then fine-tuned Qwen 3.6-27B (a 27-billion-parameter open-source model) using supervised fine-tuning, where the model learns from expert-written solutions. They also tried "repair-conditioned training"—training on pairs of broken and corrected kernels—which improved performance on some tasks but created uneven generalization. The key lesson: dataset quality and diversity matter as much as size. A well-reasoned teacher signal and balanced task coverage outperformed simply scaling training data.

Why This Matters

GPU kernel optimization is a concrete, measurable challenge where LLMs have clear value but also clear limits. PTXBench provides an auditable way to track progress as GPUs evolve and new architectures emerge. Rather than relying on anecdotal reports of LLM capability, researchers can now quantify exactly where models fail—whether it's generating syntactically correct but semantically wrong code, or optimizing poorly despite correct execution.

The benchmark also reveals that fine-tuning helps but isn't a silver bullet. Improving LLM performance on GPU kernels requires careful attention to training data quality and representation, not just parameter count. This sets a baseline for the community and establishes that systematic benchmarking and adaptation are both necessary to close the gap with hand-optimized libraries.