ModelBlaster: Closing the Performance Gap through Learning-Augmented Agents Across Models, Kernels and Hardware Backends

Shengjun Kris Dong; Dima Nikiforov; Sahil Modi; Sana Damani; Edward C Lin; Siva Kumar Sastry Hari; Christos Kozyrakis

cs.LGNov 3, 2026

ModelBlaster: Closing the Performance Gap through Learning-Augmented Agents Across Models, Kernels and Hardware Backends

Authors: Shengjun Kris Dong, Dima Nikiforov, Sahil Modi, Sana Damani, Edward C Lin, Siva Kumar Sastry Hari, Christos Kozyrakis

Organizations: NVIDIA · University of California, Berkeley

Abstract

Optimizing accelerator code across hardware architectures is difficult. ML compilers search predefined spaces of schedules and micro-kernels, while LLM-guided optimization can explore alternative algorithms, operator decompositions, and hardware-specific techniques. However, existing LLM agents typically target a single programming interface, toolchain, or accelerator family, focus on kernel-level optimization without model-level transforms, and offer limited reuse of experience across hardware targets. We present ModelBlaster, a Memory-Augmented In-Context Reinforcement Learning (MAIC-RL) framework for optimizing data-parallel kernels and whole models across heterogeneous SIMD and SIMT accelerators. Its mutable policy is a persistent knowledge base updated at inference time by Actor-Critic Textual Gradient Descent (AC-TGD). The knowledge base maps performance states to actions, outcomes, and constraints, enabling experience reuse across kernels, models, and hardware targets. A model-level optimizer uses kernel-level knowledge to inform whole-model decisions. The same optimization loop and knowledge schema support CUDA on NVIDIA GPUs, OpenCL on Qualcomm GPUs, and RVV intrinsics on RISC-V vector cores, with only lowering and measurement layers specific to each target. Compared with PyTorch, ModelBlaster achieves geometric-mean speedups of 1.85x, 5.58x, and 5.18x on KernelBench Levels 1, 2, and 3, and 1.69x on SOL-ExecBench Levels 1 and 2 on a B200 GPU. On Qualcomm Adreno GPUs, it achieves 5.8x over the vendor SDK and 1.7x over Apache TVM; on RISC-V vector cores, it achieves 2.7x over ExecuTorch with XNNPACK across 70 KernelBench problems.

Figures & tables

Panel of Figure 1 — Figure 1: Taxonomy of approaches to ML code optimization. ModelBlaster , at the bottom, combines long-term memory, textual-gradient learning, and the kernel-level and graph-level functions of ML compilers.

Figure 2

Component	Definition
Task	$\tau=(x,y,b,m)$ , where $x$ is the input program or model fragment, $y$ is a reference implementation, $b$ is the backend, and $m$ is the measurement interface.
State	$s_{t}=\phi(x_{t},c_{t},p_{t},h_{t})$ , where $x_{t}$ is the current implementation, $c_{t}$ is compiler and validation feedback, $p_{t}$ is timing or profiling evidence, and $h_{t}$ is rollout history.
Actor action	$a_{t}\sim\pi_{\theta}(a\mid s_{t},\mathcal{K}_{t})$ , where $a_{t}$ is a code transformation or implementation plan and $\mathcal{K}_{t}$ is the set of retrieved Knowledge Base entries.
Reward signal	$(r_{t},z_{t})=R(x_{t+1},y,b,m)$ , where $r_{t}$ is the correctness-gated scalar performance reward and $z_{t}$ is the text-based reward signal.
Memory update	$\mathcal{K}_{k+1}=U\!\left(\mathcal{K}_{k},\tau,\{(s_{t},a_{t},r_{t},z_{t})\}_{t=0}^{T}\right)$ , where the critic incorporates rollout outcomes into persistent long-term memory.

Table 1: Problem formulation for ModelBlaster ’s Memory-Augmented In-Context Reinforcement Learning (MAIC-RL) loop.

