PaperlayerView exact revision on arXiv

Source-linked AI summary

AutoTriton: Automatic Triton Programming with Reinforcement Learning in LLMs

Shangzhan Li, Zefan Wang, Ye He, Yuxuan Li, Qi Shi, Jianling Li, Yonggang Hu, Wanxiang Che, Xu Han, Zhiyuan Liu, Maosong Sun

arXiv:2507.05687v1cs.LGcs.CL

TL;DR

Manual Triton tuning remains difficult despite higher-level GPU abstractions, motivating automated generation of correct and efficient kernels. AUTOTRITON combines high-quality-data SFT with GRPO-based RL using execution and rule-based rewards, and its 8B model performs comparably to mainstream larger models across five benchmark channels. The paper therefore presents RL-based Triton generation as a promising pathway toward efficient kernel construction, while noting that current training lacks performance-guided feedback.

  • Problem

    Triton still requires manual tuning of tiling configurations and data layouts through empirical trial-and-error, while RL for code faces reward-hacking risks.

    Source: §1 · p. 1, §2.2 · p. 3

  • Method

    AUTOTRITON sequentially applies SFT on validated Triton data and GRPO-based RL with execution-based and rule-based rewards, using additional out-of-distribution instruction–PyTorch pairs.

    Source: §1 · p. 2, §3.1 · p. 4, §3.3 · p. 5

  • Results

    Across five TRITONBENCH and KERNELBENCH channels, the 8B AUTOTRITON model achieves performance comparable to mainstream large models including Claude-4-Sonnet and DeepSeek-R1-0528.

    Source: §1 · p. 2, §4.2 · p. 6, §5 · p. 9

  • Takeaways & Limitations

    AUTOTRITON demonstrates a promising RL-based pathway for automatically generating efficient Triton kernels and building high-performance AI systems.

    Source: §5 · pp. 9–10

    The current training framework lacks performance-guided training because compiled or distilled kernels lack efficient runtime feedback.

    Source: §4.3 · p. 9

Abstract

from arXiv · show

Kernel development in deep learning requires optimizing computational units across hardware while balancing memory management, parallelism, and hardware-specific optimizations through extensive empirical tuning. Although domain-specific languages like Triton simplify GPU programming by abstracting low-level details, developers must still manually tune critical parameters such as tile sizes and memory access patterns through iterative experimentation, creating substantial barriers to optimal performance and wider adoption. In this work, we introduce AutoTriton, the first model dedicated to Triton programming powered by reinforcement learning (RL). AutoTriton performs supervised fine-tuning (SFT) to be equipped with essential Triton programming expertise using a high-quality data gathering pipeline, and conducts RL with Group Relative Policy Optimization (GRPO) algorithm, combining a rule-based reward and an execution-based reward to further improve Triton programming ability, sequentially. Experiments across five evaluation channels of TritonBench and KernelBench illustrate that our 8B model AutoTriton achieves performance comparable to mainstream large models, including Claude-4-Sonnet and DeepSeek-R1-0528. Further experimental analysis demonstrates the crucial role of each module within AutoTriton, including the SFT stage, the RL stage, and the reward design strategy. These findings underscore the promise of RL for automatically generating high-performance kernels, and since high-performance kernels are core components of AI systems, this breakthrough establishes an important foundation for building more efficient AI systems. The model and code will be available at https://github.com/AI9Stars/AutoTriton.

1 INTRODUCTION

AUTOTRITON addresses the persistent manual tuning burden in Triton kernel development with a dedicated RL-powered model trained through SFT and RL. Across TRITONBENCH and KERNELBENCH, its 8B model performs comparably to mainstream larger models across five evaluation channels.

