자동 GPU 커널 최적화 AutoKernel 시스템

초록

AutoKernel은 PyTorch 모델을 자동으로 프로파일링하고, Amdahl 법칙에 기반해 병목 커널을 선정한 뒤, Triton·CUDA C++ 백엔드에서 수백 차례의 실험을 수행해 최적의 GPU 커널을 스스로 찾아내는 오픈소스 프레임워크이다. 5단계 검증 파이프라인을 통해 정확성을 보장하고, H100 기준 RMSNorm·Softmax·Cross‑Entropy 등 주요 트랜스포머 연산에서 기존 PyTorch eager와 torch.compile을 크게 앞선 속도를 달성한다.

상세 분석

본 논문은 GPU 커널 최적화 작업을 인간 전문가의 반복적인 “작성‑벤치마크‑보존/복구” 루프와 동일한 형태의 자동화 루프로 추상화한다. 핵심 아이디어는 (1) 모델 전체를 프로파일링해 실행 시간 비중이 큰 커널을 Amdahl 법칙으로 정량화하고, (2) 해당 커널에 대해 단일 파일을 LLM 에이전트가 수정하도록 제한함으로써 변경 이력을 명확히 관리한다는 점이다. 에이전트는 909줄에 달하는 ‘program.md’ 지침서를 읽어 6단계 최적화 플레이북(블록 사이즈, 메모리 접근, 연산, 고급 기법, 아키텍처 특화, 커널 특화)을 순차적으로 적용한다. 각 수정 후에는 5단계 검증 파이프라인—스모크 테스트, 다양한 입력 형태와 데이터 타입을 통한 Shape Sweep, 수치 안정성 검증, 결정론성 확인, 비전형 차원(edge case) 테스트—을 거쳐야만 성능 측정이 진행된다. 이 설계는 메모리 오버플로, 레이스 컨디션, 비정형 입력에 대한 회귀를 조기에 차단해 실험 비용을 최소화한다.

시스템 구현은 9,200줄 이상의 파이썬 코드와 18개의 스타터 커널(9개 Triton, 9개 CUDA C++)을 포함한다. Triton 백엔드는 빠른 JIT 컴파일(1~5초)과 고수준 DSL을 제공해 초기에 높은 탐색 효율을 보장하고, CUDA C++ 백엔드는 WMMA·warp‑shuffle·vectorized load 등 저수준 프리미티브에 직접 접근해 최종 성능 한계에 근접한다. 오케스트레이터는 Amdahl 식 S = 1‑f + f/s 를 이용해 각 커널별 기대 속도 향상을 예측하고, 5번 연속 리버트, GPU 피크 90% 도달, 2시간 제한, 2배 속도 달성 등 네 가지 이동 기준 중 하나가 충족되면 다음 병목 커널로 전환한다.

실험 결과는 H100 GPU에서 RMSNorm 5.29×, Softmax 2.82×, Cross‑Entropy 2.21× 등 주요 연산에서 기존 eager 구현을 크게 앞섰으며, torch.compile(max‑autotune) 대비 각각 2.83×, 3.44×, 2.94×의 향상을 기록한다. 또한 커뮤니티 벤치마크인 vectorsum_v2 B200에서 1위에 오르는 등 실제 배포에서도 경쟁력을 입증한다. 논문은 기존 LLM 기반 커널 자동화 연구와 차별화되는 세 가지 포인트—모델‑레벨 프로파일링, 이중 백엔드 지원, 단순 루프 기반 설계—를 강조한다. 전체 코드는 GitHub에 공개돼 재현성과 확장성을 보장한다.