Post-Training LLMs for Kernel Generation #

Overview #

This section is about training an AI model to become good at writing kernels, the small hand-tuned programs that make AI hardware run fast. We start with a model that already writes ordinary code and improve it in stages.

First the model learns by imitating a collection of good examples. Then it learns by actually running its own attempts and being rewarded when they turn out correct and fast, which lets it go beyond anything it was shown. Finally we look at how to design that reward carefully, since the model will chase whatever score we hand it, and a careless score invites exactly the corner-cutting that Section 2 warns about. The thread running through it all: copying examples only gets you so far, and the real gains come from letting the model try, run, and learn from the results.

Contents of this section #

3.1 Supervised Fine-Tuning on Kernel Corpora #

The model practices by imitating a large collection of tasks paired with good kernel solutions. Building that collection well is its own craft, and this approach teaches the model how kernels are written, but cannot push it past the quality of the examples.

3.2 Reinforcement Learning from Execution Feedback #

Instead of only copying, the model now writes a kernel, actually runs it on the hardware, and is rewarded when the result is correct and fast. Learning directly from what works lets it discover solutions better than any example it was given.

3.3 Single-Turn and Multi-Turn RL for Kernel Optimization #

Comparing a single try per task with taking several tries in a row, reading the feedback each time to fix and improve the next attempt, much like a human debugging their own code. Letting the model work over several turns is powerful but raises new questions about what to remember and which attempt to credit.

3.4 Reward Design for Kernel RL #

How to design the score that guides learning, from simple weighted sums of “does it compile, is it correct, is it fast” to more careful designs. The aim is that being genuinely correct and fast is what gets rewarded, and the model cannot win by cutting corners.