🤖 machine learning

GPU Forecasters: Language Models as Selective Surrogates for Kernel Runtime Optimization

This paper proposes using language models as selective, reinforcement-learning-enhanced surrogates to forecast GPU kernel performance, thereby reducing costly on-device evaluations and enabling the discovery of faster kernels within limited measurement budgets.

Original authors: Zaid Khan, Justin Chih-Yao Chen, Jaemin Cho, Elias Stengel-Eskin, Mohit Bansal

Published 2026-06-01

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Zaid Khan, Justin Chih-Yao Chen, Jaemin Cho, Elias Stengel-Eskin, Mohit Bansal

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to find the fastest route through a massive, complex maze. In the world of computer chips (GPUs), this "maze" is a piece of code called a kernel that tells the chip how to do math. Finding the absolute fastest version of this code is crucial for making AI run quickly, but it's incredibly expensive.

The Problem: The "Real-World" Test is Too Slow
Right now, to see if a new piece of code is fast, you have to:

Write the code.
Compile it (translate it into machine language).
Run it on a real, expensive graphics card (GPU).
Measure how long it took.

This process is like testing a new car design by actually building the car, driving it on a track, and timing it. It takes a long time and costs a lot of money. If you want to test 1,000 different designs, you have to build and drive 1,000 cars. This slows down the search for the best design.

The Solution: The "Virtual" Test Driver
The authors of this paper propose using a smart AI (a Large Language Model or LLM) to act as a predictor or a surrogate. Instead of building and driving every single car, you ask the AI: "Based on the blueprints of this new car design, how fast do you think it will go compared to the old one?"

The AI looks at the code and tries to guess the speed.

The Catch: The AI isn't perfect. Sometimes it guesses wrong.
The Fix: The paper teaches the AI to be honest about its uncertainty. It's like a weather forecaster who says, "I'm 90% sure it will rain," versus one who says, "I have no idea, but I'll guess."

How It Works: The "Selective" Strategy
The researchers created a system where the AI acts as a filter:

Confident Guesses: If the AI is very confident in its prediction (e.g., "This new code will definitely be 2x faster"), the system accepts the guess and skips the expensive real-world test.
Uncertain Guesses: If the AI is unsure (e.g., "It might be faster, or it might be slower"), the system says, "Okay, we don't trust that guess yet," and sends the code to the real GPU for a physical test.

This is called a Selective Surrogate. It's like a hiring manager who interviews 100 candidates. For the obvious top-tier and bottom-tier candidates, the manager makes a quick decision based on the resume. But for the "maybe" candidates, they schedule a full, expensive in-person interview.

The Training: Teaching the AI to be a Better Coach
The paper shows that they can make this AI predictor even better using Reinforcement Learning (a type of training where the AI gets points for being right and loses points for being overconfident when wrong).

They taught the AI not just to guess the speed, but to spread its "uncertainty" correctly.
Result: The AI learned to say, "I'm not sure about this one," more often when it was actually unsure, and to be more confident when it was right. This made the system much more reliable.

The Results: Finding Faster Kernels Faster
When they put this system to work:

Efficiency: They could look at four times as many code candidates using the same amount of time and money. Instead of testing 100 designs on the real GPU, they could test 100 on the AI and only send the top 25 to the real GPU.
Performance: Because they could look at more options, they found faster code than systems that only used the real GPU.
Discovery: The AI was surprisingly good at spotting "breakthrough" moments—times when a small change in code made a huge difference in speed.

In Summary
This paper introduces a "virtual test driver" for computer chips. By using an AI to predict which code designs are likely to be fast, and only physically testing the ones the AI is unsure about, we can search for better software much faster and cheaper. The AI doesn't replace the real test; it acts as a smart gatekeeper, saving the expensive real-world tests for the moments that matter most.

Technical Summary: GPU Forecasters: Language Models as Selective Surrogates for Kernel Runtime Optimization

Problem Statement

GPU kernels are the fundamental workhorses of modern deep learning, where minor efficiency gains translate to significant reductions in training and inference costs. Optimizing these kernels typically involves a generation-evaluation cycle where Large Language Models (LLMs) propose candidate kernels, which are then compiled and executed on physical hardware to measure performance. However, this evaluation process is a severe bottleneck. Each evaluation requires compilation (taking 60–120 seconds for CUDA) and on-device execution, making the throughput of evaluation orders of magnitude slower than generation. As search budgets scale, the cost of repeated physical measurements limits the ability to explore the vast space of potential kernel implementations.

While prior work has focused on using LLMs to generate better kernels, this paper addresses the complementary problem of reducing the evaluation cost. The authors propose using LLMs not just as generators, but as selective surrogates for the physical GPU. The goal is to forecast the relative performance of a candidate kernel against a reference implementation and to identify when the model is uncertain, thereby deferring only the uncertain cases to the expensive physical GPU.

Methodology

Task Formulation: Selective Forecasting

