CuTeGen: An LLM-Based Agentic Framework for Generation and Optimization of High-Performance GPU Kernels using CuTe

1. Problem Statement

Developing high-performance GPU kernels (e.g., for matrix multiplication and activation functions) is a bottleneck in modern AI systems. While Large Language Models (LLMs) have shown promise in code generation, they struggle with GPU kernels due to:

• Complexity of Optimization: High performance requires tightly coupled decisions regarding hardware instructions (tensor cores), memory hierarchy (tiling, shared memory), and data movement.
• Fragility of Iterative Refinement: Naive iterative editing often breaks correctness or causes performance regressions because changes in one area (e.g., tile shape) cascade into others (e.g., shared memory layout).
• Limitations of One-Shot Generation: Generating a fully optimized kernel in a single pass is unlikely to succeed given the vast, coupled search space.
• Representation Gap: Existing LLM approaches often use raw CUDA (too low-level/complex) or high-level DSLs like Triton (too abstract, limiting low-level control).

2. Methodology: The CuTeGen Framework

CuTeGen is an agentic framework that treats kernel development as a structured, iterative generate–test–refine workflow. It does not rely on large-scale search or model retraining but focuses on the progressive refinement of a single evolving kernel.

A. Core Abstraction: CuTe

A pivotal design choice is generating kernels using CuTe, a C++/CUDA template abstraction layer (part of NVIDIA CUTLASS).

• Why CuTe? It exposes performance-critical structures (tiling, layouts, data movement) while providing a stable, structured representation.
• Advantage: Unlike raw CUDA, CuTe makes dimensional relationships and tiling explicit, reducing ambiguity for the LLM. Unlike higher-level DSLs, it retains the low-level control necessary for inserting inline PTX and optimizing hardware-specific features.

B. The Agentic Workflow

The system operates in three distinct phases:

1. Initial Synthesis: The LLM generates an initial CuTe-based kernel based on a PyTorch reference implementation and a task specification.
2. Correctness Assurance (Testing & Debugging):
   • The kernel is compiled and executed against randomized inputs.
   • Two-Stage Debugging: If errors occur, the system separates Diagnosis (analyzing compiler/runtime errors) from Repair (generating localized patches).
   • Constraint: The LLM is strictly guided not to simplify the kernel (e.g., by removing tiling or falling back to library calls) but to fix the specific structural errors while preserving the optimization intent.
3. Optimization: Once correctness is established, the system enters a performance tuning loop.
   • Prompt-Guided Structural Optimization: The LLM applies domain-specific guidelines (e.g., tiling strategies for GEMM, vectorization for activations) to introduce one optimization at a time.
   • Delayed Profiling Integration: A key innovation is not exposing hardware profiling metrics (via NVIDIA Nsight Compute) immediately.
     • For Complex Kernels (e.g., MatMul): Profiling is delayed until the kernel has a stable structural design (e.g., correct tiling). Early profiling often leads to "myopic" parameter tuning (e.g., tweaking tile sizes) before the algorithmic structure is sound.
     • For Simple Kernels (e.g., Activations): Profiling is introduced earlier as the design space is less structurally complex.

3. Key Contributions

1. CuTeGen Framework: An iterative synthesis system that uses execution feedback to refine kernels, separating debugging from optimization to ensure functional correctness before performance tuning.
2. CuTe as a Target Representation: Demonstrating that using the CuTe abstraction layer guides LLMs toward a search space rich in efficient kernels while maintaining the low-level control needed for expert-level optimization (e.g., tensor cores, async copies).
3. Delayed Profiling Integration: A strategy to prevent premature convergence to local optima by withholding hardware metrics until the kernel's structural foundation is solid.
4. Structured Debugging Guidelines: A set of operational rules that prevent the LLM from "de-optimizing" the kernel (e.g., removing tiling) to fix bugs, forcing it to reason about the root cause of structural errors.

4. Experimental Results

The framework was evaluated on 12 matrix multiplication and 14 activation function kernels from the KernelBench suite, compared against PyTorch reference implementations.

• Activation Kernels: CuTeGen achieved an average speedup of 1.70× over PyTorch references.
• Matrix Multiplication:
  • Achieved competitive performance with optimized library implementations (cuBLAS) in several cases.
  • Outperformed the cuBLAS reference implementation in 2 out of 12 benchmark cases (e.g., "MatMul with Diagonal Matrices" showed a massive 17.66× speedup).
  • Demonstrated that automated synthesis can approach expert-level performance without manual engineering.
• Ablation on Profiling: Experiments confirmed that delaying profiling feedback for matrix multiplication kernels led to significantly better optimization trajectories compared to immediate feedback, which often resulted in unstable or suboptimal performance.

5. Significance and Conclusion

CuTeGen represents a shift from "one-shot" code generation to structured, agentic refinement. By combining the CuTe abstraction layer with a delayed profiling strategy, the system successfully navigates the complex trade-offs of GPU kernel optimization.

• Practical Impact: It reduces the reliance on manual, expert-driven kernel engineering for performance-critical operators.
• Generalizability: The approach is not limited to GEMM; the structured refinement loop and workload-aware strategies are applicable to a broad class of GPU operators, including attention mechanisms and reductions.
• Future Direction: The paper suggests that integrating such frameworks into standard development workflows can significantly accelerate the deployment of high-performance AI systems.

In summary, CuTeGen proves that with the right abstraction layer (CuTe) and a disciplined, feedback-driven workflow, LLMs can autonomously generate GPU kernels that rival or exceed hand-tuned expert implementations.