기술 요약: MTIA를 위한 Triton

문제 정의

머신러닝 워크로드(LLM에서 추천 시스템까지)의 급격한 다양화는 Meta의 MTIA-2i와 같은 커스텀 AI 가속기의 확산을 불러왔습니다. 이러한 커스텀 ASIC은 특정 워크로드에 최적화된 성능과 비용 이점을 제공하지만, 동시에 상당한 프로그래밍 과제를 안겨줍니다. 하드웨어 관리형 캐시와 스레드 스케줄러를 갖춘 SIMT(Single Instruction, Multiple Threads) 모델을 사용하는 GPU와 달리, MTIA-2i와 같은 커스텀 가속기는 종종 판이하게 다른 프로그래밍 모델을 채택합니다. 예를 들어, MTIA-2i는 고정 기능 유닛(FFU)에 대한 비동기 명령 발행, 순환 버퍼(Circular Buffer, CB)를 통한 소프트웨어 관리형 온칩 메모리, 그리고 레이턴시를 숨기기 위한 명시적 소프트웨어 파이프라이닝에 의존합니다.

이러한 아키텍처적 차이는 "프로그래밍 모델 격차"를 발생시킵니다. 개발자들은 각 하드웨어 세대마다 저수준 언어(예: C++)로 커널을 직접 최적화해야 하며, 이는 속도가 느리고 오류가 발생하기 쉬우며 현대적인 모델 개발에 필요한 빠른 반복(iteration)을 저해합니다. Triton과 같은 도메인 특화 언어(DSL)가 GPU의 복잡성을 성공적으로 추상화해 왔지만, 근본적으로 다른 실행 및 메모리 모델을 가진 비-GPU 아키텍처에서의 생존 가능성은 그동안 대규모 프로덕션 규모에서 입증된 바 없었습니다.

방법론

저자들은 MTIA-2i 가속기 상에서 Triton을 적용한 첫 번째 프로덕션 규모의 사례를 제시합니다. 이 접근 방식은 새로운 컴파일러 백엔드를 개발하고, TorchInductor 통합을 강화하며, 고수준 Triton 구문과 MTIA의 저수준 제약 사이의 간극을 메우기 위해 최소한의 언어 확장을 도입하는 것을 포함합니다.

1. 컴파일러 아키텍처

Triton-MTIA 컴파일러는 MLIR을 기반으로 구축되었으며 4단계 흐름을 따릅니다:

• 파싱(Parsing): 업스트림 Triton 파서를 활용하여 Triton IR(ttir)을 생성함으로써, 진화하는 Triton 언어와의 호환성을 보장합니다.
• 미들엔드(매핑, Mapping): 연산을 MTIA 하드웨어 엔진에 매핑합니다. 이는 내적 연산 엔진(DPE)이나 SIMD 엔진(SE)과 같은 고정 기능 유닛(FFU)에 적합한 연산과, RISC-V 벡터 코어를 필요로 하는 연산을 구분합니다. 결정적으로, 메모리 액세스 패턴을 분석하여 이를 DMA 전송(구조적 액세스가 필요함)으로 낮출 수 있는지, 아니면 스칼라 로드/스토어로 폴백(fallback)해야 하는지를 결정합니다.
• 백엔드(낮추기, Lowering): 텐서 버퍼화(tensor bufferization)를 포함한 아키텍처 특유의 세부 사항을 처리합니다. Triton은 온칩 메모리 계층을 네이티브로 모델링하지 않으므로, 컴파일러는 로컬 메모리(LS) 내의 순환 버퍼(CB)를 할당하고 관리하며, FIFO 동작을 시뮬레이션하기 위해 읽기/쓰기 포인터 의존성을 강제합니다.
• 코드 생성(Code Generation): LLVM IR과 사람이 읽을 수 있는 C++ 생성을 모두 지원하여, 개발자가 최종 링킹 전에 생성된 코드를 검사하고 수동으로 튜닝할 수 있도록 합니다.

2. 최적화 전략

전문가가 튜닝한 C++와 경쟁할 만한 성능을 달est하기 위해, 컴플러는 세 가지 핵심 전략을 사용합니다:

• 소프트웨어 파이프라이닝(Software Pipelining): 컴파일러는 MTIA의 RISC-V 코어와 FFU의 비동기적 특성을 활용합니다. 컴파일러는 파이프라인이 계속 채워진 상태를 유지하면서도 필요한 최소한의 온칩 메모리를 할당하기 위해 "탐욕적 CB 크기 지정(greedy CB sizing)"을 수행합니다. 또한, 마스킹된 데이터에 대해 불필나한 계산과 버퍼 초기화를 건너뛰기 위해 "효과적인 형상 전파(effective shape propagation)"를 구현합니다.
• 벡터 코드 생성(Vector Codegen): FFU에서 지원되지 않는 "롱테일(long-tail)" 연산의 경우, 컴파일러는 RISC-V 벡터 확장(RVV)을 활용합니다 불필요한 브로드캐스트를 피하기 위해 벡터-스칼라 연산을 최적화하고, LS와 레지스터 사이의 중복 메모리 전송을 제거하기 위해 백투백(back-to-back) 벡터 루프를 퓨즈(fuse)합니다.
• 부하 분산(Load Balancing): 컴파일러는 유연한 PID-to-PE(Processing Element) 매핑 전략(선형, 라운드 로빈, 지그재그)을 지원합니다. 불균형한 분포를 가진 워크로드(예: 인과적 어텐션 마스크)의 경우, 8x8 PE 그리드 전체의 부하를 균형 있게 맞추기 위해 지그재그 분포를 사용하여 상당한 속도 향상을 얻습니다.

