Making LLMs Optimize Multi-Scenario CUDA Kernels Like Experts

This paper introduces MSKernelBench, a comprehensive benchmark covering diverse multi-scenario GPU kernels, and CUDAMaster, a multi-agent, hardware-aware system that leverages this benchmark to achieve significant speedups, often matching or surpassing closed-source libraries like cuBLAS, thereby advancing general-purpose automated CUDA kernel optimization beyond current ML-focused methods.

The Gist

Imagine you have a super-fast sports car (your GPU), but it's stuck in traffic because the driver (the software code) doesn't know the best route. Usually, you need a world-class racing engineer to manually tweak the engine, change the tires, and map out the perfect path to win the race. This is what optimizing "CUDA kernels" (the code that runs on graphics cards) has always been like: hard, expensive, and done by a tiny group of experts.

This paper introduces a new way to solve this problem using AI (Large Language Models) to act as that racing engineer, but with a twist: instead of just being good at one type of race (like deep learning), this AI can handle any kind of race, from scientific simulations to complex math.

Here is the breakdown of their solution, MSKernelBench and CUDAMaster, using simple analogies.

1. The Problem: The "Deep Learning" Tunnel Vision

Previously, AI tools trying to optimize code were like driving instructors who only taught you how to drive on a specific, smooth highway (Deep Learning applications like PyTorch). They were great at that one road but got completely lost if you asked them to drive on a bumpy dirt track (Scientific Computing) or through a narrow city alley (Sparse Matrix operations).

The researchers realized that to make AI truly useful, it needed to learn how to drive every type of road, not just the highway.

2. The Solution Part 1: MSKernelBench (The "Grand Prix" Test Track)

To teach the AI properly, you need a good test track. The authors built MSKernelBench, which is like a massive, multi-terrain driving course.

• The Terrain: Instead of just one smooth road, this track has:
  • Dense Highways: Standard math operations (like multiplying big matrices).
  • Dirt Roads: Sparse operations (where data is scattered and messy, common in science).
  • City Streets: Operations used in Large Language Models (LLMs).
  • Off-Road: Scientific simulations and stencil computations.
• The Rules: They made sure the AI had to drive this track in two different weather conditions (FP32 and BF16 precision) and at different speeds (different data sizes).
• Why it matters: If an AI can win a race on this diverse track, it proves it's a true expert, not just a memorizer of one specific route.

3. The Solution Part 2: CUDAMaster (The "Pit Crew" of AI Agents)

Once they had the test track, they built CUDAMaster, a team of AI agents working together like a professional Formula 1 pit crew. Instead of one AI trying to do everything, they split the job up:

1. The Scanner (Hardware Filter): Before the AI even touches the code, this agent looks at the car's dashboard (hardware profiling data). It asks: "Is the engine overheating (Compute Bound)? Is the car waiting for fuel (Memory Latency)? Or is the road too narrow for the tires (Memory Bandwidth)?" It filters out the noise so the team only sees the real problem.
2. The Strategist (Planner Agent): Based on the dashboard, this agent comes up with a game plan. "Okay, the car is waiting for fuel. Let's change the fuel injection timing."
3. The Mechanic (Coder Agent): This agent actually rewrites the code (the engine parts) based on the plan.
4. The Race Director (Compiler Agent): This agent makes sure the new engine parts fit the car and the car can actually start (compilation and execution).
5. The Inspector (Debug Agent): If the car stalls or crashes, this agent finds the bug and fixes it immediately.

The Magic Loop: This team doesn't just try once. They run the car, check the time, fix a part, run it again, and repeat this process dozens of times until they find the absolute fastest version.

4. The Results: Beating the Pros

The researchers put their AI team against the best human-engineered libraries (like cuBLAS and cuSPARSE), which are the "Ferraris" of the current world, built by NVIDIA's top engineers over many years.

• The Outcome: The AI team didn't just keep up; in many cases, they beat the human experts.
• The Speed: On average, their AI was 35% faster than other AI optimization tools (like Astra).
• The Shock: For some specific tasks, the AI-generated code was actually faster than the official, closed-source libraries that have been perfected for decades.

The Big Picture Takeaway

Think of this paper as the moment AI went from being a "novice driver" who only knows how to drive on a straight highway to a "Grand Prix Champion" who can handle any terrain, any weather, and any car.

They proved that with the right test track (MSKernelBench) and the right pit crew strategy (CUDAMaster), AI can now automatically tune complex computer code to run as fast as, or even faster than, the best human experts in the world. This opens the door for faster scientific discoveries, better AI models, and more efficient software, all generated automatically.

Technical Summary

Here is a detailed technical summary of the paper "Making LLMs Optimize Multi-Scenario CUDA Kernels Like Experts".

1. Problem Statement

Optimizing GPU kernels manually is a labor-intensive, time-consuming process requiring deep hardware expertise. While Large Language Models (LLMs) have shown promise in automating software engineering, current automated optimization efforts (e.g., KernelBench) suffer from two critical limitations:

• Narrow Scope: Existing benchmarks focus almost exclusively on high-level deep learning frameworks (like PyTorch) and common LLM operators. They fail to cover broader High-Performance Computing (HPC) domains such as sparse matrix operations, scientific simulations, and irregular memory access patterns.
• Lack of Generalizability: Current methods often rely on the model's ability to recall known solutions for standard operators rather than testing true optimization capabilities on "open questions" without standard answers. Furthermore, existing benchmarks often use fixed data sizes and framework-level abstractions (Python/PyTorch), which obscure low-level hardware bottlenecks and limit scalability evaluation.

The core challenge is to develop a general-purpose, automated system capable of optimizing diverse CUDA kernels across multiple scenarios (dense, sparse, scientific, LLM) to match or exceed the performance of hand-tuned, closed-source libraries (e.g., cuBLAS, cuSPARSE).

