TritonGym: A Benchmark for Agentic LLM Workflows in Triton GPU Code Generation
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.
Lay Summary
Modern AI systems rely on graphics processors (GPUs) to run quickly, but writing the small programs that make GPUs efficient is difficult and time-consuming. Large language models can now write some of this code, yet most tests only ask them to produce an answer once, without checking how well they use tools such as compilers, correctness checks, and speed measurements. We built TritonGym, a benchmark for testing AI coding systems in a setting that is closer to real software development: the system can write code, run tools, learn from feedback, and try again. TritonGym gives every system the same tools, tasks, and budget, making comparisons fairer. It also includes new tasks whose solutions are not available online, so success depends more on understanding the problem than on copying known examples. Our experiments show that today’s strongest models can write useful GPU code for familiar tasks, but they still struggle on new kinds of problems and hardware settings. Tool-using workflows help, but there remains a gap between AI-generated code and carefully written expert code. By making this evaluation open and standardized, TritonGym can help researchers build more reliable AI assistants for high-performance programming.