PTXBench: Benchmarking and Adapting LLMs for GPU Kernel Optimization with Architecture-specific PTX

Genghan Zhang; Yixin Dong; Chengze Fan; Zhichen Zeng; Yueming Yuan; Shaowei Zhu; Kunle Olukotun

cs.CLAug 18, 2026

PTXBench: Benchmarking and Adapting LLMs for GPU Kernel Optimization with Architecture-specific PTX

Authors: Genghan Zhang, Yixin Dong, Chengze Fan, Zhichen Zeng, Yueming Yuan, Shaowei Zhu, Kunle Olukotun

Organizations: Stanford University · Carnegie Mellon University · Independent Researcher · RadixArk

Abstract

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.

Figures & tables

Figure 1: Overview of the PTXBench benchmark and adaptation workflow.

GPU	Model	GEMM	MHA-Fwd	MHA-Fwd-Causal	MHA-Bwd	MHA-Bwd-Causal
H100	Gemini 3.1 Pro	33.3 / 60.4 / 56.2 (33.3 / 62.5 / 59.4)	33.3 / 39.6 / 45.8 (33.3 / 39.6 / 45.8)	25.0 / 52.1 / 55.2 (25.0 / 52.1 / 55.2)	8.3 / 33.3 / 38.5 (8.3 / 33.3 / 38.5)	– / 22.9 / 25.0 (8.3 / 25.0 / 26.0)
	Claude Opus 4.8	91.7 / 95.8 / 94.8 (91.7 / 95.8 / 94.8)	50.0 / 81.2 / 90.6 (66.7 / 89.6 / 94.8)	50.0 / 77.1 / 75.0 (83.3 / 89.6 / 82.3)	8.3 / 62.5 / 79.2 (50.0 / 83.3 / 89.6)	– / 22.9 / 44.8 (25.0 / 41.7 / 59.4)
	GLM-5.2	33.3 / 60.4 / 61.5 (33.3 / 62.5 / 62.5)	16.7 / 8.3 / 8.3 (16.7 / 14.6 / 15.6)	– / 8.3 / 10.4 (8.3 / 16.7 / 17.7)	8.3 / 6.2 / 5.2 (8.3 / 18.8 / 16.7)	– (– / 22.9 / 15.6)
	Qwen3.6-27B	–	–	–	–	–
	GPT-5.6 xhigh	91.7 / 81.2 / 77.1 (91.7 / 83.3 / 79.2)	8.3 / 39.6 / 60.4 (16.7 / 58.3 / 69.8)	– / 50.0 / 56.2 (– / 64.6 / 63.5)	8.3 / 39.6 / 55.2 (58.3 / 60.4 / 65.6)	– / 31.2 / 52.1 (33.3 / 56.2 / 65.6)
B200	Gemini 3.1 Pro	8.3 / 47.9 / 45.8 (8.3 / 47.9 / 45.8)	– / 12.5 / 15.6 (– / 12.5 / 15.6)	– / 8.3 / 10.4 (8.3 / 12.5 / 12.5)	– / 2.1 / 2.1 (– / 2.1 / 2.1)	– (– / 2.1 / 1.0)

Table 1: Target-instruction and unrestricted turn correctness rates (%) on H100 and B200.

Figure 2: $\mathrm{Fast}^{\mathrm{Inst.}}_{p}$ on H100 (top) and B200 (bottom).

Model	SWE-bench Pro (%)	Model release	Knowledge cutoff	Lag after PTX release (months)
				Hopper	Blackwell
Gemini 3.1 Pro	54.2	Feb. 2026	Jan. 2025	25	$\approx 0$
Claude Opus 4.8	69.2	May 2026	Jan. 2026	37	12
GPT-5.6 Sol	64.6	July 2026	Feb. 2026	38	13
GLM-5.2	62.1	June 2026	Not disclosed	$\leq 42$	$\leq 17$
Qwen3.6-27B	53.5	Apr. 2026	Not disclosed	$\leq 40$	$\leq 15$

Table 2: Model release dates, knowledge cutoffs, vendor-reported SWE-bench Pro scores, and estimated calendar lag from the release of Hopper PTX ISA 8.0 (Dec. 2022) and Blackwell PTX ISA 8.7 (Jan. 2025) [ 25 , 28 ] . Lags use monthly granularity, from PTX release to disclosed cutoff; otherwise, model release dates provide upper bounds.

