PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2608.17379v1cs.CLcs.AI

TL;DR

LLMs lack consistent evidence of exploiting architecture-specific PTX for GPU kernel optimization, especially on complex attention workloads. PTXBench benchmarks correctness, runtime instruction execution, and performance while studying repair-conditioned adaptation, finding persistent capability gaps and uneven training gains.

  • Problem

    Evidence is limited on whether LLMs can exploit evolving architecture-specific PTX mechanisms that high-performance GPU kernels often require.

  • Method

    PTXBench combines a multi-turn CUDA–PTX generation benchmark with structured execution feedback, separate correctness, instruction-execution, and performance measures, and targeted post-training.

  • Results

    Current LLMs sometimes execute requested instructions and solve forward workloads but struggle with backward attention, while repair-conditioned training improves several tasks without uniformly outperforming direct-solution supervision.

  • Takeaways & Limitations

    Architecture-specific PTX capability and adaptation gains remain uneven across workloads, models, and training recipes rather than consistently matching frontier performance.

  • Takeaways & Limitations

    PTXBench currently focuses on BF16 GEMM and attention kernels on H100 and B200, which do not represent the full diversity of GPU operators and hardware.

Abstract

from arXiv · show

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.

1 INTRODUCTION

PTXBench evaluates whether LLMs can exploit architecture-specific PTX to produce correct, target-instruction-executing, and performant GPU kernels. It combines controlled architecture knowledge with MiniPTXAgent’s feedback-driven CUDA–PTX generation and supports characterization across models and GPUs plus repair-conditioned post-training.

  • Motivation: PTXBench addresses the need for architecture-specific optimization because portable CUDA can remain correct while failing to exploit new PTX mechanisms for tensor cores, memory units, and synchronization.PTX is the lowest-level programmable interface that CUDA developers can explicitly control, and high-performance kernels often require PTX-level optimization.
  • Benchmark design: PTXBench separately measures functional correctness, runtime execution of required target instructions, and performance relative to frontier libraries.The benchmark evaluates whether LLM-generated kernels exploit specified low-level GPU mechanisms across architectures and workloads.
  • Benchmark design: MiniPTXAgent supplies architecture-specific knowledge and structured execution feedback so models iteratively generate and revise CUDA kernels with inline PTX.The resulting CUDA–PTX kernels are produced through a multi-turn agent loop.
  • Adaptation study: The framework characterizes architecture-specific PTX capability across current models and GPUs, then tests targeted, repair-conditioned post-training.PTXBench is designed both as a controlled measurement framework and as an environment for collecting adaptation data.

2 PTXBENCH

PTXBench evaluates architecture-specific PTX kernel generation through controlled architecture knowledge, runtime verification of target instructions, functional correctness, and efficiency against frontier libraries. Its workflow generates kernels from scratch, iteratively compiles and profiles candidates, and measures correctness and speedup under fixed workloads.

  • Benchmark design: PTXBench pairs frontier-library reference operators [17] [18] with fixed workloads, target GPU architectures, and required architecture-specific PTX instruction families, while generating CUDA kernels with inline PTX from scratch.The FlashInfer-Trace schema represents workloads, solutions, and correctness checks across the benchmark workflow.
  • Benchmark design: Every trajectory receives identical architecture-specific knowledge packs containing hardware parameters, PTX wrappers, and contracts for layouts, synchronization, and memory consistency.Packs occupy 20k–30k tokens; selected well-studied operators also receive expert-validated scheduling principles.
  • Evaluation workflow: MiniPTXAgent iteratively generates, compiles, and revises kernels using execution feedback, sending only successful compilations to profiling for safety checks, functional evaluation, runtime errors, and latency measurement [19] [20] [21].Each trajectory allows a fixed number of model calls and retains prior kernels and feedback.
  • Evaluation criteria: Target-instruction correctness requires functional correctness plus runtime execution of selected tensor paths: GMMA or UTMA on Hopper and TCGEN05 on Blackwell.TMA alone does not qualify on Blackwell; execution is verified through SASS analysis.
  • Evaluation criteria: Correct candidates are compared with references using torch.allclose at atol=rtol=1e-2, while speedup is reference latency divided by candidate latency after 10 warmups and the median of 50 timed iterations.Latency measurement uses CUPTI, and cuDNN or cuBLAS headers in candidate CUDA code are considered incorrect.

