DRTriton: Large-Scale Synthetic Data Reinforcement Learning for Triton Kernel Generation
DRTriton 提出一个可扩展的学习框架,通过大规模合成数据和强化学习训练 LLM 自动将 PyTorch 代码转换为高度优化的 Triton kernel。核心包含三个组件:CSP-DAG 数据合成算法保证算子空间的全覆盖和无偏采样;课程强化学习通过解耦奖励同时优化转换成功率和推理速度;测试时计算扩展策略将成功率从 62.5% 提升至 91.2%。在 KernelBench 上达到 SOTA,成功率超越 GPT-5.2 和 Claude-Sonnet-4.5。
INSIGHT这项工作展示了 RL + 合成数据在底层系统优化领域的巨大潜力。CSP-DAG 算法解决了 kernel 训练数据稀缺和分布偏差的根本问题,课程学习策略有效平衡了正确性和性能两个目标。对于搜推系统的自定义算子优化(如注意力机制、特征交叉等),这种自动化 kernel 生成方案可以显著降低系统工程师的手工优化成本,加速从算法到部署的迭代周期。