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

이 논문은 현대 머신러닝 시스템의 핵심 병목 현상인 고성능 GPU 커널 개발의 어려움을 해결하기 위해 제안된 CuTeGen이라는 새로운 에이전트 프레임워크를 소개합니다. CuTeGen 은 대규모 언어 모델 (LLM) 을 활용하여 GPU 커널을 자동으로 생성하고, 실행 피드백을 기반으로 점진적으로 최적화하는 구조화된 워크플로우를 제공합니다.

1. 문제 정의 (Problem)

고성능 GPU 커널 (예: 행렬 곱셈, 어텐션 연산) 은 알고리즘 구조, 메모리 계층 활용, 하드웨어 특정 최적화 (Tensor Core, 공유 메모리 등) 가 긴밀하게 결합되어 있어 개발이 매우 어렵고 전문가의 수동 노력이 필요합니다.

• 기존 LLM 기반 접근법의 한계: 최근 연구들은 LLM 을 이용해 커널을 자동 생성하려 시도했으나, 생성된 코드는 정합성 (Correctness) 유지가 어렵고, 반복적인 수정 과정에서 성능이 저하되거나 오류가 발생하는 경우가 많았습니다.
• 검색 공간의 복잡성: 커널 최적화는 타일링 (tiling), 데이터 이동, 명령어 선택 등 상호 의존적인 요소들의 조합을 요구합니다. 단일 생성 (one-shot) 이나 무작위 탐색은 이 복잡한 검색 공간에서 최적의 해를 찾기 어렵습니다.

2. 방법론 (Methodology)

CuTeGen 은 커널 개발을 "생성 - 테스트 - 수정 (Generate-Test-Refine)"의 구조화된 반복 과정으로 정의하며, 다음과 같은 핵심 설계 원칙을 따릅니다.

2.1 CuTe 추상화 레이어 활용

CuTeGen 은 원시 CUDA 대신 CuTe(CUDA Template) 라는 C++/CUDA 템플릿 추상화 레이어를 사용하여 커널을 생성합니다.

• 이유: CuTe 는 타일링, 레이아웃, 데이터 이동과 같은 성능에 결정적인 구조를 명시적으로 노출하면서도, 반복적인 수정을 안정적으로 수행할 수 있는 기반을 제공합니다.
• 장점: Triton 같은 고수준 DSL 보다 낮은 수준의 제어력을 유지하면서도, 원시 CUDA 보다 LLM 이 생성하기 쉬운 구조화된 중간 표현 (IR) 역할을 합니다.

2.2 구조화된 에이전트 워크플로우

CuTeGen 은 다음 세 가지 단계를 순환적으로 수행합니다:

1. 정합성 보장 (Correctness Assurance):
   • 초기 생성된 커널을 컴파일하고 참조 구현 (PyTorch) 과 비교하여 정합성을 검증합니다.
   • 오류 발생 시, 단순한 재생성이 아닌 2 단계 디버깅 프로세스를 적용합니다:
     • 진단 (Diagnosis): 컴파일 오류, 런타임 예외, 정합성 불일치 등을 분석하여 원인을 파악합니다.
     • 수리 (Repair): 전체 재작성이 아닌, 국소적인 패치 (patch) 만을 생성하여 적용합니다. 이는 기존 최적화 구조를 유지하면서 오류만 수정하도록 유도합니다.
2. 최적화 (Optimization):
   • 정합성이 확보된 후 성능 최적화 단계로 진입합니다.
   • 구조적 최적화: 커널 유형 (행렬 곱셈 vs 활성화 함수) 에 맞는 최적화 가이드라인을 LLM 에게 제공하여 점진적인 구조 변경을 유도합니다.
   • 프로파일링 기반 정제: NVIDIA Nsight Compute 를 사용하여 성능 지표를 수집합니다.
3. 지연된 프로파일링 통합 (Delayed Profiling Integration):
   • 핵심 아이디어: 복잡한 커널 (예: 행렬 곱셈) 의 경우, 초기 단계에서 프로파일링 데이터를 즉시 제공하면 LLM 이 구조적 설계가 완성되기 전에 타일 크기 등 파라미터만 미세 조정하는 국소 최적점 (Local Optima) 에 빠질 위험이 있습니다.
   • 전략: 먼저 구조적 최적화 (타일링 전략, 메모리 계층 활용 등) 를 통해 안정적인 베이스라인을 확보한 후, 프로파일링 피드백을 도입하여 저수준 파라미터를 정밀 조정합니다.

3. 주요 기여 (Key Contributions)

1. CuTeGen 프레임워크: CuTe 추상화 레이어를 기반으로 한 최초의 반복형 GPU 커널 합성 프레임워크입니다. LLM 을 효율적인 커널이 존재하는 검색 공간으로 유도하면서도 저수준 최적화 능력을 유지합니다.
2. 지연된 프로파일링 통합 기법: 최적화 단계의 타이밍을 작업 부하 (Workload) 에 따라 조절하여, 구조적 설계가 안정화되기 전의 과도한 파라미터 튜닝을 방지하고 전역 최적점을 찾을 수 있도록 합니다.
3. 구조화된 디버깅 및 패치 전략: 전체 재작성 대신 국소적 패치를 통해 정합성과 성능을 동시에 개선하는 메커니즘을 제시합니다.

4. 실험 결과 (Results)

KernelBench 벤치마크 suite 의 12 개 행렬 곱셈 커널과 14 개 활성화 함수 커널에 대해 실험을 수행했습니다.

• 활성화 함수 (Activation Functions): PyTorch 참조 구현 대비 평균 1.70 배의 속도 향상을 달성했습니다. (예: Swish 는 2.45 배, Softsign 는 3.45 배 향상).
• 행렬 곱셈 (Matrix Multiplication): cuBLAS 와 같은 최적화된 라이브러리 구현을 참조로 하여, 2 개의 벤치마크 케이스에서 참조 구현보다 더 높은 성능을 기록했습니다. 전체 평균적으로도 경쟁력 있는 성능을 보였습니다.
• 프로파일링 타이밍 효과: 행렬 곱셈 커널에서 초기부터 프로파일링을 적용한 경우보다, 구조적 최적화 후 (Depth 11) 적용했을 때 성능 향상 곡선이 더 안정적이고 최종 성능이 우수함이 입증되었습니다.

5. 의의 및 결론 (Significance)

CuTeGen 은 고성능 GPU 커널 개발에 있어 전문가의 수동 최적화 없이도 LLM 을 통해 경쟁력 있는 성능을 달성할 수 있음을 입증했습니다.

• 구조적 안정성: CuTe 의 추상화를 통해 LLM 이 복잡한 메모리 계층과 하드웨어 특성을 올바르게 이해하고 적용할 수 있는 기반을 마련했습니다.
• 실용성: 이 프레임워크는 커널 개발을 단순한 코드 생성이 아닌, 실행 피드백에 기반한 체계적인 최적화 과정으로 재정의하여, 실제 AI 시스템 개발 워크플로우에 통합될 수 있는 가능성을 제시합니다.
• 확장성: 행렬 곱셈과 활성화 함수 외에도 어텐션 메커니즘, 리덕션 연산 등 다양한 GPU 연산으로 확장 가능한 일반적인 방법론을 제공합니다.

결론적으로, CuTeGen 은 LLM 기반 코드 생성의 정합성 문제와 성능 한계를 극복하기 위해 적절한 추상화 레이어 (CuTe) 와 전략적인 피드백 루프 (지연된 프로파일링) 를 결합한 성공적인 사례로 평가됩니다.