DRTriton: Tritonカーネル生成のための大規模合成データ強化学習
arXiv cs.CL / 2026/3/24
💬 オピニオンSignals & Early TrendsIdeas & Deep AnalysisModels & Research
要点
- 本論文は、LLMを学習させてPyTorchの参照コードを最適化されたTritonカーネルへ変換し、実行時にCUDAとしてコンパイルするためのフレームワークであるDRTritonを提案する。これは生成AIエンジニアリングにおける重要な課題、すなわち効率的なCUDAカーネル作成を対象としている。
- DRTritonは、演算子空間を偏りのない一様サンプリングで幅広くカバーし、タスクの難易度を制御することを目的とした合成データ戦略（CSP-DAG）を用いる。
- カリキュラム型の強化学習と、報酬を分離した設計（decoupled rewards）を適用し、変換成功率と推論速度を共同で改善する。さらに、ランタイム性能を一層高めるために、テスト時探索手法を追加する。
- 合成データのみで学習したにもかかわらず、DRTritonは、熟練エンジニアでも難しいケースを含む、困難な実世界のCUDAカーネルに対して良好に汎化できると報告されている。
- 実験では、DRTriton-7BがKernelBench Level 2の92%で高速化を達成し、GPT-5.2（23%）およびClaude-Sonnet-4.5（19%）を大幅に上回った。
