cs.LGOct 14, 2026

How to Train Your Little Dragon: Teaching a Small Language Model to Construct and Refine GPU Optimization Scripts

Authors: Annabelle Sujun Tang, Wei Tang, Songyao Jiang, Tianhui Li, Shengzhe Xu, Jiin Woo, Kaan Ozkara, Youngsuk Park, 

Organizations: University of California, San Diego · Amazon Annapurna Labs

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Compiler transformations interact, making GPU kernel optimization a search for combinations that preserve the computation while improving performance. We introduce HOW TO TRAIN YOUR LITTLE DRAGON, a training method that teaches a small language model both to construct complete optimization scripts and to improve existing ones. Instead of ranking candidates or generating kernel code, our model generates optimization scripts. We devise a compact, lossless representation whose deterministic decoding reconstructs the exact compiler trace, preserving its transformation parameters and dependencies. We use a compiler framework TVM MetaSchedule to build training data on 54,116 distinct workloads, and construct 108,594 examples in two training sets: 60,236 teach the model to generate an optimized script from an unoptimized kernel, while 48,358 teach it to produce a better script given the kernel and a somewhat optimized script. Supervised fine-tuning teaches the model to construct and improve optimization scripts from these search-derived examples. We then add reinforcement learning with a shaped reward that combines correctness with measured performance, awarding a latency-based bonus only to numerically correct kernels. We apply this training method to a model with 8 billion parameters. After supervised fine-tuning and reinforcement learning, the model achieves per-attempt correctness of 91.27% for construction and 92.01% for refinement, while reinforcement learning produces faster kernels than supervised training on 88 of 102 construction tasks and 145 of 170 refinement prompts. With only 8B parameters, our LITTLE DRAGON achieves higher per-attempt correctness and better kernel performance than six frontier LLMs: GPT-6 Astra, Opus 5, Kimi-K3, and GLM-5.3 among them, each with hundreds of billions to trillions of parameters, for which task-specific training is unavailable or expensive.

Figures & tables

Figure 1: How to Train Your Little Dragon : data generation and two-stage policy training. (A) Compiler search supplies verified schedules. Deterministic lifting preserves their executable edits in our Typed Schedule Edit Language ( TSEL ); Kimi-K3 writes explanations from compiler-derived facts. Static and round-trip checks admit construction and refinement labels. (B) Weighted token supervision trains Seed-Coder-8B-Instruct to produce complete optimization graphs; GAG_{A} denotes the optional prior graph for refinement. (C) Starting from SFT, GRPO updates the policy using numerical correctness and measured latency of sampled schedules.
Figure 2: A concrete optimization graph. The payload fragment and graph projection come from the same FP16 projection training label. Node text is shortened; ellipses omit other edits. Colored edges EE explain choices; gray arrows show replay order PP , and the dotted callout expands t22 . WMMA denotes tensor-core operations. The expected effect p1 concerns the whole schedule.
 Construction (%)Refinement (%)
ModelCorrect@1fast2\mathrm{fast}_{2}fast10\mathrm{fast}_{10}Correct@1Correct and improves graph ACorrect and ≥2×\geq 2\times over graph A
GPT-6 Astra9.387.656.0345.5536.4717.33
Opus 564.0158.3036.9645.8131.7014.24
Kimi-K334.9125.4314.5537.8219.594.70
GLM-5.3 (Z.ai)15.7311.757.0018.568.251.80
Qwen3.8-27B0.110.000.0019.594.571.48
Table 1: Per-output correctness and optimization success. Results cover 116 construction tasks and 194 refinement prompts spanning 39 tasks.
Figure 3: First output kernel speedups. Our Little Dragon achieves better speedups for both the construction and the refinement tasks across multiple kernel famiilies.
Figure 4: Kernel speedup and optimization cost comparisons by family between first SFT+RL output and TVM MetaSchedule search run.
Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: RL training utility. Ten-rollout means, plotted at band midpoints. The series is measured on screened, dynamically retained training samples, not the sealed set.
Figure 6: RL training correctness. Ten-rollout sample-correctness summaries.
More in this topic

Compiler Optimization

See all 27 papers

Explore similar work

CardsList
Sep 24, 2026

KernelOPT: Dispatch-Aware Agentic Search for GPU Kernel Optimization

