Selective GPU Surrogates for GPU Kernel Optimization

• Selective GPU surrogates are mechanisms that integrate language-model predictions and profiler analyses to forecast GPU kernel speedups and trigger on-device measurements when confidence is low.
• They employ selective prediction and ranking by expected speedup, balancing calibration with efficiency in budgeted kernel optimization loops.
• Reinforcement learning and expert surrogates refine calibration and ranking, enabling targeted hardware measurements and improved optimization outcomes.

Selective GPU surrogates are surrogate mechanisms inserted into GPU kernel optimization loops only when their signal is reliable or relevant. In current work, the term denotes two closely related practices. The first treats a LLM as a selective surrogate for expensive on-device runtime measurement: the model forecasts the relative speedup of a candidate kernel, reports a calibrated confidence, and defers to hardware measurement when uncertain. The second treats profiler-interpreting tools as selective expert surrogates: callable analyzers embody human bottleneck heuristics, are invoked only for bottleneck-relevant cases, and convert hardware telemetry into targeted optimization guidance. These two formulations are exemplified by "GPU Forecasters: LLMs as Selective Surrogates for Kernel Runtime Optimization" (Khan et al., 29 May 2026) and "Optimizing CUDA like a Human: Micro-Profiling Tools as Expert Surrogates for LLM-Based GPU Kernel Optimization" (Gai et al., 24 Jun 2026).

1. Definition and operating principle

GPU kernel optimization commonly relies on repeated compilation and execution on target hardware. In the runtime-forecasting formulation, this cost becomes a bottleneck because each evaluation of a kernel requires compilation and repeated execution on a GPU, whereas improvements in LLM inference reduce the cost of writing novel kernels and LLM-driven searches scale to large search budgets. Selective GPU surrogates address this asymmetry by forecasting performance for some candidates and reserving on-device measurement for cases where the forecast is uncertain or where confirmation remains necessary (Khan et al., 29 May 2026).

In the forecasting setting, the surrogate ingests a reference kernel, a candidate kernel, and a hardware description, then predicts the candidate’s speedup relative to the reference rather than its absolute runtime. The surrogate is selective in the precise sense that it does not merely output a point forecast; it also outputs a confidence that can be used as a deferral trigger. A useful surrogate should therefore be both accurate and selective, “by knowing when it could be wrong, and deferring to the GPU” (Khan et al., 29 May 2026).

In the expert-surrogate setting, selectivity refers not to abstention from prediction but to targeted invocation. KernelPro defines an expert surrogate as a programmatic embodiment of a human heuristic: a callable analyzer that recognizes bottleneck signatures in profiler telemetry, diagnoses root cause, and prescribes targeted fixes. Its micro-profiling tools are filtered by bottleneck class so that only relevant analyses run, avoiding prompt noise and irrelevant recommendations such as tensor-core diagnostics on clearly memory-bound kernels (Gai et al., 24 Jun 2026).

A recurrent misconception is that a surrogate replaces hardware validation or profiling altogether. The reported systems do not support that conclusion. In the forecasting formulation, GPU confirmation remains the ground-truth signal for kernel search. In the expert-surrogate formulation, profiler-derived guidance improves search quality, but the loop still compiles, runs, and validates kernels on device.

2. Formalization as selective prediction over relative speedups

The runtime-forecasting line formalizes GPU runtime prediction as selective prediction. For a candidate kernel $i$, the input is

$x_i = (r, k_i, h),$

where $r$ is the reference, $k_i$ the candidate code, and $h$ the hardware description. The ground-truth target is the candidate’s speedup relative to the reference,

$S_i = T_{\text{ref}} / T_i,$

where $T_{\text{ref}}$ and $T_i$ are the measured mean runtimes for the reference and candidate, respectively (Khan et al., 29 May 2026).

