TritonGym: A Benchmark for Agentic LLM Workflows in Triton GPU Code Generation

Yue Guan, Yichen Lin, Xu Zhao, Jianzhu Yao, Xinwei Qiang, Zhongkai Yu, Pramod Viswanath, Yufei Ding, Adnan Aziz
Proceedings of the 43rd International Conference on Machine Learning, PMLR 306:37457-37473, 2026.

Abstract

Large language models (LLMs) can already draft plausible Triton kernels, yet most existing evaluations still focus on single-shot generation and underplay tool use and feedback. We introduce TritonGym, a benchmark and orchestration framework for evaluating agentic workflows in GPU code generation. TritonGym standardizes access to a set of code generation tools via function-calls, separating intrinsic model capability from workflow design and enabling fair, apples-to-apples comparison. The benchmark spans a maintained operator set, community samples, out-of-distribution tasks, and DSL extensions, ensuring both generality and extensibility. By providing a common orchestration and evaluation framework, TritonGym democratizes the development of GPU coding agents, supports practical adoption of agent-generated kernels, and facilitates progress on advanced agentic systems.

Cite this Paper

BibTeX
@InProceedings{pmlr-v306-guan26d, title = {{T}riton{G}ym: A Benchmark for Agentic {LLM} Workflows in Triton {GPU} Code Generation}, author = {Guan, Yue and Lin, Yichen and Zhao, Xu and Yao, Jianzhu and Qiang, Xinwei and Yu, Zhongkai and Viswanath, Pramod and Ding, Yufei and Aziz, Adnan}, booktitle = {Proceedings of the 43rd International Conference on Machine Learning}, pages = {37457--37473}, year = {2026}, editor = {Zhang, Tong and Dudik, Miroslav and Jaggi, Martin and Agarwal, Alekh and Li, Sharon and Schuurmans, Dale and Zhu, Jerry and Berkenkamp, Felix and Dong, Hanze and Bietti, Alberto}, volume = {306}, series = {Proceedings of Machine Learning Research}, month = {06--11 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v306/main/assets/guan26d/guan26d.pdf}, url = {https://proceedings.mlr.press/v306/guan26d.html}, abstract = {Large language models (LLMs) can already draft plausible Triton kernels, yet most existing evaluations still focus on single-shot generation and underplay tool use and feedback. We introduce TritonGym, a benchmark and orchestration framework for evaluating agentic workflows in GPU code generation. TritonGym standardizes access to a set of code generation tools via function-calls, separating intrinsic model capability from workflow design and enabling fair, apples-to-apples comparison. The benchmark spans a maintained operator set, community samples, out-of-distribution tasks, and DSL extensions, ensuring both generality and extensibility. By providing a common orchestration and evaluation framework, TritonGym democratizes the development of GPU coding agents, supports practical adoption of agent-generated kernels, and facilitates progress on advanced agentic systems.} }
Endnote
%0 Conference Paper %T TritonGym: A Benchmark for Agentic LLM Workflows in Triton GPU Code Generation %A Yue Guan %A Yichen Lin %A Xu Zhao %A Jianzhu Yao %A Xinwei Qiang %A Zhongkai Yu %A Pramod Viswanath %A Yufei Ding %A Adnan Aziz %B Proceedings of the 43rd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2026 %E Tong Zhang %E Miroslav Dudik %E Martin Jaggi %E Alekh Agarwal %E Sharon Li %E Dale Schuurmans %E Jerry Zhu %E Felix Berkenkamp %E Hanze Dong %E Alberto Bietti %F pmlr-v306-guan26d %I PMLR %P 37457--37473 %U https://proceedings.mlr.press/v306/guan26d.html %V 306 %X Large language models (LLMs) can already draft plausible Triton kernels, yet most existing evaluations still focus on single-shot generation and underplay tool use and feedback. We introduce TritonGym, a benchmark and orchestration framework for evaluating agentic workflows in GPU code generation. TritonGym standardizes access to a set of code generation tools via function-calls, separating intrinsic model capability from workflow design and enabling fair, apples-to-apples comparison. The benchmark spans a maintained operator set, community samples, out-of-distribution tasks, and DSL extensions, ensuring both generality and extensibility. By providing a common orchestration and evaluation framework, TritonGym democratizes the development of GPU coding agents, supports practical adoption of agent-generated kernels, and facilitates progress on advanced agentic systems.
APA
Guan, Y., Lin, Y., Zhao, X., Yao, J., Qiang, X., Yu, Z., Viswanath, P., Ding, Y. & Aziz, A.. (2026). TritonGym: A Benchmark for Agentic LLM Workflows in Triton GPU Code Generation. Proceedings of the 43rd International Conference on Machine Learning, in Proceedings of Machine Learning Research 306:37457-37473 Available from https://proceedings.mlr.press/v306/guan26d.html.

Related Material