3 ADAPTING LLMS TO ARCHITECTURE-SPECIFIC PTX PROGRAMMING

This section presents Fixit, a supervision strategy for architecture-specific PTX programming that uses model failures, execution feedback, and teacher-generated repairs and rationales. It targets the data bottleneck caused by scarce expert knowledge and costly kernel validation.

  • Motivation: The approach addresses a durable optimization target whose main obstacle is scarce expert knowledge and iterative validation for high-quality kernels.Deployed GPUs retain fixed low-level capabilities for years, motivating post-training on architecture-specific PTX.
  • Fixit: Fixit constructs supervision from failures produced by the model being adapted, combining problem context, failed kernels, and compilation or execution feedback.Failures are sampled from the pre-adaptation model and evaluated for functional correctness.
  • Fixit: A repair teacher generates corrected kernels from each failure, retaining only repairs that pass the functional-correctness check.The repair is conditioned on the same problem, failure, and feedback as the failed kernel.
  • Fixit: A reasoning teacher then synthesizes a rationale connecting the observed failure to the retained repair.Each training example supervises the student with the teacher-generated rationale and corrected kernel.

4 BENCHMARK RESULTS

Benchmark results show that architecture-specific PTX capability varies sharply by model, workload, and GPU generation: Blackwell attention is especially difficult, while higher-level Triton often remains more reliable than direct CUDA-PTX. Explicit PTX knowledge improves target-instruction execution and correctness, but does not eliminate the performance gap.

  • Model comparison: Claude Opus 4.8 reaches 1.012× cuBLAS on Blackwell GEMM, while Gemini 3.1 Pro reaches 0.892× despite a cutoff coinciding with Blackwell PTX release.Claude also achieves a substantially higher target-instruction correctness rate on Blackwell.
  • Model comparison: GLM-5.2 is competitive with Gemini 3.1 Pro on GEMM across H100 and B200 and on unmasked H100 attention, but lags on B200 attention.GLM-5.2 and Claude Opus 4.8 often fall back to generic CUDA on Blackwell, whereas Gemini more successfully executes Blackwell instructions.
  • Implementation comparison: On Blackwell, Triton reaches 0.484× and 0.436× on two attention backward workloads, versus 0.133× and 0.015× for CUDA-PTX.On Hopper, CUDA-PTX nearly matches Triton on GEMM and reaches 0.768× versus 0.759× for causal MHA forward.
  • Architecture effects: Blackwell’s specialized mechanisms and limited Blackwell-specific training data may explain why direct CUDA-PTX generation trails higher-level languages, while CuTeDSL’s success rate remains extremely low.High-level languages help produce correct kernels on new architectures, but direct low-level implementations can sometimes match or outperform them once correct.
  • Knowledge ablation: With architecture parameters alone, the model achieves 26.0% correct turns but no target-instruction successes; template functions enable execution, and an architecture contract improves correctness.The ablation shows that explicit PTX knowledge is necessary rather than automatically inferred from architecture parameters.

5 ADAPTATION RESULTS