Rather than regress directly on $S_i$, the surrogate predicts a categorical distribution over eight ordinal speedup bins: $$p_i = (p_{i,1}, \ldots, p_{i,8}), \qquad p_{i,b} \ge 0, \qquad \sum_{b=1}^{8} p_{i,b} = 1.$$ The predicted bin is

$x_i = (r, k_i, h),$0

and the stated confidence is

$x_i = (r, k_i, h),$1

Selective use is threshold-based: trust the forecast if $x_i = (r, k_i, h),$2, and otherwise defer to hardware measurement. The decision rule is

$x_i = (r, k_i, h),$3

The paper does not fix a single $x_i = (r, k_i, h),$4 or report a formal risk-coverage frontier, but it gives the ingredients for the implied coverage

$x_i = (r, k_i, h),$5

and a corresponding selective risk over the retained set

$x_i = (r, k_i, h),$6

where $x_i = (r, k_i, h),$7 is the representative speedup for bin $x_i = (r, k_i, h),$8 (Khan et al., 29 May 2026).

The surrogate is trained and scored on relative speedups and is used by kernel search primarily as a ranking signal. For candidate $x_i = (r, k_i, h),$9, it defines an expected speedup score

$r$0

Search ranks candidates by $r$1 and measures only a top fraction under a fixed budget. This design is notable because it decouples forecast utility from exact-bin correctness: a distribution that places probability mass appropriately across adjacent bins can be more useful for budgeted ranking than a brittle argmax.

The same paper evaluates surrogate quality along three axes. “Speedup recovered” measures how much of the best-available speedup a search would recover if it only measured the top-$r$2 candidates ranked by $r$3, averaged over $r$4. Expected Calibration Error (ECE) measures agreement between stated confidence and exact-bin accuracy across ten confidence buckets. Forecast error measures the mean absolute gap between the representative speedup of the argmax bin and the measured speedup. The paper does not report NLL or Spearman/Kendall rank correlations; its emphasis is calibration, budgeted ranking utility, and absolute forecast error (Khan et al., 29 May 2026).

3. Empirical performance and reinforcement-learning calibration

On 480 held-out Triton kernels across six GPU Mode tasks—TriMul, cross-entropy, two Gated DeltaNet forward kernels, GDN recompute W/U, and FP8 quantization—all evaluated on NVIDIA A100, off-the-shelf language-model surrogates already recover a large fraction of the available speedup under limited measurement budgets (Khan et al., 29 May 2026).

Model	Recovered	ECE / Forecast error
Gemini-3 Flash	93.12 ± 0.97%	0.624 ± 0.002 / 0.421 ± 0.021
GPT-OSS-120B	86.85 ± 2.51%	0.479 ± 0.005 / 0.632 ± 0.029
GPT-OSS-20B (base)	82.88 ± 3.85%	0.441 ± 0.011 / 0.643 ± 0.011
DeepSeek-V4	79.86 ± 3.98%	0.401 ± 0.009 / 0.805 ± 0.023

These results show that high recovered speedup and good calibration are not identical objectives. Gemini-3 Flash has the highest recovered speedup, while DeepSeek-V4 has the lowest ECE among the listed off-the-shelf models. The reliability analysis adds another axis: deviations from monotone decrease of forecast error with rising confidence are summarized by $r$5, where lower is better.

The same study investigates whether reinforcement learning can improve forecast quality. It trains LoRA adapters on GPT-OSS-20B using GRPO through Tinker, with single-turn rollouts that end in a structured forecast tool call. Training rows come directly from measured search artifacts and use only successful device measurements. The reported hyperparameters are LoRA rank 32, GRPO with group size 8, 8 groups per batch, learning rate $r$6, 20 iterations, max rollout length 8192 tokens, and temperature 1.0 (Khan et al., 29 May 2026).

Three reward designs are reported: $r$7

$r$8

$r$9

$k_i$0

