pith. sign in

arxiv: 2605.28213 · v1 · pith:X6WR7TXXnew · submitted 2026-05-27 · 💻 cs.AI

Learning When to Optimize: Verified Optimization Skills from Expert GPU-Kernel Lineages

Shuoming Zhang , Qiuchu Yu , Yangyu Zhang , Ruiyuan Xu , Xiyu Shi , Guangli Li , Xiaobing Feng , Huimin Cui
show 1 more author
Jiacheng Zhao
This is my paper

Pith reviewed 2026-06-29 12:24 UTC · model grok-4.3

classification 💻 cs.AI
keywords GPU kernel optimizationLLM-based code generationexpert code lineagesoptimization skillsverified transformationsNVIDIA architecturescurriculum learning for agents
0
0 comments X

The pith

KLineage reverses expert GPU kernel lineages to extract verified optimization skills that teach LLMs precisely when optimizations apply.

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The paper establishes that tracing expert-optimized GPU kernels backward through validation-gated simplifications yields discrete, reusable skills. Each skill encodes an optimization's intent, code location, validity conditions, performance effect, and avoided failure modes. These skills then serve as a curriculum for an LLM to optimize new kernels under identical compile, correctness, and profiling gates. On five workloads across two NVIDIA architectures the approach produces higher-quality kernels and reaches them faster than memory-based LLM baselines under the same budget. A separate 22-instance held-out set confirms the gains are not from source memorization.

Core claim

KLineage learns optimization skills from expert kernels by walking backward through validation-gated simplifications, reversing each accepted step into a reusable skill that records the optimization intent, its application location in code, validity conditions, effect, and avoided failures. A downstream LLM applies these skills to new code under the same compile, correctness, and profile gates. This lineage-derived curriculum outperforms memory-based LLM-kernel baselines on five expert workloads across two NVIDIA architectures in both kernel quality and optimization efficiency under fixed budget, with a 22-instance held-out check against memorization.

What carries the argument

KLineage, the process of reversing expert kernel lineages via validation-gated simplifications to produce skills that encode optimization applicability conditions.

If this is right

  • The skills improve both final kernel quality and number of optimization steps needed relative to memory baselines.
  • The same skill set works across two different NVIDIA architectures without retraining.
  • A 22-instance held-out check shows the method avoids simple memorization of source cases.
  • Skills are applied only when the original validation gates are satisfied, preserving correctness.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

  • Similar backward extraction could be applied to other code-generation domains where expert examples exist but applicability rules are implicit.
  • The approach suggests that expert code contains explicit, transferable knowledge about optimization preconditions that forward search alone does not discover.
  • If the skills remain stable across larger sets of workloads, the method could shrink the effective search space for LLM kernel agents.
  • Testing the same lineage reversal on CPU or accelerator kernels would reveal whether the technique generalizes beyond GPU code.

Load-bearing premise

Skills extracted from expert lineages will transfer to new code surfaces while preserving soundness and performance gains when applied under the same compile, correctness, and profile gates.

What would settle it

Applying the extracted skills to a new held-out workload produces either invalid kernels or no improvement in final performance or optimization steps compared with memory-based baselines under identical budget and gates.

Figures

Figures reproduced from arXiv: 2605.28213 by Guangli Li, Huimin Cui, Jiacheng Zhao, Qiuchu Yu, Ruiyuan Xu, Shuoming Zhang, Xiaobing Feng, Xiyu Shi, Yangyu Zhang.

