Triton for MTIA: Bridging the Programming Model Gaps for Custom AI Accelerators
이 논문은 Meta의 커스텀 MTIA-2i 가속기에 대한 Triton 프로그래밍 언어의 첫 번째 프로덕션 규모 배포를 제시하며, 새로운 컴파일러 백엔드와 최소한의 언어 확장이 ML 프레임워크와 커스텀 하드웨어 사이의 간극을 메우는 고성능, 고효율 커널 개발을 가능하게 함을 입증한다.
원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기
당신이 대륙을 가로질러 승객(데이터)을 실어 나를 거대하고 고속인 열차 시스템을 구축하려고 한다고 상상해 보십시오. 수년 동안 사용 가능한 열차는 모두가 운전법을 알고 있는 표준적인 하이테크 전기 열차(GPU)뿐이었습니다. 그 열차들은 자동문, 내장형 GPS, 그리고 부드러운 승차감을 갖추고 있었습니다. 하지만 최근에 새로운 종류의 열차 엔진이 발명되었습니다(커스텀 AI 가속기). 이 새로운 엔진은 엄청나게 강력하며 무거운 화물을 운송하기 위해 특별히 제작되었지만, 매우 기묘합니다. 자동문이 없어서 문을 직접 수동으로 열고 닫아야 합니다. GPS도 없습니다. 지도는 직접 그려야 합니다. 또한 기존의 열차와는 같은 궤도를 달리지 않습니다.
문제는 열차 칸을 만드는 엔지니어들(소프트웨어 개발자)이 기존의 쉽고 운전하기 편한 전기 열차에 익숙해져 있다는 점입니다. 만약 그들이 이 새로운 초고속 엔진을 사용하고 싶다면, 완전히 새롭고 어려운 운전법을 배우거나, 매번 새로운 경로가 생길 때마다 완전히 새로운 열차 칸을 처음부터 만들어야 합니다. 이는 모든 과정을 느리게 만듭니다. 여기서 핵심적인 질문은 이것입니다. 기존의 쉽고 편리한 운전 스타일을 속도 이점을 잃지 않으면서도 이 기묘하고 새로운 엔진에서 작동하도록 가르칠 수 있을까요? 이 논문은 바로 그 간극을 메우기 위해, 표준 AI 소프트웨어의 친숙한 세계와 Meta와 같은 기업들이 AI 모델을 실행하기 위해 만든 기묘한 커스텀 하드웨어 사이의 차이를 탐구합니다.
이야기: 새로운 엔진에게 옛 언어를 가르치기
Facebook과 Instagram의 모회사인 Meta는 자신의 AI 모델을 더 빠르고 저렴하게 실행하기 위해 MTIA-2i라고 불리는 특별한 컴퓨터 칩을 구축해 왔습니다. MTIA-2i를 커스텀 제작된 레이싱 카 엔진이라고 생각하십시오. 이 엔진은 하는 일에 있어서는 놀랍지만, 대부분의 AI 프로그래머가 익숙한 표준 엔진(GPU)과는 다르게 설계되었습니다.
AI의 세계에는 Triton이라는 인기 있는 프로그래밍 언어가 있습니다. Triton을 AI를 위한 "만능 리모컨"이라고 생각할 수 있습니다. 개발자가 표준 GPU에 데이터를 어떻게 이동시킬지 지시하기 위해 복잡하고 지저분한 코드를 작성하는 대신, Triton으로 단순하고 깔기한 코드를 작성하면 스마트한 컴파일러가 이를 GPU가 이해할 수 있는 기계어로 번격해 줍니다. 이는 마치 기계와 대화할 줄 아는 통역사에게 영어로 말하는 것과 같습니다.
하지만 Triton은 원래 표준 GPU 엔진만을 위해 설계되었습니다. Meta가 이 Triton을 자신들의 커스텀 MTIA-2i 레이싱 카 엔진에 사용하려고 했을 때, 제대로 작동하지 않았습니다. "리모컨"이 새로운 엔진의 버튼에 맞지 않았던 것입니다. MTIA-2i 엔진은 비동기적으로 작동하며(명령을 받은 후 나중에 수행함), 메모리를 "순환 버퍼(circular buffer)"라는 기묘한 방식으로 스스로 관리하며, 모든 것이 원활하게 흐르도록 유지하기 위해 매우 구체적인 지침이 필요합니다. 표준 Triton 코드는 이러한 특이사항들을 처리할 수 없었습니다.
논문의 내용: 새로운 어댑터 구축하기
Meta 팀은 이를 해결하기로 했습니다. 그들은 단순히 기존의 리모컨을 억지로 작동시키려 한 것이 아니라, MTIA-2i를 위해 특별히 설계된 **새로운 컴파일러 백엔드(compiler backend)**를 구축했습니다. 이 새로운 시스템은 초스마트한 어댑터 역할을 합니다. 개발자가 작성한 단순하고 읽기 쉬운 Triton 코드를 가져와서, MTIA-2i 엔진이 필요로 하는 복잡하고 낮은 수준의 명령어로 완벽하게 번역해 줍니다.
그들이 이 작업을 어떻게 성공시켰는지 몇 가지 재미있는 비유를 들어 설명하겠습니다:
- "FIFO" 뷔페 줄: 표준 GPU에서는 컴퓨터가 스마트한 뷔페처럼 접시를 즉시 가져가고 교체하며 메모리를 자동으로 관리합니다. 반면 MTIA-2i는 "선입선출(First-In, First-Out, FIFO)" 버퍼를 사용하는 수동 조립 라인과 같습니다. 새로운 컴파일러는 이 라인을 완벽하게 관리하는 법을 학습하여, 프로그래머가 일일이 접시를 밀어 넣지 않아도 데이터가 기계가 필요로 하는 시점에 정확히 도착하도록 보장했습니다.
- 두 명의 셰프가 있는 주방: MTIA-2i 칩에는 동시에 작업할 수 있는 두 개의 "코어"(주방에 있는 두 명의 셰프와 같은 역할)가 있습니다. 컴파일러는 두 셰프가 서로 부딪히거나 기다리는 일이 없도록 요리 작업을 나누는 방법을 찾아냈습니다. 심지어 숙련된 프로그래머가 원한다면 셰프들에게 누가 양파를 썰고 누가 수프를 저을지 아주 상세하게 지시할 수도 있게 했습니다.
- "지그재그" 배달 경로: 당신이 64채의 집에 패키지를 배달해야 한다고 상상해 보십시오. 만약 길을 따라 직선으로(linear)만 간다면, 앞쪽 집들은 빠르게 완료되겠지만 마지막 집들은 영원히 기다려야 할 것입니다. 팀은 "지그재그" 배달 경로(왔다 갔다 하는 방식)가 작업을 훨씬 더 균형 있게 분배한다는 것을 발견했습니다. 그들은 프로그래머가 이 더 스마트한 경로를 쉽게 선택할 수 있도록 Triton에 기능을 추가했습니다.
결과: 빠르고, 유연하며, 실전에 투입될 준비가 된 기술
팀은 단순히 실험실에서 이 기술을 구축한 것이 아니라, 실제 현장에 적용했습니다. 그들은 이 새로운 Triton 커널을 60가지 이상의 다양한 AI 모델 유형에 걸쳐 실제 운영 환경에 성공적으로 배치했습니다.
그들이 발견한 결과는 다음과 같습니다:
- 속도: 단순한 Triton 언어로 작성된 코드는 전문가들이 복잡한 저수준 C++(칩의 "네이티브" 언어)로 작성한 코드만큼 빠르게 실행되었습니다. 실제로 행렬 곱셈(GEMM) 작업의 경우, Triton 코드는 칩의 이론적 최대 속도의 80% 이상에 도달했습니다.
- 채택률: 단 한 분기 만에 MTIA에서 Triton을 사용하는 모델 유형의 수가 3배로 증가했습니다.
- 영향력: 이제 Triton은 이 모델들의 **레이어 중 50%**와 **비-GEMM 실행 시간의 47%**를 처리하고 있습니다! 이 새로운 쉽고 편리한 시스템이 전체 작업의 거의 절반을 수행하고 있는 것입니다.
이것이 왜 중요한가
이 논문은 "쓰기 쉬운 것"과 "초고속인 것" 사이에서 하나를 선택할 필요가 없다고 주장합니다. 이전에는 MTIA와 같은 커스텀 칩을 사용하려면 저수준 코딩의 마법사가 되어 모든 것을 처음부터 작성해야 했습니다. 만약 당신이 Triton과 같은 고수준 언어를 사용하고 싶었다면, 표준 GPU에 머물러 있어야만 했습니다.
이 간극을 메움으로써, Meta는 고수준 언어가 거의 모든 커스텀 하드웨어에서 작동하도록 적응될 수 있음을 보여주었습니다. 이는 미래에 기업들이 훨씬 더 기묘하고 특화된 AI 칩을 만들더라도, 개발자들이 매번 새로운 어려운 언어를 배울 필요가 없음을 의미합니다. 그들은 이미 알고 있는 익숙하고 강력한 도구를 계속 사용할 수 있으며, 컴파일러가 새로운 하드웨어의 복잡한 세부 사항을 처리해 줄 것입니다.
또한 논문은 이전에 KNYFE라는 다른 매우 낮은 수준의 언어를 시도했지만, 너무 배우기 어려워 포기했다는 점을 언급합니다. 그들은 약간의 커스텀 수정만 거친다면, Triton과 같은 유연한 고수준 언어를 고수하는 것이 AI 개발을 빠르고 효율적으로 유지하는 데 있어 승리하는 전략임을 입증했습니다.
요약하자면, 적절한 컴파일러의 마법이 있다면 우리는 "만능 리모컨"이 아무리 기묘하게 만들어진 커스텀 TV라 할지라도, 화질을 떨어뜨리지 않고 완벽하게 작동하게 만들 수 있다는 것을 이 논문은 보여줍니다.
연구 분야의 논문에 파묻히고 계신가요?
연구 키워드에 맞는 최신 논문의 일일 다이제스트를 받아보세요 — 기술 요약 포함, 당신의 언어로.