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

Clone with submodules and install in editable mode with uv:

git clone --recurse-submodules https://github.com/codezakh/gpu-surrogates.git
cd gpu-surrogates
uv sync

If you already cloned without --recurse-submodules, run git submodule update --init --recursive from the repo root.

Requires Python 3.13+ and a Linux machine with CUDA drivers.

The runbook scripts call Modal, Tinker, and various LLM providers. Put the keys you have in a .env file at the repo root:

MODAL_TOKEN_ID=...        # kernel measurement on Modal
MODAL_TOKEN_SECRET=...
TINKER_API_KEY=...        # GPT-OSS-20B sampling and GRPO training via Tinker
GEMINI_API_KEY=...        # Gemini-3 Flash / Pro
TOGETHER_API_KEY=...      # GPT-OSS-120B via Together
DEEPSEEK_API_KEY=...      # DeepSeek-V4
HF_TOKEN=...              # read the published artifacts on HuggingFace

You only need keys for services you actually use. The per-script breakdown is in runbook/README.md.

uv run --env-file .env pytest

Tests that need Modal or a CUDA GPU are excluded by default. No API keys required, but the .env file must exist (an empty one is fine).

runbook/ has seven numbered scripts that reproduce the paper's main results. Every script reads a JSON config and writes its outputs under --output-dir. Inputs come from HuggingFace. Nothing needs to be downloaded by hand.

Every script accepts --debug for a fast test run:

uv run --env-file .env python runbook/01_score_baseline.py \
    --config runbook/configs/baseline_scoring/gemini3_flash.json \
    --output-dir runbook_output/01_smoke/ \
    --debug

Remove --debug for the full run. See runbook/README.md for details on each script.

Long-running scripts save their progress as they go and skip work that's already done when you rerun them. Use the same --config and --output-dir to continue from where you stopped. There is no resume flag.

• PUCT searches (00, 04): the script writes an event log at events.jsonl. When you rerun it, the script replays the log and continues mid-step.
• Surrogate scoring (01, 03, 05): each (surrogate, repeat, row) writes its own cache file. When you rerun, already-scored rows are skipped.
• Training (02): writes a training_artifact.json with the Tinker checkpoint URI. The checkpoint itself lives on Tinker's servers.

Datasets:

LoRA adapters for openai/gpt-oss-20b:

Surrogate and search:

• src/gpu_forecasters/landscape_map/v1/: the surrogate prompt template, the response parser, and the code that classifies predicted speedups into one of 8 categories.
• src/gpu_forecasters/max_reward_puct/v3/: the PUCT search code. Picks the next kernel to mutate, runs it, keeps the best in an archive.
• src/gpu_forecasters/gpu_mode_kernel/surrogate_search/v1/: surrogate-filtered PUCT. Each candidate is first scored by the surrogate. Only candidates the surrogate predicts will be faster are run on a real GPU.

Kernel evaluation on Modal:

• src/gpu_forecasters/gpu_mode_kernel/: code that runs candidate kernels on the six GPU Mode tasks (TriMul, cross-entropy, three Gated DeltaNet variants, FP8 quantization). One file per task in packs/. Shared evaluation code in modal_scoring.py.
• src/gpu_forecasters/kernelbench/: runs the full-network problems from KernelBench (VGG19, SwinMLP, EfficientNet, and other PyTorch models).

Training:

• src/gpu_forecasters/runbook/training.py: the GRPO trainer and the three reward functions (correctness, +Brier, +CRPS).

Data and runbook library:

• src/gpu_forecasters/eval_dataset_builder/v1/: builds the held-out evaluation set published as gpu-forecasters-eval-set from the PUCT search event logs.
• src/gpu_forecasters/runbook/: the shared library the runbook scripts use. Pydantic config types, HuggingFace dataset loaders, and the --debug override logic.

github.com/codezakh/gpu-surrogates/issues

This work uses and takes inspiration from KernelBench and ttt-discover. Thanks to Modal for an academic compute grant.

@article{khan2026gpuforecasters,
  title={GPU Forecasters: Language Models as Selective Surrogates for Kernel Runtime Optimization},
  author={Khan, Zaid and Chen, Justin Chih-Yao and Cho, Jaemin and Stengel-Eskin, Elias and Bansal, Mohit},
  journal={arXiv preprint arXiv:2605.31464},
  year={2026}
}