KernelClaw: A Multi-Agent Framework for GPU Kernel Optimization with Dual-Memory Mechanisms

Qitong Sun; Jun Han; Tianlin Li; Zhe Tang; Sheng Chen; Fei Yang; Aishan Liu; Xianglong Liu

cs.LGMar 10, 2026

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,

Organizations: School of Computer Science and Engineering, Beihang University, China · Zhejiang Lab, China · Nanyang Technological University, Singapore

Abstract

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

Method	Level 1			Level 2			Level 3
	Success $\uparrow$	Fast 1 $\uparrow$	Speedup $\uparrow$	Success $\uparrow$	Fast 1 $\uparrow$	Speedup $\uparrow$	Success $\uparrow$	Fast 1 $\uparrow$	Speedup $\uparrow$
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 $\uparrow$	Fast 1 $\uparrow$	Speedup $\uparrow$	Success $\uparrow$	Fast 1 $\uparrow$	Speedup $\uparrow$	Success $\uparrow$	Fast 1 $\uparrow$	Speedup $\uparrow$
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 $\times$	7.89 $\times$	1.94 $\times$
Audited	3.47 $\times$	3.92 $\times$	1.94 $\times$

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 $\uparrow$	Fast 1 $\uparrow$	Speedup $\uparrow$
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 $\rightarrow$ Best	$\Delta$	Relative	w/o Expert Methods $\Delta$
358	CMP_Irregular_Compute_S2 ; ML_REUSE	0.174 $\rightarrow$ 0.727	+0.553	+318%	+0.097
254	MBW_* ; ML_FALLBACK	0.521 $\rightarrow$ 0.916	+0.395	+76%	+0.171
119	GEMM ; CUBLASLT ; TILING	0.211 $\rightarrow$ 0.437	+0.226	+107%	+0.000
79	CMP_Irregular_Compute_S2	0.150 $\rightarrow$ 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 $\downarrow$	Output $\downarrow$	Total $\downarrow$
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 $\times$	7.33 $\times$	+2.62 $\times$	+55.6%
batched_vecmat_mult	Rewrite_Mulsum_to_tl_dot ; Enable_tl_dot_tf32	1.25 $\times$	6.07 $\times$	+4.82 $\times$	+385.6%
dequantize_matmul	Quant_Dequant_Matmul ; Tune_Num_Warps	0.99 $\times$	4.55 $\times$	+3.56 $\times$	+359.6%
layernorm_fwd_triton	tl_rsqrt ; Welford_Layernorm	1.28 $\times$	1.80 $\times$	+0.52 $\times$	+40.6%
Mean over 30 tasks	–	1.38 $\times$	1.90 $\times$	+0.52 $\times$	+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 $\uparrow$	Fast 1 $\uparrow$	Speedup $\uparrow$
A40	Level 1	1.00	0.68	5.86 $\times$
	Level 2	1.00	0.98	6.14 $\times$
	Level 3	1.00	0.70	1.81 $\times$
RTX 4090	Level 1	1.00	0.71	5.48 $\times$
	Level 2	1.00	0.94	5.60 $\times$
	Level 3	1.00	0.68	1.71 $\times$

Appendix

Table 11: Cross-GPU evaluation of KernelClaw on KernelBench. Speedup is measured relative to Torch Eager on the corresponding GPU.

Method	Success $\uparrow$	Geo. Mean Speedup $\uparrow$
$\mu$ CUTLASS+SOL	1.00	1.99 $\times$
KernelClaw	1.00	2.21 $\times$

Appendix

Table 12: Controlled comparison with

\mu

CUTLASS+SOL on the same 59 KernelBench tasks under GPT-5.1 and a 15-round optimization budget.

GPU Kernel Optimization

GPU kernel optimization improves the speed and memory efficiency of computations by tailoring code to GPU hardware and workloads.

See all 202 papers

daVinci-kernel: Co-Evolving Skill Selection, Summarization, and Utilization via RL for GPU Kernel Optimization
Jun 15, 2026Dayuan Fu, Mohan Jiang, Tongyu Wang +5RL for Code Generation GPU Kernel Optimization
1SII · 2SJTU · 3GAIR
EGG: An Expert-Guided Agent Framework for Kernel Generation
Jun 25, 2026Yaochen Han, Ke Fan, Hongxu Jiang +5GPU Kernel Optimization Large Language Model-Guided Optimization
Beihang University, Beijing, China · Shanghai Jiao Tong University, Shanghai, China · Tsinghua University, Beijing, China +1
Harness Engineering for LLM-Driven GPU Kernel Generation
Jul 20, 2026Yue Shui, Chenyu Ma, Hangfei Xu +2GPU Kernel Optimization Large Language Model-Guided Optimization
Baidu, Inc.

Jun 15, 2026J/K move · Enter open · S save

daVinci-kernel: Co-Evolving Skill Selection, Summarization, and Utilization via RL for GPU Kernel Optimization

Dayuan Fu, Mohan Jiang, Tongyu Wang, Dian Yang+4

1SII · 2SJTU · 3GAIR

RL for Code Generation GPU Kernel Optimization Large Language Model-Guided Optimization GPU Kernel Generation RL Post-Training

Open paper AskPDF arXiv

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 Fast

_1

threshold, outperforming the strongest prior RL-trained model, Dr. Kernel-14B.

Abstract

GPU Kernel Optimization

Explore similar work

daVinci-kernel: Co-Evolving Skill Selection, Summarization, and Utilization via RL for GPU Kernel Optimization

EGG: An Expert-Guided Agent Framework for Kernel Generation

Harness Engineering for LLM-Driven GPU Kernel Generation