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

Abstract: GPU kernels are the workhorse of modern deep learning, and optimizing them (via evolutionary search or coding agents) usually requires repeated measurement on target hardware. While these measurements provide the ground-truth signal necessary for kernel search, they are costly, because each evaluation of a kernel requires compilation and repeated execution on a GPU. As improvements in LLM inference reduce the cost of writing novel kernels and LLM-driven searches scale to large search budgets, on-device evaluation becomes a bottleneck. To address this, we study how LLMs can serve as selective GPU surrogates for kernel evaluation, by forecasting the performance of proposed kernels. A useful surrogate should be accurate, and it should be selective, by knowing when it could be wrong, and deferring to the GPU. To evaluate surrogates, we measure whether their forecasts are accurate, calibrated, and practically useful for recovering fast kernels under limited GPU-measurement budgets. Next, we study whether reinforcement learning can improve forecast accuracy and confidence calibration. Our experiments demonstrate that LLMs can accurately forecast relative kernel performance, that their utility can be improved through reinforcement learning. Used inside a kernel search, the surrogate lets the search consider several times as many candidates under the same GPU evaluation budget, and that leads to finding faster kernels than an equal-budget baseline. These results suggest that LLMs can play a broader role in kernel optimization, by acting as virtual models of a GPU rather than solely as kernel generators for search.

• The paper demonstrates that LLM-based surrogates, notably GPT-OSS-20B fine-tuned with RL, can predict discretized kernel speedup bins with calibrated confidence.
• It shows that selective prediction via thresholding confidence significantly cuts GPU measurements by focusing on high-confidence candidate kernels.
• RL-based calibration improves ranking and speedup recovery, underscoring the potential of model-driven, resource-efficient kernel optimization.

LLMs as Surrogates for GPU Kernel Runtime Optimization

Motivation and Problem Statement

Optimizing low-level GPU kernels is essential for efficient execution of modern deep learning workloads. While LLM-driven generation of candidate kernels has become increasingly efficient, the measurement of kernel performance on target hardware remains a major bottleneck due to costs associated with compilation, correctness validation, and on-device profiling. The paper "GPU Forecasters: LLMs as Selective Surrogates for Kernel Runtime Optimization" (2605.31464) addresses whether LLMs can serve as accurate and calibrated surrogates for kernel runtime prediction, thereby reducing the dependency on expensive physical device measurements during kernel search.

This work frames the problem as a selective prediction task: the LLM surrogate forecasts the relative speedup of candidate kernels and defers uncertain cases to real GPU measurement. The primary research questions are whether such a surrogate can (1) reliably prioritize candidates, (2) quantify and communicate its own uncertainty for selective use, (3) be improved via reinforcement learning (RL), and (4) improve the efficacy of actual kernel search under strict measurement budgets.

Figure 1: LLM-driven kernel optimizers rely on costly physical GPU measurements, whereas a learned surrogate forecasts kernel outcomes and flags uncertain candidates, motivating selective use of the surrogate.

Methodology

Surrogate Architecture and Task Formulation

The surrogate receives as input both a reference implementation and a candidate kernel, plus the target hardware specification, and outputs a probability distribution over discretized speedup bins, representing the candidate's runtime improvement relative to the reference. Predicting absolute runtimes is intractable from code alone; thus, the authors define eight ordinal bins on a logarithmic scale, which aggregate coarse-grained performance factors meaningful for kernel search.

