KernelClaw: A Multi-Agent Framework for GPU Kernel Optimization with Dual-Memory Mechanisms
Authors: Qitong Sun, Jun Han, Tianlin Li, Zhe Tang, Sheng Chen, Fei Yang, Aishan Liu, Xianglong Liu, Yang Liu
Organizations: School of Computer Science and Engineering, Beihang University, China · Zhejiang Lab, China · Nanyang Technological University, Singapore
Improving GPU kernel efficiency is crucial for efficient AI systems. Recent work has explored leveraging large language models (LLMs) for GPU kernel generation and optimization. However, existing LLM-based kernel optimization pipelines typically rely on opaque, implicitly learned heuristics within the LLMs to determine optimization strategies. This leads to inefficient trial-and-error and weakly interpretable optimizations. Our key insight is to replace implicit heuristics with expert optimization knowledge that is aware of task trajectories. Specifically, we present KernelClaw, a multi-agent framework with a dual-level memory architecture. KernelClaw operates by coordinating agents with long-term memory of reusable expert knowledge and short-term memory to prevent repetitive backtracking. On KernelBench Levels 1-3, KernelClaw achieves a 100% success rate and average speedups of 6.98x, 7.89x, and 1.94x over Torch Eager on Levels 1, 2, and 3, respectively, outperforming prior baselines. Code is available at https://github.com/0satan0/KernelMem/.
Figures & tables
Figure 1: Overview of KernelClaw
Method
Level 1
Level 2
Level 3
Success ↑
Fast 1 ↑
Speedup ↑
Success ↑
Fast 1 ↑
Speedup ↑
Success ↑
Fast 1 ↑
Speedup ↑
Torch Compile(default)
1.00
0.61
1.51
1.00
0.76
1.22
1.00
0.68
1.40
Torch Compile(max-autotune)
1.00
0.68
1.42
1.00
0.86
1.29
1.00
0.82
1.38
Kevin-32B
0.69
0.04
0.84
0.67
0.08
0.40
0.46
0.02
0.32
AutoTriton
0.36
0.09
1.20
0.55
0.22
0.96
0.56
0.26
0.83
Astra
0.95
0.43
1.48
0.98
0.73
0.99
0.93
0.35
0.90
Table 1: Success, Fast 1 and Speedup Results. Success is the percentage of tasks for which a method generates a kernel that compiles and passes correctness verification. Fast 1 is the percentage of tasks whose generated kernel is at least as fast as the Torch baseline. Speedup is the average runtime speedup over Torch Eager. KernelClaw achieves the best overall performance, matching the highest success rate and Fast 1 rate while delivering the largest speedup across all levels among all compared baselines.
Method
Level 1
Level 2
Level 3
Success ↑
Fast 1 ↑
Speedup ↑
Success ↑
Fast 1 ↑
Speedup ↑
Success ↑
Fast 1 ↑
Speedup ↑
w/o memory
0.96
0.46
1.52
0.98
0.75
1.05
0.94
0.36
0.95
w/o ST memory
0.96
0.51
3.53
0.98
0.85
2.42
0.94
0.70
1.78
w/o LT memory
1.00
0.55
1.77
1.00
1.00
1.18
1.00
0.79
1.04
KernelClaw
1.00
0.74
6.98
1.00
1.00
7.89
1.00
0.88
1.94
Table 2: Ablation Results. We ablate memory components: Long-term memory stores retrievable kernel optimization knowledge for reusable and explainable method selection, while Short-term memory records per-task optimization/repair trajectories to stabilize multi-round refinement. Removing either memory reduces Fast 1 and Speedup, confirming both components are critical to KernelClaw ’s improvements.
Result
Level 1
Level 2
Level 3
Original
6.98 ×
7.89 ×
1.94 ×
Audited
3.47 ×
3.92 ×
1.94 ×
Table 3: Speedup of KernelClaw before and after excluding two anomalous kernels.
Appendix figures & tables14 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 2: The short-term memory for the current repair round
Figure 3: The short-term memory for the current optimization round
Figure 4: Speedup distribution on 100 tasks at Level 1.
Figure 5: Speedup distribution on 100 tasks at Level 2.
Figure 6: Speedup distribution on 50 tasks at Level 3.
KernelBench
Success ↑
Fast 1 ↑
Speedup ↑
Level 1
0.51
0.10
0.60
Level 2
0.07
0.00
0.01
Level 3
0.04
0.00
0.03
Appendix
Table 4: Results with Qwen3-30B-A3B under 15-round optimization budget.
Stage / Outcome
Count
Rate
Sampled tasks
30
–
Valid seed kernels
21 / 30
70.0%
Matched expert methods
17 / 21
81.0%
Improved CUDABench-Score
14 / 17
82.4%
Improved both Score and speedup
12 / 17
70.6%
Failed after matched methods
3 / 17
17.6%
Appendix
Table 5: Statistics on 30 randomly sampled CUDABench tasks.
Task
Matched Method
Seed → Best
Δ
Relative
w/o Expert Methods Δ
358
CMP_Irregular_Compute_S2 ; ML_REUSE
0.174 → 0.727
+0.553
+318%
+0.097
254
MBW_* ; ML_FALLBACK
0.521 → 0.916
+0.395
+76%
+0.171
119
GEMM ; CUBLASLT ; TILING
0.211 → 0.437
+0.226
+107%
+0.000
79
CMP_Irregular_Compute_S2
0.150 → 0.351
+0.201
+134%
+0.061
Appendix
Table 6: Representative CUDABench cases improved by matched optimization methods. The last column reports the CUDABench-Score improvement obtained when expert-method retrieval is disabled.
Method
Input ↓
Output ↓
Total ↓
CudaForge
51k
20k
71k
Astra
56k
21k
77k
PRAGMA
53k
21k
74k
AdaExplore
84k
9k
93k
OpenEvolve
58k
12k
70k
EGG
55k
22k
77k
Appendix
Table 7: Average token consumption per task using GPT-5.1 under a unified budget of 15 optimization rounds. All values are reported in thousands of tokens.
Method
L1
L2
L3
Avg.
w/o short-term memory
67.9%
74.7%
70.4%
71.0%
KernelClaw
73.2%
78.2%
72.1%
74.5%
Appendix
Table 8: Percentage of rounds devoted to performance optimization under the same 15-round budget. Higher values indicate that a larger fraction of the trajectory is used to improve already correct kernels rather than repair invalid ones.
Outcome
Count
Rate
Sampled tasks
30
–
Valid seed kernels
30 / 30
100.0%
Tasks with matched expert methods
29 / 30
96.7%
Tasks improved by KernelClaw
23 / 29
79.3%
Tasks not improved by KernelClaw
6 / 29
20.7%
Tasks improved w/o expert methods
11 / 29
37.9%
Appendix
Table 9: Optimization outcomes on 30 randomly sampled TritonBench-G tasks.
Task
Matched Expert Methods
w/o Expert Methods
KernelClaw
Absolute Gain
Relative Gain
geglu_tanh_triton
Tune_Num_Warps ; Constexpr_Block_Tuning
4.71 ×
7.33 ×
+2.62 ×
+55.6%
batched_vecmat_mult
Rewrite_Mulsum_to_tl_dot ; Enable_tl_dot_tf32
1.25 ×
6.07 ×
+4.82 ×
+385.6%
dequantize_matmul
Quant_Dequant_Matmul ; Tune_Num_Warps
0.99 ×
4.55 ×
+3.56 ×
+359.6%
layernorm_fwd_triton
tl_rsqrt ; Welford_Layernorm
1.28 ×
1.80 ×
+0.52 ×
+40.6%
Mean over 30 tasks
–
1.38 ×
1.90 ×
+0.52 ×
+37.7%
Appendix
Table 10: Representative TritonBench-G tasks improved by the expert methods retrieved by KernelClaw. Speedup is measured relative to the original Triton seed kernel.
GPU
Level
Success ↑
Fast 1 ↑
Speedup ↑
A40
Level 1
1.00
0.68
5.86 ×
Level 2
1.00
0.98
6.14 ×
Level 3
1.00
0.70
1.81 ×
RTX 4090
Level 1
1.00
0.71
5.48 ×
Level 2
1.00
0.94
5.60 ×
Level 3
1.00
0.68
1.71 ×
Appendix
Table 11: Cross-GPU evaluation of KernelClaw on KernelBench. Speedup is measured relative to Torch Eager on the corresponding GPU.
Method
Success ↑
Geo. Mean Speedup ↑
μ CUTLASS+SOL
1.00
1.99 ×
KernelClaw
1.00
2.21 ×
Appendix
Table 12: Controlled comparison with μ CUTLASS+SOL on the same 59 KernelBench tasks under GPT-5.1 and a 15-round optimization budget.
GPU kernel optimization represents a paradigm where functional correctness is assumed and execution efficiency is the objective. We present daVinci-kernel, a reinforcement learning framework that couples skill discovery with skill exploitation through a dynamically evolving skill library. daVinci-kernel jointly trains three agents sharing one LLM backbone: a Skill Selection Agent that retrieves relevant techniques via BM25 and LLM reranking, a Policy Agent that generates multi-turn CUDA/Triton kernels conditioned on selected skills, and a Skill Summary Agent that distills successful rollouts into reusable skills. Candidate skills are added only after execution-based verification confirms reproducible speedups. All three agents share a single LLM backbone, are initialized via a structured SFT cold start on diversity-filtered data, and are then jointly optimized end-to-end with multi-turn REINFORCE and per-agent advantage estimation. On KernelBench, daVinci-kernel-14B achieves 37.2%, 70.6%, and 32.2% on Level 1, Level 2, and Level 3 under the Fast1 threshold, outperforming the strongest prior RL-trained model, Dr. Kernel-14B.
High-performance GPU kernels are critical for reducing the exponentially growing computational costs of large language models (LLMs), but their development heavily relies on manual tuning by domain experts. While recent advances in LLM-based approaches show promise for automating kernel generation, they still struggle to achieve both correctness and high performance. This limitation primarily arises from the lack of domain-specific optimization guidance, hindering effective exploration of the optimization space. We propose EGG, an Expert-Guided Agent Framework for Kernel Generation, which incorporates expert optimization principles to guide LLMs' decisions. Inspired by expert workflows, we decompose kernel generation into two hierarchical stages: 1) algorithmic structure design, which establishes a high-quality computational structure foundation; 2) hardware-specific tuning, which performs targeted adjustments through parallel mapping, tensor tiling, and memory optimization. This staged decomposition defines explicit optimization objectives, structuring the design space to achieve progressive refinement. To this end, a stage-aware multi-agent collaboration mechanism is designed for inter and intra-stage context management, ensuring stable optimization trajectories. Experiments on KernelBench and real-world workloads show that EGG achieves a 2.13x average speedup over PyTorch, outperforming existing agent-based and RL-based approaches.
Yaochen Han, Ke Fan, Hongxu Jiang +5
Beihang University, Beijing, China · Shanghai Jiao Tong University, Shanghai, China · Tsinghua University, Beijing, China +1
Large language models (LLMs) can assist GPU kernel generation, but their practical effectiveness depends on whether generated code can be reliably constrained, validated, profiled, and selected. This paper presents a harness-centered system for LLM-driven GPU kernel optimization in the MLSys 2026 FlashInfer AI Kernel Generation Contest on NVIDIA Blackwell B200 GPUs. The system separates an evaluation harness from a profile-backed optimization controller: the harness enforces compilation, correctness, official-aligned timing, and artifact archival, while the controller turns profiler and workload evidence into bounded candidate-generation decisions. Human-authored skills capture operator constraints, references, profiling procedures, and promotion rules, while Codex and Claude Code agents generate candidate kernels inside those constraints. Across five operator definitions, the retained official-aligned artifacts achieved mean-latency speedups over supplied FlashInfer baselines of 1.62x, 18.05x, 29.68x, 1.12x, and 13.70x. The Agent-Assisted kernels outperform the Full-Agent artifacts across the evaluated definitions, indicating that expert-provided optimization directions, high-quality references, and workload context remain critical for reliable AI-driven kernel optimization.