The empirical outcome is asymmetric. Correctness-only reward improves ECE by $k_i$1 but reduces recovered speedup by $k_i$2 percentage points, increases forecast error by $k_i$3, and yields $k_i$4, meaning confidence no longer aligns well with error. Correctness + CRPS improves ECE by $k_i$5 but reduces recovered speedup by $k_i$6 percentage points and increases latency by $k_i$7 s. Correctness + Brier improves ECE by $k_i$8, increases recovered speedup by $k_i$9 percentage points, increases forecast error by $h$0, yields $h$1, shortens inference chains by $h$2 tokens, and reduces latency by $h$3 seconds per call (Khan et al., 29 May 2026).

This combination of results is important because it demonstrates a nontrivial trade-off: proper scoring rules can improve calibration and ranking even when argmax-bin forecast error worsens. The paper’s mechanism analysis attributes this to probability-mass redistribution across bins. The base surrogate overpredicts slow kernels and underpredicts the fastest ones; correctness-only RL mainly shifts argmax decisions, whereas Brier-shaped RL spreads probability appropriately across adjacent bins, improving calibration and ranking.

4. Integration into kernel search under fixed GPU budgets

The forecasting surrogate is embedded in a max-reward tree search with a PUCT selection rule. Without a surrogate, every generated candidate is measured. With a surrogate, the system increases proposal multiplicity and filters proposals down to a budgeted subset: parent nodes are selected via PUCT, a mutator LLM proposes $h$4 mutations per parent, the surrogate produces $h$5 for each candidate, candidates are ranked by expected speedup $h$6, and only a top fraction is measured under a fixed GPU budget (Khan et al., 29 May 2026).

For a candidate batch $h$7 at step $h$8, with per-candidate forecast $h$9 and budget $S_i = T_{\text{ref}} / T_i,$0, the paper writes the budgeted measurement policy as

$S_i = T_{\text{ref}} / T_i,$1

It also describes an optional confidence threshold $S_i = T_{\text{ref}} / T_i,$2: measure any $S_i = T_{\text{ref}} / T_i,$3 with $S_i = T_{\text{ref}} / T_i,$4, and trust the surrogate for $S_i = T_{\text{ref}} / T_i,$5 with $S_i = T_{\text{ref}} / T_i,$6. In the reported extrinsic experiment, the search proposes $S_i = T_{\text{ref}} / T_i,$7 more candidates per parent and measures only the top 25% ranked by $S_i = T_{\text{ref}} / T_i,$8, maintaining equal GPU measurements per step with and without the surrogate.

The resulting equal-budget comparison is mixed but favorable overall. With the Brier-trained GPT-OSS-20B surrogate and Gemini-3 Flash as mutator, the surrogate-enabled search finds faster kernels than the equal-budget baseline on FP8 quantization, GDN ChunkFwd-o, and GDN Recompute W/U, and both procedures reach essentially the same best speedup on TriMul. The baseline wins by about 5% on cross-entropy and 7% on GDN ChunkFwd-h; in both cases the search saturates early, with near-optimal kernels appearing within approximately 15 measurements. Overall, the surrogate translates intrinsic ranking improvements into extrinsic search gains on four of six tasks (Khan et al., 29 May 2026).

The measurement protocol matters because the surrogate’s utility is defined against actual device timings. Core measurements use NVIDIA A100-80GB. Triton kernels, with Triton 3.3.1, and CUDA kernels are compiled and timed via a standardized Modal A100 worker protocol with seeds set for Python, NumPy, PyTorch CPU, and PyTorch CUDA to 42; correctness is checked first; CUDA-event timing uses pre/post synchronizations; and an adaptive loop with at least 3 and at most 100 iterations stops early if standard error over mean is below 0.001, if total runtime exceeds 10 seconds, or if wall-clock exceeds 120 seconds. The geometric mean across cases defines scalar speedup reward. Appendix latency numbers further quantify the budget trade-off: Triton kernel benchmark evaluations average 20.16 s, CUDA kernel benchmark evaluations 152.55 s, and LLM surrogate forecasts 14.69 s, with surrogate throughput scaling with concurrency on a single A100 (Khan et al., 29 May 2026).

