[Paper] tritonBLAS: 트리톤 기반 GEMM 커널 파라미터 선택을 위한 분석적 접근법

Source: arXiv - 2512.04226v1

개요

이 논문은 tritonBLAS라는 결정론적 분석 모델을 소개한다. 이 모델은 GPU에서 GEMM(일반 행렬 곱) 커널에 대한 고성능 파라미터를 자동으로 선택한다. 캐시 계층 구조와 데이터 배치와 같은 아키텍처 세부 정보를 활용함으로써, tritonBLAS는 대부분의 라이브러리가 의존하는 비용이 큰 런타임 자동 튜닝 없이도 거의 최적에 가까운 커널을 생성할 수 있다.

주요 기여

• 분석 성능 모델: GPU 마이크로아키텍처(캐시 크기, 공유 메모리 레이아웃 등)를 GEMM 차단 파라미터에 매핑한다.
• Triton‑전용 구현: 경량 GEMM 프레임워크를 Triton만으로 구현해 손으로 작성한 CUDA 커널이나 외부 라이브러리가 필요 없게 한다.
• 런타임 자동 튜닝 제로: 모델이 컴파일 시 최적 구성을 예측하여 최신 자동 튜닝 솔루션 성능의 95 % 이상을 달성한다.
• 광범위한 평가: 다양한 행렬 형태와 최신 GPU(NVIDIA Ampere, Hopper 등)에서 일관된 속도 향상과 낮은 오버헤드를 보여준다.
• 오픈소스 가능성: 기존 Triton 기반 프로젝트에 통합하거나 생산 파이프라인에서 cuBLAS/rocBLAS를 대체하는 드롭‑인 방식으로 사용할 수 있다.

방법론

1. 아키텍처 프로파일링 – 저자들은 대상 GPU에서 핵심 하드웨어 파라미터(L1/L2 캐시 용량, SM당 공유 메모리, 워프 수 등)를 추출한다.
2. 분석 차단 모델 – 이러한 파라미터와 GEMM 타일링 선택(블록 크기, 스레드 블록 형태, 레지스터 사용량) 사이의 관계를 나타내는 방정식 집합을 만든다. 모델은 세 단계 차단을 포괄한다:
   • 레지스터 타일링(내부 마이크로 커널)
   • 공유 메모리 타일링(중간 단계)
   • 캐시 레벨 타일링(외부)
3. 파라미터 선택 알고리즘 – 가벼운 탐색(대부분 폐쇄형)으로 가능한 타일링 구성을 평가하고, 메모리 대역폭 및 점유율 제약을 만족하면서 추정 연산 강도를 최대화하는 구성을 선택한다.
4. Triton에서 커널 생성 – 선택된 파라미터를 일반적인 Triton GEMM 템플릿에 전달하면, Triton 컴파일러가 이를 구체적인 GPU 커널로 특수화한다.
5. 검증 – 생성된 커널을 자동 튜닝된 라이브러리(cuBLAS, CUTLASS 등)와 비교 벤치마크한다. 행렬 크기는 64×64와 같은 작은 경우부터 8192×8192와 같은 큰 경우까지 다양하게 테스트한다.

결과 및 발견

• 일관된 성능: 정방형, 길쭉한, 짧은 행렬 모두에서 동일하게 우수한 결과를 보인다.
• 자동 튜닝 시간 제로: 분석 모델은 몇 마이크로초 내에 실행되며, 전통적인 자동 튜너가 수분에 걸쳐 수행하는 검색 단계와 비교할 수 없을 정도로 빠르다.
• 메모리 효율성: 선택된 타일링이 L2 캐시 재사용을 보장해, 단순 커널보다 DRAM 트래픽을 감소시킨다.
• 확장성: 모델이 각 디바이스에 국한돼 있기 때문에 다중 GPU 환경에서도 별도의 디바이스 간 프로파일링 없이 적용 가능하다.

실용적 함의

• 배포 주기 가속 – 팀은 새로운 GEMM‑중심 워크로드(예: 트랜스포머 추론, 과학 시뮬레이션)를 각 대상 머신에서 자동 튜닝을 기다릴 필요 없이 바로 배포할 수 있다.
• 예측 가능한 성능 – 모델이 결정론적이므로 성능 퇴보를 추적하고 디버깅하기가 자동 튜닝의 확률적 결과보다 쉽다.
• 클라우드 비용 절감 – 긴 자동 튜닝 단계가 사라짐으로써 온디맨드 GPU 인스턴스 사용 시 직접적인 컴퓨팅 비용 감소 효과가 있다.
• 통합 용이성 – TritonBLAS는 기존 Triton 코드베이스(예: diffusion 모델용 커스텀 커널)에 단일 import만으로 삽입할 수 있어, 즉시 고성능 GEMM 프리미티브를 제공한다.
• 이식성 – 분석 모델은 새로운 GPU 세대에 자동으로 적응한다. 하드웨어 파라미터 추출 단계만 업데이트하면 되므로 향후 아키텍처에도 미래 지향적이다.

제한 사항 및 향후 연구

• 모델 정확도 – 자동 튜닝 성능의 95 % 이상을 달성하는 것은 인상적이지만, 매우 불균형한 행렬 형태와 같은 극단적인 경우에는 여전히 작은 격차가 존재한다.
• 비‑GEMM 커널 – 현재 프레임워크는 밀집 행렬 곱에만 초점을 맞추고 있다. 컨볼루션이나 희소 커널에 분석 접근법을 확장하는 것은 아직 해결되지 않은 과제이다.
• 동적 워크로드 – 런타임에 행렬 크기가 변하는 경우, 모델을 가볍게 재평가해야 한다. 저자들은 이전 선택을 캐시하고 재사용하는 방안을 계획하고 있다.
• 하드웨어 다양성 – 연구는 NVIDIA GPU에 집중했으며, AMD나 Intel GPU에 적용하려면 아키텍처 모델을 별도로 조정해야 한다.

전반적으로 tritonBLAS는 잘 설계된 분석 모델이 GEMM에 대한 비용이 큰 자동 튜닝을 대체할 수 있음을 보여준다. 이를 통해 개발자는 빠르고 신뢰할 수 있으며 이식 가능한 방법으로 GPU 성능에 근접할 수 있다.

저자

• Ryan Swann
• Muhammad Osama
• Xiaohu Guo
• Bryant Nelson
• Lixun Zhang
• Alex Brown
• Yen Ong
• Ali Yazdani
• Sean Siddens
• Ganesh Dasika
• Alex Underwood

논문 정보

• arXiv ID: 2512.04226v1
• Categories: cs.DC
• Published: December 3, 2025
• PDF: Download PDF