AutoTriton: Automatic Triton Programming with Reinforcement Learning in LLMs
- URL: http://arxiv.org/abs/2507.05687v1
- Date: Tue, 08 Jul 2025 05:38:24 GMT
- Title: AutoTriton: Automatic Triton Programming with Reinforcement Learning in LLMs
- Abstract summary: We introduce AutoTriton, the first model dedicated to Triton programming powered by reinforcement learning (RL)<n>AutoTriton performs supervised fine-tuning (SFT) to be equipped with essential Triton programming expertise using a high-quality data gathering pipeline.<n> Experiments across five evaluation channels of TritonBench and KernelBench illustrate that our 8B model AutoTriton achieves performance comparable to mainstream large models.
- Score: 87.8306870967343
- License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
- Abstract: 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.
Related papers
- RealisticTritonBench: A Benchmark for Triton-Kernel Generation in Real-World AI Frameworks [7.371299443009948]
We introduce RealisticTritonBench, the first benchmark to derive Triton kernel generation tasks from real-world pull requests in popular AI frameworks.<n> RealisticTritonBench integrates PRs that modify Triton kernels from popular open-source AI frameworks and transforms them into generation tasks with concrete engineering contexts.<n>We evaluate leading LLMs on RealisticTritonBench and find that they still struggle with real-world Triton kernel generation tasks.
arXiv Detail & Related papers (2026-08-12T12:38:17Z) - SLAI T-Rex: Full-Parameter Post-training of the DeepSeek-V4 Family on Ascend SuperPOD [91.71050339479261]
This report presents an end-to-end optimization practice on the Ascend NPU SuperPOD.<n>We develop a hierarchical optimization framework spanning model-level parallelism, computation-communication orchestration, and low-level kernel execution.<n>The resulting system achieves 34.22% Model FLOPs Utilization (MFU) with a 2.93x improvement over the open-source baseline recipe.
arXiv Detail & Related papers (2026-07-22T13:49:17Z) - DRTriton: Large-Scale Synthetic Data Reinforcement Learning for Triton Kernel Generation [35.68487613748886]
Large Language Models (LLMs) automatically convert PyTorch reference implementations to kernels, significantly reducing engineering efforts.<n>We propose DRTriton, a scalable learning framework for training LLMs to convert PyTorch codes into highly optimized Triton kernels.<n> DRTriton consists of three key components: (i) a synthetic algorithm CSP-DAG that guarantees full coverage and unbiased uniform sampling over the operator space with difficulty; (ii) a curriculum reinforcement learning with decoupled reward efficiently optimize conversion success rate and 19% speed simultaneously; and (iii) a test-time search algorithm that further improves inference speed.
arXiv Detail & Related papers (2026-03-23T00:59:35Z) - Towards Cold-Start Drafting and Continual Refining: A Value-Driven Memory Approach with Application to NPU Kernel Synthesis [68.7701048879757]
EvoKernel is a self-evolving agentic framework that automates the lifecycle of kernel synthesis.<n>It learns stage-specific Q-values that prioritize experiences based on their contribution to the current objective.<n>It improves models' correctness from 11.0% to 83.0% and a median speedup of 3.60x over initial drafts.
arXiv Detail & Related papers (2026-03-11T14:57:06Z) - RollArt: Scaling Agentic RL Training via Disaggregated Infrastructure [49.88201789074532]
Agentic Reinforcement Learning (RL) enables Large Language Models (LLMs) to perform autonomous decision-making and long-term planning.<n>We present RollArc, a distributed system designed to maximize throughput for multi-task agentic RL on disaggregated infrastructure.
arXiv Detail & Related papers (2025-12-27T11:14:23Z) - TritonForge: Profiling-Guided Framework for Automated Triton Kernel Optimization [24.065109818256605]
TritonForge is a profiling-guided framework for automated GPU kernel optimization.<n>It integrates kernel analysis, runtime profiling, and iterative code transformation to streamline the process.<n>It achieves up to 5x performance improvement over baseline implementations and on average 1.76x of the cases are successful.
arXiv Detail & Related papers (2025-12-09T23:44:35Z) - CarBench: A Comprehensive Benchmark for Neural Surrogates on High-Fidelity 3D Car Aerodynamics [5.362190317243247]
CarBench is the first comprehensive benchmark dedicated to large-scale 3D car aerodynamics.<n>We evaluate state-of-the-art models on DrivAerNet++, the largest public dataset for automotive aerodynamics.<n>To accelerate progress in data-driven engineering, we open-source the benchmark framework.
arXiv Detail & Related papers (2025-11-25T21:53:43Z) - Eliminating Multi-GPU Performance Taxes: A Systems Approach to Efficient Distributed LLMs [61.953548065938385]
We introduce the ''Three Taxes'' (Bulk Synchronous, Inter- Kernel Data Locality, and Kernel Launch Overhead) as an analytical framework.<n>We propose moving beyond the rigid BSP model to address key inefficiencies in distributed GPU execution.<n>We observe a 10-20% speedup in end-to-end latency over BSP-based approaches.
arXiv Detail & Related papers (2025-11-04T01:15:44Z) - TritonRL: Training LLMs to Think and Code Triton Without Cheating [14.859089778654933]
TritonRL is trained with a novel training framework that enables robust and automated kernel synthesis.<n>Our experiments on KernelBench demonstrate that TritonRL achieves state-of-the-art correctness and speedup, surpassing all other Triton-specific models.
arXiv Detail & Related papers (2025-10-18T21:36:10Z) - Ring-lite: Scalable Reasoning via C3PO-Stabilized Reinforcement Learning for LLMs [51.21041884010009]
Ring-lite is a Mixture-of-Experts (MoE)-based large language model optimized via reinforcement learning (RL)<n>Our approach matches the performance of state-of-the-art (SOTA) small-scale reasoning models on challenging benchmarks.
arXiv Detail & Related papers (2025-06-17T17:12:34Z) - TritonBench: Benchmarking Large Language Model Capabilities for Generating Triton Operators [59.625889531331815]
Triton is a high-level Python-like language designed for building efficient GPU kernels.<n>Despite advances in large language models (LLMs) for conventional code generation, these models struggle to generate accurate, performance-optimized Triton code.<n>In this work, we introduce TritonBench, the first comprehensive benchmark for Triton operator generation.
arXiv Detail & Related papers (2025-02-20T17:21:27Z) - TD3: Tucker Decomposition Based Dataset Distillation Method for Sequential Recommendation [50.23504065567638]
This paper introduces textbfTD3, a novel textbfDataset textbfDistillation method within a meta-learning framework.<n> TD3 distills a fully expressive emphsynthetic sequence summary from original data.<n>An augmentation technique allows the learner to closely fit the synthetic summary, ensuring an accurate update of it in the emphouter-loop.
arXiv Detail & Related papers (2025-02-05T03:13:25Z) - AutoSculpt: A Pattern-based Model Auto-pruning Framework Using Reinforcement Learning and Graph Learning [32.10443611442628]
AutoSculpt is a pattern-based automated pruning framework for deep neural networks (DNNs)<n>It automatically identifies and prunes regular patterns within DNN architectures that can be recognized by existing inference engines.<n>It achieves pruning rates of up to 90% and nearly 18% improvement in FLOPs reduction, outperforming all baselines.
arXiv Detail & Related papers (2024-12-24T02:05:51Z) - Quantized Distillation: Optimizing Driver Activity Recognition Models
for Resource-Constrained Environments [34.80538284957094]
This paper introduces a lightweight framework for resource-efficient driver activity recognition.
The framework enhances 3D MobileNet, a neural architecture optimized for speed in video classification.
It achieves a threefold reduction in model size and a 1.4-fold improvement in inference time.
arXiv Detail & Related papers (2023-11-10T10:07:07Z) - A Learned Performance Model for Tensor Processing Units [5.733911161090224]
We demonstrate a method of learning performance models from a corpus of graph programs for Processing Unit (TPU) instances.
We show that our learned model outperforms a heavily-optimized analytical performance model on two tasks.
It helps an autotuner discover faster programs in a setting where access to TPUs is limited or expensive.
arXiv Detail & Related papers (2020-08-03T17:24:52Z) - Auto-PyTorch Tabular: Multi-Fidelity MetaLearning for Efficient and
Robust AutoDL [53.40030379661183]
Auto-PyTorch is a framework to enable fully automated deep learning (AutoDL)
It combines multi-fidelity optimization with portfolio construction for warmstarting and ensembling of deep neural networks (DNNs)
We show that Auto-PyTorch performs better than several state-of-the-art competitors on average.
arXiv Detail & Related papers (2020-06-24T15:15:17Z) - PolyScientist: Automatic Loop Transformations Combined with Microkernels
for Optimization of Deep Learning Primitives [55.79741270235602]
We develop a hybrid solution to the development of deep learning kernels.
We use the advanced polyhedral technology to automatically tune the outer loops for performance.
arXiv Detail & Related papers (2020-02-06T08:02:34Z)
This list is automatically generated from the titles and abstracts of the papers in this site.
This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.