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

GPU Forecasters: Using LLMs as Selective Surrogates to Accelerate Kernel Optimization

Optimizing GPU kernels is notoriously slow. Every new version must be compiled and tested on real hardware to verify its speed. A new paper suggests we can stop wasting GPU cycles on bad ideas. Researchers can use LLMs as "virtual GPUs" to predict performance. By treating the LLM as a selective surrogate (a model that stands in for an expensive system), engineers can filter out unpromising code. They can then reserve actual hardware for candidates that look truly promising.

The Problem

In modern deep learning, GPU kernels are the fundamental unit of performance. These are low-level programs that execute tensor operations. Optimizing them typically involves an iterative search. This search uses evolutionary algorithms or coding agents to propose new versions. Each candidate must be compiled, checked for correctness, and profiled on a physical device.

As LLM-driven search methods scale, this "generation-evaluation" cycle hits a massive bottleneck. Compilation alone can take 60 to 120 seconds per kernel. Furthermore, the on-device evaluation throughput is often an order of magnitude slower than the LLM's generation speed. Current paradigms, as shown in, force the optimizer to wait on the hardware for every single proposal.

This creates a massive wall-clock penalty. The core tension is that LLMs are great at writing code. However, they are not inherently aware of microarchitectural nuances. Factors like memory hierarchy or register pressure dictate real-world runtime.

How It Works

The authors propose moving from a direct measurement paradigm to a surrogate-assisted one. Instead of sending every candidate to the GPU, they use an LLM to act as a selective surrogate. This model forecasts the relative speedup of a candidate compared to a reference implementation.

The mechanism works in three primary stages:

1. Discretized Forecasting: The model does not attempt to predict absolute runtime in seconds. Such a task is nearly impossible due to compiler and hardware unpredictability. Instead, the model predicts a speedup ratio ($S_i$) categorized into eight logarithmic bins. These bins range from "severe slowdown" to "extreme speedup" [Table 3]. The LLM outputs a probability distribution over these bins. This provides a measure of its own certainty.
2. Selective Prediction: This is the "selective" part of the surrogate. The system uses the predicted probability distribution to decide whether to trust the forecast. If the model's confidence is high, the search proceeds based on the forecast. If the model is uncertain, the candidate is sent to the physical device for ground-truth measurement.
3. Reinforcement Learning (RL) Refinement: The authors apply RL to an open-weights model (GPT-OSS-20B). They use specific reward functions to shape the model's behavior. They move beyond simple "correctness" rewards. They incorporate proper scoring rules like the Brier score and CRPS (Continuous Ranked Probability Score). These rules train the model to produce well-calibrated probability distributions. This ensures that high confidence actually correlates with accuracy.

As shown in, the training loop leverages the byproduct of the search itself.

Every time a kernel is measured on a GPU, that data point becomes a training example for the surrogate.

Numbers

The headline result is a significant boost in search efficiency. The authors report that using the surrogate inside a kernel search allows the optimizer to consider several times as many candidates. This happens under the same GPU evaluation budget. Specifically, the Brier-trained surrogate helps recover more speedup at lower measurement budgets.

The paper demonstrates that off-the-shelf models already have surprising utility. The base GPT-OSS-20B model recovers 82.88% of the available speedup. This holds even when only 1% to 50% of candidates are measured [Table 1]. However, the real wins come from RL. The authors find that Brier-shaped training improves the "speedup recovered" metric. It provides an average gain of 2.10 percentage points over the base model .

Regarding calibration, the authors note that training on correctness alone can be destructive. It might improve the accuracy of the most likely bin. However, it can destroy the model's ability to signal uncertainty. In contrast, adding a Brier calibration term improves the Expected Calibration Error (ECE). ECE measures how well reported confidence matches empirical accuracy. This training makes the relationship between stated confidence and forecast error nearly monotonic .

What's Missing

While the results are compelling, there are gaps that a production engineer should note:

• Discovery Moment Precision: The paper explores whether surrogates can identify "discovery moments." These are steps where a mutation produces a significantly faster kernel. The authors show that surrogates can detect these better than chance . However, the precision is still quite low. This means a surrogate-led search will still trigger many "false positive" GPU measurements. These false positives could eat into the theoretical cost savings.
• Complexity of the Reasoning Chain: The effectiveness of the surrogate relies heavily on Chain-of-Thought (CoT) reasoning. As seen in the failure case analysis (Appendix A), reasoning can collapse. If the model commits to an unverified technical prior, the entire forecast fails. The paper does not deeply investigate how to guard against these "hallucinated" architectural insights.
• Generalization to New Architectures: The experiments focus on A100 GPUs and specific tasks. It is unclear how much retraining would be required to move this surrogate to different architectures. Hardware like H100s or specialized AI accelerators have different memory latencies and instruction sets.

Should You Prototype This

Depends on your search budget.

If you are running small-scale kernel optimizations, the overhead may not be worth it. If you only evaluate a few dozen candidates, the complexity is likely too high. However, if you run large-scale, automated kernel discovery pipelines, this is worth a prototype. This is especially true if you spend hundreds of GPU-hours on compilation and profiling.

The authors have released a collection of 12,388 LLM-generated GPU kernels with measured runtimes. This provides a solid foundation for testing your own surrogate. Code is reportedly available; see the paper for the canonical link at github.com/codezakh/gpu-forecasters. If you can afford the initial cost of generating a training set, the payoff in reduced hardware dependency is substantial.

Figures from the paper