🤖 AI

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

The paper introduces KLineage, a framework that extracts verified optimization skills from expert GPU kernels by walking backward through validation-gated simplifications, enabling LLM-based agents to learn not just what optimizations to apply but when they are sound, thereby outperforming memory-based baselines in kernel quality and efficiency.

Original authors: Shuoming Zhang, Qiuchu Yu, Yangyu Zhang, Ruiyuan Xu, Xiyu Shi, Guangli Li, Xiaobing Feng, Huimin Cui, Jiacheng Zhao

Published 2026-05-28

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Shuoming Zhang, Qiuchu Yu, Yangyu Zhang, Ruiyuan Xu, Xiyu Shi, Guangli Li, Xiaobing Feng, Huimin Cui, Jiacheng Zhao

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a very talented but inexperienced apprentice how to build a high-performance race car engine.

The Problem: Knowing "What" vs. Knowing "When"
Currently, AI agents (the apprentices) are great at knowing what optimizations to try. They know they should use "vectorized loads" or "software pipelining." It's like the apprentice knows the names of all the fancy tools in the toolbox.

However, they don't know when to use them.

They might try to use a specialized tool on a part that isn't the right shape.
They might try to speed up a process that requires a specific setup they haven't built yet.
The result? The engine crashes, the code won't compile, or it runs slower than before.

The AI knows the recipe (the steps), but it lacks the timing (the conditions under which those steps actually work).

The Solution: KLINEAGE (The "Reverse-Engineered" Cookbook)
The paper introduces a new system called KLINEAGE. Instead of letting the AI guess its way forward (trial and error), KLINEAGE learns from expert engines that are already built and working perfectly.

Here is how it works, using a creative analogy:

1. The "Backward Walk" (Deoptimization)

Imagine you have a perfectly tuned, high-speed race car (the Expert Kernel).

Traditional AI: Tries to build a car from scratch, hoping to stumble upon the right design.
KLINEAGE: Takes the perfect race car and asks, "What if we took this one specific high-tech part out?"
- It removes a part (like a turbocharger).
- It checks: "Does the car still run? Is it safe?"
- If the car still runs (but is slower), it records that step.
- It keeps removing parts one by one until the car is just a basic, slow, "naive" engine.

By walking backward from the expert to the beginner, KLINEAGE creates a map of every single step that must be taken to get from "slow" to "fast."

2. Creating "Verified Skills"

When KLINEAGE walks backward, it doesn't just write down "remove turbo." It creates a Skill Card for the reverse journey (going from slow to fast).

Each Skill Card is like a detailed instruction sticker that says:

The Action: "Add the turbocharger."
The Condition: "Only add this if the engine block is already reinforced." (This is the "When" knowledge).
The Risk: "If you add this without the reinforcement, the engine will explode."
The Evidence: "We tested this on 5 different cars, and it worked every time."

These skills are code-anchored, meaning they aren't just vague advice; they are specific instructions tied to specific parts of the code.

3. The "Safety Gate" (Verification)

Before an AI can use these Skill Cards on a new project, it has to pass a test.

The AI tries to apply the skill to a new piece of code.
A "Safety Gate" checks three things:
1. Compile: Does the code even run?
2. Correctness: Does it give the right answer?
3. Profile: Is it actually faster?
If the skill fails any of these checks, it is rejected. Only the skills that pass become part of the library.

4. The Result: A Smarter Apprentice

When the AI faces a new, difficult coding problem, it doesn't guess. It looks up the Skill Cards that match the current situation.

It sees: "Oh, I need to speed up this memory access. The Skill Card says I can use 'vectorized loads,' but only if the data is aligned."
It checks the alignment. It's aligned! It applies the skill.
It moves to the next step, checking the conditions for the next skill.

The Findings
The paper tested this on five complex tasks (like matrix multiplication and image processing) on two different types of powerful computer chips (NVIDIA H100 and RTX PRO 6000).

Speed & Success: The KLINEAGE system produced working, fast code 1.12 times faster on average than the best previous AI methods.
Efficiency: It reached its best results much faster (spending less "money" on computer time) because it didn't waste time trying things that were guaranteed to fail.
No Cheating: The system wasn't just memorizing the answers. When tested on 22 completely new problems it had never seen before, it still performed well, proving it learned the principles of optimization, not just the specific answers.

In Summary
KLINEAGE teaches AI to optimize code by reverse-engineering expert solutions to understand the hidden rules and conditions. It turns "guessing what to do" into "following a verified map of when to do it," resulting in faster, more reliable, and smarter code generation.

Technical Summary: Learning When to Optimize (KLINEAGE)

Problem Statement

While Large Language Model (LLM) agents are increasingly capable of generating GPU kernels, they frequently struggle with the "when" of optimization. Current agents often possess the vocabulary of optimizations (e.g., tiling, vectorization, pipelining) but lack the conditional knowledge to determine when these optimizations are valid and effective. Applying a familiar optimization to an incorrect state often results in compilation errors, correctness failures, or performance degradation.

