우리가 아직 공개할 생각이 없었던 놀라울 정도로 빠른 AI 생성 커널들

• AI가 생성한 CUDA-C 커널들이 PyTorch의 전문가 최적화 커널과 비슷하거나 더 나은 성능을 보임
• Conv2D, LayerNorm 등에서 PyTorch 대비 최대 ~179.9% 에서 484.4% 까지 성능 향상 확인
• 기존 순차 수정 방식 대신, 자연어 최적화 아이디어 도출 및 병렬 브랜칭 탐색 도입
• 여러 메모리 최적화, 비동기 연산, 데이터 타입 전략 등 고급 컴퓨팅 기법을 활용하여 경쟁력 확보
• FP16 등 최신 ML 연산자에 대한 성능 한계는 있으나, 지속적인 탐색과 데이터 축적으로 개선 가능성 확인

개요

Stanford CRFM 연구팀은 AI가 생성한 고성능 CUDA-C 커널이 기존 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%

GPU는 Nvidia L40S 사용. 성능(%): PyTorch 기준 실행 시간 ÷ 생성 커널 실행 시간

커널 생성 방법론

• 자체 벤치마크인 KernelBench(2024년 12월 공개) 태스크 적용
  • 주어진 torch 코드에 대해, LLM이 최적 속도 커널을 새로 작성
  • 입력/출력 결과의 수치 동치 여부로 정확성 검증
• 기존 방식: 커널을 단계별로 조금씩 고치며 점차 개선하는 순차적 수정 루프
  • 단점: 아이디어 다양성 부족, 같은 최적화 반복, 지역 최적점 수렴
• 개선 아이디어
  1. 최적화 아이디어를 자연어로 발상 및 기록한 뒤, 그 아이디어들의 코드 구현분기를 여러 개 동시에 생성
  2. 각 라운드마다 여러 최적화 시도 병렬 진행 → 최고 성능 커널로 다음 라운드 씨앗(시드) 설정
  • 이렇게 하면 한정된 탐색 반복 안에서 다양한 최적화 전략 탐구와 병렬 탐색이 가능

주된 최적화 유형

• 메모리 구조 최적화: 전역/공유/레지스터 간 효율적 데이터 이동
• 비동기 연산·지연 시간 숨김: 느린 메모리 작업과 연산을 중첩해 실행
• 데이터 타입·정밀도 전략: FP16, BF16 등 저정밀 데이터 활용으로 대역폭 저감 및 전용 하드웨어 활용
• 산술·명령어 최적화: 연산 카운트 줄이기, 특수 명령 활용
• 병렬화·SM 점유율 극대화: 더 많은 워프 투입으로 대기시간 은폐
• 분기·루프 최적화: 반복/분기/인덱싱 계산 오버헤드 감소

커널 최적화 예시 (Conv2D)

• 최초(0라운드): 단순 CUDA 변환 – PyTorch 대비 20% 속도
• 다음 라운드:
  • 읽기 전용 캐시 활용, Tensor-Core GEMM 전환, double buffer 구조, 인덱스 메모리화, 벡터화 연산, K-루프 동시버퍼링 등 고도의 GPU 아키텍처 활용
• 최종(13라운드): PyTorch 대비 179.9% 성능 확보
• 실제 코드는 고급 CUDA 프로그래밍 기법 대거 활용

시사점 및 한계

• AI 기반 커널 생성이 강력한 탐색 루프와 자연어 기반 다양한 최적화 아이디어 조합으로 인간 전문가 수준을 뛰어넘을 수 있음
• 일부 최신 연산자(FP16 matmul, Flash Attention)는 현 시점에서 PyTorch 대비 아직 성능 낮음
• FP32 계산이 최근 하드웨어에서는 FP16/BF16에 비해 상대적으로 최적화 덜 됨 → 해당 영역에서 성능 우위 가능
• 제한된 탐색 토큰(입출력 합 700만 토큰) 상황에서도 지속적인 성능 개선 확인
• 궁극적으로 지속적 개선과 자기개선 AI 시스템으로 발전 가능성 높음

감사의 글

Christopher Rinard, Saman Amarasinghe, Allen Nie 외 논의 참여자, Standard Kernel Co.·Prime Intellect의 지원에 감사 표명

부록: 최적화 Conv2D CUDA 코드 샘플

• 실제 함수와 커널 전체 소스코드가 포함됨(상세 코드 생략)
• 코드 내 주요 특징:
  • 공유 메모리 벡터화, K-dim 계층적 double-buffering, CUDA WMMA, warp-level output buffer 등
  • 실시간 동적 인덱스 계산, bias 처리, K 루프 안의 지연 데이터 동시 로드 등
• 완전한 샘플 코드 및 예시 커널은 깃허브 저장소에서 확인 가능