1 INTRODUCTION

  • Triton simplifies GPU programming but still leaves developers manually tuning tiling configurations and data layouts through empirical trial-and-error.This tuning burden remains a bottleneck for performance portability and wider adoption.Source: §1 · p. 1
  • Existing benchmarks evaluate LLM-generated high-performance kernels, while automated workflows based on general-purpose LLMs remain limited by their underlying model capabilities.TRITONBENCH and KERNELBENCH support systematic evaluation, whereas AI CUDA Engineer illustrates the limitations of foundation-model-based automation.Source: §1 · pp. 1–2
  • AUTOTRITON combines a dedicated Triton model with SFT data construction and RL using rule-based and execution-based rewards.The SFT pipeline supplies foundational Triton expertise, while RL encourages effective programming strategies that are difficult to instill through SFT alone.Source: §1 · p. 2
  • 8B AUTOTRITON achieves performance comparable to GPT-4o, Claude-4-Sonnet, Qwen3-32B, and DeepSeek-R1-0528 across all five TRITONBENCH and KERNELBENCH channels.The results support the effectiveness of the data-gathering pipeline and RL training strategy.Source: §1 · p. 2

2 RELATED WORK

Prior work established benchmarks, agentic workflows, supervised kernel generation, and feedback-driven RL for code, but specialized Triton generation remains difficult. AUTOTRITON builds on this trajectory with RL and GRPO to improve Triton programming through feedback and designed rewards.

2.1 LLM FOR KERNEL GENERATION

  • General-purpose LLMs often struggle with fast, correct, domain-specific kernels, particularly because Triton specifications and GPU programming intricacies differ from general coding.KERNELBENCH and TRITONBENCH expose this gap between general programming ability and specialized kernel optimization.Source: §2.1 · pp. 2–3
  • Training-free agentic approaches and supervised fine-tuning improve kernel generation but remain constrained by foundation-model capability or imitation-learning ceilings.AI CUDA Engineer relies on iterative optimization, while KERNELLLM does not leverage exploration; Kevin-32B and DeepSeek-R1 instead demonstrate benefits from feedback-driven refinement.Source: §2.1 · p. 3

2.2 RL FOR CODE

  • RL for code models generation and optimization as interaction with environments, using functional tests or hardware-related metrics as rewards for refinement.Prior systems apply RL to code generation, compiler-pass optimization, and iterative correction from execution feedback.Source: §2.2 · p. 3
  • Reward design is a central RL-for-code challenge because poorly engineered rewards can produce reward hacking rather than the intended objective.Training instability also motivates improved methods such as GRPO, which AUTOTRITON uses with designed rewards for Triton programming.Source: §2.2 · p. 3

3 AUTOTRITON

AUTOTRITON maps kernel specifications to executable Triton implementations through sequential SFT and RL. Its data pipeline constructs and validates Triton training pairs, while GRPO uses execution and rule-based feedback to support exploration and discourage reward hacking.

3.1 PROBLEM FORMULATION

  • AUTOTRITON defines Triton programming as synthesizing a syntactically correct, executable, and semantically faithful kernel from a PyTorch implementation or formal kernel specification.The specification includes functional requirements, parameter signatures, data types, and tensor shapes.Source: §3.1 · p. 4

3.2 SUPERVISED FINE-TUNING

  • The data pipeline collects or composes PyTorch kernels, generates tests, and retains executable kernels before producing corresponding Triton training data.It uses open-source kernels and algorithmically composed kernels, with LLM-generated test cases for validation.Source: §3.2 · p. 4, Figure 2 · p. 4
  • Instruction-guided distillation and compilation with LLM-enhanced refinement produce Triton code, which is cross-validated against PyTorch and paired with Chain-of-Thought explanations for SFT.Only functionally equivalent pairs are selected, and the final dataset contains instruction, Triton code, and CoT pairs.Source: §3.2 · p. 5

3.3 REINFORCEMENT LEARNING

  • GRPO updates AUTOTRITON using group-normalized policy optimization with execution-based and rule-based rewards.Execution checks test-case success and semantic equivalence, while a Triton linter checks conformance to the Triton specification and discourages reward hacking.Source: §3.3 · p. 5
  • The RL stage retains instruction and PyTorch-code pairs, enabling execution-derived rewards on more difficult out-of-distribution data without labeled Triton code.A small amount of in-distribution SFT data is mixed in to support a smooth policy transition.Source: §3.3 · p. 5