Component	RL counterpart	Functionality
Knowledge Base	Parameters $\theta$	Store of states, techniques, hypotheses, claims, outcomes, and verified implementations; with the prompt it forms the policy parameters.
Operator Selector		Picks the next kernel slot by time share, stall decay, and exploration.
State Classifier	State $s_{t}$	Assigns the deepest matching performance state from the profiling signature, with the program, feedback, and rollout history.
Technique Selector	Policy $\pi_{\theta}(a_{t}\mid s_{t})$	Picks a technique for the state with a UCB rule over reward statistics.
Actor CodeGen Agent	Action $a_{t}$	Rewrites the selected slot using claims, similar kernels, and past attempts.
Fix Agent	Action $a_{t}$	Repairs candidates that fail the correctness check, revising the action.

Table 2: Components of the ModelBlaster agentic workflow and their counterparts in the RL formulation.

Hardware / Benchmark	Level	Method	Share	Mean	GeoM	Med	Min	Max	% $>$ 1x	% $<$ 1x
NVIDIA L40S KernelBench	1	AI CUDA Engineer	82%	1.728	1.101	1.202	0.0659	32.35	74.39%	25.61%
		StitchCUDA	98%	1.718	1.335	1.116	0.0250	9.77	86.73%	11.22%
		Ours	96%	3.972	1.850	1.758	0.2621	158.51	85.42%	14.58%
	2	AI CUDA Engineer	83%	3.865	1.695	1.574	0.328	123.39	79.52%	20.48%
		StitchCUDA	79%	2.888	2.177	1.665	0.296	9.48	94.94%	5.06%
		Ours	100%	22.823	5.576	4.934	0.8084	314.64	85.00%	15.00%

Table 3: ModelBlaster ’s speedup across NVIDIA, Qualcomm, and RISC-V backends. Baselines are PyTorch eager Ansel et al. (2024) (NVIDIA), QNN SDK (Adreno), and scalar kernels or ExecuTorch/XNNPACK (RISC-V). Share denotes the percentage of solved benchmark problems.

Figure 3: Distribution of 3972 CUDA optimization applications across performance states.

System	Language	Agent’s Models	Hardware Setup	Dataset	Hyperparameter	Test Harness
ModelBlaster (Ours)	CUDA	GPT-5.0-mini, GPT-4.1, GPT-5.0, Sonnet 4.6, Opus 4.6	NVIDIA A6000, A100, H100, L40S; RISC-V Vector core (Saturn)	KernelBench L1–3, SOL-ExecBench	10 iterations, 10 rollout steps per iteration	Functionality Check : passing kernels are compiled and executed on GPUs; outputs are compared to PyTorch baselines with multiple randomized seeds to ensure correctness and prevent overfitting.
AI CUDA Engineer	CUDA	O4-mini, Sonnet 3.7, Gemini 2.5 Pro, GPT-4.1, O3	NVIDIA H100	KernelBench L1–2	10 generations; 10 proposals/generation; top 4 evaluated	LLM Soft-Verification : an LLM scans each kernel for likely compilation, memory, or numerical errors, filtering bad code before GPU execution.
StitchCUDA	CUDA	Opus 4.6 (our run)	NVIDIA L40S (our run)	KernelBench L1–2	15 attempts per problem with 2 replans	Verifier agent : checks correctness and profiles with Nsys and NCU; speedup over PyTorch is measured as wall-clock time in its own harness.
QNN SDK Compiler	OpenCL	N/A	Qualcomm Adreno 650	KernelBench L1	N/A	(SDK-internal validation)

Table 4: Comparison among ModelBlaster (Ours), AI CUDA Engineer, and StitchCUDA.

Figure 4:

\mathit{fast}_{p}

curves on an L40S for KernelBench Levels 1 and 2: the share of problems with a correct result whose speedup over PyTorch eager is at least

r

. All three systems are measured as end-to-end wall-clock time.

n

is the number of problems with valid results; AI CUDA Engineer has kernels that directly call external libraries excluded.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Comparison	All overlapping	Compute-bound
ModelBlaster over TVM	1.69 $\times$ ( $n=39$ )	2.30 $\times$ ( $n=16$ )
ModelBlaster over TVM, optimized pipeline only	1.50 $\times$ ( $n=32$ )
ModelBlaster over QNN	5.80 $\times$ ( $n=42$ )
TVM over QNN	3.36 $\times$ ( $n=37$ )	1.21 $\times$ ( $n=8$ )

Appendix

