Improving Efficiency of GPU Kernel Optimization Agents using a Domain-Specific Language and Speed-of-Light Guidance

Siva Kumar Sastry Hari; Vignesh Balaji; Sana Damani; Qijing Huang; Christos Kozyrakis

cs.LGMar 30, 2026

Improving Efficiency of GPU Kernel Optimization Agents using a Domain-Specific Language and Speed-of-Light Guidance

Authors: Siva Kumar Sastry Hari, Vignesh Balaji, Sana Damani, Qijing Huang, Christos Kozyrakis

Organizations: NVIDIA

Abstract

Optimizing GPU kernels with LLM agents is an iterative process over a large design space. Every candidate must be generated, compiled, validated, and profiled, so fewer trials will save both runtime and cost. We make two key observations. First, the abstraction level that agents operate at is important. If it is too low, the LLM wastes reasoning on low-impact details. If it is too high, it may miss important optimization choices. Second, agents cannot easily tell when they reach the point of diminishing returns, wasting resources as they continue searching. These observations motivate two design principles to improve efficiency: (1) a compact domain-specific language (DSL) that can be learned in context and lets the model reason at a higher level while preserving important optimization levers, and (2) Speed-of-Light (SOL) guidance that uses first-principles performance bounds to steer and budget search. We implement these principles in $μ$ CUTLASS, a DSL with a compiler for CUTLASS-backed GPU kernels that covers kernel configuration, epilogue fusion, and multi-stage pipelines. We use SOL guidance to estimate headroom and guide optimization trials, deprioritize problems that are near SOL, and flag kernels that game the benchmark. On 59 KernelBench problems with the same iteration budgets, switching from generating low-level code to DSL code using GPT-5-mini turns a 0.40x geomean regression into a 1.27x speedup over PyTorch. Adding SOL-guided steering raises this to 1.56x. Across model tiers, $μ$ CUTLASS + SOL-guidance lets weaker models outperform stronger baseline agents at lower token cost. SOL-guided budgeting saves 19-43% of tokens while retaining at least 95% of geomean speedup, with the best policy reaching a 1.68x efficiency gain. Lastly, SOL analysis helps detect benchmark-gaming cases, where kernels may appear fast while failing to perform the intended computation.

Figures & tables

Built-in: relu, gelu, silu, sigmoid, tanh, mish, hardswish, leaky_relu, elu, clip, clamp, bias, per_channel_scale, per_row_scale, per_col_scale, scale, aux_store, aux_load

Custom: custom(’expr’, inputs={…}) (SM90a † )

Table 1: Coverage and feature support in

\mu

CUTLASS. ∗ On SM90+, depthwise convolutions route to a CuTe-based backend with additional restrictions (e.g., limited stride/dilation and epilogue support). † Custom epilogue expressions are supported on SM90a. ‡ .with_operand_swap(true) applies to SM90+ FP32 GEMM and requires

M{=}N

Panel of Figure 1 — Figure 1: $\mu$ CUTLASS compilation pipeline. The compiler parses and lowers a kernel.dsl program to a typed configuration IR, validates architecture and other constraints, and emits the CUTLASS C++ code in a wrapper header ( ucutlass_hash.h ). Optionally, it emits a CUTLASS Python API plan ( kernel.py ) and invokes cutlass_cppgen to produce kernel.cu , which is then wrapped into the final header.

Panel of Figure 2 — Figure 2: SOL-guided budget scheduling. The scheduler combines SOL gap, progress, and speedup signals to determine problem eligibility, then selects the next work item: an attempt for MI agents or an iteration for orchestrated MANTIS.

SOL-guided steering
Variant	$\mu$ CUTLASS	SOL-Guidance	Total Attempts
MI w/o $\mu$ CUTLASS	$\times$	–	40
MI + $\mu$ CUTLASS	✓	–	40
In-prompt steering w/o $\mu$ CUTLASS	$\times$	In-Prompt	40
In-prompt steering + $\mu$ CUTLASS	✓	In-Prompt	40
Orchestrated steering w/o $\mu$ CUTLASS	$\times$	Orchestrated	40 ( $5\times 2\times 4$ )

Table 2: Experimental variants and default budgets. Each of three controller types (MI, in-prompt steering, orchestrated steering) is paired with and without

\mu

CUTLASS, isolating the effects of code representation and SOL-guided steering under matched per-problem budgets.

