💬 NLP

Kernel-Smith: A Unified Recipe for Evolutionary Kernel Optimization

Kernel-Smith is a unified framework that combines a stable, evaluation-driven evolutionary agent with a specialized post-training recipe to generate high-performance GPU kernels, achieving state-of-the-art speedups on both NVIDIA and MetaX platforms while demonstrating practical deployment in production systems.

Original authors: 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 W

Published 2026-03-31

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: 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, Qipeng Guo, Kai Chen

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 build the fastest possible race car engine. You have a blueprint (the standard code), but you know that with some clever tweaks, you could make it go twice as fast.

For a long time, AI models trying to do this were like novice mechanics. They would look at the blueprint, guess a few changes, and say, "Here is my new engine!" If it worked, great. If it didn't, they'd try again from scratch. They were bad at iterating—they couldn't learn from their mistakes to make the next version slightly better.

Kernel-Smith is a new framework that changes the game. It treats AI not as a one-shot guesser, but as a team of expert mechanics working in a workshop, using a specific recipe to evolve the perfect engine.

Here is how it works, broken down into simple concepts:

1. The Workshop: Evolutionary Search (The "Darwinian" Garage)

Instead of one mechanic guessing a solution, Kernel-Smith sets up a garage full of prototype engines.

The Population: The AI generates many different versions of the code (the engines).
The Test Track: Every single prototype is put on a test track (the GPU) to see how fast it runs and if it breaks.
The Archive: The system keeps a "Hall of Fame" for the fastest engines and a "Museum of Diversity" for engines that tried weird, unique things.
The Evolution: The AI looks at the winners and the interesting failures, mixes their ideas, and builds a new generation of engines. It doesn't just guess; it evolves. Over 40 rounds, the engines get faster and faster, just like species evolving over millions of years.

2. The Referee: Stable Evaluation (The "Fair Judge")

In a real race, if the stopwatch is broken or the track is bumpy, you can't tell who is actually faster.

The Problem: Computer chips are noisy. Sometimes a program looks slow just because the computer was busy doing something else, not because the code is bad.
The Solution: Kernel-Smith built a super-strict referee system. It runs the tests many times, warms up the engine first, and uses special tools to ignore the "noise." It ensures that if a code is declared "faster," it is actually faster, not just lucky. This prevents the AI from getting confused by bad data.

3. The Teacher: The "Step-by-Step" Training (The "Coach")

This is the secret sauce. Usually, AI models are trained to write a perfect essay in one go. But writing a perfect engine is hard.

The Old Way: "Here is a bad engine. Write a good one." (The AI often fails).
The Kernel-Smith Way: The researchers watched the "garage" evolve. They looked at the specific moments where an engine went from "okay" to "amazing." They taught the AI only those specific steps.
The Analogy: Imagine a chess coach. Instead of showing the AI a whole game and saying "Win," the coach says, "Look at this specific move where you captured the Queen. That was the winning move. Learn that move."
By training the AI to be a master of small improvements rather than a master of one giant leap, the AI becomes incredibly good at refining code over and over again.

4. The Results: Beating the Giants

The paper tested this system against the biggest, most expensive AI models in the world (like Google's Gemini and Anthropic's Claude).

The Scoreboard: On a standard test called "KernelBench," Kernel-Smith didn't just keep up; it won. It produced code that was significantly faster than the others.
Real World Proof: It's not just a lab experiment. The team took their AI-generated code and actually merged it into real software used by millions of people (like SGLang and LMDeploy).
- In one case, they made a specific part of a video game or AI chat system run 4.7 times faster.
- In another, they made a complex research tool run 14.5 times faster.

Why Does This Matter?

Think of the AI chips in your phone or the servers running your favorite apps as high-performance sports cars.

Before: The AI was like a mechanic who could only change the oil.
Now: Kernel-Smith is a mechanic who can redesign the entire engine to squeeze out every drop of speed.

This means that in the future, the AI systems we use every day could run much faster, use less battery, and cost less to operate, all because this new "recipe" taught the AI how to evolve its own code rather than just guessing it.

1. Problem Statement

High-performance GPU kernel generation is critical for optimizing large language model (LLM) training/inference and scientific computing workloads. While Large Language Models (LLMs) have shown promise in code generation, they face significant hurdles in this domain:

One-Shot Limitations: Current LLM approaches struggle to move beyond single-pass code generation. They often fail to sustain reliable optimization through iterative refinement.
Search Complexity: Efficient kernels require exploring a vast space of implementation choices (fusion patterns, tiling strategies, rewrite directions). Standard multi-turn agent loops often get "anchored" to early decisions, limiting exploration diversity.
Evaluation Instability: Performance optimization is highly sensitive to measurement noise. In GPU environments, timing fluctuations can lead evolutionary agents to preserve suboptimal kernels or discard promising ones.
Correctness vs. Performance: Generating a kernel that is both functionally correct and significantly faster than the baseline is a coupled challenge that standard models struggle to solve iteratively.

2. Methodology

Kernel-Smith introduces a unified framework combining a stable evaluation-driven evolutionary agent with an evolution-oriented post-training recipe.