Figure 3: Gemini 3.1 Pro $\mathrm{Fast}_{p}$ distributions for Triton and CUDA-PTX on H100 and B200.

Method or ratio	GEMM	MHA-Fwd	MHA-Fwd-Causal	MHA-Bwd	MHA-Bwd-Causal
Multi-turn refinement	$0.505\times$	$0.584\times$	$0.560\times$	$0.240\times$	$0.166\times$
Antigravity	$0.583\times$	$0.425\times$	$0.413\times$	$0.053\times$	$0.000\times$
Agent / refinement total tokens	$3.32\times$	$5.06\times$	$3.90\times$	$3.73\times$	$3.52\times$

Table 3: Antigravity versus multi-turn refinement after eight H100 evaluation requests per trajectory. Speedups are medians of four per-trajectory best instruction-qualified values. Token ratios divide the agent’s median cumulative total tokens by the refinement median within each workload.

Prompt knowledge	Target inst. correctness (%) (turn correctness)	Target inst. best speedup (best speedup)
Architecture parameters	– (50.0 / 29.2 / 26.0)	– (0.056 / 0.119 / 0.273)
Architecture parameters + PTX template functions	– / 20.8 / 19.8 (– / 20.8 / 19.8)	– / 0.542 / 0.542 (– / 0.542 / 0.542)
Architecture parameters + PTX template functions + architecture contract	8.3 / 33.3 / 38.5 (8.3 / 33.3 / 38.5 )	0.206 / 0.375 / 0.515 (0.206 / 0.375 / 0.515)

Table 4: Ablation of architecture-specific prompt knowledge for Gemini 3.1 Pro.

Figure 4: Base versus s1 generalization on the s1 training problem classes and held-out workloads. Top: Qwen3.6-27B before adaptation. Bottom: Qwen3.6-27B-s1, trained with shuffle seed 0. Runtime-instruction rows require dynamically observed Hopper SASS execution.

Figure 5: Six-recipe data ablation. Pie area is proportional to record count, slices show the fraction drawn from each problem, and the bottom line identifies the reasoning synthesizer.

Figure 6: Six-recipe SFT data ablation at eight turns (complete results in Appendix Tables 10 and 11 ).

Figure 7: How Fixit SFT changes reasoning length and error types across turns.

Figure 8: Cross-language transfer of Fixit SFT from CUDA-PTX to Triton on Hopper.

Figure 9: Turn-level error-state transitions for GEMM.

Condition	MHA-Fwd	MHA-Fwd-Causal	MHA-Bwd	MHA-Bwd-Causal
Qwen3.6-27B w/o expert guidance	–	–	–	–
Qwen3.6-27B	–	–	–	–
Qwen3.6-27B-s1 w/o expert guidance	–	– / – / 4.2 (– / – / 4.2)	– / – / 4.2 (– / – / 4.2)	– / 4.2 / 3.1 (– / 4.2 / 3.1)
Qwen3.6-27B-s1	16.7 / 16.7 / 19.8 (16.7 / 16.7 / 19.8)	– / – / 3.1 (– / – / 3.1)	– / 2.1 / 4.2 (– / 2.1 / 4.2)	– / 2.1 / 4.2 (– / 2.1 / 5.2)
Qwen3.6-27B + retrieved repair notes	–	–	–	–
Qwen3.6-27B + retrieved repair notes and fixed kernel	– / 29.2 / 29.2 (– / 29.2 / 29.2)	– / 20.8 / 27.1 (– / 20.8 / 27.1)	– / 14.6 / 18.8 (– / 14.6 / 20.8)	– / 20.8 / 33.3 (– / 25.0 / 37.5)

Table 5: Target-instruction and unrestricted turn correctness rates (%) under SFT and prompt-time supervision.

Figure 10: $\mathrm{Fast}^{\mathrm{Inst.}}_{p}$ under SFT and prompt-time supervision.

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

Mechanism	Indicator	PTX interface	SASS execution criterion
Tensor compute	$C_{\mathrm{Gen4}}$	wgmma.mma_async	$N_{\mathrm{HGMMA}}>0$
Tensor compute	$C_{\mathrm{Gen5}}$	tcgen05.mma , tcgen05.ld	$(N_{\mathrm{UTCHMMA}}>0)\land(N_{\mathrm{LDTM}}>0\lor N_{\mathrm{LDT}}>0)$
Data movement	$D_{\mathrm{TMA}}$	cp.async.bulk.tensor	$N_{\mathrm{UTMALDG}}>0\lor N_{\mathrm{UTMASTG}}>0$
Data movement	$D_{\mathrm{TMEM}}$	tcgen05.cp , tcgen05.st	$N_{\mathrm{UTCCP}}>0\lor N_{\mathrm{STTM}}>0\lor N_{\mathrm{STT}}>0$