Fixit SFT with Qwen3.6-27B improves selected PTX kernel outcomes, but gains depend on training format, problem balance, reasoning-teacher strength, and workload or language transfer. The adapted model broadens executable-kernel search and improves some peak speedups, while leaving substantial correctness and generalization gaps.

  • Training format: Fixit conditioning improves correctness on GEMM, MHA-Fwd-Causal, and MHA-Bwd but trails direct-solution supervision on MHA-Fwd and MHA-Bwd-Causal.The comparison uses the same eight problem classes, GLM-5.2 reasoning, and similar record counts; thus, failure-conditioned supervision helps selectively rather than uniformly.
  • Coverage and balance: Balanced problem coverage, rather than record count alone, determines adaptation breadth: s1 and s5 solve all five problems, while larger imbalanced recipes still fail on MHA-Bwd-Causal.Increasing s4 to balanced s5 by 1.5× improves eight-turn correctness on four problems, ties MHA-Fwd, and improves peak speedup on four; the best recipe still varies by problem.
  • Reasoning synthesizer: A stronger reasoning teacher is important: with identical Fixit examples, s5 solves all five problems, whereas s6 using Qwen3.6-27B as synthesizer solves only GEMM.Target-model failures are useful training inputs, but generating repair rationales benefits from a stronger teacher.
  • Transfer and generalization: S1 transfers across some related workloads but not uniformly, producing no correct kernels on either d96 backward task or GQA after training on four d128 MHA tasks.The d96 backward tasks are especially difficult because d96 does not align with H100 WGMMA’s m64 tile.
  • Transfer and generalization: Under Triton transfer, s1 lowers turn-level correctness on every workload but raises best correct speedup from 0.238× to 0.632× on MHA-Fwd-Causal and from 0.043× to 0.331× on MHA-Bwd-Causal.The CUDA-PTX-trained recipe can improve peak performance despite harming cross-language correctness.
  • Search process and prompt-time supervision: SFT lengthens reasoning and shifts failures from compilation toward runtime and numerical errors, while expert guidance cannot make the base model write correct MHA kernels.S1 produces correct GEMM kernels at turn 0 and in at least one of seven later turns, whereas the base model produces none; retrieval helps sharply only when it supplies a fixed kernel.

6 RELATED WORK

Prior GPU-kernel benchmarks emphasize functional correctness and overall efficiency but do not isolate execution of specified architecture mechanisms [34] [35] [36] [38] [39] [40]. Existing adaptation methods use execution feedback to optimize correctness and speed, whereas this work targets a harder constrained regime of generating kernels from scratch [41] [42] [43] [44] [45] [46] [47] [48].

  • 6 RELATED WORK: GPU-kernel benchmarks span translation, Triton generation, device portability, serving traces, deployment, and hardware limits, but primarily assess correctness and overall efficiency rather than specified architecture-mechanism execution [34] [35] [36] [38] [39] [40].The benchmarks therefore do not isolate whether a generated kernel executes the target mechanism instead of relying on generic approaches.
  • 6 RELATED WORK: SFT and RL adaptation methods use execution feedback for Triton and CUDA, often through DSLs, compilers, ordinary CUDA, or libraries; this work targets generating kernels from scratch in a harder constrained regime [41] [42] [43] [44] [45] [46] [47] [48].CUDA Agent permits existing cuDNN functions (cf. Figure 15 in [50]), while CUDA-L2 permits CUTLASS and CuTe [51].

7 LIMITATIONS

The study is limited by modest adaptation experiments using one 27B model and by PTXBench’s current focus on BF16 GEMM and attention kernels on H100 and B200 GPUs.

  • Adaptation scope: Adaptation results may not transfer unchanged to industry-scale post-training or other model families because experiments use modest LoRA datasets and a single 27B base model.This limits conclusions about repair conditioning, data balance, and teacher quality beyond the evaluated setting.
  • Benchmark scope: PTXBench does not represent the full diversity of GPU operators and hardware because it currently evaluates BF16 GEMM and attention kernels on H100 and B200 GPUs.The benchmark stresses recent architecture-specific tensor cores and asynchronous memory, while its separated workload, architecture-context, and evaluation workflow can extend to broader workloads, models, and future GPUs.

8 CONCLUSION

PTXBench is an auditable benchmark and adaptation environment for architecture-specific PTX programming. It reveals that current LLMs may execute requested instructions and solve forward workloads, yet struggle with backward attention and competitive performance across H100 and B200.

  • 8 CONCLUSION: PTXBench provides an auditable benchmark and adaptation environment for architecture-specific PTX programming.
  • 8 CONCLUSION: The benchmark separately measures functional correctness, runtime execution of target instructions, and speedup over frontier libraries.
  • 8 CONCLUSION: Current LLMs can sometimes execute requested instructions and solve forward workloads, but struggle with backward attention.
  • 8 CONCLUSION: Across H100 and B200, current LLMs do not consistently achieve competitive performance.

