Making LLMs Optimize Multi-Scenario CUDA Kernels Like Experts

AI summary

Overview

• Research area: Automated GPU (CUDA) kernel optimization using large language model agents, and benchmark design for evaluating it across multiple computational domains.
• Technical level: Advanced. The paper assumes familiarity with GPU programming, hardware profiling tools such as NVIDIA Nsight Compute, and agentic LLM pipelines.
• Scope: The paper introduces MSKernelBench, a 50-task multi-scenario CUDA optimization benchmark with FP32 and BF16 variants, and CUDAMaster, a four-agent system that uses filtered hardware profiling to automatically generate, compile, debug, and optimize CUDA kernels.

What This Paper Is About

Existing automated kernel-optimization work with LLMs is largely limited to deep-learning operators written for frameworks such as PyTorch, which biases evaluation toward regular, computationally intense workloads and toward solutions the model may already have memorized. The authors argue that this narrow scope ignores sparse linear algebra and scientific computing kernels with irregular memory access, where no standard answer exists for the model to recall. Their goal is therefore twofold: build a systematic, multi-scenario benchmark that measures optimization ability at the level of raw C/CUDA rather than framework abstractions, and build an agent system that can reach or exceed hand-tuned library performance across those diverse scenarios.

Key Contributions

1. MSKernelBench, a multi-scenario CUDA optimization benchmark of 50 tasks spanning foundational algebraic operations, common LLM kernels, sparse matrix operators, scientific computing routines, and numerical methods. Each task is implemented in both FP32 and BF16 precision (giving 100 optimization tasks total), with support for scalable data sizes, complexity-weighted scoring, and hardware profiling. Unlike KernelBench, it is written in C/C++ rather than Python and includes sparse matrix and scientific kernels.
2. CUDAMaster, a multi-agent, filtered-profiling-guided, end-to-end framework that generates optimized kernels together with the full toolchain needed to compile and execute them. It comprises four agents (Planner, Coder, Compiler, Debug) and is controlled by iteration rounds R and debug rounds D.
3. A hardware analysis filter that classifies each kernel into one of three bottleneck types using data-driven thresholds derived with Otsu's method, then forwards only the relevant profiling metrics to the LLM agents.
4. Empirical evidence that the approach delivers substantial speedups on most operators, outperforming Astra by about 35%, and in some cases matching or exceeding closed-source libraries such as cuBLAS, cuDNN, and cuSPARSE.

Main Findings

• Benchmark composition: Of the 50 benchmark tasks, 11 are Compute Bound, 24 are Memory Latency Bound, and 15 are Memory Bandwidth Bound, as determined by the automated classification. Compute Bound tasks cluster around dense linear algebra, stencil computation, and core LLM operators; Memory Latency Bound tasks around normalization, loss, and activation functions; Memory Bandwidth Bound tasks around sparse matrix computations such as SpMV and tensor processing routines.
• Threshold calibration: Using NVIDIA Nsight Compute (v2024.1) profiling data and Otsu's method, the optimal classification thresholds for Compute (SM) Throughput, DRAM Throughput, and Memory Throughput all converged near 30%, yielding rules of Compute Bound (Compute SM Throughput > 30%), Memory Latency Bound (Compute < 30%, DRAM < 30%, Memory < 30%), and Memory Bandwidth Bound as the default otherwise.
• Model comparison: On the cumulative success rate across thresholds, OpenAI o4-mini consistently led DeepSeek-V3.2: 100% vs. 95% at τ=0, 94% vs. 80% at τ=1, and 60% vs. 49% at τ=2.
• Precision effects: BF16 gave o4-mini a slight edge at moderate-to-high thresholds (τ ≤ 8), while FP32 provided greater stability for both models under stringent optimization targets (τ ≥ 32).
• Against closed-source libraries: Both models outperformed cuSPARSE on SpMV CSR, and achieved up to 1.8× speedup over cuDNN in 2D Convolution and over cuBLAS in Dot Product.
• Against Astra: On Astra-optimized SGLang kernels, o4-mini led by around 35% on RMSNorm and matched or surpassed Astra on the other two fused operators tested (SiLU & Mul, and Merge Attention States).
• Ablation on iterations and debugging: The Full configuration (R=3, D=3) achieved the highest success rates. With o4-mini at τ=1, Full reached 94%, versus 77% for No Debug, 74% for Single Iteration, and 77% for Single Run. At τ ≥ 8, Full maintained 25% versus 17–19% for the ablated variants. Degrading R had a more pronounced effect than eliminating debugging.
• Ablation on profiling strategy: For o4-mini, the Filtered configuration matched Full Profile performance across all thresholds while reducing overhead; at τ=1 both reached 94–95%, above the 90% of No Profile. At τ ≥ 8, Filtered maintained 25% versus 22% for No Profile.
• Cost and token efficiency: With o4-mini, Filtered Profile cost $0.2707 and 123,036 tokens per task on average, versus $0.2120 and 96,373 tokens for No Profile and $0.3804 and 172,917 tokens for Full Profile. With DeepSeek-V3.2, Filtered cost $0.0310 and 223,294 tokens, versus $0.0277 and 167,971 for No Profile and $0.0687 and 387,432 for Full Profile. The authors report filtering cuts cost by up to 32% and token usage by 30–40% relative to full profiling.

Methodology in Plain English

The authors first built a benchmark by hand. They curated 50 tasks from authoritative sources including official NVIDIA documentation, common cuBLAS and cuSPARSE use cases, the open-source KernelBench, and the LeetGPU cloud learning platform. Because well-optimized industrial libraries are closed-source, they wrote each task in pure C or CUDA themselves and validated numerical correctness against a trusted reference, such as a closed-source optimized library or a serial CPU version, before adopting it as a performance baseline. Each task ships with distinct FP32 and BF16 implementations behind a unified C wrapper and a shared test template, so a researcher only has to replace a clearly demarcated optimization segment.