Deep learning inference and training performance depends critically on GPU kernel efficiency. Modern compilers such as PyTorch Inductor automatically generate GPU kernels from high-level model code, but frequently underperform expert-written implementations by wide margins. Recent LLM-assisted kernel optimizers can close this gap for standalone kernels, yet treat compiled models as black boxes, generally optimizing individual standalone kernels without respecting the compiler's structural decisions or verifying the model end-to-end. We present KernelOPT, a multi-agent system that treats compiled models as structured artifacts. It preserves vendor library calls (cuBLAS, cuDNN) and exclusively targets generated Triton sub-kernels using five profiling-guided LLM agents. A four-gate verification cascade applies static validation, multi-seed correctness checking, model-level float64-fallback verification, and performance gating (γ=1.03γ{=}1.03) to filter candidates and verify the re-stitched model end-to-end. When candidates fail verification, the system preserves the compiler baseline. The system accepts PyTorch nn Modules, standalone Triton kernels, and Helion kernels. Evaluated on 250 KernelBench problems (100 Level 1, 100 Level 2 and 50 Level 3) on NVIDIA H200, KernelOPT achieves geometric mean speedups over torch compile of 1.40×\times (L1), 1.15×\times (L2), and 1.07×\times (L3) across all kernels, including fallback cases. Optimized-only geomeans (excluding cases where verification gates preserve the compiler baseline) are substantially higher: 2.54×\times (L1: 36/100), 1.84×\times (L2: 23/100), and 1.37×\times (L3: 11/50), reflecting where the optimizer achieves meaningful leverage.

Aheli Poddar, Sanskar Prasad, Arindam Samanta +3

Red Hat

GPU Kernel OptimizationAgentic SearchLarge Language Model-Guided OptimizationCompiler OptimizationNeural Network Hardware Acceleration
AskPDF
Mar 22, 2026

AutoKernel: Autonomous GPU Kernel Optimization via Iterative Agent-Driven Search

Writing high-performance GPU kernels is among the most labor-intensive tasks in machine learning systems engineering. We present AutoKernel, an open-source framework that applies an autonomous agent loop to GPU kernel optimization for arbitrary PyTorch models. Given a model, AutoKernel profiles it to identify computational bottlenecks, ranks them by Amdahl's law impact, and iteratively refines Triton or CUDA C++ kernel implementations through hundreds of experiments without human intervention. A five-stage correctness harness covering smoke tests, shape sweeps, numerical stability, determinism verification, and edge-case coverage ensures every candidate kernel is validated before any speedup is recorded. The system comprises over 9,000 lines of Python, 18 starter kernel implementations across two backends, a six-tier optimization playbook, and integration with the KernelBench benchmark suite. AutoKernel covers nine kernel types spanning the dominant operations in modern transformer architectures. On an NVIDIA H100, our Triton kernels outperform both PyTorch eager and torch.compile (max-autotune) on the majority of tested configurations: 5.29x over eager on RMSNorm, 2.82x on softmax, and 2.21x on cross-entropy, while beating torch.compile by 2.83x, 3.44x, and 2.94x respectively. In community deployment, an AutoKernel-optimized kernel achieved first place on the vectorsum_v2 B200 leaderboard. The full system is available at https://github.com/RightNow-AI/autokernel.

Jaber Jaber, Osama Jaber

RightNow AI

GPU Kernel OptimizationGPU Kernel Generation
AskPDF
May 27, 2026

KLineage: Recovering the Missing When of Kernel Optimization by Deoptimizing Experts

LLM-based agents are increasingly used to generate GPU kernels, but they often struggle to determine when an optimization is sound because its required code state and dependencies are implicit in expert implementations. We introduce KLineage, which learns this missing "when" knowledge from expert kernels: instead of relying on forward rollouts, KLineage walks expert implementations backward through validation-gated simplifications and reverses each accepted step into a reusable optimization skill. Each skill records not only the optimization intent, but also when to apply the optimization technique, including where it applies in code, what conditions made it valid, what effect it has, and what failures its assumptions avoid. A downstream LLM materializes these skills on new code surfaces under the same compile/correctness/profile gate. This guidance on when to apply each optimization can help downstream models to generate higher-performance kernels. On five expert workloads across two NVIDIA architectures, these lineage-derived skills serve as an effective optimization curriculum, exceeding recent memory-based LLM-kernel baselines in both final kernel quality and optimization efficiency under the same fixed budget. We also demonstrate that the KLineage framework extends beyond NVIDIA GPUs to Ascend NPUs. Our code is publicly available at https://github.com/ict-agent/klineage.

Shuoming Zhang, Qiuchu Yu, Ruiyuan Xu +11

SKLP, Institute of Computing Technology, Chinese Academy of Sciences · University of Chinese Academy of Sciences · University of New South Wales

Agent Skill LearningGPU AccelerationHigh-Performance ComputingLarge Language Model-Guided OptimizationLearning to OptimizeGPU Kernel Generation+2
AskPDF