Triton (인공지능)

Triton은 Python 기반으로 GPU 커널을 작성할 수 있게 해주는 오픈소스 언어 및 컴파일러로, 고성능 CUDA 수준의 연산을 훨씬 간단한 코드로 구현할 수 있도록 설계되었다. 원래 Harvard 및 OpenAI에서 개발되었으며, 현재는 PyTorch의 컴파일 스택에 통합되어 딥러닝 연산 최적화를 위해 널리 사용되고 있다.

개요

Triton은 연구자와 개발자가 복잡하고 오류가 발생하기 쉬운 CUDA C++ 코드를 직접 작성하지 않고도, Python 스타일의 간결한 문법으로 고성능 GPU 커널을 작성할 수 있게 한다. Triton 컴파일러는 이러한 코드를 GPU 실행 가능한 저수준 커널 코드로 자동 변환하며, 연산 최적화, 자동 타일링, 벡터라이제이션 등 다양한 성능 향상 기능을 제공한다.

특징

Python 기반 GPU 커널 작성

Triton은 GPU 커널을 Python 코드처럼 작성하며, 이를 컴파일하여 CUDA 수준의 고성능을 제공한다.

• Python과 유사한 문법
• 저수준 메모리 접근 및 병렬 처리 구조를 자동 최적화
• 복잡한 CUDA boilerplate 코드 불필요

자동 타일링 및 벡터라이제이션

Triton 컴파일러는 matmul, convolution 같은 구조화된 연산에 대해:

• 자동 타일 크기 선택
• Vectorized memory load/store
• Shared memory 활용
• 워프/스레드 구조 최적화

를 처리하여, GPU 아키텍처에 맞춘 성능 이식성을 제공한다.

커널 Fusion 친화적 구조

Triton은 PyTorch 2.0의 TorchInductor와 결합하여 연산 그래프의 연속된 연산들을 하나의 fused kernel로 압축할 수 있다. 이는 메모리 왕복 및 커널 호출 오버헤드를 크게 감소시키며, 트레이닝 및 추론 성능 향상에 기여한다.

간결한 문법

예시:

import triton
import triton.language as tl

@triton.jit
def matmul_kernel(a_ptr, b_ptr, c_ptr, M, N, K):
    pid = tl.program_id(0)
    row = pid * 16 + tl.arange(0, 16)
    col = tl.arange(0, 16)
    a = tl.load(a_ptr + row[:, None] * K + tl.arange(0, K))
    b = tl.load(b_ptr + tl.arange(0, K)[:, None] * N + col)
    c = tl.dot(a, b)
    tl.store(c_ptr + row[:, None] * N + col, c)

위 코드는 전형적인 16×16 타일 기반 행렬 곱셈 커널이며, CUDA와 비교하면 훨씬 짧고 직관적이다.

필요성 및 배경

GPU 프로그래밍은 전통적으로 CUDA C++을 사용해야 했지만:

• boilerplate 코드가 많고
• 메모리 접근 패턴 설계가 어렵고
• 오류 발생률이 높으며
• cuDNN/cuBLAS가 제공하지 않는 커스텀 fused kernel이 필요할 때 문제가 발생했다.

Triton은 이러한 문제를 해결하기 위해 등장한 언어로, 프레임워크 개발자 및 연구자들이 보다 쉽게 커스텀 고성능 GPU 연산을 구축할 수 있게 한다.

PyTorch와의 통합

PyTorch 2.0에서 Triton은 중요한 역할을 한다.

TorchInductor Backend

• PyTorch eager graph를 분석
• 연산을 그룹화(fusion)
• Triton kernel로 자동 생성
• GPU에서 실행

이를 통해 기존 PyTorch 대비 상당한 성능 향상을 제공한다.

예시 활용 분야

• LayerNorm, Softmax 같은 element-wise 연산 최적화
• Attention 연산의 fused kernel
• Custom matmul / conv 변형
• Transformer block 내 다양한 연산 통합

장점

• CUDA 수준의 성능을 Python 기반으로 달성 가능
• 학습曲線이 CUDA보다 훨씬 낮음
• 유지보수 및 커스텀 연산 개발에 유리
• PyTorch와의 긴밀한 통합으로 실질적 속도 향상 제공

한계

• CUDA, cuDNN, cuBLAS만큼 성숙하지는 않음
• 매우 복잡한 커널이나 SM 구조에 최적화된 커널은 직접 CUDA로 작성해야 할 수도 있음
• NVIDIA GPU 중심 생태계 (다른 GPU 아키텍처 지원은 제한적)

함께 보기

• PyTorch
• CUDA
• GPU
• 자동 미분
• 딥 러닝 컴파일러
• XLA

참고 문헌

• T. Tillet, H.-T. Kung, D. Cox, Triton: An intermediate language and compiler for tiled neural network computations. Proceedings of ACM SIGPLAN International Workshop on Machine Learning and Programming Languages (2019).
• PyTorch Triton 공식 문서