Evaluation has three gates. The submitted code must compile against a fixed C function interface, its output must match the ground-truth kernel within tolerance across multiple predefined data sizes, and its runtime is measured after 3 warm-up runs followed by 50 executions averaged together. The final score is a complexity-weighted average of speedups, where each data size's weight is proportional to the theoretical computational complexity of the unoptimized baseline at that size — so larger, more expensive workloads count more, and genuine algorithmic improvements are amplified.

For the optimization system, the authors mirrored how a human expert works but removed the two main bottlenecks: manually sifting through enormous profiling output, and repeatedly context-switching between debugging, analysis, and coding. Their hardware analysis filter collects full profiling data once, uses Otsu's method to calibrate objective bottleneck thresholds, classifies each kernel, and forwards only the metrics relevant to that bottleneck class. A loop then runs for R=3 rounds: the Planner Agent proposes a high-level strategy from the distilled profiling, the Coder Agent implements it inside the test harness, and the Compiler Agent handles build and execution commands. If correctness fails, a Debug Agent gets up to D=3 rounds to diagnose and repair the code. Only valid kernels that improve on the best score are kept, and the benchmark's core test logic is never modified by any agent, keeping comparisons fair.

Why This Matters

The paper challenges the assumption that LLM kernel optimization is a solved-ish problem for deep learning operators and reframes it as a general-purpose, multi-scenario engineering challenge. It shows that agent systems can reach expert-level tuning on irregular, non-LLM workloads where no memorized answer exists — a stronger test of genuine optimization capability — and it provides an open benchmark and framework for others to build on.

Real-world applications:

• Scientific computing and simulation: Sparse matrix-vector multiplication, stencil computation, and numerical methods underpin fluid dynamics, climate modeling, and physics simulation, where the benchmark's sparse and scientific tasks are directly relevant.
• LLM inference serving: Optimized RMSNorm, SiLU & Mul, and Merge Attention States kernels directly affect the throughput and latency of deployed language models.
• AI hardware and library development: The ability to generate kernels competitive with closed-source libraries (cuBLAS, cuDNN, cuSPARSE) points toward faster generation of tuned libraries for new hardware generations, where per-operator, per-generation engineering is currently expensive.
• Cost reduction in GPU-accelerated pipelines: Both the runtime speedups and the reported API cost and token savings (up to 32% cost, 30–40% token reduction from profiling filtering) matter for organizations running large-scale automated tuning.

Industry relevance centers on the tension the paper identifies between hand-tuned vendor libraries, which set the performance ceiling but lack flexibility and cost enormous engineering effort per operator, and compiler-based approaches like TVM and Triton, which improve productivity but struggle to match expert performance across diverse domains. CUDAMaster's claim to match or exceed closed-source libraries in several cases speaks directly to that gap.

Future Directions

• Scaling to more scenarios and hardware: The benchmark's bottleneck distribution (11 Compute Bound, 24 Memory Latency Bound, 15 Memory Bandwidth Bound) suggests uneven coverage; extending task coverage and validating the agent across hardware beyond the single NVIDIA RTX 4090 used here is an open question.
• Systematic evaluation for multi-scenario settings is a stated motivation but the paper reports results on 50 tasks and two LLMs; whether the approach generalizes to other models, other precisions, and other GPU architectures is not established.
• Reducing dependence on expensive profiling: The Full Profile configuration consumes the most tokens and cost, and the paper reports that iteration count (R) has a stronger effect than debugging (D). Finding cheaper ways to obtain equivalent diagnostic signal remains an open problem.
• Bridging toward algorithmic, not just implementation, improvements: The scoring metric is designed to reward reducing runtime complexity (for example from O(N²) to O(N log N)), but the paper does not report whether agents actually produced such algorithmic changes rather than implementation-level speedups.

Target Audience

This paper is most useful to researchers and engineers working on automated code generation and GPU kernel optimization, including those building LLM agent systems for systems-level programming tasks. It also serves practitioners in high-performance computing and scientific computing who want to know whether automated tuning can handle irregular, sparse, and non-deep-learning workloads. Benchmark designers will find the comparison against KernelBench and the details of the complexity-weighted scoring protocol particularly relevant, while ML systems engineers evaluating whether agent-generated kernels can replace or augment hand-tuned library calls will find the cuBLAS, cuDNN, cuSPARSE, and Astra comparisons most directly applicable.

Authors’ abstract

Optimizing GPU kernels manually is a challenging and time-consuming task. With the rapid development of LLMs, automated GPU kernel optimization is gradually becoming a tangible reality. However, current LLM-driven automated optimization methods narrowly focus on machine learning applications, such as PyTorch operator optimization, while overlooking broader domains like sparse matrix operations in scientific computing. Extending to these broader applications brings new challenges for the benchmark and algorithm. Therefore, developing a general-purpose automated kernel optimization method becomes our primary focus. In this paper, we address the absence of systematic evaluation for multi-scenario settings by introducing MSKernelBench, which spans multiple scenarios, including fundamental algebraic operations, common LLM kernels, sparse matrix operators, and scientific computing routines, each supporting both FP32 and BF16 precision. Building on this benchmark, we introduce CUDAMaster, a multi-agent, hardware-aware system for kernel optimization that leverages profiling information and automatically constructs the full compilation and execution toolchain. Experimental results demonstrate that CUDAMaster achieves significant speedups across most operators, outperforming Astra by about 35%. In several cases, its performance matches or surpasses that of highly optimized, closed-source libraries such as cuBLAS. A demo showcasing the original and optimized code for each operator is available at https://hanyx2021.github.io/MSKernelBenchDemo/.