MaxKernel, 멀티 에이전트로 TPU 커널 생성과 최적화

가속기에서 높은 성능을 내는 커널을 작성하는 일은 일반적인 코드 생성보다 훨씬 어렵다. 연산의 의미뿐 아니라 메모리 접근, 병렬 실행, 데이터 배치, 컴파일러 제약, 특정 칩의 실행 특성까지 함께 고려해야 하기 때문이다. MaxKernel은 대규모 언어 모델을 컴파일러와 하드웨어의 피드백을 받을 수 있는 개발 루프 안에 배치한다. 한 번 코드를 생성하고 끝내는 대신, 구현을 제안하고 컴파일하고 테스트하고 프로파일링한 뒤 다음 시도를 수정하는 방식이다.

세 가지 개발 방식

MaxKernel은 여러 에이전트가 역할을 나누는 구조를 사용한다. 계획, 구현, 자체 디버깅, 정확성 테스트, 하드웨어 프로파일링을 전문 서브 에이전트가 담당하며, 이 공통 기반 위에서 세 가지 방식을 제공한다.

• Human-in-the-Loop: 사람과 에이전트가 단계별로 설계를 진행한다. 개발자가 중요한 방향을 직접 정할 수 있어 초기 탐색이나 협업 개발에 적합하다.
• Autonomous: 성능 지표와 하드웨어 trace를 바탕으로 에이전트가 구현을 반복적으로 개선한다. 반복적인 실험과 수작업을 줄이는 데 초점이 있다.
• 그래프 기반 자율 탐색: 여러 구현 후보와 최적화 경로를 탐색 그래프로 관리한다. 하나의 수정 경로만 따라가는 대신 더 넓은 설계 공간을 살펴보려는 접근이다.

이 세 가지 방식은 커널 개발이 단순한 코드 자동완성이 아니라 실험의 연속이라는 점을 반영한다. 사람이 설계의 방향을 잡아야 하는 경우와, 동일한 작업을 반복하며 성능을 측정해야 하는 경우를 하나의 시스템에서 구분해 처리할 수 있다.

실제 측정 결과를 다음 시도의 근거로 활용

TPU 커널에서는 같은 계산을 수행하는 코드라도 타일링, 데이터 배치, 메모리 접근 방식에 따라 실행 성능이 달라질 수 있다. 따라서 사전 학습된 지식만으로 최적의 구현을 선택하기는 어렵다. MaxKernel은 컴파일 결과, 테스트 결과, 프로파일링 정보를 에이전트의 판단 근거로 사용해 다음 후보를 조정한다.

평가에는 50개의 다양한 TPU 커널 과제가 포함된 JaxBench와, 최신 오픈소스 모델에서 가져온 복잡한 실제 워크로드가 사용됐다. 논문은 MaxKernel이 여러 과제에서 높은 수준으로 최적화된 구현을 생성하고, 전문가가 수작업으로 튜닝한 기준 구현에 가까운 결과를 낸다고 설명한다. 제공된 자료에는 구체적인 가속 배수가 없으므로, 이를 특정 커널의 최고 기록이라기보다 전체 개발 방식의 가능성을 보여주는 결과로 해석하는 편이 적절하다.

의미와 남은 과제

MaxKernel이 보여주는 핵심은 에이전트가 실험을 구성하고, 컴파일러가 즉각적인 기술 피드백을 제공하며, 하드웨어가 최종 성능을 판단하는 협업 모델이다. 이 방식은 TPU뿐 아니라 GPU 등 다른 AI 가속기의 자동 튜닝에도 참고가 될 수 있고, 전용 커널 개발에 필요한 전문 지식의 진입 장벽을 낮출 가능성이 있다.

물론 탐색 비용, 제한적인 하드웨어 관측 범위, 테스트의 완전성은 여전히 해결해야 할 문제다. 새로운 아키텍처, 익숙하지 않은 연산 조합, 실제 운영 환경에서도 안정적으로 작동하는지는 추가 검증이 필요하다. MaxKernel은 오픈소스로 공개됐기 때문에 연구자와 엔지니어가 에이전트 역할 분담, 탐색 절차, 평가 구성을 직접 살펴보고 확장할 수 있다.

출처: Hugging Face Daily Papers