KernelSkill: A Multi-Agent Framework for GPU Kernel Optimization

KernelSkill is a multi-agent framework that enhances GPU kernel optimization by replacing opaque LLM heuristics with a knowledge-driven, dual-level memory architecture of expert skills, achieving state-of-the-art speedups and a 100% success rate on KernelBench.

The Gist

Imagine you are trying to build the most efficient, high-speed race car engine possible. This engine is the "GPU Kernel," the tiny piece of code that makes your computer's graphics card do the heavy lifting for things like training AI or rendering video games.

For decades, building these engines has been like trying to fix a Ferrari with a blindfold on. You need a world-class expert mechanic who knows exactly which wrench to use, but they have to guess, test, break it, fix it, and guess again. It's slow, expensive, and requires a genius-level understanding of the machine.

Recently, we tried using AI (Large Language Models) to act as that mechanic. The AI is smart and can write code, but it often acts like a student who has read a textbook but hasn't actually driven the car. It tries random fixes based on "gut feeling" (implicit heuristics), often making the same mistakes over and over or choosing the wrong tool for the job.

Enter KernelSkill.

KernelSkill is a new framework that turns the AI mechanic into a Master Team of Engineers with a super-organized memory system. Here is how it works, using some everyday analogies:

1. The Problem: The "Forgetful" Mechanic

Imagine an AI mechanic trying to fix a car.

• The Old Way: The mechanic looks at the engine, guesses a fix, tries it, and if it fails, they guess again. They don't remember that they tried "tightening the bolt" five minutes ago, so they try it again. They also don't remember that "tightening the bolt" worked on a different car last week. They are reinventing the wheel every time.
• The Result: They waste time, get stuck in loops, and the car never gets truly fast.

2. The Solution: The "Two-Level Memory" System

KernelSkill gives the AI a brain upgrade with two distinct types of memory, like a Library and a Daily Logbook.

📚 The Long-Term Memory: The "Expert Library"

Think of this as a massive, organized library containing the best repair manuals ever written by human experts.

• How it works: Instead of guessing, the AI first checks the "Library." It looks up the specific problem (e.g., "The engine is overheating because of poor airflow").
• The Magic: It doesn't just guess; it pulls out a proven strategy: "When you see X, do Y." This is Knowledge-Driven. It ensures the AI picks the right tool for the job based on real-world expertise, not just a random guess.
• Analogy: It's like a chef who doesn't just taste the soup and guess what's missing; they consult a master recipe book that says, "If the soup is salty, add acid, not more salt."

📝 The Short-Term Memory: The "Daily Logbook"

Think of this as a sticky note on the mechanic's dashboard that says, "I already tried changing the spark plugs, and it didn't work. Don't do that again."

• How it works: As the AI works on a specific car (a specific code task), it writes down every step it takes. If it tries a fix and it fails, the Logbook records it.
• The Magic: This prevents the AI from getting stuck in a loop where it keeps trying the same bad fix. It also helps the AI see the "big picture" of the repair process, ensuring that fixing one part doesn't break another.
• Analogy: It's like a detective solving a mystery. They write down every clue and every suspect they've already cleared. This stops them from going back to the same dead end.

3. The Team: A Multi-Agent Workflow

KernelSkill isn't just one AI; it's a team of specialists working together, like a pit crew at a race track:

• The Generator: Builds the initial engine (writes the code).
• The Reviewer: Checks if the engine runs without exploding (compilation) and if it drives the same way as the original (correctness).
• The Profiler: Takes the car for a test drive to see exactly where it's slow (performance metrics).
• The Planner: Looks at the test drive data and the "Expert Library" to decide the best fix.
• The Repairer: Actually makes the changes to the code.

4. The Results: Speeding Up the World

The researchers tested this system on KernelBench, which is like a giant obstacle course for these engines.

• Success Rate: KernelSkill got a 100% score. It fixed every single problem it was given. Other AI systems failed on the harder levels.
• Speed:
  • On easy tasks, it made the code 5.4 times faster.
  • On medium tasks, 2.8 times faster.
  • On hard tasks, 1.9 times faster.
• Efficiency: It didn't just get faster; it got there faster. It needed fewer attempts to find the perfect solution because it wasn't wasting time guessing or repeating mistakes.

Summary

KernelSkill is like upgrading a chaotic, forgetful apprentice mechanic into a highly organized, expert-led pit crew. By giving the AI a Library of Expert Knowledge (so it knows what to do) and a Daily Logbook (so it remembers what it already tried), it stops wasting time and starts delivering lightning-fast results.

This means the AI systems we use every day—from chatbots to self-driving cars—can run much more efficiently, saving energy and time, all thanks to a smarter way of organizing the AI's memory.

Technical Summary

Here is a detailed technical summary of the paper "KernelSkill: A Multi-Agent Framework for GPU Kernel Optimization."

1. Problem Statement

Efficient GPU kernels are fundamental to modern AI systems, yet optimizing them remains a highly expert-driven, costly, and slow process involving iterative profiling, bottleneck diagnosis, and manual code edits. While Large Language Models (LLMs) have been applied to automate this process, existing approaches suffer from two critical limitations:

