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

2026-05-29

AI Summary

This work addresses the high cost of repeated empirical timing measurements in GPU kernel optimization, which constitutes a major bottleneck in search-based tuning. To overcome this limitation, the paper proposes the first approach that leverages a large language model (LLM) as a self-reflective, selective performance surrogate to predict the execution time of candidate kernels. The LLM actively defers to ground-truth measurements when its prediction confidence is insufficient, and it is further enhanced through reinforcement learning to improve both prediction accuracy and confidence calibration. Under a fixed GPU measurement budget, this method enables the evaluation of several times more candidate implementations than conventional approaches, substantially increasing the efficiency of discovering high-performance kernels.

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.

Problem

• GPU kernel optimization
• performance forecasting
• surrogate models
• measurement bottleneck
• language models

Innovation

• LLM surrogate
• kernel performance forecasting
• selective evaluation
• reinforcement learning
• GPU kernel optimization