Table of Contents
Fetching ...

Fine-Tuning GPT-5 for GPU Kernel Generation

Ali Tehrani, Yahya Emara, Essam Wissam, Wojciech Paluch, Waleed Atallah, Łukasz Dudziak, Mohamed S. Abdelfattah

TL;DR

This work tackles GPU kernel generation, a domain where data scarcity and hardware-specific optimization hinder supervised fine-tuning. It introduces reinforcement learning from verifiable rewards (RLVR) in the Makora environment, leveraging GPT-5 as a base model to learn correct and efficient Triton kernels. Through careful dataset curation, tool-augmented multi-turn refinement, and robust defenses against reward hacking, the approach achieves state-of-the-art kernel generation: 77.0% functional correctness in a single attempt and 97.4% correctness with MakoraGenerate, along with a 2.12x geometric mean speedup over TorchInductor. The results demonstrate that targeted post-training RL can unlock LLM capabilities in highly specialized domains, enabling practical AI-assisted accelerator programming and scalable kernel optimization.

Abstract

Developing efficient GPU kernels is essential for scaling modern AI systems, yet it remains a complex task due to intricate hardware architectures and the need for specialized optimization expertise. Although Large Language Models (LLMs) demonstrate strong capabilities in general sequential code generation, they face significant challenges in GPU code generation because of the scarcity of high-quality labeled training data, compiler biases when generating synthetic solutions, and limited generalization across hardware generations. This precludes supervised fine-tuning (SFT) as a scalable methodology for improving current LLMs. In contrast, reinforcement learning (RL) offers a data-efficient and adaptive alternative but requires access to relevant tools, careful selection of training problems, and a robust evaluation environment. We present Makora's environment and tools for reinforcement learning finetuning of frontier models and report our results from fine-tuning GPT-5 for Triton code generation. In the single-attempt setting, our fine-tuned model improves kernel correctness from 43.7% to 77.0% (+33.3 percentage points) and increases the fraction of problems outperforming TorchInductor from 14.8% to 21.8% (+7 percentage points) compared to baseline GPT-5, while exceeding prior state-of-the-art models on KernelBench. When integrated into a full coding agent, it is able to solve up to 97.4% of problems in an expanded KernelBench suite, outperforming the PyTorch TorchInductor compiler on 72.9% of problems with a geometric mean speedup of 2.12x. Our work demonstrates that targeted post-training with reinforcement learning can unlock LLM capabilities in highly specialized technical domains where traditional supervised learning is limited by data availability, opening new pathways for AI-assisted accelerator programming.

Fine-Tuning GPT-5 for GPU Kernel Generation

TL;DR

This work tackles GPU kernel generation, a domain where data scarcity and hardware-specific optimization hinder supervised fine-tuning. It introduces reinforcement learning from verifiable rewards (RLVR) in the Makora environment, leveraging GPT-5 as a base model to learn correct and efficient Triton kernels. Through careful dataset curation, tool-augmented multi-turn refinement, and robust defenses against reward hacking, the approach achieves state-of-the-art kernel generation: 77.0% functional correctness in a single attempt and 97.4% correctness with MakoraGenerate, along with a 2.12x geometric mean speedup over TorchInductor. The results demonstrate that targeted post-training RL can unlock LLM capabilities in highly specialized domains, enabling practical AI-assisted accelerator programming and scalable kernel optimization.

Abstract

Developing efficient GPU kernels is essential for scaling modern AI systems, yet it remains a complex task due to intricate hardware architectures and the need for specialized optimization expertise. Although Large Language Models (LLMs) demonstrate strong capabilities in general sequential code generation, they face significant challenges in GPU code generation because of the scarcity of high-quality labeled training data, compiler biases when generating synthetic solutions, and limited generalization across hardware generations. This precludes supervised fine-tuning (SFT) as a scalable methodology for improving current LLMs. In contrast, reinforcement learning (RL) offers a data-efficient and adaptive alternative but requires access to relevant tools, careful selection of training problems, and a robust evaluation environment. We present Makora's environment and tools for reinforcement learning finetuning of frontier models and report our results from fine-tuning GPT-5 for Triton code generation. In the single-attempt setting, our fine-tuned model improves kernel correctness from 43.7% to 77.0% (+33.3 percentage points) and increases the fraction of problems outperforming TorchInductor from 14.8% to 21.8% (+7 percentage points) compared to baseline GPT-5, while exceeding prior state-of-the-art models on KernelBench. When integrated into a full coding agent, it is able to solve up to 97.4% of problems in an expanded KernelBench suite, outperforming the PyTorch TorchInductor compiler on 72.9% of problems with a geometric mean speedup of 2.12x. Our work demonstrates that targeted post-training with reinforcement learning can unlock LLM capabilities in highly specialized technical domains where traditional supervised learning is limited by data availability, opening new pathways for AI-assisted accelerator programming.
Paper Structure (49 sections, 9 equations, 7 figures, 3 tables)

This paper contains 49 sections, 9 equations, 7 figures, 3 tables.

Figures (7)

  • Figure 1: Reinforcement Learning fine tuning with Makora's evaluation infrastructure and tools.
  • Figure 2: Workflow of dataset processing and cleaning.
  • Figure 3: Runtime distribution of measured training problems and KernelBench.
  • Figure 4: Predicted difficulty level distribution across sampled training datasets, illustrating the coverage of problem complexity for the 100, 1000, and larger reference sets.
  • Figure 5: One attempt benchmark results where higher is better. GPT-5 trained with reinforcement learning outperforms all baselines. MakoraGenerate is an agentic system that is not restricted to a single attempt and instead leverages multiple strategies and search procedures to generate and optimize kernels.
  • ...and 2 more figures