Variant	M	A	N	T	I	S	Cross-prob
MANTIS	✓	✓	✓	✓	✓	✓	✓
MNTIS (no Analyze)	✓	$\times$	✓	✓	✓	✓	✓
MANIS (no Triage)	✓	✓	✓	$\times$	✓	✓	✓
MANTI (no Summarize)	✓	✓	✓	✓	✓	$\times$	$\times$
MANTIS-noXmem	✓	✓	✓	✓	✓	✓	$\times$

Table 3: MANTIS component ablations. Each row removes one mechanism. “Cross-prob” indicates whether summaries are persisted across problems.

Figure 3: Geomean speedup over PyTorch for the four main variants across three model tiers with same attempt budget is shown here.

Panel of Figure 4 — Figure 4: Four main variants across three model tiers are shown here. Left column shows Fast- $p$ curves (percentage of problems with speedup above $r$ ). Right column shows Attempt-Fast- $p(2)$ (percentage of problems reaching $\geq 2\times$ at different attempts). Dashed gray lines show the MI baseline of the next stronger model.

Panel of Figure 5 — Figure 5: Orchestrated vs. in-prompt SOL guidance implementation is shown here. Dotted = orchestrated, solid = in-prompt. Same color = same setting.

Panel of Figure 6 — Figure 6: Component ablation results are shown here. Fast- $p$ curves for full MANTIS and four single-component ablations on configurations where SOL guidance provides benefits.

Figure 7: Independent parameter sweeps for GPT-5.2 $\mu$ CUTLASS + SOL-guided variant are shown here. (a) SOL-headroom threshold $\epsilon$ with no-progress window off ( $w{=}0$ ). (b) No-progress window $w$ with SOL-headroom threshold fixed at $\epsilon{=}100\%$ . Bars show token and attempts savings. Lines show median and geomean speedup retention.

Figure 8: Scheduler policy Pareto frontiers: normalized dollar cost vs. geomean speedup for nine variants across three per model tier are shown here. Each point is one $(\epsilon,w)$ combination. Marker shape and size encode $\epsilon$ and $w$ , respectively. Star markers show the exhaustive baseline. Hexagons show MI agent reference points. Roofline envelopes trace the upper convex hull of each variant’s frontier.

Figure 9: Best scheduler policy per variant is shown here. For each variant, the $(\epsilon,w)$ combination maximizing efficiency gain is selected subject to $\geq 95\%$ geomean retention. Annotations show the chosen parameters and token savings.

Figure 10: Review outcome composition for each variant. Six-band stacked bars show No Issues (accepted), Minor Issues (accepted), SOL Ceiling (rejected), PyTorch-only (rejected), Original Gaming (rejected), and Inherited Gaming (rejected).

Figure 11: LGD category breakdown for each variant is shown here. Original Gaming and Minor Issues subcategories are stacked.

Figure 12: speedup inflation without the integrity pipeline. Green bars show the filtered results. Orange, light red, and red bars additionally allow PyTorch-only, LGD-labeled gaming solutions, and SOL-ceiling detected solutions, respectively.

	PyTorch-only		Gaming (LGD)
Variant	Run 1	Run 2	Run 1	Run 2
MI	113	8	9	1
$\mu$ CUTLASS + MI	345	51	50	95
SOL-Guided	15	10	3	11
$\mu$ CUTLASS + SOL-Guided	63	26	35	19

Table 4: Prompt-level integrity guardrails on GPT-5-mini. Run 1 used the original prompt. Run 2 added anti-PyTorch-only and anti-gaming instructions.

Figure 13: Run-to-run variation across model tiers. CV is the coefficient of variation ( $\sigma/\mu$ ). GPT-5.2 groups contain $N{=}5$ (w/o $\mu$ CUTLASS) and $N{=}4$ (+ $\mu$ CUTLASS, fewer ablations were run) ablation variants. GPT-5-mini groups contain $N{=}6$ each: 5 ablation variants plus one independent repeat run.

Figure 14: Fast- $p$ curves comparing our $\mu$ CUTLASS + SOL-guided variants across three model tiers against Sakana AI CUDA Engineer (dashed gray). The FP16 SOL curve (dashed purple) shows the theoretical speedup limit.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Level	ID	Rationale for inclusion
L1	1	Basic GEMM baseline.
L1	2	LLM-like GEMM shapes (e.g., M=2048, K=8192, N=4096).
L1	3	Batched matmul (BMM) used in attention score/value products.
L1	4	Matrix–vector multiply representative of single-token decode.
L1	6	Matmul with large $K$ (common in MLP projections).
L1	7	Matmul with small $K$ (e.g., attention head dimension).

