[Paper] 타일 설계: 타일 기반 다중 PE 가속기에서 GEMM 배포 자동화

Source: arXiv - 2512.13638v1

Overview

이 논문은 Design in Tiles (DiT) 라는 자동화 프레임워크를 소개합니다. DiT는 고수준 GEMM(General Matrix Multiplication) 코드와 오늘날 AI 칩을 구동하는 저수준 타일 기반 다중 PE(Processing Element) 가속기 사이의 격차를 메워줍니다. 배포 툴체인과 구성 가능한 실행 모델을 결합함으로써, DiT는 현재 하드웨어‑특화 라이브러리에서 요구되는 수작업 튜닝 없이도 거의 최적에 가까운 GEMM 커널을 생성할 수 있게 합니다.

주요 기여

• 엔드‑투‑엔드 자동화: GEMM 설명을 받아 타일 기반 가속기 구성에 맞는 완전 튜닝된 구현을 생성하는 툴체인.
• 구성 가능한 실행 모델: PE 수, 타일 차원, 메모리 계층, 타일 간 대역폭을 포착하는 파라메트릭 성능 모델로, 빠른 설계 공간 탐색을 가능하게 함.
• 확장 가능한 매핑 전략: 2‑D 타일 그리드(예: 32 × 32 타일) 전반에 행렬을 분할하는 체계적인 방법으로, PE 활용도를 극대화하고 통신 지연을 숨김.
• 성능 동등성 및 그 이상: 4 TB/s 대역폭을 가진 시뮬레이션 1979 TFLOPS FP8 가속기에서 NVIDIA GH200의 전문가 튜닝 GEMM 라이브러리 대비 1.2–2.0× 속도 향상 입증.
• 오픈‑소스 프로토타입: 저자들은 DiT 프레임워크와 벤치마크 스크립트 세트를 공개하여 커뮤니티 채택 및 추가 연구를 장려.

방법론

1. Hardware Abstraction: 저자들은 타일 기반 가속기를 동일한 PE들로 구성된 격자(grid)로 모델링합니다. 각 PE는 로컬 레지스터, 작은 스크래치패드, 고대역폭 인터커넥트를 갖습니다. 타일 크기, PE 수, 메모리 대역폭과 같은 매개변수들을 조정 가능한 knob으로 노출합니다.
2. Tile‑Level Scheduling: GEMM의 세 중첩 루프(i, j, k)를 하드웨어 격자에 맞게 재배열하고 타일링합니다. DiT는 자동으로 각 타일이 출력 행렬의 몇 행/열을 계산할지, 그리고 축소 차원(k)이 시간 단계에 걸쳐 어떻게 분할될지를 결정합니다.
3. Data Movement Planning: 실행 가능한 모델을 사용해 DiT는 공유 메모리 계층에서 타일을 로드하는 비용을 예측하고, 더블 버퍼링을 스케줄링하며, PE가 바쁘게 유지되도록 프리패치를 삽입합니다.
4. Code Generation: 고수준 스케줄은 C++/CUDA와 유사한 커널로 낮춰져 특정 PE ISA(예: RISC‑V 기반 또는 맞춤 SIMD)를 목표로 합니다. 생성된 코드는 벤더 백엔드로 컴파일되어 대상 가속기에 바로 사용할 수 있는 바이너리를 생성합니다.
5. Iterative Optimization: 가벼운 자동 튜너가 모델의 성능 추정치를 기반으로 소수의 타일 크기 후보를 탐색하여, 예측된 PE 활용도가 가장 높은 구성을 선택합니다.

결과 및 발견

• PE 활용도: DiT는 불규칙한 형태에서 손수 튜닝한 GH200 커널이 약 70 %에 머물렀던 것에 비해, PE를 95 % 이상 지속적으로 활성화했습니다.
• 메모리 대역폭: 모델 기반 프리패치 스케줄은 이론적인 4 TB/s 대역폭의 90 % 이상을 달성하여, 단순 데이터 이동으로 인한 정체를 없앴습니다.
• 컴파일 시간: 자동 튜닝을 포함한 엔드‑투‑엔드 생성은 표준 워크스테이션에서 5 분 미만에 완료되었으며, 맞춤 ASIC 라이브러리의 수주에 걸친 수동 튜닝 사이클보다 훨씬 빠릅니다.

실용적인 시사점

• Accelerator Vendors: DiT는 새로운 AI 칩과 함께 “플러그‑앤‑플레이” GEMM 라이브러리를 제공할 수 있는 경로를 제공하여, 비용이 많이 들고 전문가 전용인 소프트웨어 팀의 필요성을 줄여줍니다.
• Framework Integrators: 딥러닝 스택(예: PyTorch, TensorFlow)은 커스텀 하드웨어의 백엔드로 DiT를 호출할 수 있어, 커널을 다시 작성하지 않고도 성능 향상을 얻을 수 있습니다.
• Edge & Cloud Deployments: GEMM을 다양한 행렬 형태에 자동으로 맞출 수 있는 능력은 비정방형 텐서(예: 트랜스포머 어텐션 헤드)를 자주 사용하는 추론 워크로드에 대해 더 나은 활용도를 제공합니다.
• Hardware‑Software Co‑Design: 설계자는 모델에서 타일 수, 인터커넥트 대역폭, 메모리 크기를 실험하면서 GEMM 성능에 미치는 영향을 즉시 확인할 수 있어, 설계 주기의 초기 단계에서 더 나은 실리콘 결정을 내릴 수 있습니다.

제한 사항 및 향후 작업

• 모델 충실도: 실행 가능한 모델이 주요 지연 원인을 포착하지만, 캐시 라인 충돌 및 전압‑주파수 스케일링과 같은 미세 효과는 추상화되어 실제 실리콘 결과에 영향을 줄 수 있습니다.
• GEMM을 넘어선 범위: DiT는 현재 조밀한 행렬 곱셈에 초점을 맞추고 있으며, 프레임워크를 컨볼루션, 희소 커널 또는 혼합 정밀도 워크로드로 확장하는 것은 아직 해결되지 않은 과제입니다.
• 하드웨어 다양성: 프로토타입은 특정 타일 기반 다중 PE 가속기에 초점을 맞추고 있으며, 타일과 벡터 유닛을 결합하는 이종 설계에 적용하려면 추가 모델링 작업이 필요합니다.
• 자동 튜닝 오버헤드: 경량이지만, 자동 튜너가 매우 큰 설계 공간에서 전역 최적 타일 크기를 놓칠 수 있으며, 보다 정교한 탐색 알고리즘을 통합하는 것이 계획된 방향입니다.

Design in Tiles는 올바른 추상화와 도구를 통해 전통적으로 어려웠던 GEMM을 최첨단 AI 가속기에 매핑하는 작업을 자동화할 수 있음을 보여주며, 성능 향상과 차세대 하드웨어의 시장 출시 시간을 단축시킵니다.

저자

• Aofeng Shen
• Chi Zhang
• Yakup Budanaz
• Alexandru Calotoiu
• Torsten Hoefler
• Luca Benini

논문 정보

• arXiv ID: 2512.13638v1
• 분류: cs.DC, cs.AR
• 출판일: 2025년 12월 15일
• PDF: PDF 다운로드