사카나 AI, 파이토치 코드 CUDA로 자동 변환·최적화하는 에이전트 발표

PyTorch로 딥러닝 모델을 짜는 일은 비교적 간단하지만, GPU 연산 성능을 극한까지 끌어올리려면 결국 저수준의 CUDA 커널을 직접 손봐야 합니다.

이 작업은 하드웨어에 대한 깊은 이해가 필요하고 많은 시간이 듭니다. Sakana AI는 이러한 병목을 해결하기 위해 CUDA 커널을 자동으로 발견하고 검증하며 최적화하는 에이전트 프레임워크를 공개했습니다.

고수준 코딩을 넘어 저수준 커널로

최근 대형 언어 모델들은 테스트 타임 컴퓨트 확장을 통해 소프트웨어 개발에서 뛰어난 역량을 보여주고 있지만, 대부분 고수준 애플리케이션 개발에 치우쳐 있었습니다.

기존의 커널 생성 벤치마크 역시 허점이 존재하고 테스트 환경의 다양성이 부족해 진정한 일반화 성능을 평가하기 어려웠습니다.

이에 따라 Sakana AI는 다양한 시나리오에서 커널의 성능과 정확성을 엄격하게 평가할 수 있는 새로운 벤치마크 robust-kbench를 먼저 도입했습니다.

이 벤치마크를 바탕으로 최첨단 LLM이 PyTorch 코드를 기능적으로 동등한 CUDA 커널로 변환한 뒤, 성능을 점진적으로 개선해 나가는 순차적 워크플로를 구축했습니다.

진화적 탐색과 검증기를 통한 성능 개선

변환된 커널은 CUDA 생태계에 맞춰 설계된 새로운 진화적 메타 생성 절차를 거치며 실행 속도가 최적화됩니다.

잘못된 코드가 하드웨어 자원을 낭비하지 않도록, 파이프라인 내부에는 LLM 기반 검증기가 배치되어 코드의 정확성을 판별하고 비정상적인 커널을 사전에 걸러냅니다.

robust-kbench 평가 결과, 이 시스템이 생성한 CUDA 커널은 순전파와 역전파를 포함한 실제 응용 작업에서 기본 PyTorch 구현체보다 뛰어난 실행 속도를 기록했습니다.

또한 여러 연산을 하나로 합치는 연산자 융합을 비롯해 다양한 런타임 최적화 기법을 자율적으로 적용하는 능력도 입증했습니다.

왜 중요한가

인공지능 모델을 최적화하고 배포하는 엔지니어에게 저수준 CUDA 작업은 가장 까다롭고 번거로운 영역입니다. PyTorch 프로토타입 코드에서 곧바로 고성능 커널을 뽑아낼 수 있다면 모델 서빙과 연구의 반복 주기가 비약적으로 단축될 수 있습니다.

업계 전반으로 보아도 GPU 아키텍처를 극한까지 다루는 시스템 엔지니어의 수는 매우 부족합니다. Sakana AI의 이번 연구는 LLM의 추론과 탐색 능력이 고난도 시스템 최적화 영역에서도 전문 인력의 역할을 일정 부분 대신할 수 있음을 보여줍니다.

하드웨어 테스트 이전에 LLM 검증기를 배치해 오류를 미리 차단한 구조는 연산 자원을 아끼는 매우 현실적이고 영리한 선택이라고 봅니다. 다만 현재 성과가 robust-kbench 환경을 중심으로 입증된 만큼, 규격화되지 않은 복잡한 실제 프로덕션 환경에서도 동일한 안정성을 보여줄지는 지켜볼 필요가 있어 보입니다.