Appendix

Table 6: PTX interfaces and SASS execution criteria for the selected BF16 mechanisms.

Figure 11: Detailed execution infrastructure supporting PTXBench. MiniPTXAgent compiles generated CUDA–PTX locally and sends successfully compiled candidates to an isolated GPU profiling service for sanitization, evaluation, and optional diagnostics. A reliability monitor pauses dispatch, restarts unhealthy service state, and excludes affected turns from trajectories. After a restart, it also excludes thermally abnormal GPUs because throttling distorted latency by over 15% in our measurements [ NVIDIA Corporation, 2026f ] .

Figure 12: Workload-state caching (a), rollout concurrency with generation–profiling pipelining (b), and profiling-GPU scaling (c).

GPU	Component	Tokens
H100	Architecture parameter	259
	Template functions	18,601
	Architecture contract	3,454
	Total	22,314
B200	Architecture parameter	413
	Template functions	18,871

Appendix

Table 7: Token counts for architecture-specific prompt components.

Figure 13: GPT-5.6 Sol $\mathrm{Fast}_{p}$ distributions for Triton and CUDA-PTX on H100 and B200.

Method or ratio	GEMM	MHA-Fwd	MHA-Fwd-Causal	MHA-Bwd	MHA-Bwd-Causal
Multi-turn refinement	$0.641\times$	$0.682\times$	$0.805\times$	$0.383\times$	$0.214\times$
Codex	$0.725\times$	$0.565\times$	$0.581\times$	$0.274\times$	$0.116\times$
Agent / refinement total tokens	$5.78\times$	$4.50\times$	$3.47\times$	$5.60\times$	$4.20\times$

Appendix

Table 8: Codex versus multi-turn refinement after eight H100 evaluation requests per trajectory. Speedups are medians of four per-trajectory best instruction-qualified values; token ratios compare workload-level medians of cumulative total tokens.

Figure 14: Median best instruction-qualified speedup at shared per-trajectory total-token budgets on H100 for Antigravity/Gemini 3.1 Pro (top) and Codex/GPT-5.6 Sol (bottom), each compared with multi-turn refinement using the same model. Each curve includes four trajectories per method and workload; bands show pointwise 95% bootstrap intervals. The horizontal axis is logarithmic.

Figure 15: H100 $\mathrm{Fast}_{p}$ and $\mathrm{Fast}^{\mathrm{Inst.}}_{p}$ .

Figure 16: B200 $\mathrm{Fast}_{p}$ and $\mathrm{Fast}^{\mathrm{Inst.}}_{p}$ .

GPU	Model	GEMM	MHA-Fwd	MHA-Fwd-Causal	MHA-Bwd	MHA-Bwd-Causal
H100	Gemini 3.1 Pro	0.687 / 0.934 / 0.962 (0.687 / 0.934 / 0.962)	0.555 / 0.730 / 0.730 (0.555 / 0.730 / 0.730)	0.614 / 0.651 / 0.768 (0.614 / 0.651 / 0.768)	0.206 / 0.375 / 0.515 (0.206 / 0.375 / 0.515)	– / 0.634 / 0.639 (0.065 / 0.634 / 0.639)
	Claude Opus 4.8	0.770 / 0.968 / 0.976 (0.770 / 0.968 / 0.976)	0.759 / 0.770 / 0.839 (0.759 / 0.770 / 0.839)	0.758 / 0.806 / 0.806 (0.758 / 0.806 / 0.806)	0.300 / 0.440 / 0.489 (0.300 / 0.440 / 0.489)	– / 0.499 / 0.499 (0.058 / 0.499 / 0.499)
	GLM-5.2	0.447 / 0.692 / 0.692 (0.447 / 0.692 / 0.692)	0.407 / 0.470 / 0.607 (0.407 / 0.470 / 0.607)	– / 0.471 / 0.533 (0.015 / 0.471 / 0.533)	0.316 / 0.437 / 0.437 (0.316 / 0.437 / 0.437)	– (– / 0.101 / 0.101)
	Qwen3.6-27B	–	–	–	–	–
	GPT-5.6 xhigh	0.788 / 0.870 / 0.870 (0.788 / 0.870 / 0.870)	0.282 / 0.774 / 0.872 (0.282 / 0.774 / 0.872)	– / 0.699 / 0.865 (– / 0.699 / 0.865)	0.230 / 0.699 / 0.703 (0.230 / 0.699 / 0.703)	– / 0.650 / 0.746 (0.094 / 0.650 / 0.746)