Appendix

Table 5: KernelBench problems in our LLM-relevant subset (Levels 1–3).

GPU Kernel Optimization

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

See all 209 papers

Jun 24, 2026

Optimizing CUDA like a Human: Micro-Profiling Tools as Expert Surrogates for LLM-Based GPU Kernel Optimization

We present KernelPro, a closed-loop multi-agent system that automatically generates, profiles, and iteratively optimizes GPU kernel code by integrating large language model (LLM) code generation with hardware profiler feedback and pluggable bottleneck detection tools. KernelPro introduces four contributions: (1) a semantic feedback operator that encodes expert heuristics as pluggable micro-profiling tools, transforming raw hardware metrics into actionable natural language guidance; (2) a two-stage tool invocation architecture where roofline-based bottleneck classification filters which specialized analysis tools execute, combining kernel-level (ncu), instruction-level (SASS), and system-level (nsys) profiling; (3) a domain-adapted MCTS with progressive widening, asymmetric branching, log-reward calibration, dead-end pruning, and search memory for cross-iteration learning; and (4) direct CuTe source-level code generation via autonomous code search over the CUTLASS/CuTe codebase. On KernelBench, KernelPro achieves geometric mean speedups of 2.42x/4.69x/5.30x on Levels 1/2/3, establishing state-of-the-art performance across all difficulty levels. On VeOmni's expert-optimized MoE training kernels, KernelPro achieves 1.23x over hand-tuned Triton by generating a from-scratch raw-CUDA+CuTe Hopper WGMMA kernel. Ablation studies demonstrate that each design component independently and significantly improves optimization quality: micro-profiling tools (p < 0.0001 vs raw metrics), MCTS search (26% higher geometric mean vs greedy, p = 0.004), and proactive tool orchestration (23% improvement, p = 0.035). Finally, KernelPro is the first CUDA kernel coding agent to optimize energy efficiency beyond the speed-only focus of prior systems, demonstrating an 11.6% measured energy reduction at matched speed.

Jiading Gai, Shuai Zhang, Kaj Bostrom +6

Amazon · Siemens · University of Minnesota

GPU Kernel Optimization High-Performance Computing Large Language Model-Guided Optimization Monte Carlo Tree Search GPU Kernel Generation

Ask PDF

Oct 2, 2026

D2K-Bench: Can LLM Agents Turn Expert Designs into Efficient GPU Kernels?

GPU kernels generated by large language model (LLM) agents can remain less efficient than expert implementations, but runtime alone does not reveal how the gap relates to design discovery and implementation. We introduce D2K-Bench, a diagnostic benchmark of 26 tasks and 85 workloads that measures how effectively agents translate expert design guidance into efficient GPU kernels. The guidance covers L1: high-level algorithmic insights, L2: dataflow design, and L3: low-level optimization tricks, including dependencies among these levels. Pairwise runs with and without guidance share task descriptions, workloads, tools, hardware, and a 350-turn budget. Complementary assessments examine independently proposed designs and the design properties implemented in generated code. Across five models on NVIDIA B200 GPUs, guidance raises correctness over 130 model-task pairs from 93.1% to 98.5% and increases the Performance Score over all 26 tasks from 1.46 to 1.95. For the three frontier models with correct submissions on all 26 tasks in both runs (GPT-6-Astra, Claude-Opus-4.8, and GPT-5.6-Sol), geometric mean speedup increases from

1.69\times

2.49\times

. Across all five models, the mean combined implementation score increases from 57 to 70 out of 100. These results show the value of expert design guidance while identifying design properties that remain unimplemented.

Daifeng Li, Huiqiang Jiang, Chengruidong Zhang +6

HKUST · Alibaba Group · USTC

GPU Kernel Generation GPU Kernel Optimization Large Language Model-Guided Optimization AI Agent Benchmarks

Ask PDF

Mar 10, 2026

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

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/.

Qitong Sun, Jun Han, Tianlin Li +6

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

GPU Kernel Optimization GPU Kernel Generation Large Language Model-Guided Optimization Agent Skill Learning

Ask PDF

Abstract

GPU Kernel Optimization

Explore similar work

Optimizing CUDA like a Human: Micro-Profiling Tools as Expert Surrogates for LLM-Based GPU Kernel Optimization

D2K-Bench: Can LLM Agents Turn Expert Designs into Efficient GPU Kernels?

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