A. Agent Framework: Evolutionary Search with Stable Evaluation

Instead of refining a single trajectory, Kernel-Smith maintains a population of executable candidates and evolves them over multiple rounds.

Population Management: The system uses an archive (inspired by MAP-Elites) to store top-performing and diverse programs.
Structured Feedback: At each step, the agent receives structured feedback including compilation status, correctness checks, speedup metrics, runtime measurements, and error logs.
Evaluation Backends: To ensure reliability, the authors built backend-specific evaluation services for NVIDIA GPUs (Triton) and MetaX GPUs (Maca).
- Stability Measures: They employ warm-up executions, multiple measurements with outlier removal, and CUDAGraph technology to constrain execution time fluctuations to within 1%.
- Hacking Detection: The system includes mechanisms to detect "lazy optimization" (e.g., trivial rewrites) or "advanced hacking" (calling native PyTorch operators to fake speedups), ensuring only genuine kernel improvements are accepted.

B. Training Recipe: Step-Centric Supervision and RL

The core innovation is training the LLM not as a one-shot generator, but as a strong local improver within the evolutionary loop.

Data Synthesis:
- Source: A curated dataset of 59k high-quality PyTorch modules extracted from diverse GitHub repositories.
- Cold-Start: Using an open-source teacher model (DeepSeek-V3.2-Speciale) to generate initial rollout trajectories, filtered for correctness and speedup.
- Cluster-Seeded Expert Data: Using HDBSCAN clustering to identify representative operators for manual expert annotation, feeding these back into the system for higher-fidelity trajectory generation.
Step-Centric Filtering: Instead of training on full evolution trajectories (which may contain redundant or misleading steps), the system filters for high-gain revisions.
- Correctness-Oriented: Filters for functional correctness in the initial translation (PyTorch $\to$ Triton).
- Performance-Oriented: Filters for speedup ratios $>1.0$ in iterative evolution steps.
Reinforcement Learning (RL): The model is fine-tuned using Group Relative Policy Optimization (GRPO). The reward signal is the speedup ratio relative to the parent code. Crucially, the training data consists only of the "Best Steps" from evolution trajectories, teaching the model to recognize and execute atomic improvements that compound over time.

3. Key Contributions

Kernel-Smith Framework: A unified system that integrates a stable, noise-resistant evolutionary agent with a specialized training pipeline tailored for iterative kernel optimization.
Evolution-Oriented Training Recipe: A novel approach that converts long-horizon search trajectories into step-centric supervision signals, optimizing the model to be a "local improver" rather than a one-shot generator.
Robust Evaluation Infrastructure: The development of backend-specific evaluation services for NVIDIA (Triton) and MetaX (MACA) that include rigorous noise reduction and anti-hacking mechanisms.
Real-World Validation: Successful upstream contributions to production inference engines (SGLang and LMDeploy) and optimization of cutting-edge research modules (DeepSeek Engram), proving the framework's transferability from benchmarks to production.

4. Results

The framework was evaluated on KernelBench (NVIDIA backend) and a custom MetaX MACA backend.

NVIDIA Backend (KernelBench):
- Kernel-Smith-235B-RL achieved State-of-the-Art (SOTA) performance with an average speedup ratio (Avg AMSR) of 3.70, outperforming frontier proprietary models like Gemini-3.0-pro (2.83) and Claude-4.6-opus (3.33).
- It demonstrated superior performance on medium-difficulty tasks (Level 2), achieving an Avg AMSR of 7.77 compared to Claude-4.6-opus's 5.83.
- It maintained high correctness (96.33%) while significantly outperforming open-weight baselines.
- Growth Curve: The model's best-score growth curve formed the "upper envelope" of all competitors, indicating it benefits more effectively from additional test-time compute (evolutionary steps).
MetaX Backend:
- Kernel-Smith-MACA-30B surpassed large-scale counterparts like DeepSeek-V3.2-think and Qwen3-235B-2507-think, demonstrating seamless adaptation to heterogeneous platforms.
Real-World Deployments:
- SGLang: A fused Triton kernel for normal_decode_set_metadata achieved a 4.78x local speedup and was merged into the production codebase.
- LMDeploy: A fused MoE routing kernel achieved 1.36x local speedup and improved end-to-end throughput by 1.85%–3.00%.
- DeepSeek Engram: Optimizing a recent research module resulted in a 14.59x speedup, which was subsequently merged into the DLBlas library.

5. Significance

Paradigm Shift: The paper shifts the focus from "one-shot code generation" to "iterative, evolution-driven optimization," proving that LLMs can act as effective local optimizers when guided by stable feedback.
Practical Impact: Unlike many research benchmarks that remain theoretical, Kernel-Smith has successfully contributed code to major production systems (SGLang, LMDeploy), bridging the gap between LLM research and industrial deployment.
Generalizability: The framework's backend-decoupled design allows it to adapt to new hardware (e.g., Huawei NPUs) without retraining the core agent logic, only requiring new evaluation backends.
Methodological Insight: It highlights that stable evaluation (suppressing noise) and step-centric training (focusing on high-gain atomic steps) are the critical ingredients for making LLM-driven systems optimization reliable and scalable.