Figure 1
Figure 1. Figure 1: Motivation of KLINEAGE. (a) Forward search often knows which optimizations to try but not when their preconditions hold. (b) KLINEAGE walks expert kernels backward to recover validated forward transitions from simpler states toward expert states. (c) The recovered skills are reused only as code-anchored candidates, with compile/test/profile gates deciding admission on new targets. 2025; Li et al., 2025d,c)… view at source ↗
Figure 2
Figure 2. Figure 2: Overview of KLINEAGE. Offline, validation-gated deoptimization recovers forward transitions from expert kernels and admits lifted skills only after held-out roundtrip materialization. Online, retrieved skills are materialized on the target code surface and filtered by the same compile/correctness/profile gate. where ι is the intent, anchor identifies the code or IR surface on which it acts, and carrier is … view at source ↗
Figure 3
Figure 3. Figure 3: Generated CUDA-kernel speedup over the platform-specific vendor reference. Higher is better; red [PITH_FULL_IMAGE:figures/full_fig_p007_3.png] view at source ↗
Figure 4
Figure 4. Figure 4: Running-best speedup vs. Torch as a function of cumulative LLM-API cost on SM120, under the shared [PITH_FULL_IMAGE:figures/full_fig_p007_4.png] view at source ↗
Figure 5
Figure 5. Figure 5: GDN deoptimization traces. The back￾ward edges remove expert optimizations; KLINEAGE stores the inverse directions as forward skill candidates. The largest drops isolate the optimization knowledge that is hardest for unconditioned generation: shared￾memory transpose for cumsum, phase decomposition for fused_fwd, and MMA/TMA layout for kkt_solve. Three sub-kernels, three dominant skills. The largest drop in… view at source ↗
Figure 6
Figure 6. Figure 6: AdaExplore SM120 per-step trajectories on GEMM (top) and Conv2d (bottom). Green circles are [PITH_FULL_IMAGE:figures/full_fig_p013_6.png] view at source ↗
read the original abstract

LLM-based agents are increasingly used to generate GPU kernels, but they often know what optimizations to try without knowing when those optimizations are sound. 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 where it applies in code, what conditions made it valid, what effect it had, and what failures its assumptions avoid. A downstream LLM materializes these skills on new code surfaces under the same compile/correctness/profile gate. 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 additionally use a separate 22-instance held-out check as a sanity test against source-case memorization.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Referee Report

1 major / 0 minor

Summary. The paper introduces KLineage, which extracts reusable optimization skills from expert GPU-kernel lineages by walking backward through validation-gated simplifications and reversing each accepted step. Each skill records the optimization intent, applicability location in code, validity conditions, performance effect, and avoided failure modes. A downstream LLM then applies these skills to new code surfaces under identical compile/correctness/profile gates. On five expert workloads across two NVIDIA architectures the lineage-derived skills are reported to outperform recent memory-based LLM-kernel baselines in both final kernel quality and optimization efficiency under a fixed budget; a separate 22-instance held-out set serves only as a sanity check against source-case memorization.

Significance. If the extracted skills prove transferable while preserving soundness and gains, the approach would supply a concrete mechanism for learning the missing 'when' component of optimization from verified expert lineages rather than forward rollouts. The fixed-budget comparison and explicit recording of conditions and failure modes are methodological strengths that could be reused beyond the current setting.

major comments (1)
  1. [Abstract] Abstract: the central claim that the lineage-derived skills constitute an effective optimization curriculum that exceeds baselines on new code surfaces is not supported by any quantitative results on the 22 held-out instances. These instances are described solely as a sanity test against source-case memorization; no performance, efficiency, or correctness numbers are supplied for them, leaving the transfer claim unverified.

Simulated Author's Rebuttal

1 responses · 0 unresolved

We thank the referee for the careful review and for highlighting the methodological strengths of the work. We address the single major comment below.

read point-by-point responses

  1. Referee: [Abstract] Abstract: the central claim that the lineage-derived skills constitute an effective optimization curriculum that exceeds baselines on new code surfaces is not supported by any quantitative results on the 22 held-out instances. These instances are described solely as a sanity test against source-case memorization; no performance, efficiency, or correctness numbers are supplied for them, leaving the transfer claim unverified.

    Authors: We agree that the abstract's phrasing could be read as implying quantitative support for transfer specifically on the 22 held-out instances, yet the manuscript supplies none. The five expert workloads constitute the primary evaluation of skill transfer to new code surfaces, while the 22-instance set functions only as a memorization sanity check. To resolve the ambiguity and strengthen the transfer claim, the revised manuscript will include the corresponding performance, efficiency, and correctness metrics for the 22 instances. revision: yes

Circularity Check

0 steps flagged

No significant circularity detected in derivation chain

full rationale

