- AI가 생성한 CUDA-C 커널들이 PyTorch의 전문가 최적화 커널과 비슷하거나 더 나은 성능을 보임
- 단일 LLM(대형언어모델)이 자연어 최적화 아이디어 생성과 다양한 코드 브랜칭을 반복, 기존 방법보다 최적화 다양성과 병렬 탐색에서 뛰어난 성능 달성
- 대표 벤치마크 결과, Matmul(101%), Conv2D(179.9%), Softmax(111.8%), LayerNorm(484.4%), Conv2D+ReLU+MaxPool(290.1%) 등에서 PyTorch 대비 압도적
- 기존 “순차적 커널 개선”의 한계를 넘기 위해 자연어 추론 + 브랜칭 구조 적용, 빠른 속도로 고성능 커널을 생성
- FP16, Flash Attention 등 최신 ML 워크로드에서도 진보 중이며, 미래에는 AI가 자체적으로 더 빠른 커널을 탐색·개선하는 패러다임이 주류가 될 가능성을 보여줌
TL;DR 주요 성과
- Stanford CRFM 연구팀은 AI가 생성한 고성능 CUDA-C 커널이 기존 PyTorch의 전문가 설계 커널과 유사하거나 더 나은 속도를 내는 예를 발견함
- 원래는 합성 데이터를 통해 더 나은 커널 자동생성 모델을 학습시키려고 했으나, 합성 데이터 생성만으로도 놀랄 만한 수준의 빠른 커널이 자동 생성되는 현상을 관찰함
- 이에 따라 방법, 성능 벤치마크, 최적화 전략, 앞으로의 방향 등을 조기에 공개함
- 벤치마크: Nvidia L40S GPU 기준. 성능(%): PyTorch 기준 실행 시간 ÷ 생성 커널 실행 시간
- Matmul (FP32): PyTorch 대비 101.3% (4096x4096 행렬)
- Conv2D: 179.9% (입력: 100, 3, 224, 224; AlexNet Conv1 규격)
- Softmax: 111.8% (4096x65536 텐서)
- LayerNorm: 484.4% (16, 64, 256, 256 텐서)
- Conv2D + ReLU + MaxPool: PyTorch reference 대비 290.1%, torch.compile() 대비 189.0%
연구 동기 및 방법
- 원래 목적은 커널 생성 모델 학습용 합성 데이터 생성이었으나, 실험 과정에서 자체적으로 생성된 커널이 예상을 뛰어넘는 고성능을 달성
- KernelBench(Stanford 공개 벤치마크, 2024년 12월 공개) 활용
- 주어진 torch 코드에 대해, LLM이 최적 속도 커널을 새로 작성
- 입력/출력 결과의 수치 동치 여부로 정확성 검증
- 기존 방식: 커널을 단계별로 조금씩 고치며 점차 개선하는 순차적 수정 루프
- 단점: 아이디어 다양성 부족, 같은 최적화 반복, 지역 최적점 수렴
- 개선 아이디어
- 최적화 아이디어를 자연어로 발상 및 기록한 뒤, 그 아이디어들의 코드 구현분기를 여러 개 동시에 생성
- 각 라운드마다 여러 최적화 시도 병렬 진행 → 최고 성능 커널로 다음 라운드 씨앗(시드) 설정
- 이렇게 하면 한정된 탐색 반복 안에서 다양한 최적화 전략 탐구와 병렬 탐색이 가능
최적화 아이디어 예시
- 메모리 접근 최적화: global/shared/register 메모리 계층 효율 개선
- 비동기 처리 및 레이턴시 은닉: 연산과 데이터 이동을 오버랩
- 데이터타입/정밀도 최적화: FP16/BF16 활용 및 하드웨어 특화 연산
- 계산 및 명령어 최적화: 연산량, 명령어 수, 하드웨어 특수 명령 최적 활용
- 병렬성 및 오큐펀시: SM(Streaming Multiprocessors) 활용 극대화
- 루프/분기/인덱싱 최적화: 루프 최소화, 불필요한 인덱스 연산 제거
커널 최적화 실제 과정(Conv2D 예시)
- 라운드별 최적화 아이디어 및 성능 개선 흐름
- 최초(0라운드): 단순 CUDA 포팅(PyTorch 대비 20% 속도)
- 다음 라운드: → 읽기 캐시 활용 → FP16 Tensor Core 연산(GEMM 변환) → 이중 버퍼링/파이프라인 → 인덱스 사전계산/공유메모리 → 벡터화 → K-루프 동시버퍼링 등 고도의 GPU 아키텍처 활용
- 최종(13라운드): PyTorch 대비 179.9% 성능 확보
- 실제 코드는 고급 CUDA 프로그래밍 기법 대거 활용
- 각 라운드마다 새로운 아이디어를 자연어로 생성하고, 여러 구현을 병렬 시도하여 최적 코드를 선택
Takeaways 및 시사점
- AI 기반 커널 생성이 강력한 탐색 루프와 자연어 기반 다양한 최적화 아이디어 조합으로 인간 전문가 수준을 뛰어넘을 수 있음
- 일부 최신 연산자(FP16 matmul, Flash Attention)는 현 시점에서 PyTorch 대비 아직 성능 낮음
- FP32 계산이 최근 하드웨어에서는 FP16/BF16에 비해 상대적으로 최적화 덜 됨 → 해당 영역에서 성능 우위 가능
- 제한된 탐색 토큰(입출력 합 700만 토큰) 상황에서도 지속적인 성능 개선 확인
- AlphaEvolve, Gemini 2.5 Pro 등 최근 연구도 대량 브랜칭+검증 기반 탐색이 재학습 없이도 획기적 성능 개선 가능함을 시사
- 앞으로는 이런 방식으로 합성 데이터 및 고성능 커널을 대량 생성하며, AI가 자체 개선하는 루프(자기개선형 AI)로 발전할 것임
결론
- AI 기반 커널 자동 생성·최적화는 이미 전문가 핸드코딩 수준에 도달했으며, 빠른 미래에 모델+브랜칭 탐색+검증 조합으로 자가개선 AI 시스템이 가능할 전망
- PyTorch·TensorFlow와 같은 프레임워크의 성능 한계를 AI가 자체적으로 넘어서게 될 가능성 대두
부록: Conv2D 커널 전체 코드
- 실제 함수와 커널 전체 소스코드가 포함됨(상세 코드 생략)
- 코드 내 주요 특징:
- 공유 메모리 벡터화, K-dim 계층적 double-buffering, CUDA WMMA, warp-level output buffer 등
- 실시간 동적 인덱스 계산, bias 처리, K 루프 안의 지연 데이터 동시 로드 등
- 완전한 샘플 코드 및 예시 커널은 깃허브 저장소에서 확인 가능