The authors formulate the problem as forecasting the relative speedup ( $S_i = T_{ref} / T_i$ ) of a candidate kernel compared to a reference kernel, given the hardware specifications. Predicting absolute runtime in seconds is deemed intractable due to the complexity of modeling compiler behavior and micro-architectural details. Instead, the target is discretized into eight ordinal speedup bins (ranging from "Severe Slowdown" to "Extreme Speedup").

The surrogate (an LLM) takes the reference kernel, candidate kernel, and hardware specs as input. It generates a chain of reasoning and outputs a probability distribution over the eight bins. The model's stated confidence is defined as the maximum probability assigned to any single bin.

Selective Prediction Strategy

The system operates under a selective prediction framework. The surrogate is used to rank candidates. If the surrogate's confidence exceeds a certain threshold, the forecast is accepted, and the candidate is not measured on the GPU. If the confidence is low (indicating uncertainty), the candidate is deferred to the physical GPU for ground-truth measurement. This allows the search process to consider many more candidates within a fixed GPU measurement budget.

Reinforcement Learning (RL) for Calibration

To improve the surrogate's accuracy and, crucially, its calibration (the alignment between stated confidence and actual accuracy), the authors apply Reinforcement Learning to an open-weights model (GPT-OSS-20B). They define rewards based on:

Correctness: Whether the predicted bin matches the measured bin.
Categorical Calibration (Brier Score): Penalizing the distribution based on the distance from the true one-hot target.
Distance-Aware Calibration (CRPS): Penalizing the cumulative distribution, accounting for the ordinal nature of speedup bins (e.g., predicting a 2× speedup when the truth is 1.5× is less severe than predicting 4×).

The training data is constructed as a by-product of the kernel search itself, using successful candidates from previous search runs as training examples.

Key Contributions

LLMs as GPU Surrogates: The paper demonstrates that LLMs can achieve non-trivial accuracy in forecasting how a physical GPU responds to candidate kernels, recovering a substantial fraction of the available speedup by ranking candidates based on forecasts rather than measurements.
Selective Prediction Framework: The authors cast GPU forecasting as a selective prediction problem, showing that surrogates can report calibrated uncertainty. This enables a search strategy that trusts high-confidence forecasts and reserves physical measurements for uncertain cases.
Impact of RL on Surrogate Behavior: Through experiments, the paper characterizes how RL training shifts the surrogate's properties. Specifically, adding calibration rewards (Brier or CRPS) improves the reliability of the model's confidence estimates, whereas training on correctness alone can degrade the relationship between confidence and accuracy.
End-to-End Search Improvement: The study shows that integrating the surrogate into a kernel search loop (specifically a PUCT-based search) allows the search to explore significantly more candidates per GPU measurement. On most tasks, this leads to finding faster kernels than a baseline that measures every proposal under the same GPU budget.
Dataset Release: The authors release a reusable collection of 12,388 LLM-generated GPU kernels with measured runtimes, spanning 118 problems, multiple GPU types (A100, L40S, H100), and various generation methods. This dataset supports broader research in GPU performance prediction and kernel search.

Results

Forecast Accuracy: Off-the-shelf LLMs (e.g., GPT-OSS-20B) recover approximately 83% of the best available speedup when only the top 1% of candidates are measured on the GPU. This performance improves to 94% when measuring the top 50%.
Calibration and RL: RL training significantly alters the surrogate's behavior.
- Correctness-only training improves exact-bin accuracy but degrades the correlation between confidence and error (the model becomes overconfident on wrong predictions).
- Calibration training (Brier/CRPS) improves the monotonicity of forecast error vs. confidence. Notably, Brier-calibrated training improves the "speedup recovered" metric by ~2.1 percentage points over the base model, suggesting that better-calibrated distributions lead to better ranking signals for the search algorithm.
Search Efficiency: In end-to-end kernel search, the surrogate-enabled search finds faster kernels than the baseline on four out of six tasks, often reaching the baseline's best speedup with substantially fewer GPU measurements. On tasks where the baseline wins, the search appears saturated early, indicating the surrogate is most beneficial when the search space is large and unexplored.
Discovery Moments: The surrogate can identify "discovery moments" (mutations where a child kernel is significantly faster than its parent) better than random chance, though precision remains a challenge, necessitating GPU confirmation for high-stakes decisions.

Significance and Claims

The paper claims that LLMs can play a broader role in kernel optimization beyond mere code generation. By acting as virtual models of a GPU, they can serve as high-throughput, low-cost filters that manage the scarcity of physical hardware resources.

The authors emphasize that the primary value of this approach is not just speed, but the ability to perform model-based optimization without querying the original system at every step. The surrogate offers higher throughput and lower per-query cost than a GPU, trading off forecast noise that can be managed through calibrated uncertainty estimates. The work connects kernel optimization to the broader field of building neural simulators for complex physical systems, suggesting a path toward optimizers that treat hardware measurement as a scarce resource to be used selectively.