Selective use of the surrogate is enabled by thresholding on the maximum predicted probability (i.e., the model's stated confidence). When the surrogate's forecast is insufficiently confident, device measurement is triggered instead.

Figure 2: Kernel evaluation as selective forecasting over discrete speedup bins; the surrogate outputs a forecast distribution, defers low-confidence cases, and is trained via RL to optimize correctness and calibration rewards.

Reinforcement Learning for Surrogate Calibration

Open-weight models, specifically GPT-OSS-20B, are fine-tuned with RL to improve both exact-bin accuracy and the quality of the forecast distribution. RL rewards include: (1) correctness for predicting the true bin, (2) Brier score for categorical calibration, and (3) CRPS for distance-aware calibration across bin orderings. These shape the model's ability to quantify uncertainty and to calibrate its output probabilities, impacting its value in selective prediction settings.

Results

Surrogate Forecasting Performance

The evaluation uses 480 held-out kernels across six realistic tasks, targeting competitive and non-trivial heterogenous operations on NVIDIA A100s. The metric of interest is the fraction of best-achievable speedup captured by only GPU-measuring the surrogate's top-k% of candidates. GPT-OSS-20B achieves 82.9% of maximal speedup averaged across small budgets (1%-50%), while Gemini-3 Flash achieves 93.1%. Calibration quality varies across models, with ECE ranging from 0.401 to 0.624.

Notably, off-the-shelf LLM surrogates enable recovery of a substantial amount of available speedup while measuring only a fraction of candidates on the GPU.

Effect of Reinforcement Learning

RL tuning impacts both raw forecast distributions and their calibration. The base surrogate tends to overestimate the performance of slow kernels and underestimate the fastest. RL redistributes probability mass and tunes the confidence signal.

Figure 3: The surrogate's error distribution before and after RL; RL spreads uncertainty more effectively, particularly for fast kernels.

Calibration-shaped RL (Brier score) improves both calibration and speedup-recovery utility, with a +2.06% gain in speedup found under budgeted search, at the expense of a slight increase in forecast error for the argmax prediction. Correctness-only RL improves exact-bin accuracy but correlates less well with confidence and weakens ranking utility.

Figure 4: Calibration rewards cause reported confidence to correlate with forecast error, whereas correctness-only RL disrupts this relationship.

Figure 5: Brier calibration increases the speedup found under most measurement budgets in comparison to the base surrogate and other RL variants.

Surrogate Integration in Kernel Search

The surrogate mechanism is evaluated in a real search loop using a PUCT-based max-reward tree search. When candidate selection is guided by surrogate forecasts (with Brier-calibrated GPT-OSS-20B), the search finds matching or superior kernels on four of six tasks compared to a baseline that directly evaluates all proposed candidates, but with fewer required measurements. On the remaining two tasks, the traditional baseline finds slightly faster kernels, coinciding with scenarios in which the search landscape is already saturated.

Figure 6: In real search, the surrogate-moderated procedure matches or exceeds the baseline in most tasks, requiring substantially fewer GPU measurements.

Detecting Discovery Moments

The surrogate can also identify "discovery moments"—mutations where a child kernel is significantly faster than its parent—better than random, especially for large speedup thresholds. However, the precision remains moderate and does not eliminate the need for confirmation by GPU measurement.

Figure 7: Precision-recall tradeoff for surrogate identification of discovery moments; performance is above chance but not sufficient for standalone use.

Surrogate Efficiency

A single A100 can serve substantially more surrogate queries per hour than direct kernel measurement, especially when concurrent inference is exploited, highlighting the throughput advantage of model-based forecasting.

Figure 8: Evaluations per hour for kernel benchmarking vs. surrogate: surrogates enable 1-2 orders of magnitude higher throughput at scale.

Implications and Future Directions

This study provides evidence that LLMs can function as calibrated, selective surrogates for GPU kernel performance prediction, making them directly useful in budget-constrained optimization loops. The results suggest a model-based search paradigm, where hardware measurement is treated as a scarce resource and surrogate forecasts are exploited wherever reliable.

The findings impact two major research fronts: (1) kernel optimization tooling—enabling more expansive, risk-aware search over candidate programs, and (2) the construction of learned "neural simulators" of physical computing systems, relevant to broader efforts in model-based RL and surrogate-assisted optimization. Future investigations could target (a) incorporation of native abstention and risk control in surrogate design, (b) domain adaptation for new hardware or kernel families, and (c) self-improving surrogates via search-bootstrapped data generation, further tightening the loop between generative and evaluative components.

Conclusion

The integration of LLM-derived surrogates into kernel optimization search demonstrates significant reductions in the number of expensive hardware measurements needed to recover competitive performance. Selective prediction with calibrated surrogates, especially those improved by RL-driven calibration, enhances both practical search performance and the interpretability of surrogate forecasts. These insights establish a path toward highly efficient, model-based optimization pipelines in compiler and systems ML, with wider applicability to domains where physical system measurements are bottlenecked or costly.