5. Profiling tools as selective expert surrogates

KernelPro generalizes selectivity from forecasting to diagnosis. Its micro-profiling tools are expert surrogates formalized as a 5-tuple

$S_i = T_{\text{ref}} / T_i,$9

where $T_{\text{ref}}$0 denotes required metrics, $T_{\text{ref}}$1 thresholds, and $T_{\text{ref}}$2 emits severity, root cause, ranked recommendations, and expected improvement. A representative example is WarpStallTool, which computes the dominant stall reason

$T_{\text{ref}}$3

over $T_{\text{ref}}$4, triggers at 40%, and prescribes stall-specific remedies such as coalescing and shared memory for memory dependency, $T_{\text{ref}}$5 over $T_{\text{ref}}$6 for barrier, and $T_{\text{ref}}$7 for long-latency loads (Gai et al., 24 Jun 2026).

Selective invocation is organized as a two-stage process. Stage 1 performs roofline-based bottleneck classification using

$T_{\text{ref}}$8

If $T_{\text{ref}}$9, the kernel is classified as memory-bound; otherwise compute-bound. Stall and occupancy cues further mark latency-bound or mixed cases, with dominant warp stall $T_i$0 and achieved occupancy $T_i$1 being key thresholds. Stage 2 executes only bottleneck-relevant tools: memory-bound kernels receive coalescing, vectorization, shared-memory tiling, DRAM throughput, and system-level launch or transfer analyses; compute-bound kernels receive tensor-core utilization, SASS instruction-mix, register spill, and occupancy/scheduler analyses; mixed cases receive a combination (Gai et al., 24 Jun 2026).

The central rationale is empirical rather than merely architectural. KernelPro’s ablation shows that “raw ncu only” underperforms “no feedback” at 1.77x versus 3.35x geometric mean, while tool-interpreted feedback reaches 4.00x; the difference between raw metrics and tools is reported as $T_i$2. Proactive tool orchestration yields 23% better speedups than reactive tool calls, with $T_i$3. The system attributes this to reduced prompt noise, removal of irrelevant diagnoses, and conversion of counters into ranked natural-language directives (Gai et al., 24 Jun 2026).

KernelPro combines these surrogates with a domain-adapted MCTS. Each node is a compiled-and-profiled CUDA kernel. Selection uses

$T_i$4

with $T_i$5. Progressive widening constrains a node with $T_i$6 visits to at most $T_i$7 children, with $T_i$8, $T_i$9 for ROOT and FAILED nodes, and $S_i$0 for SUCCESSFUL nodes. Reward calibration is solution-level: $S_i$1 for correct kernels, $S_i$2 for incorrect output, and $S_i$3 for crash or compile error. Search memory carries “Search Findings,” “Errors & Corrections,” “Successful Patterns,” “Key Files,” and “Learnings” into later prompts (Gai et al., 24 Jun 2026).

The performance figures place selective expert surrogates within a broader optimization agenda than runtime forecasting alone. On KernelBench, KernelPro reports geometric mean speedups of 2.42x, 4.69x, and 5.30x on Levels 1, 2, and 3, respectively, with 100% task coverage. On VeOmni’s expert-optimized MoE training kernels, it achieves 1.23x over hand-tuned Triton by generating a from-scratch raw-CUDA+CuTe Hopper WGMMA kernel. It also reports an 11.6% measured energy reduction at matched 2.52x speed on Swish, attributed to a lower-energy instruction stream with 56 rather than 216 SASS instructions and the use of read-only $S_i$4 and fast reciprocal $S_i$5 (Gai et al., 24 Jun 2026).

6. Scope, limitations, and practical design patterns