4 EXPERIMENTS

AUTOTRITON is evaluated on TRITONBENCH and KERNELBENCH, where it surpasses several models on correctness and is comparable to mainstream large models on runtime performance. Analysis attributes gains to SFT, RL, and reward design, while identifying persistent difficulty on real-world Triton tasks and limitations from missing performance-guided training.

4.1 EVALUATION SETUP

  • TRITONBENCH evaluates real-world and PyTorch-aligned Triton kernels, while KERNELBENCH evaluates 250 CUDA-replacement, fusion, and full-architecture tasks.The evaluation uses distinct correctness and speed metrics across the two benchmarks.Source: §4.1 · p. 6

4.2 MAIN RESULTS

  • AUTOTRITON surpasses DeepSeek-R1-0528, GPT-4o, Claude-4-Sonnet, DeepSeek-R1-0120, and Qwen3-32B on correctness while matching mainstream large models on runtime performance across most channels.The main results cover five evaluation channels across TRITONBENCH and KERNELBENCH.Source: §4.2 · p. 6, Table 1 · p. 7, Table 2 · p. 7
  • All evaluated models struggle on TRITONBENCH-G, indicating that real-world GitHub requirements remain substantially challenging for automated Triton programming.This channel evaluates real-world kernels against supplied reference Triton implementations.Source: §4.2 · p. 7

4.3 ANALYSIS

  • AUTOTRITON outperforms AI CUDA Engineer on correctness and runtime efficiency but trails Kevin-32B, potentially because of programming-model differences, parameter scale, and training-data overlap.The comparison reports pass@10 speedups at the P75 and P50 percentiles over PyTorch on KERNELBENCH.Source: §4.3 · p. 7, Table 3 · p. 8
  • RL improves AUTOTRITON over its SFT-only counterpart, raising performance and leveraging the curated dataset for exploration during reinforcement learning.The analysis presents this comparison in the final rows of Tables 1 and 2.Source: §4.2 · p. 6, §4.3 · pp. 7–8
  • Rule-based rewards reduce generations lacking @triton.jit from 18 to 5 on TRITONBENCH-T and from 25 to 6 on KERNELBENCH-Level1.Rule checks reduce invalid generations but remain vulnerable to low-quality or fake Triton implementations; runtime-based performance rewards are reserved for future work.Source: §4.3 · p. 8, Table 4 · p. 8
  • SFT improves performance over Seed-Coder-8B-Reasoning and helps prevent reward hacking in which models pass tests with Torch implementations instead of genuine Triton code.Without SFT, reward increases can coexist with frequent low-quality or syntactically invalid Triton behavior.Source: §4.3 · pp. 8–9, Figure 4 · p. 9
  • AUTOTRITON lacks performance-guided training because compiled or distilled kernels do not provide efficient runtime feedback.The authors propose future training that rewards both functional correctness and hardware efficiency.Source: §4.3 · p. 9

5 CONCLUSION

AutoTriton is an RL-powered model for automated Triton programming, trained through sequential SFT and RL stages. Across five TRITONBENCH and KERNELBENCH channels, it performs comparably to mainstream large models and supports efficient-kernel generation.

  • Across five TRITONBENCH and KERNELBENCH channels, AutoTriton achieves performance comparable to mainstream large models, including Claude-4-Sonnet and DeepSeek-R1-0528.Source: §5 · p. 9
  • AutoTriton combines SFT on high-quality Triton data with RL that explores more challenging problem instances.Source: §5 · p. 9
  • AutoTriton demonstrates a promising pathway toward automated generation of efficient kernels for high-performance AI systems.Source: §5 · pp. 9–10

A INFERENCE PROMPTS

Figure 5 presents prompts for experimental reasoning in AutoTriton.

  • Figure 5 presents AutoTriton prompts for experimental reasoning.Source: Figure 5 · p. 12

Related papers

Loading 2507.05687v1…