A APPENDIX · A.1 EXPERIMENT SETUP

The experiment evaluates generated CUDA kernels on H100 and B200 GPUs using architecture-specific compilation and fixed workload profiling. Each optimization trajectory retains prior kernels and feedback, allows eight model calls, and can terminate early upon reaching 1.2× speedup or exhausting context.

  • A.1 EXPERIMENT SETUP: Kernel evaluation spans NVIDIA H100 80 GB (Hopper) and B200 (Blackwell) GPUs, targeting sm 90a and sm 100a, respectively.MiniPTXAgent compiles each candidate with nvcc -O3 and evaluates it against the fixed FlashInfer-Trace workload.
  • A.1 EXPERIMENT SETUP: Each trajectory permits eight model calls while retaining preceding kernels and execution feedback.The retained history enables subsequent calls to use earlier candidate kernels and profiling feedback.
  • A.1 EXPERIMENT SETUP: Trajectories terminate early when speedup reaches 1.2× or the LLM runs out of context.

A.2 TARGET INSTRUCTION EXECUTION MEASUREMENT … A.6 SUPPLEMENTARY RESULTS

The appendix details auditable target-instruction measurement, profiling-service optimizations, architecture-specific prompt construction, and corrected FlashAttention-3 context. Supplementary figures and tables extend evaluation across speedup, reasoning length, SFT, training recipes, and expert guidance.

  • A.2 TARGET INSTRUCTION EXECUTION MEASUREMENT: Runtime qualification replays statically positive candidates with Nsight Compute counters, while excluding unrelated instructions and retaining raw counts for audit.The paper uses a Boolean runtime-execution indicator and excludes Hopper TMA cache-control and prefetch instructions unrelated to tensor-computation paths.
  • A.2 TARGET INSTRUCTION EXECUTION MEASUREMENT: Dynamic profiling rejects candidates whose selected instructions appear only in unlaunched kernels and conservatively marks undefined-behavior cases as unknown.A use-after-free during stream-ordered execution explains the 4.2% target-instruction rate versus 5.2% unrestricted correctness rate for Qwen3.6-27B-s1 on MHA-Bwd-Causal at eight turns.
  • A.3 PROFILING SERVICE: Pipelining generation with profiling increases rollout throughput 2.78×, while GPU-resident workload-state reuse improves /evaluate throughput 2.24×.Throughput saturates as concurrency rises, and four profiling GPUs provide approximately the same end-to-end throughput as eight.
  • A.4 ARCHITECTURE-SPECIFIC PROMPT TOKEN COUNTS: Architecture-specific prompt packs combine parameters, PTX wrappers, and contracts for layouts and memory consistency, with B200 prompts adding descriptors for shared memory, peer CTAs, and tensor memory.The packs are validated against documentation, reports, and execution because PTX documentation can be uneven, underspecified, or wrong.
  • A.5 PUBLISHED AND CORRECTED FLASHATTENTION-3 PSEUDOCODE: The controlled context corrects two internal inconsistencies in published FlashAttention-3 Algorithm 2 [62]: an incomplete loop bound and a missing probability-state update.The corrected version iterates through j = Tc −1, carries pipeline states forward, rescales using the old row maximum, and normalizes in the epilogue.
  • A.6 SUPPLEMENTARY RESULTS: Supplementary results report H100 and B200 Fastp/FastInst, best speedups with target-instruction execution, and five-problem SFT correctness and speedups.The supplied materials identify these evaluations through Figures 13–14 and Tables 7–9, without providing their cell values.
  • A.6 SUPPLEMENTARY RESULTS: Additional supplementary analyses cover reasoning lengths, training-data recipes, and expert guidance used in MHA evaluation.Figure 15 compares the base model, KernelGen s0, Fixit s3, and Gemini 3.1 Pro; Figure 16 documents expert guidance.
Loading 2608.17379v1…