The paper derives optimization skills by walking expert GPU-kernel implementations backward through validation-gated simplifications and reversing accepted steps into reusable skills that record intent, applicability conditions, effects, and avoided failures. These skills are then materialized by a downstream LLM on code surfaces under the same compile/correctness/profile gates. No step reduces by construction to its inputs via self-definition, fitted parameters renamed as predictions, or self-citation load-bearing; the reported outperformance on the five expert workloads is an empirical comparison against external baselines rather than a tautological restatement of the lineage extraction process itself. The 22 held-out instances serve only as a memorization sanity check and introduce no circular reduction.

Axiom & Free-Parameter Ledger

0 free parameters · 1 axioms · 0 invented entities

Abstract implies but does not detail assumptions; review is abstract-only so ledger is minimal.

axioms (1)
  • domain assumption Expert kernels contain sequences of simplifications that remain valid when reversed into general skills.
    Central to extracting reusable skills from lineages.

pith-pipeline@v0.9.1-grok · 5726 in / 1007 out tokens · 26510 ms · 2026-06-29T12:24:03.657578+00:00 · methodology

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Reference graph

Works this paper leans on

29 extracted references · 24 canonical work pages · 8 internal anchors

  1. [1]

    write newline

    " write newline "" before.all 'output.state := FUNCTION n.dashify 't := "" t empty not t #1 #1 substring "-" = t #1 #2 substring "--" = not "--" * t #2 global.max substring 't := t #1 #1 substring "-" = "-" * t #2 global.max substring 't := while if t #1 #1 substring * t #2 global.max substring 't := if while FUNCTION word.in bbl.in capitalize " " * FUNCT...

  2. [2]

    Martin Andrews and Sam Witteveen. 2025. https://arxiv.org/abs/2506.20807 GPU kernel scientist: An LLM -driven framework for iterative kernel optimization . Preprint, arXiv:2506.20807

    work page arXiv 2025

  3. [3]

    Carlo Baronio, Pietro Marsella, Ben Pan, Simon Guo, and Silas Alberti. 2025. https://arxiv.org/abs/2507.11948 Kevin: Multi-turn RL for generating CUDA kernels . Preprint, arXiv:2507.11948

    work page arXiv 2025

  4. [4]

    Gonzalez, and Ion Stoica

    Shiyi Cao, Ziming Mao, Joseph E. Gonzalez, and Ion Stoica. 2026. https://arxiv.org/abs/2602.19128 K-search: Llm kernel generation via co-evolving intrinsic world model . Preprint, arXiv:2602.19128

    work page arXiv 2026

  5. [5]

    Ruifan Chu, Anbang Wang, Xiuxiu Bai, Shuai Liu, and Xiaoshe Dong. 2025. Gpu kernel optimization beyond full builds: An llm framework with minimal executable programs. arXiv preprint arXiv:2512.22147

    work page arXiv 2025

  6. [6]

    Joshua H Davis, Klaudiusz Rydzy, Srinivasan Ramesh, Aadit Nilay, Daniel Nichols, Swapna Raj, Nikhil Jain, and Abhinav Bhatele. 2026. Keet: Explaining performance of gpu kernels using llm agents. arXiv preprint arXiv:2605.04467

    work page internal anchor Pith review Pith/arXiv arXiv 2026

  7. [7]

    Kris Shengjun Dong, Sahil Modi, Dima Nikiforov, Sana Damani, Edward Lin, Siva Kumar Sastry Hari, and Christos Kozyrakis. 2026. https://arxiv.org/abs/2602.14293 Kernelblaster: Continual cross-task cuda optimization via memory-augmented in-context reinforcement learning . Preprint, arXiv:2602.14293

    work page arXiv 2026

  8. [8]

    He Du, Qiming Ge, Jiakai Hu, Aijun Yang, Zheng Cai, Zixian Huang, Sheng Yuan, Qinxiu Cheng, Xinchen Xie, Yicheng Chen, Yining Li, Jiaxing Xie, Huanan Dong, Yaguang Wu, Xiangjun Huang, Jian Yang, Hui Wang, Bowen Zhou, Bowen Li, and 2 others. 2026 a . https://arxiv.org/abs/2603.28342 Kernel-smith: A unified recipe for evolutionary kernel optimization . Prep...

    work page internal anchor Pith review Pith/arXiv arXiv 2026

  9. [9]

    Weihua Du, Jingming Zhuo, Yixin Dong, Andre Wang He, Weiwei Sun, Zeyu Zheng, Manupa Karunaratne, Ivan Fox, Tim Dettmers, Tianqi Chen, Yiming Yang, and Sean Welleck. 2026 b . https://arxiv.org/abs/2604.16625 Adaexplore: Failure-driven adaptation and diversity-preserving search for efficient kernel generation . Preprint, arXiv:2604.16625

    work page internal anchor Pith review Pith/arXiv arXiv 2026

  10. [10]

    FlagOpen Project . 2026. Flaggems: A high-performance triton operator library for large language models. https://github.com/FlagOpen/FlagGems. V5.0.2, accessed 2026-05

    2026

  11. [11]

    Siva Kumar Sastry Hari, Vignesh Balaji, Sana Damani, Qijing Huang, and Christos Kozyrakis. 2026. Improving efficiency of gpu kernel optimization agents using a domain-specific language and speed-of-light guidance. arXiv preprint arXiv:2603.29010

    work page arXiv 2026

  12. [12]

    Robert Tjarko Lange, Qi Sun, Aaditya Prasad, Maxence Faldor, Yujin Tang, and David Ha. 2025. https://arxiv.org/abs/2509.14279 Towards robust agentic CUDA kernel benchmarking, verification, and optimization . Preprint, arXiv:2509.14279

    work page arXiv 2025

  13. [13]

    Haonan Li, Keyu Man, Partha Kanuparthy, Hanning Chen, Wei Sun, Sreen Tallam, Chenguang Zhu, Kevin Zhu, and Zhiyun Qian. 2025 a . Tritonforge: Profiling-guided framework for automated triton kernel optimization. arXiv preprint arXiv:2512.09196

    work page arXiv 2025

  14. [14]

    Jianling Li, Shangzhan Li, Zhenye Gao, Qi Shi, Yuxuan Li, Zefan Wang, Jiacheng Huang, Haojie Wang, Jianrong Wang, Xu Han, Zhiyuan Liu, and Maosong Sun. 2025 b . https://arxiv.org/abs/2502.14752 Tritonbench: Benchmarking large language model capabilities for generating Triton operators . Preprint, arXiv:2502.14752

    work page arXiv 2025

  15. [15]

    Shangzhan Li, Zefan Wang, Ye He, Yuxuan Li, Qi Shi, Jianling Li, Yonggang Hu, Wanxiang Che, Xu Han, Zhiyuan Liu, and Maosong Sun. 2025 c . https://arxiv.org/abs/2507.05687 Autotriton: Automatic Triton programming with reinforcement learning in LLM s . Preprint, arXiv:2507.05687

    work page arXiv 2025

  16. [16]

    Xiaoya Li, Xiaofei Sun, Albert Wang, Jiwei Li, and Chris Shum. 2025 d . https://arxiv.org/abs/2507.14111 CUDA-L1 : Improving CUDA optimization via contrastive reinforcement learning . Preprint, arXiv:2507.14111

    work page arXiv 2025

  17. [17]

    Mark Lou and Stefan K Muller. 2024. Automatic static analysis-guided optimization of cuda kernels. In Proceedings of the 15th International Workshop on Programming Models and Applications for Multicores and Manycores, pages 11--21

    2024

  18. [18]

    Xing Ma, Yangjie Zhou, Wu Sun, Zihan Liu, Jingwen Leng, Yun Lin, Shixuan Sun, Minyi Guo, and Jin Song Dong. 2026. Cubridge: An llm-based framework for understanding and reconstructing high-performance attention kernels. arXiv preprint arXiv:2605.05023

    work page internal anchor Pith review Pith/arXiv arXiv 2026

  19. [19]

    Gabriele Oliaro, Yichao Fu, May Jiang, Owen Lu, Junli Wang, Zhihao Jia, Hao Zhang, and Samyam Rajbhandari. 2026. https://arxiv.org/abs/2605.23215 Fastkernels: Benchmarking gpu kernel generation in production . Preprint, arXiv:2605.23215

    work page internal anchor Pith review Pith/arXiv arXiv 2026

  20. [20]

    KernelBench: Can LLMs Write Efficient GPU Kernels?

    Anne Ouyang, Simon Guo, Simran Arora, Alex L. Zhang, William Hu, Christopher R \'e , and Azalia Mirhoseini. 2025. https://arxiv.org/abs/2502.10517 Kernelbench: Can llms write efficient GPU kernels? Preprint, arXiv:2502.10517

    work page internal anchor Pith review Pith/arXiv arXiv 2025

  21. [21]

    Qiuyi Qu, Yicheng Sui, Yufei Sun, Rui Chen, Xiaofei Zhang, Yuzhi Zhang, Haofeng Wang, and Ge Lan. 2026. A two-stage gpu kernel tuner combining semantic refactoring and search-based optimization. arXiv preprint arXiv:2601.12698

    work page arXiv 2026

  22. [22]

    Tara Saba, Anne Ouyang, Xujie Si, and Fan Long. 2026. Cutegen: An llm-based agentic framework for generation and optimization of high-performance gpu kernels using cute. arXiv preprint arXiv:2604.01489

    work page internal anchor Pith review Pith/arXiv arXiv 2026

  23. [23]

    Spector, Simran Arora, Aaryan Singhal, Arjun Parthasarathy, Daniel Y

    Benjamin F. Spector, Simran Arora, Aaryan Singhal, Arjun Parthasarathy, Daniel Y. Fu, and Christopher R \'e . 2025. https://openreview.net/forum?id=0fJfVOSUra ThunderKittens : Simple, fast, and adorable AI kernels . In The Thirteenth International Conference on Learning Representations (ICLR)

    2025

  24. [24]

    Qitong Sun, Jun Han, Tianlin Li, Zhe Tang, Sheng Chen, Fei Yang, Aishan Liu, Xianglong Liu, and Yang Liu. 2026. https://arxiv.org/abs/2603.10085 Kernelskill: A multi-agent framework for GPU kernel optimization . Preprint, arXiv:2603.10085

    work page arXiv 2026

  25. [25]

    Jianghui Wang, Vinay Joshi, Saptarshi Majumder, Xu Chao, Bin Ding, Ziqiong Liu, Pratik Prabhanjan Brahma, Dong Li, Zicheng Liu, and Emad Barsoum. 2025. https://arxiv.org/abs/2507.23194 Geak: Introducing Triton kernel AI agent and evaluation benchmarks . Preprint, arXiv:2507.23194

    work page arXiv 2025

  26. [26]

    Anjiang Wei, Tianran Sun, Yogesh Seenichamy, Hang Song, Anne Ouyang, Azalia Mirhoseini, Ke Wang, and Alex Aiken. 2025. https://openreview.net/forum?id=IZKZIcPaHz Astra: A multi-agent system for GPU kernel performance optimization . In NeurIPS 2025 Fourth Workshop on Deep Learning for Code

    2025

  27. [27]

    Nina Wiedemann, Quentin Leboutet, Michael Paulitsch, Diana Wofk, and Benjamin Ummenhofer. 2026. Kernelfoundry: Hardware-aware evolutionary gpu kernel optimization. arXiv preprint arXiv:2603.12440

    work page arXiv 2026

  28. [28]

    Genghan Zhang, Shaowei Zhu, Anjiang Wei, Zhenyu Song, Allen Nie, Zhen Jia, Nandita Vijaykumar, Yida Wang, and Kunle Olukotun. 2025. Accelopt: A self-improving llm agentic system for ai accelerator kernel optimization. arXiv preprint arXiv:2511.15915

    work page internal anchor Pith review Pith/arXiv arXiv 2025

  29. [29]

    Xinguo Zhu, Shaohui Peng, Jiaming Guo, Yunji Chen, Qi Guo, Yuanbo Wen, Hang Qin, Ruizhi Chen, Qirui Zhou, Ke Gao, Yanjun Wu, Chen Zhao, and Ling Li. 2026. https://doi.org/10.1609/aaai.v40i34.40155 QiMeng-Kernel : Macro-thinking micro-coding paradigm for LLM -based high-performance GPU kernel generation . Proceedings of the AAAI Conference on Artificial In...

    work page doi:10.1609/aaai.v40i34.40155 2026