드니얼 박사: 트리톤 커널 생성을 위한 올바른 강화 학습 접근법

초록

고품질 커널은 확장 가능한 AI 시스템에 필수적이며, LLM이 이러한 코드를 생성할 수 있게 되면 AI 개발이 크게 진전될 수 있습니다. 하지만 이 작업을 위해 LLM을 훈련시키기 위해서는 충분한 데이터와 강력한 환경이 필요하며, 이 과정은 종종 보상 해킹과 게으른 최적화에 취약합니다. 이러한 경우 모델은 훈련 보상을 해킹하여 의미 있는 속도 향상보다 사소한 정확성에 우선순위를 둘 수 있습니다. 본 논문에서는 커널 생성을 위한 강화학습(RL)을 체계적으로 연구합니다. 먼저 보상 해킹 검사, 다중 턴 상호작용에서의 데이터 수집, 장기 RL 훈련을 지원하는 강력한 분산 GPU 환경인 KernelGYM을 설계합니다. KernelGYM을 기반으로 효과적인 다중 턴 RL 방법을 조사하고 GRPO에서 자기 포함으로 인한 편향된 정책 경사 문제를 확인합니다. 이를 해결하기 위해 다중 턴 RL에 대한 편향되지 않은 어드밴티지 추정을 제공하는 턴 단위 Reinforce-Leave-One-Out(TRLOO)을 제안합니다. 게으른 최적화를 완화하기 위해 훈련 안정성을 위한 불일치 수정을 통합하고, 해당 문제를 극복하기 위해 프로파일링 기반 보상(PR)과 프로파일링 기반 반응 샘플링(PRS)을 도입합니다. 훈련된 모델인 Dr.Kernel-14B는 Kernelbench에서 Claude-4.5-Sonnet에 버금가는 성능에 도달합니다. 마지막으로 Dr.Kernel-14B에 대한 순차적 테스트 타임 스케일링을 연구합니다. KernelBench Level-2 하위 집합에서 생성된 커널의 31.6%가 Torch 참조 대비 최소 1.2배의 속도 향상을 달성하여 Claude-4.5-Sonnet(26.7%) 및 GPT-5(28.6%)를 능가합니다. 모든 턴에서 최상의 후보를 선택할 경우, 이 1.2배 속도 향상 비율은 47.8%로 더욱 증가합니다. 환경, 훈련 코드, 모델, 데이터셋을 포함한 모든 리소스는 https://www.github.com/hkust-nlp/KernelGYM에서 확인할 수 있습니다.

English

High-quality kernel is critical for scalable AI systems, and enabling LLMs to generate such code would advance AI development. However, training LLMs for this task requires sufficient data, a robust environment, and the process is often vulnerable to reward hacking and lazy optimization. In these cases, models may hack training rewards and prioritize trivial correctness over meaningful speedup. In this paper, we systematically study reinforcement learning (RL) for kernel generation. We first design KernelGYM, a robust distributed GPU environment that supports reward hacking check, data collection from multi-turn interactions and long-term RL training. Building on KernelGYM, we investigate effective multi-turn RL methods and identify a biased policy gradient issue caused by self-inclusion in GRPO. To solve this, we propose Turn-level Reinforce-Leave-One-Out (TRLOO) to provide unbiased advantage estimation for multi-turn RL. To alleviate lazy optimization, we incorporate mismatch correction for training stability and introduce Profiling-based Rewards (PR) and Profiling-based Rejection Sampling (PRS) to overcome the issue. The trained model, Dr.Kernel-14B, reaches performance competitive with Claude-4.5-Sonnet in Kernelbench. Finally, we study sequential test-time scaling for Dr.Kernel-14B. On the KernelBench Level-2 subset, 31.6% of the generated kernels achieve at least a 1.2x speedup over the Torch reference, surpassing Claude-4.5-Sonnet (26.7%) and GPT-5 (28.6%). When selecting the best candidate across all turns, this 1.2x speedup rate further increases to 47.8%. All resources, including environment, training code, models, and dataset, are included in https://www.github.com/hkust-nlp/KernelGYM.

드니얼 박사: 트리톤 커널 생성을 위한 올바른 강화 학습 접근법

Dr. Kernel: Reinforcement Learning Done Right for Triton Kernel Generations

초록

Support