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.


翻译:大型语言模型(LLMs)的快速发展推动了对自动化高性能系统内核的日益增长的需求,以加速机器学习工作负载。我们提出了TritonRL,这是一个针对Triton编程的领域专用8B规模LLM,通过一种新颖的强化学习(RL)框架进行训练。尽管Triton合成面临独特挑战,包括数据稀缺性和对奖励黑客攻击的高度敏感性,但我们的方法通过两项主要创新实现了稳健的内核生成。首先,我们实现了一个多层验证系统,提供高保真度的奖励信号,确保生成的内核在语法和功能上均有效。其次,我们提出了分层奖励分解(HRD),将高层推理和底层实现的强化解耦,以解决长序列生成中的信用分配问题。在KernelBench上的综合评估表明,TritonRL在正确性和运行时加速方面达到了最先进的水平,超越了同期针对Triton的专用模型,并与超过100B参数的前沿模型性能相当。我们的结果突显了硬件感知的强化学习范式在专业领域适应中的有效性。

0
下载
关闭预览

相关内容

代码(Code)是专知网的一个重要知识资料文档板块,旨在整理收录论文源代码、复现代码,经典工程代码等,便于用户查阅下载使用。
什么是后训练?大语言模型训练后优化方法综述,87页pdf
通过强化学习增强代码生成中的代码大语言模型:综述
专知会员服务
29+阅读 · 2025年1月1日
《大型语言模型(LLMs): 训练到推理》全面概述技术细节
哈工大讯飞联合实验室发布中文XLNet预训练模型
哈工大SCIR
13+阅读 · 2019年8月20日
自然语言处理中的语言模型预训练方法
PaperWeekly
14+阅读 · 2018年10月21日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
1+阅读 · 2015年12月31日
国家自然科学基金
2+阅读 · 2015年12月31日
国家自然科学基金
1+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
国家自然科学基金
4+阅读 · 2014年12月31日
VIP会员
相关VIP内容
什么是后训练?大语言模型训练后优化方法综述,87页pdf
通过强化学习增强代码生成中的代码大语言模型:综述
专知会员服务
29+阅读 · 2025年1月1日
《大型语言模型(LLMs): 训练到推理》全面概述技术细节
相关基金
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
1+阅读 · 2015年12月31日
国家自然科学基金
2+阅读 · 2015年12月31日
国家自然科学基金
1+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
国家自然科学基金
4+阅读 · 2014年12月31日
Top
微信扫码咨询专知VIP会员