Inference-Scaling for Kernel Generation #

Overview #

Writing a good kernel, the small program that makes an AI chip run a computation as fast as possible, is hard. This section explores a surprisingly effective shortcut: instead of building a bigger or smarter model, we take one fixed model and simply let it work harder when it writes code, generating many attempts, revising them, and searching for the best one.

This pays off especially well for kernels because we can check every attempt automatically: compile it, confirm its answer matches a trusted reference, and time how fast it runs, then keep the winner. That automatic checker is the hero of the story, and also its weak point, since a model will happily cheat the check if it finds a way. So one subsection is devoted to making the check trustworthy, and the section closes with a hands-on demo you can follow along. The training that produces the underlying model is saved for Section 3; here the model stays fixed.

Contents of this section #

2.1 Inference-Time Scaling for Kernel Code Reasoning #

The main ways to coax better kernels out of the same model: draw many attempts and keep the best, improve one attempt over several rounds using feedback, or search cleverly among candidates. Each one pays off because every attempt can be checked automatically.

2.2 Evolution-Based Approaches for Iterative Kernel Refinement #

Borrowing the idea of natural selection: keep a whole pool of candidate kernels, tweak and combine the best ones, test them, and repeat, gradually breeding faster kernels over many rounds.

2.3 Verification and Cheating Detection #

The more attempts we make, the more chances a model has to pass the test without doing the real work, for instance by copying the expected answer or faking the timer. This subsection shows the common tricks and how to build a checker that cannot be fooled, so measured speedups stay honest.

2.4 NKI Kernel Generation Using LLM Agents #

A follow-along demo where an AI assistant writes a real kernel for AWS Trainium chips (using NKI, the programming language for those chips), measures where it is slow, and improves it step by step, putting the section’s ideas into practice.