Codex로 GPU 커널 232배 가속 달성
원제: Auto-research with codex: How I achieved a 232x Faster Kernel
왜 중요한가
AI 코딩 에이전트를 활용한 GPU 커널 자동 최적화가 전문 연구 수준의 성능 개선을 실현할 수 있음을 실증한 사례로, HPC 및 AI 인프라 분야에서 자동 연구의 가능성을 보여준다.
개발자 Sankalp이 OpenAI의 Codex를 활용한 자동 연구 방식으로 GPU Mode의 QR 분해 커널 최적화 콘테스트에서 베이스라인 대비 232배 속도 향상을 달성했다. 183명 참가자 중 12위를 기록했으며, FP32 배치 행렬 QR 분해 문제를 대상으로 했다.
GPU Mode와 Core Automation이 공동 주최한 'Linear Algebra Kernels in the Age of Research' 시리즈의 자동 연구 테마 콘테스트에서, 개발자 Sankalp이 OpenAI Codex를 활용해 베이스라인 대비 232배 빠른 GPU 커널을 구현하고 12위(183명 중)를 기록했다.
문제는 배치 단위의 정방 FP32 CUDA 행렬에 대해 compact Householder QR 분해를 구현하는 것으로, torch.geqrf(A)와 동일한 결과를 반환해야 했다. 512×512, 1024×1024, 2048×2048, 4096×4096 등 다양한 행렬 크기와 컨디셔닝 케이스에 걸친 실행 시간의 기하 평균으로 순위가 결정됐다. 내부적으로 FP16, FP8, NVFP4 사용이 허용됐으나, 최종 반환값은 FP32 수준의 QR 검증을 통과해야 했다.
이 문제가 자동 연구에 적합했던 이유는 GPU Mode가 제공하는 popcorn CLI 덕분이다. 에이전트가 CLI를 통해 직접 테스트, 벤치마크, 리더보드 제출을 자동화할 수 있었다. Sankalp은 Codex를 활용해 blocked Householder 알고리즘으로 직렬 연산을 줄이고, 아이디어 다양성을 도입해 로컬 최적점 탈출을 시도하는 등 반복적인 커널 개선을 진행했다. 그는 수학적 배경 이해를 바탕으로 더 나은 질문을 AI에게 던지는 방식이 핵심이었다고 설명했다.