B200	Gemini 3.1 Pro	0.273 / 0.680 / 0.892 (0.273 / 0.680 / 0.892)	– / 0.280 / 0.280 (– / 0.280 / 0.280)	– / 0.206 / 0.248 (0.013 / 0.206 / 0.248)	– / 0.087 / 0.133 (– / 0.087 / 0.133)	– (– / 0.015 / 0.015)

Appendix

Table 9: Best speedup with target instruction execution on H100 and B200.

Figure 17: Generalization of the 158-record 4ops s1 recipe under additional SFT shuffle seeds. Top: seed 1. Bottom: seed 2. The seed-1 GEMM cell contains 90 of 96 expected turn rows; all other plotted cells contain the complete 96 turns. Runtime-instruction rows require dynamically observed Hopper SASS execution.

SFT-ed Model Label	GEMM	MHA-Fwd	MHA-Fwd-Causal	MHA-Bwd	MHA-Bwd-Causal
Qwen3.6-27B-s1	25.0 / 14.6 / 13.5 (25.0 / 14.6 / 13.5)	16.7 / 16.7 / 19.8 (16.7 / 16.7 / 19.8)	– / – / 3.1 (– / – / 3.1)	– / 2.1 / 4.2 (– / 2.1 / 4.2)	– / 2.1 / 4.2 (– / 2.1 / 5.2)
Qwen3.6-27B-s2	– / 2.1 / 2.1 (– / 2.1 / 2.1)	– / 2.1 / 5.2 (– / 2.1 / 5.2)	– / – / 1.0 (– / – / 1.0)	– / – / 1.0 (– / – / 1.0)	–
Qwen3.6-27B-s3	– / 2.1 / 3.1 (– / 2.1 / 3.1)	– / 4.2 / 2.1 (– / 4.2 / 2.1)	– / 6.2 / 4.2 (– / 6.2 / 4.2)	– / 4.2 / 4.2 (– / 4.2 / 4.2)	–
Qwen3.6-27B-s4	8.3 / 8.3 / 10.4 (8.3 / 8.3 / 10.4)	– / – / 4.2 (– / – / 4.2)	– / 4.2 / 2.1 (– / 4.2 / 2.1)	– / – / 1.0 (– / – / 1.0)	–
Qwen3.6-27B-s5	16.7 / 14.6 / 12.5 (16.7 / 14.6 / 12.5)	– / 2.1 / 4.2 (– / 2.1 / 4.2)	– / 2.1 / 3.1 (– / 2.1 / 3.1)	– / – / 5.2 (– / – / 5.2)	– / – / 4.2 (– / – / 4.2)
Qwen3.6-27B-s6	8.3 / 2.1 / 1.0 (8.3 / 2.1 / 1.0)	–	–	–	–

Appendix

Table 10: Turn correctness rates for the five-problem SFT evaluation.

SFT-ed Model Label	GEMM	MHA-Fwd	MHA-Fwd-Causal	MHA-Bwd	MHA-Bwd-Causal
Qwen3.6-27B-s1	0.303 / 0.303 / 0.340 (0.303 / 0.303 / 0.340)	0.556 / 0.556 / 0.565 (0.556 / 0.556 / 0.565)	– / – / 0.395 (– / – / 0.395)	– / 0.380 / 0.389 (– / 0.380 / 0.389)	– / 0.192 / 0.199 (– / 0.192 / 0.199)
Qwen3.6-27B-s2	– / 0.209 / 0.211 (– / 0.209 / 0.211)	– / 0.548 / 0.548 (– / 0.548 / 0.548)	– / – / 0.315 (– / – / 0.315)	– / – / 0.296 (– / – / 0.296)	–
Qwen3.6-27B-s3	– / 0.274 / 0.280 (– / 0.274 / 0.280)	– / 0.465 / 0.465 (– / 0.465 / 0.465)	– / 0.388 / 0.388 (– / 0.388 / 0.388)	– / 0.494 / 0.494 (– / 0.494 / 0.494)	–
Qwen3.6-27B-s4	0.276 / 0.373 / 0.373 (0.276 / 0.373 / 0.373)	– / – / 0.452 (– / – / 0.452)	– / 0.383 / 0.383 (– / 0.383 / 0.383)	– / – / 0.199 (– / – / 0.199)	–
Qwen3.6-27B-s5	0.325 / 0.446 / 0.446 (0.325 / 0.446 / 0.446)	– / 0.425 / 0.573 (– / 0.425 / 0.573)	– / 0.241 / 0.246 (– / 0.241 / 0.246)	– / – / 0.295 (– / – / 0.295)	– / – / 0.246 (– / – / 0.246)
Qwen3.6-27B-s6	0.073 / 0.073 / 0.073 (0.073 / 0.073 / 0.073)	–	–	–	–

