MTIA를 위한 Triton: 맞춤형 AI 가속기를 위한 프로그래밍 모델 격차 해소
요약
Meta의 맞춤형 AI 가속기 MTIA-2i에서 Triton을 프로덕션 규모로 적용한 사례를 소개합니다. 새로운 컴파일러 백엔드와 언어 확장을 통해 Triton이 다양한 가속기 아키텍처의 프로그래밍 모델 격차를 해소할 수 있음을 입증했습니다.
핵심 포인트
- MTIA-2i를 위한 새로운 Triton 컴파일러 백엔드 개발
- Triton 커널이 전문가가 튜닝한 C++ 구현과 대등한 성능 달성
- 60여 가지 모델 유형에 대해 Triton 커널의 성공적인 프로덕션 배포
- DSL을 통한 맞춤형 가속기와 ML 프레임워크 간의 격차 해소
머신러닝 (ML) 워크로드의 급격한 성장은 맞춤형 가속기 아키텍처의 확산을 가속화했습니다. 처음부터 설계된 이러한 가속기들은 종종 GPU와는 구별되는 프로그래밍 모델을 노출합니다. 하이퍼스케일러 (Hyperscalers)와 AI 칩 스타트업들이 이 분야에서 혁신을 지속하고 있지만, 다양한 모델을 지원하기 위한 광범위한 연산자 (Operator) 커버리지를 달성하는 것은 여전히 주요 과제로 남아 있습니다. 또한, 모델과 커널 (Kernel)의 빠른 반복을 위해서는 사용하기 쉬운 고수준 커널 프로그래밍 언어가 중요합니다. Triton은 TorchInductor와 함께 GPU에서 이러한 문제들을 해결하지만, 서로 다른 프로그래밍 모델을 가진 가속기에서의 생존 가능성은 아직 입증되지 않았습니다. 본 연구에서는 Meta가 개발한 맞춤형 ML 가속기인 MTIA-2i에서 Triton을 프로덕션 규모로 적용한 첫 번째 사례를 제시합니다. MTIA-2i를 지원하기 위해, 우리는 이를 타겟으로 하는 새로운 컴파일러 백엔드 (Compiler backend)를 개발하고, TorchInductor 코드 생성 (Code generation)에 대한 향상을 도입하며, MTIA 특유의 아키텍처 특징을 노출하는 최소한의 언어 확장 기능을 제안합니다. 우리는 Triton-MTIA 커널이 전문가가 튜닝한 C++ 구현과 경쟁할 만한 성능을 달성함을 입증합니다. 이러한 개발 효율성 이점을 활용하여, 우리는 수동으로 작성된 커널과 Inductor가 생성한 Triton 커널을 약 60가지의 서로 다른 모델 유형에 걸쳐 프로덕션에 성공적으로 배포하였으며, 이는 해당 모델들의 레이어 (Layer) 50%와 non-GEMM 실행 시간의 47%를 차지합니다. 우리의 결과는 Triton과 같은 DSL (Domain-Specific Languages)이 ML 프레임워크, 커널, 그리고 맞춤형 가속기 사이의 프로그래밍 모델 격차를 메울 수 있으며, 이를 통해 빠른 혁신과 대규모의 효율적인 배포를 가능하게 한다는 강력한 증거를 제공합니다.
AI 자동 생성 콘텐츠
본 콘텐츠는 arXiv cs.PL (Programming Languages)의 원문을 AI가 자동으로 요약·번역·분석한 것입니다. 원 저작권은 원저작자에게 있으며, 정확한 내용은 반드시 원문을 확인해 주세요.
원문 바로가기