Existing approaches rely on forward rollouts (searching from a naive implementation toward an optimized one) or memory systems built from failure rules and profile states. However, these methods are limited by the trajectories the agent can discover and often fail to capture the specific preconditions, risks, and code anchors required for an optimization to be sound. The paper posits that the missing knowledge is not the intent of an optimization, but the verified conditional structure under which it applies.

Methodology: KLINEAGE

The authors introduce KLINEAGE, a framework that induces reusable, code-anchored optimization skills by walking backward through expert GPU kernels. Unlike forward-search methods, KLINEAGE treats expert implementations as evidence from which to extract a curriculum.

1. Guided Deoptimization

The core induction process operates in reverse:

Backward Walk: Starting from a high-performance expert kernel, the system proposes simplifications (e.g., removing software pipelining, devectorizing, or removing shared-memory staging).
Validation-Gated: An LLM rewriter patches the kernel to remove the optimization. The patch is accepted only if it passes a strict compile/correctness/profile gate.
Inverse Transition: Accepted simplifications are stored as their inverse forward transitions ( $\tau_i$ ). This creates a validated executable curriculum from a naive state to the expert state.
Evidence Collection: Rejected simplifications serve as evidence for precondition violations, recording the specific risks and failure modes associated with applying an optimization in the wrong context.

2. Skill Induction and Abstraction

Concrete transitions are aggregated and lifted into Conditioned Procedural Skills. A skill is not merely a code snippet but a structured specification containing:

Intent ( $\iota$ ): The abstract optimization goal (e.g., "staging reused tiles").
Anchor: The specific code or IR surface where the intent acts.
Carrier: An actionable representation (diff sketch, pseudocode, or annotated snippet) to guide implementation.
Preconditions, Effects, and Risks: Explicit conditions required for validity, expected performance deltas, and potential failure modes.
Scope: Constraints over case, surface language (CUDA, TileLang, etc.), platform (SM90, SM120), and prior actions.

3. Verification-Gated Materialization

To ensure reliability, a skill is only admitted to the library after a held-out roundtrip succeeds: an agent must materialize the skill on a fresh state and reproduce the expected effect. During online reuse, the downstream LLM retrieves skills based on the target's scope (language, platform, current state) and materializes them under the same compile/correctness/profile gate used for all submissions.

Key Contributions

The paper makes three primary technical contributions:

Conditioned Procedural Skills: A formalization of optimization knowledge as verified, code-anchored specifications that explicitly encode the "when" (preconditions, risks, scope) alongside the "what" (intent).
Guided Deoptimization: A novel procedure for inducing skills by walking expert kernels backward through dependency-aware steps, storing accepted simplifications as inverse forward transitions to create an executable curriculum.
Verification-Gated Reuse: A mechanism where skills are only admitted after roundtrip validation and are reused as code-anchored candidates filtered by compile/correctness/profile gates, keeping the LLM in the role of instantiating validated intents rather than searching blindly.

Experimental Results

The framework was evaluated on five expert workloads (GEMM, Conv2d, FMHA, GDN, Top-K) across two NVIDIA architectures (H100/SM90 and RTX PRO 6000/SM120), with a fixed optimization cost budget of $10 per workload.

Performance: KLINEAGE achieved a 1.12× success-only geometric mean speedup over vendor references on 10/10 main-tier pairs. In contrast, baselines AdaExplore and AccelOpt achieved 0.25× (6/10 pairs) and 0.54× (8/10 pairs), respectively.
Conditional Structure: The performance gap was most pronounced on workloads with non-local preconditions (Conv2d, GDN, Top-K), where KLINEAGE reached 1.22× speedup while baselines dropped to 0.13× and 0.34×.
Cross-Language Transfer: When materializing retrieved skills as TileLang programs (without seeing expert TileLang source), KLINEAGE achieved a 1.06× geometric mean speedup, demonstrating that the skills capture transferable intent rather than language-specific syntax.
Ablations:
- Roundtrip Recovery: 41/50 trials successfully recovered expert-level latency using only retrieved skills, confirming the preservation of code anchors and ordering constraints.
- Generated-Only: Removing lineage-derived fields (anchors, preconditions, risks) resulted in slowdowns of 5–14× on Conv2d and up to 105× on GDN, proving that the conditional structure is critical for success.
Sanity Check: A held-out test on 22 FlagGems operators (not used in training) showed the library was not memorized to the source experts, reaching the Triton ceiling on 16/22 instances.

Significance and Claims

The paper claims that KLINEAGE is the first LLM-kernel memory construction method that automatically derives reusable lineages by deoptimizing expert implementations rather than accumulating forward-rollout states.

The authors argue that the significance lies in shifting the focus from discovering optimizations via search to externalizing the conditional structure of optimization from existing expert artifacts. By treating expert kernels as evidence for "when" an optimization is sound, KLINEAGE provides a curriculum that allows agents to avoid common pitfalls (compilation errors, correctness failures) and efficiently reach expert-level performance. The work suggests that for complex, constraint-heavy domains like GPU kernel optimization, verified procedural memory is more effective than unconditioned search or static rule-based memory.

The paper remains modest in its broader claims, noting that the method is currently restricted to compatible backends (NVIDIA SM90/SM120) and that the induced lineages represent validation-backed curricula rather than historical provenance of how the expert kernels were originally written.