Appendix

Table 11: Best speedups for the five-problem SFT evaluation.

SFT-ed Model Label	Config	Reasoning Synthesizer	Record Count
Qwen3.6-27B-s1	4ops	GLM-5.2	158
Qwen3.6-27B-s2	4ops-Extended	GLM-5.2	259
Qwen3.6-27B-s3	8ops-Extended	GLM-5.2	406
Qwen3.6-27B-s4	8ops-Post-balanced	GLM-5.2	170
Qwen3.6-27B-s5	8ops-Pre-balanced	GLM-5.2	258
Qwen3.6-27B-s6	8ops-Pre-balanced	Qwen3.6-27B	258

Appendix

Table 12: Training data recipes.

Explore similar work

Are LLM-Generated GPU Kernels Production-Ready? A Trace-Driven Benchmark and Optimization Agent
Jul 16, 2026Lingyun Yang, Yuxiao Wang, Shenghao Liang +8Graphics Processing Unit Kernels Agentic Optimization
ATREX Team†, Alibaba Group
KernelBenchX: A Comprehensive Benchmark for Evaluating LLM-Generated GPU Kernels
May 6, 2026Han Wang, Jintao Zhang, Kai Jiang +3Graphics Processing Unit Kernels Cuda
∗Equal contribution.
Xe-Forge: Multi-Stage LLM-Powered Kernel Optimization for Intel GPU
Apr 16, 2026Marcin Spoczynski, Daniel Fleischer, Moshe Berchansky +5Graphics Processing Unit Kernels Cuda
Intel Corporation

Jul 16, 2026, cs.AIJ/K move · Enter open · S save

Are LLM-Generated GPU Kernels Production-Ready? A Trace-Driven Benchmark and Optimization Agent

Lingyun Yang, Yuxiao Wang, Shenghao Liang, Linfeng Yang+7

ATREX Team†, Alibaba Group

Graphics Processing Unit Kernels Agentic Optimization Pytorch

Open paper AskPDF arXiv

Existing GPU kernel generation benchmarks draw problems from synthetic or curated sources that diverge from deployed workloads. We present Atrex-Bench, a benchmark whose 30 operators and 440 shapes are sampled directly from full-cluster production inference traces of compute-limited, memory-rich GPUs. Each problem carries an importance weight derived from its share of observed GPU time, weighted by application card-hours and computed separately for the serving phases in which it runs, together with a per-problem roofline ceiling, so the aggregate score emphasizes the kernels that consume the most serving time. Evaluating six frontier coding agents on Atrex-Bench shows that even the best vanilla model reaches only

{\sim}10\%

of the hardware roofline on production operators; and correctness alone overstates capability, since much of the apparent pass rate comes from PyTorch fallbacks rather than kernels the model wrote. To close this gap, we co-release Atrex-Kernel-Agent (AKA), a profile-driven kernel-optimization agent that combines iterative measure-revise search, optimization dropout for escaping stalled search contexts, and a layered GPU-optimization knowledge base (298 reference-kernel files and 244 optimization-knowledge documents, plus external upstream reference projects for API/ISA lookup). In a controlled case study, the agent converts zero-FlyDSL fallbacks into real kernels that match or exceed hand-tuned production baselines.

Abstract

Explore similar work

Are LLM-Generated GPU Kernels Production-Ready? A Trace-Driven Benchmark and Optimization Agent

KernelBenchX: A Comprehensive Benchmark for Evaluating LLM-Generated GPU Kernels

Xe-Forge: Multi-Stage LLM-Powered Kernel Optimization for Intel GPU