Table 5: Pairwise geometric-mean speedups on the Adreno 650 on KernelBench Level 1 (FP16).

n

is the number of operators covered by both systems. Compute-bound operators are a subset of the overlap.

Operator class	ModelBlaster faster	Speedup
Matrix multiplication	14 of 16	3.13 $\times$
Reductions	6 of 6	251.7 $\times$
Normalizations	8 of 8	47.2 $\times$
Activations	11 of 13	10.0 $\times$
Conv2d	1 of 11	0.09 $\times$
Conv1d	0 of 5	0.06 $\times$

Appendix

Table 6: ModelBlaster against ExecuTorch with XNNPACK on RVV by operator class: the number of problems where ModelBlaster is faster, and the geometric-mean speedup over ExecuTorch.

Figure 5: Geometric-mean speedup of ModelBlaster over ExecuTorch with XNNPACK on RVV by operator class (log scale; values above 1 mean ModelBlaster is faster). Labels give the speedup and the number of problems where ModelBlaster is faster out of the class total.

Panel of Figure 6 — Figure 6: Reusable persistent memory: (a) discovery and application of new optimizations with a pretrained vs. empty Knowledge Base ; (b) the same Knowledge Base trained on an A6000 transfers across other NVIDIA GPU platforms.

Panel of Figure 7 — Figure 7: Memory contents and text-based reward ablation. (a) The learned memory does not collapse to a single transformation family: no state exceeds 20% of applications, justifying state-conditioned retrieval. (b) Replacing text-based critiques with cycle-only scalar feedback degrades optimization across most kernels (1.22 $\times$ vs. 1.57 $\times$ on Level 2), confirming that language carries information scalar timing cannot.

Arm	vs eager	vs CUDA graph
A0 baseline	2.79 $\times$	1.59 $\times$
A1 round controller	2.79 $\times$	1.59 $\times$
A2 graph-first round	3.46 $\times$	1.97 $\times$
A3 sibling-conv merge	3.47 $\times$	1.98 $\times$
A4 region windows	4.08 $\times$	2.33 $\times$
A5 region loop	2.90 $\times$	1.65 $\times$

Appendix

Table 7: Flow-component ablation on eight KernelBench Level 3 problems on an L40S: geometric-mean speedup of the final model over PyTorch eager and over a CUDA graph.

Panel of Figure 8 — Figure 8: The impact of hyperparameter tuning on search dynamics: shaded regions show the inter-quartile range of optimized kernels. Breadth and depth both exhibit diminishing returns.

Figure 9: Distribution of Qualcomm optimization applications across performance states (KernelBench L1, Adreno 650 GPU).

Figure 10: Synthetic Clustering of Top Optimizations

Model	Speedup
LeNet5	10.58 $\times$
SwinMLP	4.07 $\times$
ResNetBasicBlock	3.19 $\times$
DenseNet121DenseBlock	2.72 $\times$
SqueezeNetFireModule	2.41 $\times$
MLP	1.07 $\times$

Appendix

Table 8: Speedup of ModelBlaster on representative KernelBench Level 3 models on an L40S, over PyTorch eager. The last row is the geometric mean over all 46 problems.

GPU Kernel Optimization

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

See all 227 papers

Sep 24, 2026

KernelOPT: Dispatch-Aware Agentic Search for GPU Kernel Optimization

Deep learning inference and training performance depends critically on GPU kernel efficiency. Modern compilers such as PyTorch Inductor automatically generate GPU kernels from high-level model code, but frequently underperform expert-written implementations by wide margins. Recent LLM-assisted kernel optimizers can close this gap for standalone kernels, yet treat compiled models as black boxes, generally optimizing individual standalone kernels without respecting the compiler's structural decisions or verifying the model end-to-end. We present KernelOPT, a multi-agent system that treats compiled models as structured artifacts. It preserves vendor library calls (cuBLAS, cuDNN) and exclusively targets generated Triton sub-kernels using five profiling-guided LLM agents. A four-gate verification cascade applies static validation, multi-seed correctness checking, model-level float64-fallback verification, and performance gating (

