GPU Forecasters: Language Models as Selective Surrogates for Kernel Runtime Optimization
This work addresses the bottleneck of expensive on-device GPU kernel evaluation for researchers and developers optimizing deep learning workloads, by proposing LLMs as a cost-effective alternative.
This paper explores using Large Language Models (LLMs) as selective surrogates to forecast GPU kernel performance, aiming to reduce the costly on-device measurements required for kernel optimization. The LLM surrogate allows a kernel search to consider several times more candidates within the same GPU evaluation budget, leading to the discovery of faster kernels compared to an equal-budget baseline.
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.