Technical Summary: KernelPro – Optimizing CUDA like a Human

1. Problem Statement

Optimizing GPU kernels requires a disciplined, iterative cycle of profiling, bottleneck recognition, root cause diagnosis, and targeted code transformation. While experienced engineers excel at this by interpreting raw hardware metrics through refined heuristics, recent Large Language Model (LLM)-based optimization systems have struggled to replicate this expertise.

Current approaches typically feed raw or lightly summarized hardware profiler metrics (e.g., Nsight Compute counters) directly into the LLM prompt, relying on the model to implicitly infer diagnostic patterns and optimization strategies. This approach conflates two distinct capabilities: interpreting structured hardware telemetry (a rule-governed task) and generating optimized code (a creative task). Consequently, existing agents often miss optimizations, produce inconsistent diagnoses, and fail to incorporate new analysis patterns without retraining. Furthermore, prior systems focus almost exclusively on latency/speed, neglecting energy efficiency, which is increasingly critical for AI scaling.

2. Methodology: The KernelPro Architecture

KernelPro is a closed-loop multi-agent system designed to separate the concerns of hardware interpretation and code generation. It integrates LLM code generation with pluggable micro-profiling tools that act as expert surrogates, transforming raw metrics into actionable natural language guidance.

2.1 Core Agents and Pipeline

The system operates through three primary agents coordinated by a search orchestrator:

1. Benchmarking Agent (Stage 1): Performs a one-time analysis of the reference implementation. It identifies mathematical operations, computes a roofline estimate based on compute intensity and hardware specs, and classifies the kernel's bottleneck type (compute-bound, memory-bound, latency-bound, or mixed).
2. Profiling Agent (Stage 2): Executes a suite of 15 micro-profiling tools filtered by the Stage 1 bottleneck classification. These tools analyze metrics from three sources:
   • ncu (Nsight Compute): Kernel-level metrics (throughput, stalls, occupancy).
   • nsys (Nsight Systems): System-level timeline data (launch overhead, synchronization).
   • SASS: Instruction-level analysis via disassembly of the compiled binary.
     Instead of forwarding raw data, the Profiling Agent synthesizes a semantic feedback prompt containing severity ratings, root cause explanations, and prioritized recommendations.
3. Programmer Agent: Generates optimized CUDA code in three modes: initial generation, debugging (fixing compilation/correctness errors), and iterative optimization based on the semantic feedback. It supports multi-language inputs (PyTorch, Triton, CUDA) and integrates directly with CUTLASS/CuTe libraries for source-level code generation.

2.2 Key Architectural Innovations

• Semantic Feedback Operator: Tools are formalized as 5-tuples (name, required metrics, thresholds, trigger, analysis) that encode expert heuristics. For example, a tool might trigger if Tensor Core utilization is below 10%, analyze the cause, and recommend switching to a CUTLASS GEMM implementation.
• Proactive Tool Orchestration: Unlike reactive agents where the LLM decides which tools to call (risking stochastic omission), KernelPro deterministically executes all relevant tools based on the bottleneck classification before invoking the LLM. This ensures comprehensive coverage of the bottleneck.
• Domain-Adapted MCTS: The system uses Monte Carlo Tree Search (MCTS) with specific adaptations for the high-cost nature of kernel compilation and profiling:
  • Progressive Widening: Gates node expansion to ensure exploration grows sublinearly with visit count.
  • Asymmetric Branching: Allocates more children to successful nodes (optimization refinements) than failed nodes (repair attempts).
  • Log-Reward Calibration: Uses a logarithmic speedup reward to prevent outliers from dominating the search tree.
  • Search Memory: A persistent store that aggregates findings (file paths, API patterns, error corrections) across iterations to prevent the LLM from repeating past mistakes.
• Autonomous Code Search: The Programmer Agent can search the CUTLASS/CuTe codebase directly to retrieve layout algebras and copy atoms, mimicking the workflow of expert engineers writing raw CUDA+CuTe kernels.

2.3 Energy-Aware Optimization

KernelPro introduces a lexicographic reward function where speedup is the primary objective and energy efficiency is a secondary objective used to break ties. Energy is estimated via a deterministic proxy model based on DRAM bytes, uncoalesced loads, register spills, and instruction counts, validated against NVML measurements.

3. Key Contributions

The paper identifies four primary technical contributions:

1. Semantic Feedback Operator: A mechanism to encode expert heuristics as pluggable micro-profiling tools, transforming raw hardware metrics into natural language guidance before the LLM sees them.
2. Two-Stage Tool Invocation: A pipeline where roofline-based bottleneck classification filters which specialized analysis tools (ncu, nsys, SASS) execute, reducing prompt noise.
3. Domain-Adapted MCTS: A search strategy featuring progressive widening, asymmetric branching, log-reward calibration, dead-end pruning, and search memory for cross-iteration learning.
4. Direct Source-Level Generation: The ability to generate raw CUDA+CuTe code by autonomously searching the CUTLASS/CuTe codebase, rather than relying on high-level template instantiation or library stitching.
5. Energy-Aware Agent: The first CUDA kernel coding agent to optimize for energy efficiency alongside speed, demonstrating measurable energy reductions at matched speeds.

4. Experimental Results

Evaluated on KernelBench (250 tasks across three difficulty levels) and a production VeOmni MoE training kernel:

• KernelBench Performance: KernelPro achieves geometric mean speedups of 2.42× (Level 1), 4.69× (Level 2), and 5.30× (Level 3), surpassing the prior state-of-the-art (KernelBlaster) across all levels. It achieves 100% task coverage (100/100 on L1/L2, 50/50 on L3).
• Production Validation: On a hand-tuned Triton MoE weight-gradient kernel (H100), KernelPro generated a from-scratch raw-CUDA+CuTe WGMMA kernel achieving a 1.23× speedup over the expert baseline.
• Ablation Studies:
  • Micro-profiling tools: Provide a 125% higher speedup compared to using raw metrics alone (p < 0.0001). Notably, raw metrics alone performed worse than no feedback at all, confirming that unstructured data degrades optimization quality.
  • MCTS: Outperforms greedy search by 26% in geometric mean speedup (p = 0.004), particularly on tasks requiring multi-step optimization chains.
  • Proactive Orchestration: Yields a 23% improvement over reactive tool invocation (p = 0.035).
  • Energy Efficiency: In a matched-speed A/B test on the Swish kernel, the energy-aware agent reduced dynamic energy by 11.6% by selecting lower-energy instruction streams (e.g., using __fdividef instead of IEEE-accurate division) that speed-only agents ignored.

5. Significance and Claims

The paper argues that the primary bottleneck in LLM-based GPU optimization is not the code generation capability of the LLM, but how profiling data reaches the model. By replacing implicit reasoning with explicit, systematic diagnostic workflows via micro-profiling tools, KernelPro grounds the LLM in authoritative library sources and expert heuristics.

The authors claim that KernelPro establishes a new paradigm where:

• Interpretation is more critical than data access: The translation of raw counters into actionable directives is the key driver of performance.
• Source-level generation is feasible: LLMs can successfully generate complex, raw CUDA+CuTe kernels by autonomously searching library codebases, moving beyond simple template instantiation.
• Energy is a viable optimization target: Optimization agents can effectively reduce energy consumption at fixed speeds by selecting instruction streams that minimize hardware activity without impacting latency.

The work positions KernelPro as a step toward "optimizing like a human," where the system mimics the disciplined, heuristic-driven cycle of expert engineers rather than relying on the LLM to implicitly rediscover these patterns from raw data.