3. 언어 확장

자동 추론이 불충분한 경우를 해결하기 위해, 저자들은 최소한의 확장을 도입했습니다:

• 아키텍처 노출: tl.mtia.get_pe_id 및 tl.mtia.re_reduce와 같은 새로운 프리미티브를 통해 커널 작성자가 하드웨어 토폴로지를 명시적으로 쿼리하고 리덕션 엔진 네트워크를 활용할 수 있게 합니다.
• 메모리 재사용: tl.extract_slice 확장은 정밀한 메모리 재사용을 가능하게 하여, 작성자가 데이터를 LS에 한 번 로드한 후 메인 메모리에서 다시 가져오지 않고 슬라이스를 재사용할 수 있게 합니다.
• 커스텀 타입: ctypes.Structure를 통한 C++ 구조체와의 통합은 Triton 커널이 기존 MTIA 소프트웨어 스택 메타데이터(예: TensorView)와 상호작용할 수 있도록 합니다.

4. Inductor 통합

저자들은 MTIA를 위해 TorchInductor 코드 생성을 맞춤화했습니다. 여기에는 DMA 로워링을 보장하기 위해 스칼라 인덱스를 타일링하는 분석 가능한 메모리 액세스 패턴 생성과, LayerNorm과 같은 복잡한 커널을 위한 Triton 템플릿 사용 확장이 포함됩니다. 이를 통해 활성화 함수(activation function)의 자동 퓨전을 가능하게 합니다.

주요 기여

1. 프로덕션 규모 배포: 본 논문은 고수준 DSL이 비-GPU 프로그래밍 모델을 가진 커스텀 가속기에 성공적으로 적용될 수 있음을 입증합니다. Triton 커널은 약 60개의 모델 유형에 걸쳐 프로덕션에 배치되었으며, 이는 레이어의 50%와 non-GEMM 실행 시간의 47%를 커버합니다.
2. 컴파일러 및 언어 혁신: MTIA로 Triton을 낮추는 새로운 컴파일러 백엔드와 함께, 하드웨어 기능을 노출하면서도 이식성을 해치지 않는 특정 최적화(소프트웨어 파이프라이닝, CB 관리) 및 최소한의 언어 확장을 도입했습니다.
3. 성능 검증: 연구는 Triton 생성 커널이 전문가가 튜닝한 C++ 구현과 경쟁할 만한 성능을 달성하며, GEMM 연산의 경우 아키텍처 루프라인(roofline)의 80% 이상에 도달하는 경우가 많다는 실증적 증거를 제공합니다.

결과

• 성능: Triton GEMM 커널은 일관되게 아키텍처 루프라인의 80% 이상을 달성하여, 전문가가 튜닝한 C++ 커널과 대등한 성능을 보여주었습니다. 퓨즈된 FlashAttention 커널은 특히 긴 시퀀스 길이를 가진 생성형 모델에서 상당한 속도 향상을 보였습니다.
• 개발 속도: 단 한 분기 만에 팀은 Triton을 사용하는 모델 유형의 수를 3배로 늘렸고, MTIA-2i에서 Triton 커널에 소요되는 런타임 시간을 두 배로 늘렸습니다.
• 커버리지: Inductor-Triton 흐름은 내부 모델 세트 전반에 걸쳐 추론에서 94.5%, 훈련에서 94.7%의 성공률을 달성했습니다.
• 이식성: 포팅된 16개 GPU 커널 중 11개가 최소한의 변경만으로 직접 포팅 가능했으며, 이는 정의되지 않은 동작(undefined behaviors)을 피한다면 교차 아키텍처 코드 재사용이 가능하다는 것을 보여줍니다.

의의

본 논문은 다양한 ML 하드웨어와 모델이 개발자 생산성에 큰 장벽이 된다고 주장합니다. ML 프레임워크(PyTorch), 고수준 DSL(Triton), 그리고 커스텀 가속기(MTIA) 사이의 프로그래밍 모델 격차를 성공적으로 메움으로써, 저자들은 성능을 희생하지 않으면서도 높은 개발 속도를 유지할 수 있음을 보여주었습니다. 이 작업은 타겟팅된 언어 확장과 컴파일러 최적화로 보강된 유연한 DSL이 이질적인 AI 하드웨어를 위한 통합 계층 역할을 할 수 있으며, 이를 통해 신속한 혁신과 효율적인 대규모 배포를 가능하게 한다는 점을 시사합니다. 저자들은 이를 향후 커스텀 가속기의 도입 장벽을 낮추기 위한 표준화된 DSL 중심의 강력한 커뮤니티를 구축하기 위한 단계로 규정합니다.