The runtime-forecasting study is narrow in hardware concentration but broad in archive construction. Its main evaluation and RL training target six GPU Mode tasks and a 480-kernel Triton hold-out set binned across eight speedup bins, all measured on NVIDIA A100 GPUs. The released archive contains 12,388 LLM-generated GPU kernels with measured runtimes across 118 problems, three GPU types—A100, L40S, H100—and four generation methods, consuming 400M LLM tokens and 600 GPU-hours; the subset supporting the GPU Mode held-out evaluation uses approximately 47 A100-hours and approximately 62M LLM tokens. A separate “discovery moments” analysis uses 1,347 parent-child pairs from 27 completed PUCT searches spanning 21 problems across the six GPU Mode tasks and KernelBench Level 3 problems, measured on A100 and L40S GPUs, and reports mean average precision of 0.20 at a 1.25x cutoff, 0.39 at 1.5x, and 0.31 at 2x, indicating usable but noisy signal for surfacing large improvements (Khan et al., 29 May 2026).

The most important limitations are also precisely reported. In the forecasting line, the base surrogate overpredicts speedups for slow kernels and underpredicts for the fastest ones; correctness-only RL worsens reliability; precision-recall curves are not strong enough to eliminate GPU confirmation; speculative priors such as “einsum is slow” can induce four-bin errors; and the paper does not provide cross-architecture calibration curves or per-architecture thresholds. Gains are budget-sensitive: when a baseline saturates very quickly, surrogate-enabled exploration has little room to help. Forecast error and ranking utility can also move in opposite directions, which is why the paper prioritizes ranking under a budget rather than argmax-bin accuracy (Khan et al., 29 May 2026).

KernelPro has an analogous set of constraints. Misclassification can occur at low signal-to-noise, especially for short kernels or noisy counters. Tensor-core utilization may be misleading when library calls hide instruction paths, which motivates SASS disassembly and $S_i$6 symbol checks. Portability across architectures requires SM-specific shapes, flags, and tool versions. Profiling overhead must be controlled by caching metric snapshots keyed by binary hash, skipping unchanged binaries, capping per-iteration profiler invocations, batching $S_i$7 runs, using adaptive gap thresholds, and limiting SASS disassembly to generated $S_i$8 kernels (Gai et al., 24 Jun 2026).

The practical guidance emerging from these studies converges on a common design pattern. For runtime surrogates, use $S_i$9 as the deferral trigger, inspect the reliability curve to choose $$p_i = (p_{i,1}, \ldots, p_{i,8}), \qquad p_{i,b} \ge 0, \qquad \sum_{b=1}^{8} p_{i,b} = 1.$$0, prefer correctness + Brier over correctness-only or CRPS, rank by expected speedup $$p_i = (p_{i,1}, \ldots, p_{i,8}), \qquad p_{i,b} \ge 0, \qquad \sum_{b=1}^{8} p_{i,b} = 1.$$1 rather than argmax bin alone, and retrain under drift by harvesting new $$p_i = (p_{i,1}, \ldots, p_{i,8}), \qquad p_{i,b} \ge 0, \qquad \sum_{b=1}^{8} p_{i,b} = 1.$$2 rows. For expert surrogates, perform roofline classification first, filter tools by bottleneck type, turn profiler outputs into semantic guidance, and run the optimization loop under MCTS with progressive widening and search memory. In both formulations, the selective component is not an auxiliary detail; it is the mechanism that reconciles surrogate use with the continued necessity of hardware-grounded validation.

Selective GPU surrogates therefore define a broader shift in GPU optimization: hardware measurement and expert interpretation become scarce resources to be budgeted rather than mandatory steps applied uniformly to every candidate. One branch uses calibrated language-model forecasts as virtual models of GPU execution; the other uses filtered micro-profiling tools as virtualized expert diagnosis. Together they establish a common principle: surrogate intervention is most effective not when it is ubiquitous, but when it is explicitly conditioned on confidence, bottleneck relevance, and downstream search utility (Khan et al., 29 May 2026, Gai et al., 24 Jun 2026).