트라이톤으로 HPC GPU 커널 이식성 해결하기: 파이썬 의존을 걷어낸 AOT 컴파일 실험

핵심 메시지

• GPU 커널은 벤더와 세대가 바뀔 때마다 손으로 다시 써야 해서, 실행 시간의 5%밖에 차지하지 않는 코드가 코드베이스의 95%를 차지하는 비효율이 생긴다.
• AI 업계가 만든 타일 기반 파이썬 언어 트라이톤은 메모리 배치와 아키텍처별 최적화를 컴파일러가 알아서 처리해, 새 하드웨어 기능도 코드 수정 없이 활용할 수 있게 해 준다.
• 문제는 트라이톤이 파이썬과 파이토치에 묶여 있다는 점인데, 연구진은 사전 컴파일과 C 스타일 래퍼로 이 의존성을 걷어내고 일반 오브젝트 파일처럼 링크되게 만들었다.
• 유체동역학 프록시 앱과 양자화학 텐서 축약에서 네이티브 구현과 비슷하거나 더 나은 성능이 나왔고, 특히 벤더 라이브러리로는 불가능한 커널 융합이 큰 이득을 냈다.
• 반면 스레드마다 작업량이 제각각인 몬테카를로 중성자 수송 벤치마크에서는 성능이 크게 떨어져, 타일 균일성을 전제로 한 추상화의 한계가 드러났다.

쉽게 이해하기

슈퍼컴퓨터는 점점 더 GPU에 기대고 있지만, 정작 GPU 커널은 이식성이 나쁘다. 벤더가 다르고 아키텍처 세대가 바뀔 때마다 새 프로그래밍 인터페이스를 익혀 코드를 다시 써야 하기 때문이다. 노스캐롤라이나 주립대의 얀보 자오는 ARRAY 2026 발표에서 이 오래된 생산성 문제를 AI 생태계의 도구로 풀어 보자고 제안했다.

손으로 다시 쓰는 비용이 특히 아까운 쪽은 이른바 '차가운 커널'이다. 실행 시간의 95%를 잡아먹는 5%의 뜨거운 커널은 수작업 튜닝을 감수할 가치가 있지만, 코드베이스의 95%를 차지하면서 실행 시간은 5%뿐인 나머지도 이식할 때는 똑같이 다시 써야 한다. 기존의 CUDA 변환기나 C++ 기반 이식성 라이브러리들은 타입이 강하게 묶여 있어 메모리 배치를 과감하게 바꾸지 못한다는 것이 발표자의 진단이다.

대안으로 지목된 트라이톤은 알고리즘의 의도를 높은 수준에서 기술하면 메모리 관리와 아키텍처별 코드 생성을 컴파일러가 계층별로 처리하는 구조다. 텐서 코어나 저비트 연산 같은 새 하드웨어 기능이 세대마다 추가돼도 사용자가 코드를 고칠 필요가 없고, 벤더는 자기 백엔드만 붙이면 기존 최적화를 그대로 쓸 수 있다. 커널을 CPU에서 표준 파이썬 디버거로 돌려 볼 수 있다는 점도 실무적 이점으로 꼽혔다.

걸림돌은 트라이톤이 파이썬 런타임과 파이토치에 강하게 결합돼 있다는 점이다. 권한과 환경 관리가 엄격한 슈퍼컴퓨팅 환경에서는 파이토치를 불러오는 것 자체가 부담이다. 연구진은 커널 시그니처와 텐서 레이아웃 정보를 받아 사전 컴파일한 뒤 C 스타일 래퍼로 감싸, 결과물을 하나의 오브젝트 파일처럼 링크하도록 만들었다. 사람이 읽을 수 있는 중간 표현을 남길지, 재컴파일 비용을 없앤 바이너리를 쓸지, 둘 다 담아 파일 크기를 감수할지는 선택지로 열어 뒀다.

평가는 세 갈래로 진행됐다. 라그랑지안 유체동역학 프록시 앱의 커널에서는 격자가 작을 때 불필요한 제어 흐름 탓에 손해를 봤지만 격자가 커지면 네이티브 CUDA와 비슷해졌고, 타일 균일성이 높은 축약 커널은 오히려 더 빨랐다. 양자화학의 4차 텐서 축약에서는 벤더 텐서 라이브러리와 대체로 비슷한 성능(±10% 수준)이 나왔다. 반면 스레드마다 필요한 데이터와 반복 횟수가 다른 몬테카를로 중성자 수송 벤치마크에서는 성능이 크게 떨어졌다.

주요 인사이트

• 이 연구의 핵심 주장은 새 언어를 만들자는 것이 아니라, AI 업계가 막대한 투자로 다듬어 놓은 컴파일러 자산을 과학 계산이 그대로 빌려 쓰자는 것이다. 벤더가 새 하드웨어 기능을 백엔드에 반영하면 HPC 코드도 자동으로 혜택을 본다.
• 벤더 텐서 라이브러리 대비 가장 큰 이득은 개별 커널의 속도가 아니라 커널 융합이었다. 여러 커널을 새 커널 안에서 호출하고 표시만 해 주면 컴파일러가 하나로 합쳐 주는데, 닫힌 라이브러리로는 불가능한 최적화다.
• 성능이 무너진 지점이 오히려 정직한 정보다. 타일 단위 균일 연산을 전제로 한 추상화는 정적 루프 경계를 요구하고 스레드 단위 저수준 제어를 마스크로만 허용하기 때문에, 불규칙한 워크로드에서는 결과는 맞아도 속도가 안 나온다.
• 질의응답에서 발표자는 오늘날 진짜 병목이 연산이 아니라 메모리라고 답했다. 연산 능력이 메모리보다 훨씬 빠르게 늘어난 만큼, 어떤 메모리 객체가 중요한지 어디에 둘지를 판단하는 것이 다음 과제라는 것이다.
• 코딩 에이전트로 기존 CUDA 커널을 트라이톤으로 옮기는 시나리오가 발표에서 자연스럽게 언급된 점도 눈에 띈다. 다만 발표자는 에이전트를 쓰더라도 토큰 비용과 사람의 확인 작업이 남는다고 덧붙였다.

자주 묻는 질문

관련 AI 소식