2. Methodology

The authors propose a two-pronged approach: a comprehensive benchmark (MSKernelBench) and a multi-agent optimization framework (CUDAMaster).

A. MSKernelBench: A Multi-Scenario Benchmark

To address the lack of systematic evaluation, the authors introduced MSKernelBench, a benchmark designed to strip away framework dependencies and evaluate optimization at the fundamental level.

• Scope: It comprises 50 distinct tasks spanning four categories:
  1. Dense Linear Algebra: (e.g., Dot Product, Matrix Multiplication).
  2. Sparse Matrix Operations: (e.g., SpMV/SpMM in CSR, CSC, COO, ELL formats).
  3. LLM Operators: (e.g., RMS Norm, SiLU, Attention mechanisms).
  4. Scientific Computing: (e.g., Stencil computations, 2D/3D convolutions, numerical integration).
• Implementation: Written in pure C/CUDA to ensure portability and low-level control, avoiding PyTorch overhead.
• Precision & Scale: Supports both FP32 and BF16 precisions. It evaluates kernels across multi-scale data sizes (ranging from $2^{10}$to$2^{22}$) to test scalability.
• Evaluation Metric: Uses a complexity-weighted speedup metric. Unlike simple average speedups, this metric weights performance gains by the theoretical computational complexity ($T(N)$) of the baseline, ensuring that optimizations on larger, more complex workloads contribute more significantly to the final score.

B. CUDAMaster: A Multi-Agent Optimization System

CUDAMaster is an end-to-end, hardware-aware multi-agent system that automates the kernel tuning process. It mimics the workflow of a human expert through four specialized agents:

1. Hardware Analysis Filter:
   • Collects profiling data using NVIDIA Nsight Compute.
   • Classifies kernels into three bottleneck types using Otsu's method on throughput metrics: Compute Bound, Memory Latency Bound, and Memory Bandwidth Bound.
   • Key Innovation: It filters profiling data to extract only the metrics relevant to the specific bottleneck type, reducing noise and context size for the LLM.
2. Planner Agent: Analyzes the distilled profiling insights and historical performance to propose high-level optimization strategies (e.g., loop unrolling, tiling, shared memory usage).
3. Coder Agent: Implements the proposed strategies into executable CUDA code, adhering to strict interface requirements.
4. Compiler Agent: Manages the compilation toolchain, generating optimized nvcc commands and handling linking.
5. Debug Agent: Invoked if execution fails; it diagnoses compilation errors, CUDA runtime errors, or numerical correctness issues and suggests fixes.

The system operates in an iterative loop ($R$ rounds) with a maximum of $D$ debugging attempts per candidate, continuously refining the code until the best-performing, correct kernel is found.

3. Key Contributions

1. MSKernelBench: The first comprehensive benchmark for multi-scenario CUDA optimization, covering dense, sparse, LLM, and scientific kernels with FP32/BF16 support and multi-scale evaluation.
2. CUDAMaster: A novel multi-agent framework that integrates hardware profiling filtering with iterative planning, coding, and debugging to autonomously generate optimized toolchains.
3. Hardware-Aware Filtering: A data-driven approach to classify bottlenecks and filter profiling metrics, significantly improving the efficiency and accuracy of LLM-guided optimization.
4. Open-Source Release: The authors open-sourced the benchmark, framework, and a demo of optimized code to facilitate future research.

4. Experimental Results

The system was evaluated on 50 tasks (100 total with FP32/BF16 variants) using OpenAI o4-mini and DeepSeek-V3.2 on an NVIDIA RTX 4090.

• Overall Performance: CUDAMaster achieved significant speedups across most operators, outperforming the state-of-the-art Astra framework by approximately 35%.
• Comparison with Closed-Source Libraries:
  • In several cases, the generated kernels matched or surpassed highly optimized, closed-source libraries like cuBLAS and cuSPARSE.
  • Example: For Dot Product, the system achieved a 26.09x speedup over the naive baseline, compared to cuBLAS's 26.09x (effectively matching it) and significantly outperforming other baselines.
  • Example: For SpMV (CSR), it achieved 2.96x speedup vs. cuSPARSE's 2.23x.
• Ablation Studies:
  • Iterative Planning & Debugging: Removing debugging or limiting to single iterations drastically reduced success rates (e.g., dropping from 94% to 77% at $\tau=1$ for o4-mini), proving the necessity of the multi-round refinement loop.
  • Profile Filtering: The "Filtered" profile strategy achieved performance comparable to "Full Profile" (94% success rate at $\tau=1$) but reduced API costs by 32% and token usage by 30–40%, demonstrating an optimal trade-off between guidance and cost.
• Bottleneck Transformation: The system successfully transformed many memory-bound tasks into compute-bound or bandwidth-bound tasks, effectively alleviating the most severe performance constraints (e.g., reducing memory latency-bound tasks by 67%).

5. Significance

This work represents a paradigm shift in automated GPU programming:

• Beyond Deep Learning: It proves that LLMs can optimize kernels beyond the narrow scope of deep learning operators, tackling the irregular and complex patterns found in scientific computing and sparse linear algebra.
• Expert-Level Automation: It demonstrates that with the right environment (filtered hardware data) and iterative refinement, LLM agents can approximate or even exceed the performance of human experts and proprietary libraries.
• Scalable Framework: By decoupling optimization from high-level frameworks and using a complexity-weighted metric, the approach provides a more reliable and scalable method for evaluating automated kernel engineering.
• Future Impact: The open-sourcing of MSKernelBench and CUDAMaster sets a new standard for benchmarking and provides a foundation for developing adaptive, general-purpose high-performance code generation systems.