γ{=}1.03

) to filter candidates and verify the re-stitched model end-to-end. When candidates fail verification, the system preserves the compiler baseline. The system accepts PyTorch nn Modules, standalone Triton kernels, and Helion kernels. Evaluated on 250 KernelBench problems (100 Level 1, 100 Level 2 and 50 Level 3) on NVIDIA H200, KernelOPT achieves geometric mean speedups over torch compile of 1.40

\times

(L1), 1.15

\times

(L2), and 1.07

\times

(L3) across all kernels, including fallback cases. Optimized-only geomeans (excluding cases where verification gates preserve the compiler baseline) are substantially higher: 2.54

\times

(L1: 36/100), 1.84

\times

(L2: 23/100), and 1.37

\times

(L3: 11/50), reflecting where the optimizer achieves meaningful leverage.

Aheli Poddar, Sanskar Prasad, Arindam Samanta +3

Red Hat

GPU Kernel Optimization Agentic Search Large Language Model-Guided Optimization Compiler Optimization Neural Network Hardware Acceleration

Ask PDF

Oct 28, 2026

A Case Study of a Multi-Agent Kernel-Optimization Loop That Learns a New AI Accelerator

Frontier large language models write correct CUDA kernels for most single-operator benchmark tasks, yet their accuracy collapses on accelerators whose architecture and programming language they have rarely seen. We present a case study of an LLM agent loop built for such accelerators. Our example is the FuriosaAI RNGD NPU, programmed in a recent, sparsely documented Rust DSL, where the loop optimized three Gemma-4-12B decode kernels over 12 days on a shared, scarce device. The loop builds its own target knowledge: it writes and checks a summary of the DSL documentation, measures the chip with microbenchmarks, records compiler-derived legality rules, and revisits this knowledge when the toolchain changes, rechecking rules and measurements where it can. A role-based search over per-kernel Design Trees explores whole-kernel designs and accumulates small accepted changes in the best version of each kernel. Because device time is scarce and readings are noisy, cheap checks screen new designs before any device run, and late in the run a change was accepted only after repeated comparisons with the best kernel on the device. The final kernels reach a 9.293x geomean speedup over the official baseline, each within 1.25x of the cycles needed just to move its weights. Our prompts, profiler and acceptance test are available at https://github.com/bgchun-fs/multi-agent-kernel-gen-loop.

Byonggon Chun

Fluidstack

AI Accelerator Inference Large Language Model-Guided Optimization Agent Harness Optimization Multi-Agent LLMs

Ask PDF

Feb 27, 2026

CUDA Agent: Large-Scale Agentic RL for High-Performance CUDA Kernel Generation

GPU kernel optimization is fundamental to modern deep learning but remains a highly specialized task requiring deep hardware expertise. Despite strong performance in general programming, large language models (LLMs) remain uncompetitive with compiler-based systems such as torch.compile for CUDA kernel generation. Existing CUDA code generation approaches either rely on training-free refinement or fine-tune models within fixed multi-turn execution-feedback loops, but both paradigms fail to fundamentally improve the model's intrinsic CUDA optimization ability, resulting in limited performance gains. We present CUDA Agent, a large-scale agentic reinforcement learning system that develops CUDA kernel expertise through three components: a scalable data synthesis pipeline, a skill-augmented CUDA development environment with automated verification and profiling to provide reliable reward signals, and reinforcement learning algorithmic techniques enabling stable training. CUDA Agent achieves state-of-the-art results on KernelBench, delivering 100%, 100%, and 92% faster rate over torch.compile on KernelBench Level-1, Level-2, and Level-3 splits, outperforming the strongest proprietary models such as Claude Opus 4.5 and Gemini 3 Pro by about 40% on the hardest Level-3 setting.

Weinan Dai, Hanlin Wu, Qiying Yu +13

ByteDance Seed · Institute for AI Industry Research (AIR), Tsinghua University · SIA-Lab of Tsinghua AIR and ByteDance Seed

GPU Kernel Generation GPU Kernel Optimization RL for Code Generation Agentic RL

Ask PDF

Abstract

GPU Kernel Optimization

Explore similar work

KernelOPT: Dispatch-Aware Agentic Search for GPU Kernel Optimization

A Case Study of a Multi-Agent Kernel-Optimization Loop That Learns a New AI Accelerator

CUDA Agent: Large-Scale Agentic RL for High-Performance CUDA Kernel Generation