用强化学习训练80亿参数模型，让LLM学会写高效Triton代码
TritonRL: Training LLMs to Think and Code Triton Without Cheating
- 基于强化学习和分层奖励分解，解决长序列生成中的奖励分配难题
- 在KernelBench上实现最高正确率和速度提升，超越同类小模型
- 适合对高性能编程、硬件感知模型感兴趣的开发者与研究者
大型语言模型的快速发展推动了自动化高性能系统内核的需求，以加速机器学习任务。我们提出TritonRL，一个针对Triton编程的80亿参数领域专用大模型，通过一种新颖的强化学习框架进行训练。尽管Triton代码生成面临数据稀缺和奖励欺骗等挑战，我们的方法通过两项核心创新实现了鲁棒的内核生成：首先，采用多层次验证系统，提供高保真奖励信号，确保生成代码在语法和功能上均正确；其次，提出分层奖励分解（HRD），将高层推理与底层实现的强化学习解耦，解决长序列生成中的信用分配问题。在KernelBench上的全面评估表明，TritonRL在正确性和运行时加速方面达到领先水平，优于同期的Triton专用模型，并接近参数量超1000亿的前沿模型性能。结果凸显了硬件感知强化学习范式在特定领域适配中的有效性。
原文摘要 · Abstract (English)
The rapid evolution of Large Language Models (LLMs) has driven a growing demand for automated, high-performance system kernels to accelerate machine learning workloads. We introduce TritonRL, a domain-specialized 8B-scale LLM for Triton programming, trained via a novel reinforcement learning (RL) framework. While Triton synthesis faces unique challenges, including data scarcity and a high susceptibility to reward hacking, our approach enables robust kernel generation through two primary innovations. First, we implement a multi-layered verification system that provides high-fidelity reward signals, ensuring that generated kernels are both syntactically and functionally valid. Second, we propose Hierarchical Reward Decomposition (HRD), which decouples reinforcement for high-level reasoning and low-level implementation to resolve the credit assignment problem in long-sequence generation. Comprehensive evaluations on KernelBench demonstrate that TritonRL achieves state-of-the-art correctness and runtime speedup, outperforming concurrent Triton-specific models and matching the performance of frontier models with over 100B parameters. Our results highlight the effectiveness of hardware-aware RL paradigms in specialized domain adaptation.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。