다양한 AI 가속기를 위한 에이전트 기반 커널 자동생성 프레임워크 KernelEvolve

초록

KernelEvolve는 메타가 제안한 에이전트형 커널 코딩 시스템으로, 트리톤·CuTe·저수준 하드웨어 언어 등 여러 추상화를 활용해 DLRM 학습·추론에 필요한 커널을 자동 생성·최적화한다. 그래프 기반 탐색, 적합도 함수, 동적 프롬프트 합성을 통해 하드웨어 특성을 실시간 반영하고, 지식베이스에 저장된 제약조건을 활용해 NVIDIA·AMD GPU와 메타 자체 AI 가속기(MTIA v3) 등 이기종 환경에서도 100% 정확도와 최대 17배 성능 향상을 달성한다.

상세 분석

KernelEvolve는 오늘날 AI 시스템이 직면한 세 가지 주요 다양성—모델 아키텍처, 커널 프리미티브, 그리고 하드웨어 세대·아키텍처—을 통합적으로 해결하기 위해 설계된 에이전트형 자동 커널 생성 프레임워크이다. 첫 번째 핵심은 “커널 사양 → 코드” 파이프라인을 다중 추상화 레이어로 나눈 점이다. 사용자는 고수준 연산(예: 행렬 곱, 집계)과 목표 하드웨어를 명시하면, 시스템은 Triton, CuTe DSL, 그리고 각 가속기의 저수준 진단 언어(예: PTX, AMD GCN, MTIA 전용 IR) 중 최적의 표현을 자동 선택한다. 이는 기존에 각각의 언어에 대해 전문가가 수작업으로 구현하던 과정을 크게 단축한다.

두 번째로, 커널 최적화는 그래프 기반 탐색(search graph)으로 모델링된다. 각 노드는 특정 코드 변형(스레드 블록 크기, 메모리 스케줄링, 레지스터 재배치 등)을 나타내며, 엣지는 변형 적용 순서를 의미한다. 탐색 정책은 강화학습 기반 선택기와 휴리스틱 규칙을 결합해, 탐색 공간이 수천에서 수십만 개에 달하는 경우에도 효율적으로 수렴한다. 적합도 함수는 실행 시간, 메모리 사용량, 전력 효율 등을 다중 목표로 평가하며, 하드웨어 제약(예: 레지스터 한계, 공유 메모리 용량)과 일치하도록 페널티를 부여한다. 종료 규칙은 목표 성능 도달, 탐색 예산 소진, 혹은 수렴 기준을 만족할 때 트리거된다.

세 번째 특징은 “런타임 컨텍스트 적응”이다. KernelEvolve는 실행 시점에 수집된 프로파일링 데이터(예: 캐시 미스 비율, 워프 스케줄링 효율)를 기반으로 프롬프트를 재구성하고, 대형 언어 모델(LLM)에게 새로운 코드 스니펫을 생성하도록 요청한다. 이를 “retrieval‑augmented prompt synthesis”라 부르며, 기존 LLM이 학습되지 않은 최신 가속기(MTIA v3)의 특수 명령어 집합도 정확히 반영한다.

또한 시스템은 영구적인 지식베이스(KB)를 유지한다. KB는 하드웨어 사양, 이전 최적화 사례, 그리고 제약조건을 구조화된 형태(그래프 DB)로 저장한다. 새로운 커널을 생성할 때 KB를 조회해 유사한 사례를 찾아 초기 탐색 방향을 설정하고, 최적화 과정에서 발견된 새로운 제약은 자동으로 KB에 삽입된다. 이 피드백 루프는 시간이 지날수록 시스템의 탐색 효율과 정확도를 지속적으로 향상시킨다.

실험 결과는 세 가지 주요 지표에서 눈에 띈다. 첫째, 공개된 KernelBench 250문제 전부를 100% 통과했으며, 3단계 난이도 구분에서도 동일한 성공률을 보였다. 둘째, 160개의 PyTorch ATen 연산을 3가지 이기종 하드웨어(NVIDIA A100, AMD MI250, MTIA v3)에서 테스트했을 때 480개의 연산‑플랫폼 조합 모두 정확도가 100%였으며, 평균 4.2배, 최고 17배의 성능 향상을 기록했다. 셋째, 기존 수작업 커널 개발에 소요되던 주 단위 시간을 몇 시간으로 단축했으며, 엔지니어링 비용을 크게 절감했다.

이러한 설계는 특히 “프로그래머블 장벽”을 낮추는 데 기여한다. 기존에는 새로운 AI 가속기의 ISA가 공개되지 않으면 커널 개발이 거의 불가능했지만, KernelEvolve는 KB와 LLM 기반 프롬프트 합성을 통해 비공개 ISA에도 대응한다. 따라서 메타 내부뿐 아니라 외부 기업이 차세대 가속기를 도입할 때도 빠른 포팅이 가능해진다.

요약하면, KernelEvolve는 다중 추상화, 그래프 기반 탐색, 런타임 적응형 프롬프트, 그리고 지속적인 지식베이스 축적이라는 네 가지 핵심 메커니즘을 결합해, 이기종 AI 가속기 환경에서 DLRM 워크로드의 커널 개발·최적화를 자동화하고, 성능·생산성·확장성 모두에서 현저한 이점을 제공한다.