CuTeGen: CuTe를 활용한 고성능 GPU 커널 생성 및 최적화를 위한 LLM 기반 에이전트 프레임워크

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

고성능 GPU 커널은 현대 머신 러닝 시스템에서 매우 중요하지만, 알고리즘 구조, 메모리 계층 구조 사용, 그리고 하드웨어 특정 최적화 간의 복잡한 상호 작용으로 인해 효율적인 구현을 개발하는 것은 여전히 어렵고 전문가의 지식을 필요로 하는 과정입니다. 최근 연구에서는 대규모 언어 모델(LLM)을 사용하여 GPU 커널을 자동으로 생성하는 방법을 탐구했지만, 생성된 구현은 종종 정확성을 유지하고 반복적인 개선을 통해 경쟁력 있는 성능을 달성하는 데 어려움을 겪습니다. 본 논문에서는 CuTeGen을 제안합니다. CuTeGen은 GPU 커널의 자동 생성 및 최적화를 위한 에이전트 프레임워크이며, 커널 개발을 구조화된 생성-테스트-수정 워크플로우로 처리합니다. CuTeGen은 단일 생성 또는 후보 구현에 대한 대규모 검색에 의존하는 기존 방식과 달리, 실행 기반 검증, 구조화된 디버깅, 그리고 단계별 최적화를 통해 단일 커널을 점진적으로 개선하는 데 중점을 둡니다. 핵심 설계 요소는 CuTe 추상화 레이어를 사용하여 커널을 생성하는 것입니다. CuTe는 타일링 및 데이터 이동과 같은 성능에 중요한 구조를 노출하면서 반복적인 수정에 대한 더 안정적인 표현을 제공합니다. CuTeGen은 성능 향상을 위해 워크로드 기반 최적화 프롬프트를 사용하고 프로파일링 피드백을 지연 통합합니다. 행렬 곱셈 및 활성화 워크로드에 대한 실험 결과는 CuTeGen이 기능적으로 올바른 커널을 생성하고 최적화된 라이브러리 구현과 경쟁력 있는 성능을 달성한다는 것을 보여줍니다.

High-performance GPU kernels are critical to modern machine learning systems, yet developing efficient implementations remains a challenging, expert-driven process due to the tight coupling between algorithmic structure, memory hierarchy usage, and hardware-specific optimizations. Recent work has explored using large language models (LLMs) to generate GPU kernels automatically, but generated implementations often struggle to maintain correctness and achieve competitive performance across iterative refinements. We present CuTeGen, an agentic framework for automated generation and optimization of GPU kernels that treats kernel development as a structured generate--test--refine workflow. Unlike approaches that rely on one-shot generation or large-scale search over candidate implementations, CuTeGen focuses on progressive refinement of a single evolving kernel through execution-based validation, structured debugging, and staged optimization. A key design choice is to generate kernels using the CuTe abstraction layer, which exposes performance-critical structures such as tiling and data movement while providing a more stable representation for iterative modification. To guide performance improvement, CuTeGen incorporates workload-aware optimization prompts and delayed integration of profiling feedback. Experimental results on matrix multiplication and activation workloads demonstrate that the framework produces functionally correct kernels and achieves competitive performance relative to optimized library implementations.