OpenAI Triton 1.0 출시

OpenAI는 연구자들이 광범위한 CUDA 경험 없이도 매우 효율적인 GPU 코드를 작성할 수 있도록 설계된 오픈 소스 Python 스타일 프로그래밍 언어이자 컴파일러인 Triton 1.0을 출시했습니다. Triton을 사용하면 개발자는 훨씬 적은 양의 코드로 cuBLAS 성능과 일치하는 FP16 행렬 곱셈과 같이 전문가가 튜닝한 CUDA 커널에 필적하는 하드웨어 성능을 달성할 수 있습니다.

GPU 프로그래밍의 복잡성 해결

네이티브 프레임워크 연산자에서 임시 텐서를 생성하고 이동할 때 발생하는 성능 저하를 피하기 위해 특화된 GPU 커널을 작성하는 것이 필요한 경우가 많습니다. 그러나 수동 GPU 최적화는 개발자가 세 가지 주요 아키텍처 구성 요소를 관리해야 하므로 어렵습니다:

DRAM: 현대적인 인터페이스의 버스 폭을 활용하기 위해 메모리 전송이 대규모 트랜잭션으로 병합(coalesced)되어야 합니다.
SRAM: 데이터는 재사용 전에 SRAM에 수동으로 저장되어야 하며, 공유 메모리 뱅크 충돌(bank conflicts)을 피하기 위해 관리되어야 합니다.
ALUs: 명령어/스레드 수준의 병렬성을 활용하고 텐서 코어를 사용하기 위해 계산은 스트리밍 멀티프로세서(SMs) 내부 및 전체에 걸쳐 신중하게 분할 및 스케줄링되어야 합니다.

Triton은 타일링(tiling) 및 SM 간 동기화와 같은 고수준 알고리즘 고려 사항은 개발자에게 맡기면서, 메모리 병합, 공유 메모리 관리 및 SM 내의 스케줄링을 자동화합니다.

Triton 프로그래밍 모델

Triton은 커널이 인스턴스 그리드에서 실행되는 데코레이터가 지정된 Python 함수로 정의되는 프로그래밍 모델을 사용합니다. CUDA나 Numba에서 사용하는 단일 명령어 다중 스레드(SIMT) 모델과 달리, Triton은 블록(2의 거듭제곱 차원을 가진 작은 배열)에 대한 연산을 통해 인스턴스 내 병렬성을 노출합니다.

주요 기술적 차이점

블록 기반 실행: 개별 스레드가 아닌 블록 단위로 작동함으로써, Triton은 공유 메모리 동기화 및 텐서 코어 스케줄링을 포함하여 CUDA 스레드 블록 내의 동시성 문제를 추상화합니다.
포인터 산술 연산: Triton JIT는 입력을 텐서가 아닌 포인터로 취급하여, block-sparse 텐서와 같은 복잡한 데이터 구조에 필요한 메모리 액세스에 대한 저수준 제어를 제공합니다.
융합 커널(Fused Kernels): Triton은 융합 커널 생성을 단순화합니다. 예를 들어, Triton의 융합된 softmax 구현은 정규화 과정 내내 행을 SRAM에 유지하여 데이터 재사용을 극대화할 수 있습니다. 이 접근 방식은 동일한 PyTorch 구현보다 최대 2배 더 효율적일 수 있습니다.

행렬 곱셈 성능

Triton은 신경망의 핵심 연산인 행렬 곱셈에 매우 효과적입니다. 약 25줄의 Python 코드만으로 V100 텐서 코어에서 피크 성능을 달성할 수 있습니다. 이러한 접근성은 개발자가 뛰어난 GPU 프로그래밍 전문 지식 없이도 융합된 변환(예: slicing 또는 Leaky ReLU)을 포함하도록 행렬 곱셈 커널을 맞춤 설정할 수 있게 해줍니다.

시스템 아키텍처 및 컴파일러 백엔드

Triton의 성능은 다차원 블록이 일급 시민(first-class citizens)인 LLVM 기반 중간 표현인 Triton-IR을 중심으로 하는 모듈식 아키텍처에서 비롯됩니다.

컴파일 파이프라인

Python AST: @triton.jit 데코레이터는 Python 함수의 추상 구문 트리(Abstract Syntax Tree)를 탐색합니다.
Triton-IR: AST는 정적 단일 할당(SSA) 구축 알고리즘을 사용하여 Triton-IR로 변환됩니다.
LLVM-IR & PTX: 컴파일러 백엔드는 IR을 단순화하고 최적화하며, 자동으로 병렬화하고, NVIDIA GPU에서 실행하기 위해 LLVM-IR 및 최종적으로 PTX로 변환합니다.

컴파일러 최적화

자동 메모리 저장(Automatic Memory Stashing): 컴파일러는 계산 집약적인 블록 수준 연산(예: tl.dot)의 피연산자를 분석하여 liveness 분석을 사용하여 데이터를 공유 메모리에 자동으로 저장합니다.
자동 병렬화: Triton은 서로 다른 커널 인스턴스를 동시에 실행하여 SM 간의 실행을 병렬화하고, SIMD 유닛에 걸쳐 블록 수준 연산의 반복 공간을 분할하여 SM 내의 실행을 병렬화합니다.