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

arXiv: https://arxiv.org/abs/2602.05885 · HTML: https://mlsanity.com/paper/2602.05885/

## Key facts

- Cited by 22 papers in the map (data through 14 Oct 2026).
- Placed in RL training of language and agent models, one of 15 directions in Deep reinforcement learning methods.
- Builds on 9 of its 17 references that are in the map.
- First posted on arXiv 5 Feb 2026.

## Placed in

- Deep reinforcement learning methods / [RL training of language and agent models](https://mlsanity.com/deeprl/rl-training-of-language-and-agent-models/)

## Builds on (most cited first, 9 of 9)

- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Search-R1: Training LLMs to Reason and Leverage Search Engines with Reinforcement Learning](https://mlsanity.com/paper/2503.09516/) (2025)
- [BrowseComp: A Simple Yet Challenging Benchmark for Browsing Agents](https://mlsanity.com/paper/2504.12516/) (2025)
- [FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving](https://mlsanity.com/paper/2501.01005/) (2025)
- [KernelBench: Can LLMs Write Efficient GPU Kernels?](https://mlsanity.com/paper/2502.10517/) (2025)
- [Kevin: Multi-Turn RL for Generating CUDA Kernels](https://mlsanity.com/paper/2507.11948/) (2025)
- [TileLang: A Composable Tiled Programming Model for AI Systems](https://mlsanity.com/paper/2504.17577/) (2025)
- [AutoTriton: Automatic Triton Programming with Reinforcement Learning in LLMs](https://mlsanity.com/paper/2507.05687/) (2025)
- [TritonRL: Training LLMs to Think and Code Triton Without Cheating](https://mlsanity.com/paper/2510.17891/) (2025)

## Built on by (most cited first, 22 of 22)

- [KernelBench-Verified: Do LLM-Generated Kernels Actually Beat PyTorch?](https://mlsanity.com/paper/2607.16241/) (2026)
- [Kernel-Smith: A Unified Recipe for Evolutionary Kernel Optimization](https://mlsanity.com/paper/2603.28342/) (2026)
- [AdaExplore: Failure-Driven Adaptation and Diversity-Preserving Search for Efficient Kernel Generation](https://mlsanity.com/paper/2604.16625/) (2026)
- [FastKernels: Benchmarking GPU Kernel Generation in Production](https://mlsanity.com/paper/2605.23215/) (2026)
- [daVinci-kernel: Co-Evolving Skill Selection, Summarization, and Utilization via RL for GPU Kernel Optimization](https://mlsanity.com/paper/2606.16497/) (2026)
- [Scaling Participation in Modular AI Systems](https://mlsanity.com/paper/2606.07812/) (2026)
- [RealisticTritonBench: A Benchmark for Triton-Kernel Generation in Real-World AI Frameworks](https://mlsanity.com/paper/2608.12004/) (2026)
- [Agentic Kernel Optimization: Generating State-of-the-Art GPU Kernels Without Hand-Written CUDA](https://mlsanity.com/paper/2608.14560/) (2026)
- [Surprisal-Guided Selection: Compute-Optimal Test-Time Strategies for Execution-Grounded Code Generation](https://mlsanity.com/paper/2602.07670/) (2026)
- [CUDAHercules: Benchmarking Hardware-Aware Expert-level CUDA Optimization for LLMs](https://mlsanity.com/paper/2605.08467/) (2026)
- [From Tokens to Regions: CUDA-Sensitive Instruction Tuning for GPU Kernel Generation](https://mlsanity.com/paper/2606.16231/) (2026)
- [CANN Bench: Benchmarking Agent Generated Kernels against Real NPU and Algorithmic Limits](https://mlsanity.com/paper/2607.20518/) (2026)
- [CUDABeaver: Benchmarking LLM-Based Automated CUDA Debugging](https://mlsanity.com/paper/2605.08455/) (2026)
- [KernelGenBench: Can LLMs and Agents Write Efficient Kernels Across Operator Sources and Hardware Platforms?](https://mlsanity.com/paper/2607.27231/) (2026)
- [LEAP: Lean Environment-Feedback via Adaptive Pruning for Code RL in GPU Kernel Generation](https://mlsanity.com/paper/2608.01804/) (2026)
- [A Contract-Grade Verifier for LLM-Generated GPU Kernels, and a Native Blackwell Backward for the Gated-Linear-Recurrence Family](https://mlsanity.com/paper/2608.12700/) (2026)
- [PTXBench: Benchmarking and Adapting LLMs for GPU Kernel Optimization with Architecture-specific PTX](https://mlsanity.com/paper/2608.17379/) (2026)
- [AsmEvo: Agentic Assembly-Level Optimization of AMD GPU Kernels with Functional Equivalence Verification](https://mlsanity.com/paper/2608.20711/) (2026)
- [FlashVector: Agent for Hierarchical Model Serving Stack Optimization](https://mlsanity.com/paper/2609.17391/) (2026)
- [KernelZero: Co-Evolving Proposer and Coder for Continuously Improved GPU Kernel Generation](https://mlsanity.com/paper/2609.33074/) (2026)
- [From Experience to Expertise: Adoption-Aware Memory Learning for Data-Scarce NPU Kernel Synthesis](https://mlsanity.com/paper/2609.35568/) (2026)
- [AI as a Compiler: Compiling Triton kernels without the Triton compiler](https://mlsanity.com/paper/2609.36800/) (2026)

## Abstract (arXiv)

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.