• Implicit Heuristics: Current LLM-based pipelines rely on opaque, internally learned heuristics to select optimization strategies. This leads to inefficient trial-and-error and weak interpretability.
• Instability and Myopia: Multi-round refinement often fails to stabilize. Systems may select mismatched strategies for the actual bottleneck, fail to update strategies as profiling signals evolve, or oscillate between ineffective edits due to a lack of historical state tracking.

The core challenge is to create an automated system that can make traceable, evidence-grounded decisions and maintain stable refinement across multiple rounds without relying solely on the LLM's internal "black box" reasoning.

2. Methodology: KernelSkill

KernelSkill is a memory-augmented multi-agent framework designed to replace implicit heuristics with explicit, knowledge-driven optimization skills. It operates via a closed-loop refinement process involving several specialized agents and a dual-level memory architecture.

A. Multi-Agent Pipeline

The framework coordinates the following agents in an iterative loop:

1. Generator: Translates a PyTorch reference into initial CUDA seed kernels, prioritizing correctness over speed.
2. Reviewer: Validates kernels via three checks:
   • Compiler: Ensures successful build.
   • Verifier: Checks output consistency against PyTorch.
   • Profiler: Collects runtime metrics (latency, resource usage) via tools like Nsight Compute (ncu).
3. Diagnoser: Triggered on compilation or correctness failures. It analyzes error logs and short-term repair history to propose fixes, avoiding cyclic repair loops.
4. Feature Extractor: Derives static code features (e.g., memory access patterns, loop structures) to complement runtime profiling data.
5. Planner: Selects optimization methods based on retrieved knowledge and profiling feedback, generating a step-by-step optimization plan.
6. Optimizer/Repairer: Executes the plans to modify the kernel code.

B. Dual-Level Memory Architecture

The core innovation of KernelSkill is its two-tier memory system, which addresses the instability and lack of traceability in prior works:

1. Long-Term Memory (Expert Knowledge Base):

   • Source: Distilled from expert GPU optimization literature and engineering experience.
   • Structure: Contains a Deterministic Decision Policy (rule-based gating, prioritization, and veto rules) and Method Knowledge (rationales and implementation cues).
   • Function: It externalizes expert heuristics. When a kernel is analyzed, the system uses static features and profiling metrics to query this memory. This ensures method selection is auditable and traceable (based on explicit evidence-to-rule matches) rather than relying on LLM intuition.

2. Short-Term Memory (Task Trajectory):

   • Function: Records the per-task optimization and repair history, including attempted plans, outcomes, and feedback.
   • Role in Repair: Prevents "cyclic repair" by conditioning the Diagnoser on past failed attempts, ensuring the system does not revisit known-failing edits.
   • Role in Optimization: Tracks the "base kernel" and previous optimization attempts to prevent oscillation. It updates the base kernel only if specific speedup thresholds (relative or absolute) are met, ensuring stable progression.

3. Key Contributions

• KernelSkill Framework: A novel multi-agent architecture that coordinates generation, verification, profiling, planning, and repair for iterative kernel optimization.
• Two-Level Memory Design:
  • Long-term: Encodes reusable, explainable expert optimization skills, enabling traceable method selection grounded in profiling evidence.
  • Short-term: Maintains explicit per-task state to stabilize multi-round refinement and prevent redundant or oscillating edits.
• Deterministic Decision Policy: Moves away from purely LLM-driven strategy selection to a hybrid approach where rules filter candidates, and LLMs synthesize plans based on retrieved expert knowledge.

4. Experimental Results

The framework was evaluated on KernelBench (Levels 1–3), covering single operators, multi-operator workloads, and full model architectures.

• Success Rate: KernelSkill achieved a 100% success rate across all three levels (generating kernels that compile and pass correctness checks), outperforming or matching the best baselines (e.g., STARK).
• Speedup Performance: KernelSkill delivered significant average speedups over the unoptimized PyTorch (Torch Eager) baseline:
  • Level 1: 5.44× speedup (vs. 3.03× for the next best, STARK).
  • Level 2: 2.82× speedup.
  • Level 3: 1.92× speedup.
• Efficiency: KernelSkill achieved these results in fewer refinement rounds (15 rounds) compared to baselines that required more iterations to reach lower performance ceilings.
• Ablation Studies: Removing either memory component degraded performance. Removing short-term memory reduced success rates (due to instability/oscillation), while removing long-term memory significantly reduced speedups (due to inefficient trial-and-error).

5. Significance

• Interpretability & Trust: By externalizing optimization knowledge into a retrievable, rule-based memory, KernelSkill makes optimization decisions auditable. Engineers can understand why a specific strategy was chosen based on explicit profiling evidence.
• Robustness: The separation of reusable knowledge (long-term) from task-specific state (short-term) solves the problem of "forgetting" past attempts or getting stuck in local optima, making the system robust across diverse and complex kernel tasks.
• Scalability: The framework demonstrates that combining structured expert knowledge with LLM flexibility can outperform purely training-based or inference-time agentic approaches, offering a viable path toward fully automated, high-performance GPU kernel development.