PTXBench: Benchmark and Adapt LLMs for GPU Kernel Optimization with Architecture-specific PTX
Genghan Zhang, Yixin Dong, Chengze Fan, Zhichen Zeng, Yueming Yuan, Shaowei Zhu, Kunle Olukotun
cs.CL, cs.AI
2026-08-18
PTXBench scores correctness, runtime PTX use, and speedup vs cuBLAS/cuDNN. Opus 4.8 is strong on H100 forward, weak on backward; no model matches vendor libraries.
KernelBench-style suites ask whether a model can emit a correct GPU kernel that beats PyTorch. That win can come from ordinary CUDA, or from calling cuBLAS. New GPUs earn their keep on GMMA, TMA, and Blackwell's TCGEN05. The question a kernel author actually has is whether a model will write those instructions for a named architecture and whether they run on the evaluated workload.
A PTXBench item pins down a reference (cuBLAS, cuDNN, or FlashInfer), a fixed workload, a target GPU, and a required PTX family. The model writes CUDA with inline PTX from scratch. Including vendor headers counts as wrong. Every trajectory gets the same 20k–30k-token architecture pack: parameters, instruction wrappers, and memory/sync contracts. MiniPTXAgent runs up to eight turns: nvcc in a CPU container, then a numeric check (atol=rtol=1e-2) and CUPTI timing.
Target-instruction correctness is stricter than functional correctness. The selected family must appear in SASS, and Nsight Compute must report a positive predicate-enabled thread count, so dead code does not pass. Hopper counts GMMA or UTMA. Blackwell counts TCGEN05; inherited TMA is not enough. Metrics are per-turn correctness and Fastp, the fraction of turns that are correct and faster than threshold p. Each problem uses 3 prompts × 4 eight-turn trajectories (N=32). Models are Gemini 3.1 Pro, Claude Opus 4.8, GLM-5.2, and Qwen3.6-27B, on H100 and B200.
Capability is uneven. At eight turns on H100, Opus 4.8 reaches 94.8% target-instruction correctness on GEMM and 90.6% on multi-head forward, then 44.8% on causal backward. Qwen3.6-27B produces no correct Hopper kernel. On B200 the split flips: Opus is 91.7% functionally correct on backward attention and only 10.4% target-instruction correct. The model can write a right kernel that never takes the new tensor path. No model matches vendor libraries across the suite. Gemini hits 0.892× cuBLAS on Blackwell GEMM, Opus 1.012×; attention stays far worse than on Hopper.
Triton is the stabler choice on a new chip. On Hopper, CUDA-PTX and Triton are close (causal forward peaks 0.768× versus 0.759×). On two Blackwell backward workloads Triton reaches 0.484× and 0.436×; raw PTX reaches 0.133× and 0.015×. Without the knowledge pack, Gemini still has 26.0% correct turns at eight steps and zero target-instruction hits. Adding templates plus the architecture contract lifts target-instruction correctness to 38.5%.
LoRA on Qwen3.6-27B with Gemini as the repair teacher is mixed. Fixit does not uniformly beat direct-solution supervision. Coverage and class balance matter more than raw record count. Swap the reasoning teacher for the student itself and only GEMM survives. The smallest all-task recipe, s1, transfers to some d=64 and d=96 forward shapes, and to none of the d=96 backward or GQA tasks. On Triton, turn correctness falls, but the causal-forward peak rises from 0.238× to 0.632×.
A correct kernel and a kernel that uses this generation of hardware are different claims. Kernel-agent leaderboards should split the runtime instruction probe from vendor-library speedup, or generic CUDA and quiet library calls will inflate the ranking. The knowledge pack is not optional: without the contract, models almost never emit the new PTX. Repair-conditioned SFT can unlock a 27B, but an unbalanced mix drops tasks. It is not a general training recipe.
Adaptation uses one 27B and modest LoRA sets. Repair conditioning, balance, and teacher quality may change sign at industry scale or on other families. The suite is BF16 GEMM and attention on H100 and B200, not the full operator surface. atol=1e-2 is loose for numerically fragile kernels. The instruction probe shows that an opcode ran, not that it caused the speedup.