ICML 2026poster0 citations

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

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.

LLMAgentsRobustnessBenchmark
BibTeX
@inproceedings{
guan2026tritongym,
title={TritonGym: A Benchmark for Agentic {LLM} Workflows in Triton {GPU} Code Generation},
author={Yue Guan and Yichen Lin and Xu Zhao and Jianzhu Yao and Xinwei Qiang and Zhongkai Yu and Pramod Viswanath and Yufei Ding and Adnan Aziz},
booktitle={Forty-third International Conference on Machine Learning},
year={2026},
url={https://openreview.net/forum?id=Htq2loTQTL}
}