원문 캡처 · junupark.xyz
원문 캡처 · junupark.xyz

Triton은 오픈AI가 처음 개발한 프로그래밍 언어이자 컴파일러다. CUDA보다 사용자 정의 GPU 커널을 쉽게 작성하도록 돕고, 특히 딥러닝 작업에서 효율적인 GPU 코드를 표현하기 쉽게 한다. 이 노트는 벡터 덧셈과 행렬 곱셈을 예로 들어 Triton의 추상화와 커널 실행·검증·튜닝 방법을 설명한다.

CUDA에서는 보통 각 스레드가 무엇을 할지 정하고 저수준 세부 사항을 직접 관리한다. Triton에서는 데이터의 블록, 즉 타일을 하나의 프로그램이 어떻게 처리할지 기술한다. 컴파일러가 블록 연산을 GPU 스레드에 배치하고 GPU 코드를 생성한다. Triton은 PyTorch와 NumPy에서 쓰는 배열 프로그래밍 방식도 활용한다. 배열의 모양과 브로드캐스팅이 익숙하지 않다면 NumPy 브로드캐스팅 시각 안내와 반복문 없이 하는 NumPy 배열 프로그래밍이 참고 자료가 된다.

GPU 프로그래밍은 DSL(도메인 특화 언어)과 관계없이 크게 CPU에서 실행되는 호스트 코드와 GPU에서 실행되는 디바이스 코드로 나뉜다. 호스트 코드는 입력 텐서를 준비하고 병렬 작업량을 정해 커널을 실행한다. 디바이스 코드는 각 GPU 프로그램이 수행할 계산을 정의한다. 전체 작업은 호스트와 디바이스 메모리를 만들고, 데이터를 디바이스로 복사한 뒤, 온칩 메모리로 불러와 계산하고 결과를 디바이스 메모리에 저장하는 흐름이다.

커널 실행과 타일 배정

호스트(CPU)는 디바이스 작업이 시작되기 전에 커널을 실행한다. 한 번의 실행으로 커널 인스턴스 여러 개가 만들어지며, Triton에서는 이를 프로그램이라고 부른다. 커널 실행을 준비할 때는 입력·출력 텐서와 크기, 스트라이드, 컴파일 시점 인수 등을 전달하고 그리드 크기를 정한다.

커널은 변수와 메모리를 유연하게 다루므로 모양, 스트라이드, 자료형, Blackwell에서 도입된 기능처럼 기기별로 달라지는 조건을 검사하는 것이 유용하다. 이런 검사는 오류를 일찍 드러내고 정확성을 확인하는 한편, 조건에 따라 별도의 최적화 경로를 선택할 수 있게 한다.

python
def vector_add(a: torch.Tensor, b: torch.Tensor):
    assert a.is_cuda and b.is_cuda
    assert a.device == b.device
    assert a.ndim == b.ndim == 1
    assert a.shape == b.shape
    assert a.dtype == b.dtype
    # ...

그리드는 각 축에서 실행할 프로그램 수를 지정하는 1차원·2차원·3차원 튜플이다. 커널 안에서 tl.program_id(axis)를 호출하면 해당 축에서 현재 프로그램의 위치를 얻는다. 그리드 배치가 좋지 않으면 병렬성이나 캐시 지역성이 떨어질 수 있으므로 작업에 맞게 정해야 한다.

일반적인 비영속 타일 실행에서는 출력을 BN개 원소의 타일로 나누고 타일마다 프로그램을 하나씩 실행한다. 따라서 N개 원소를 덮는 데 필요한 타일 수가 그리드 크기가 된다. triton.cdiv(a, b)는 ceil(a / b)를 계산해 마지막에 남는 부분 타일에도 프로그램을 배정한다. 커널은 텐서의 논리적 범위를 벗어나는 레인을 마스크해야 한다. 그리드를 함수로 지정하면 Triton은 BN 같은 컴파일 시점 메타 매개변수가 담긴 딕셔너리를 전달하므로, 자동 튜닝으로 선택된 블록 크기에 따라 그리드도 정할 수 있다.

python
# ...
grid = lambda meta: (triton.cdiv(N, meta["BN"]),)
_vec_add_kernel[grid](
        # ...
    )

영속 커널은 제한된 수의 프로그램을 실행하고, 각 프로그램이 커널 안에서 여러 타일을 반복 처리한다. 프로그램 수는 출력 타일 수를 넘지 않도록 하고 GPU의 SM 수에 맞춰 정하는 경우가 많다. 이 방식에서는 호스트가 출력 타일 수보다 적은 프로그램을 실행하도록 그리드를 직접 정해야 한다.

커널에는 읽을 입력 텐서와 미리 할당해 쓸 출력 텐서, 각 차원의 논리적 크기, 스트라이드, 런타임 스칼라 값, tl.constexpr로 표시한 컴파일 시점 인수를 전달한다. 스트라이드는 텐서의 한 차원에서 이웃한 위치로 이동할 때 메모리에서 건너뛰는 원소 수다. 벡터가 연속 메모리에 있으면 보통 스트라이드는 1이지만, 비연속 벡터에서는 다를 수 있다. tl.constexpr 인수는 컴파일 시점에 알려져 있어 코드 특수화, 블록 텐서 크기 지정, 컴파일 시점 분기 제거에 활용할 수 있다. 타일 크기는 자동 튜닝 대상으로 자주 쓰이며, 알고리즘 변형이나 기능 플래그에도 사용한다.

python
def vector_add(a: torch.Tensor, b: torch.Tensor):
    # ...
    _vec_add_kernel[grid](
        a, b, c, # a, b as input tensor, c as output tensor
        N,
        a.stride(0), b.stride(0), c.stride(0), # all 1
        BN=1024, # known at compile time; we can autotune it
    )

벡터 덧셈을 실행하는 호스트 코드는 입력이 같은 GPU에 있고 차원·모양·자료형이 일치하는지 검사한다. 출력 텐서를 만들고 N을 구한 뒤, BN에 맞는 그리드를 계산해 커널을 실행하고 결과를 반환한다.

python
def vector_add(a: torch.Tensor, b: torch.Tensor):
    assert a.is_cuda and b.is_cuda
    assert a.device == b.device
    assert a.ndim == b.ndim == 1
    assert a.shape == b.shape
    assert a.dtype == b.dtype
    c = torch.empty_like(a)
    N = c.numel()
    grid = lambda meta: (triton.cdiv(N, meta["BN"]),)
    _vec_add_kernel[grid](
        a, b, c, # a, b as input tensor, c as output tensor
        N,
        a.stride(0), b.stride(0), c.stride(0), # all 1
        BN=1024, # known at compile time; we can autotune it
    )
    return c

Triton 커널의 타일과 프로그램

타일은 함께 처리하는 데이터 영역이다. 예를 들어 벡터 덧셈 프로그램 하나가 출력 C의 연속된 원소 1024개를 처리할 수 있다. 타일은 프로그램 안에서 블록 형태의 오프셋·포인터·값으로 표현된다. Triton에서 타일과 블록은 흔히 비슷한 뜻으로 쓰지만, 타일은 전체 텐서에서 논리적으로 맡은 영역을, 블록은 프로그램 안에서 다루는 블록 모양의 오프셋·포인터·값을 가리킨다.

text
tile 0: C[   0 : 1024]
tile 1: C[1024 : 2048]
tile 2: C[2048 : 3072]
...

프로그램은 컴파일된 커널의 실행 인스턴스로, CUDA 스레드 블록(CTA)과 대략 대응한다. 호스트가 그리드에 지정한 수만큼 프로그램을 실행한다. N = 4,096, BN = 1024라면 네 프로그램이 각기 다른 타일을 맡는다. 모든 프로그램은 같은 커널 코드를 실행하지만 서로 다른 데이터를 처리한다. 이를 SPMD(단일 프로그램 다중 데이터) 방식이라고 한다.

text
program 0 → tile 0 → C[   0 : 1024]
program 1 → tile 1 → C[1024 : 2048]
program 2 → tile 2 → C[2048 : 3072]
program 3 → tile 3 → C[3072 : 4096]

일반적인 커널 프로그램은 프로그램 ID를 읽어 맡은 출력 타일을 찾고, 타일의 전역 원소 오프셋을 만든다. 그 오프셋으로 입력·출력 포인터를 구성하고 입력 블록을 읽어 계산한 뒤 결과를 저장한다. 텐서는 커널 실행 때 GPU 저장소를 가리키는 포인터로 전달된다. 논리적 차원 크기는 반복과 마스크에, 스트라이드는 다차원 인덱스를 메모리 위치로 바꾸는 데 쓴다. 선택적으로 드롭아웃 확률이나 스케일 값 같은 런타임 스칼라를 전달할 수 있고, 타일 크기나 기능 플래그는 컴파일 시점 메타 매개변수로 전달한다.

python
@triton.jit
def _vec_add_kernel(
    a_ptr, b_ptr, c_ptr,
    N,
    stride_a, stride_b, stride_c,
    BN: tl.constexpr,
):
    # ...

tl.program_id(axis)로 얻는 프로그램 ID(pid)는 그리드의 한 축에서 현재 프로그램을 식별한다. 1차원 그리드를 실행했다면 axis=0을 사용한다. 프로그램 ID와 타일 크기로 전역 오프셋을 만들면 프로그램마다 서로 다른 타일을 처리한다. tl.arange는 np.arange나 torch.arange처럼 타일 안의 지역 오프셋을 만든다.

python
@triton.jit
def _vec_add_kernel(
    # ...
):
    pid = tl.program_id(axis=0)
python
local_offs_n = tl.arange(0, BN)
python
offs_n = pid * BN + tl.arange(0, BN)

예를 들어 BN = 1024, pid = 3이면 오프셋은 3072부터 4095까지다.

text
3 * 1024 + tl.arange(0, 1024) == [3072, 3073, 3074, 3075, ..., 4095]

이 오프셋은 출력 좌표를 블록 모양 텐서로 만든다. 좌표에 기본 포인터와 스트라이드를 적용해 A, B, C의 포인터를 구성한다. N이 BN의 배수가 아닌 경우에도 마지막 부분 타일을 안전하게 처리하도록 전역 오프셋에 마스크를 적용한다.

python
@triton.jit
def _vec_add_kernel(
    # ...
):
    # ...
    offs_n = pid * BN + tl.arange(0, BN)
    n_mask = offs_n < N

전역 오프셋에 각 차원의 스트라이드를 곱하고 기본 포인터에 더하면 메모리 주소가 된다. 스트라이드는 해당 축에서 한 칸 이동할 때 메모리 원소를 몇 개 건너뛰는지 나타낸다.

python
@triton.jit
def _vec_add_kernel(
    a_ptr, b_ptr, c_ptr,
    # ...
):
    # ...
    a_ptrs = a_ptr + offs_n * stride_a
    b_ptrs = b_ptr + offs_n * stride_b
    c_ptrs = c_ptr + offs_n * stride_c

입력은 보통 커널 실행 전 GPU 전역 메모리에 있다. tl.load로 읽은 값을 온칩 자원(대개 레지스터)에서 계산한 뒤 tl.store로 결과를 쓴다. 캐시 계층을 거치는 저수준 데이터 이동은 컴파일러와 하드웨어가 관리한다. 마스크로 범위 밖 원소를 걸러내며, 마스크된 로드에 other를 지정해 대체 값을 준다. 벡터 덧셈에서는 0.0을 쓸 수 있고, 상황에 따라 -float("inf")를 사용하기도 한다.

python
@triton.jit
def _vec_add_kernel(
    a_ptr, b_ptr,
    # ...
):
    # ...
    a_ptrs = a_ptr + offs_n * stride_a
    b_ptrs = b_ptr + offs_n * stride_b
    a = tl.load(a_ptrs, mask=n_mask, other=0.0)
    b = tl.load(b_ptrs, mask=n_mask, other=0.0)

벡터 덧셈처럼 단순한 계산은 기본 산술 연산으로 표현할 수 있고, 더 복잡한 커널은 추가 연산이나 행렬 곱셈 tl.dot을 사용할 수 있다.

python
@triton.jit
def _vec_add_kernel(
    a_ptr, b_ptr,
    # ...
):
    # ...
    a_ptrs = a_ptr + offs_n * stride_a
    b_ptrs = b_ptr + offs_n * stride_b
    a = tl.load(a_ptrs, mask=n_mask, other=0.0)
    b = tl.load(b_ptrs, mask=n_mask, other=0.0)
    c = a + b

계산한 결과는 출력 포인터와 마스크를 사용해 디바이스 메모리에 저장한다.

python
@triton.jit
def _vec_add_kernel(
    a_ptr, b_ptr, c_ptr,
    # ...
):
    # ...
    c = a + b
    c_ptrs = c_ptr + offs_n * stride_c
    tl.store(c_ptrs, c, mask=n_mask)

앞서 설명한 단계를 합친 벡터 덧셈 커널은 다음과 같다.

python
@triton.jit
def _vec_add_kernel(
    a_ptr, b_ptr, c_ptr,
    N,
    stride_a, stride_b, stride_c,
    BN: tl.constexpr
):
    pid = tl.program_id(axis=0)
    offs_n = pid * BN + tl.arange(0, BN)
    n_mask = offs_n < N
    a_ptrs = a_ptr + offs_n * stride_a
    b_ptrs = b_ptr + offs_n * stride_b
    a = tl.load(a_ptrs, mask=n_mask, other=0.0)
    b = tl.load(b_ptrs, mask=n_mask, other=0.0)
    c = a + b
    c_ptrs = c_ptr + offs_n * stride_c
    tl.store(c_ptrs, c, mask=n_mask)

행렬 곱셈과 다차원 인덱싱

행렬 곱셈 같은 딥러닝 연산은 2차원 이상의 타일을 사용한다. 호스트 코드는 입력이 GPU에 있는지, 두 텐서가 같은 장치에 있는지, 차원이 2인지, A의 열 수와 B의 행 수가 같은지 검사한다. 예제는 두 입력의 자료형이 모두 torch.float16인 경우를 대상으로 한다. 출력 크기와 그리드를 계산한 뒤 BM=32, BN=32, BK=32 설정으로 커널을 실행한다.

python
def matmul(a: torch.Tensor, b: torch.Tensor):
    assert a.is_cuda and b.is_cuda
    assert a.device == b.device
    assert a.ndim == 2 and b.ndim == 2
    assert a.shape[1] == b.shape[0] # K dimensions must match
    assert a.dtype == b.dtype == torch.float16
    M, K = a.shape
    _, N = b.shape
    c = torch.empty((M, N), device=a.device, dtype=a.dtype)
    grid = lambda meta: (
        triton.cdiv(M, meta["BM"]),
        triton.cdiv(N, meta["BN"]),
    )
    _matmul_kernel[grid](
        a, b, c,
        M, N, K,
        a.stride(0), a.stride(1),
        b.stride(0), b.stride(1),
        c.stride(0), c.stride(1),
        BM=32, BN=32, BK=32,
    )
    return c

2차원 텐서에는 각 텐서마다 두 스트라이드가 필요하다. 행 우선(row-major)으로 저장된 모양 [M, K] 행렬 A는 열 방향으로 한 칸 이동하면 원소 하나, 행 방향으로 이동하면 K개 원소를 건너뛴다. 연속된 행렬의 스트라이드는 (K, 1)이다. 전치 텐서나 다른 뷰는 메모리 배치가 다를 수 있으므로 a.stride(0)과 a.stride(1)을 전달하면 행과 열 방향의 이동량을 반영할 수 있다.

행렬 곱셈 커널은 호스트에서 전달한 포인터, 행렬 크기와 스트라이드, 컴파일 시점 타일 크기를 받는다. 2차원 그리드의 각 축에서 프로그램 ID를 구하고 BM, BN 크기의 출력 타일 좌표와 BK 크기의 축 좌표를 만든다.

python
@triton.jit
def _matmul_kernel(
    a_ptr, b_ptr, c_ptr,
    M, N, K,
    stride_a_m, stride_a_k,
    stride_b_k, stride_b_n,
    stride_c_m, stride_c_n,
    BM: tl.constexpr,
    BN: tl.constexpr,
    BK: tl.constexpr,
):
python
    pid_m = tl.program_id(axis=0)
    pid_n = tl.program_id(axis=1)

영속 실행에서는 출력 타일마다 프로그램을 하나씩 실행하는 대신 더 적은 프로그램을 띄우고 각 프로그램이 반복문으로 여러 타일을 처리할 수 있다. tl.num_programs(axis)는 해당 축에서 실행된 프로그램 수를 돌려준다. 2차원 출력 타일 격자를 1차원 순서로 펼치면 각 프로그램이 타일 목록을 일정 간격으로 순회한다.

python
start_tile = tl.program_id(axis=0)
num_programs = tl.num_programs(axis=0)
num_pid_m = tl.cdiv(M, BM)
num_pid_n = tl.cdiv(N, BN)
num_tiles = num_pid_m * num_pid_n
for tile_id in tl.range(start_tile, num_tiles, num_programs):
    pid_m = tile_id // num_pid_n
    pid_n = tile_id % num_pid_n
    # Compute the output tile C[pid_m, pid_n].
    # ...

예를 들어 16개 출력 타일에 프로그램 네 개를 실행하면 프로그램 0은 타일 0, 4, 8, 12를 처리한다. tl.num_programs()는 실행된 프로그램 수를 알려줄 뿐, 그 자체로 커널을 영속 방식으로 만들지는 않는다. 더 작은 그리드를 호스트에서 명시적으로 지정해야 한다.

각 축의 전역 좌표는 프로그램 ID와 타일 크기에 지역 좌표를 더해 만든다.

python
    offs_m = pid_m * BM + tl.arange(0, BM)
    offs_n = pid_n * BN + tl.arange(0, BN)
    offs_k = tl.arange(0, BK)

포인터 계산에서는 차원별 좌표에 스트라이드를 곱해 기본 포인터에 더한다. [:, None]과 [None, :]는 배열 차원을 늘려 브로드캐스팅에 맞는 모양을 만든다. 브로드캐스팅은 모양이 다른 텐서끼리 원소별 연산을 할 때 크기 1인 차원의 값이 큰 차원에 맞춰 확장된 것처럼 다루는 방식이다. 이 확장은 논리적이며, 더 큰 텐서를 메모리에 실제로 할당할 필요는 없다. NumPy, PyTorch, Triton은 가장 오른쪽 차원부터 모양을 비교한다. 차원 크기가 같거나, 어느 한쪽이 1이거나, 한 모양의 차원 수가 적어 빠진 앞쪽 차원을 1로 간주할 수 있으면 호환된다.

모양 [3, 4]와 [3, 4], [1, 4], [1, 1]은 호환된다. 크기 1인 차원이 각각 확장될 수 있기 때문이다. [1, 2]는 가장 오른쪽 크기 4와 2가 다르고 어느 쪽도 1이 아니므로 호환되지 않는다. 행렬 곱셈에서는 [BM, 1] 모양의 행 좌표와 [1, BK] 모양의 K 좌표가 함께 [BM, BK] 포인터 타일을 만든다. 마찬가지로 K 좌표와 열 좌표가 [BK, BN] 타일을 만든다.

text
offs_m[:, None] : (BM, 1)
offs_k[None, :] : (1, BK)
combined        : (BM, BK)
text
offs_k[:, None] : (BK, 1)
offs_n[None, :] : (1, BN)
combined        : (BK, BN)

Triton에서는 tl.expand_dims(offs_m, axis)나 tl.unsqueeze(offs_m, axis)를 쓸 수도 있지만 보통 :와 None 표기를 사용한다.

python
offs_m[:, None]  # column shape (BM, 1)
python
offs_k[None, :]  # row shape (1, BK)

다차원 포인터 연산은 축마다 따로 이해할 수 있다. 인덱스는 이동할 칸 수를, 스트라이드는 한 칸을 이동할 때 메모리에서 건너뛸 원소 수를 나타낸다. 둘을 곱한 값이 그 축의 메모리 오프셋이며, 브로드캐스팅이 모든 축의 기여분을 합쳐 타일의 각 좌표에 대응하는 포인터를 만든다.

python
a_ptrs = (
    a_ptr
    + offs_m[:, None] * stride_a_m
    + offs_k[None, :] * stride_a_k
)

행 좌표에 행 스트라이드를 곱한 값은 모든 열에 브로드캐스팅된다. 열 좌표에 열 스트라이드를 곱한 값은 모든 행에 적용된다. 두 기여분을 a_ptr에 더하면 M×K 포인터 격자가 된다. B의 포인터도 같은 방식으로 만든다.

python
    a_ptrs = (
        a_ptr
        + offs_m[:, None] * stride_a_m
        + offs_k[None, :] * stride_a_k
    )
    b_ptrs = (
        b_ptr
        + offs_k[:, None] * stride_b_k
        + offs_n[None, :] * stride_b_n
    )

프로그램은 출력의 BM × BN 타일을 맡는다. K 차원 전체를 한꺼번에 읽는 대신 BK 원소씩 나눠 처리하고 부분 결과를 accumulator에 누적한다. 누산기는 정밀도를 위해 tl.float32로 만들고 반복이 끝난 뒤 tl.float16으로 변환한다. 각 반복에서 A의 BM × BK 타일과 B의 BK × BN 타일을 읽으며, 마스크는 포인터와 같은 브로드캐스팅 방식으로 구성한다.

python
    accumulator = tl.zeros((BM, BN), dtype=tl.float32)
    for k_tile in range(0, tl.cdiv(K, BK)):
        k = k_tile * BK + offs_k
        a = tl.load(
            a_ptrs,
            mask=(offs_m[:, None] < M) & (k[None, :] < K),
            other=0.0,
        )
        b = tl.load(
            b_ptrs,
            mask=(k[:, None] < K) & (offs_n[None, :] < N),
            other=0.0,
        )

tl.dot은 현재 K 청크의 계산 결과를 누산기에 더한다. 이어 포인터를 다음 K 청크 위치로 이동한다. 마지막에는 누산기를 float16으로 바꾸고 출력 포인터와 마스크를 만들어 저장한다.

python
        accumulator = tl.dot(a, b, acc=accumulator)
        a_ptrs += BK * stride_a_k
        b_ptrs += BK * stride_b_k
python
    c = accumulator.to(tl.float16)
    c_ptrs = (
        c_ptr
        + offs_m[:, None] * stride_c_m
        + offs_n[None, :] * stride_c_n
    )
    c_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
    tl.store(c_ptrs, c, mask=c_mask)

위 과정을 합친 행렬 곱셈 커널은 Triton 행렬 곱셈 튜토리얼의 L2 캐시 최적화를 포함하지 않는 기본 형태다.

python
@triton.jit
def _matmul_kernel(
    a_ptr, b_ptr, c_ptr,
    M, N, K,
    stride_a_m, stride_a_k,
    stride_b_k, stride_b_n,
    stride_c_m, stride_c_n,
    BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr,
):
    pid_m = tl.program_id(axis=0)
    pid_n = tl.program_id(axis=1)
    offs_m = pid_m * BM + tl.arange(0, BM)
    offs_n = pid_n * BN + tl.arange(0, BN)
    offs_k = tl.arange(0, BK)
    a_ptrs = (
        a_ptr
        + offs_m[:, None] * stride_a_m
        + offs_k[None, :] * stride_a_k
    )
    b_ptrs = (
        b_ptr
        + offs_k[:, None] * stride_b_k
        + offs_n[None, :] * stride_b_n
    )
    accumulator = tl.zeros(
        (BM, BN),
        dtype=tl.float32,
    )
    for k_tile in range(0, tl.cdiv(K, BK)):
        k = k_tile * BK + offs_k
        a = tl.load(
            a_ptrs,
            mask=(offs_m[:, None] < M) & (k[None, :] < K),
            other=0.0,
        )
        b = tl.load(
            b_ptrs,
            mask=(k[:, None] < K) & (offs_n[None, :] < N),
            other=0.0,
        )
        accumulator = tl.dot(a, b, acc=accumulator)
        a_ptrs += BK * stride_a_k
        b_ptrs += BK * stride_b_k
    c = accumulator.to(tl.float16)
    c_ptrs = (
        c_ptr
        + offs_m[:, None] * stride_c_m
        + offs_n[None, :] * stride_c_n
    )
    c_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
    tl.store(c_ptrs, c, mask=c_mask)

설정 선택과 추가 연산

Triton은 하드웨어의 세부 사항을 직접 다루지 않도록 추상화하지만, 성능은 작업과 GPU에 맞는 타일 및 실행 설정에 좌우된다. 흔히 조정하는 값은 BM, BN, BK, num_warps, num_stages다.

첫 번째 방법은 @triton.autotune과 triton.Config 목록으로 자동 튜닝하는 것이다. 각 설정은 블록 크기 같은 컴파일 시점 키워드 인수와 num_warps, num_stages, num_ctas 같은 실행 옵션을 지정할 수 있다. 전체 항목은 triton.Config 문서를 참고한다.

python
autotune_configs = [
    triton.Config({"BM": 128, "BN": 256, \
    "BK": 64}, num_stages=3, num_warps=8),
    triton.Config({"BM": 64, "BN": 256, \
    "BK": 32}, num_stages=4, num_warps=4),
    # ...
]
@triton.autotune(configs=autotune_configs, key=["M", "N", "K"])
@triton.jit
def _matmul_kernel(
    a_ptr, b_ptr, c_ptr,
    M, N, K,
    stride_a_m, stride_a_k,
    stride_b_k, stride_b_n,
    stride_c_m, stride_c_n,
    # meta-parameters
    BM: tl.constexpr, BN: tl.constexpr, \
    BK: tl.constexpr,
):
    ...

가능한 설정을 모두 시험하면 시간이 오래 걸리므로 보통 현실성 있는 후보를 제한해 제공한다. 검색 공간이 크면 prune_configs_by로 후보를 줄일 수 있다. 컴파일과 콜드 스타트 지연이 중요하다면 @triton.heuristics나 수동 조회 테이블로 실행 중 튜닝을 피할 수도 있다. @triton.autotune()의 key는 튜닝 문제를 구분하는 런타임 값이다. 키 값이 하나라도 바뀌면 Triton은 새 키에 맞춰 후보를 평가하고 결과를 캐시한다.

두 번째 방법은 모양·자료형·장치·기능 플래그를 기준으로 수동 조회 테이블이나 휴리스틱을 사용하는 것이다. 자동 튜닝은 여러 모양과 자료형에 유연하게 대응하지만, 휴리스틱은 런타임 오버헤드를 더 줄일 때 유용할 수 있다. 실행 중 튜닝과 컴파일 오버헤드가 중요한 지연 민감형 추론 서비스에서 고려할 수 있다.

행렬 곱셈에 ReLU 같은 활성화 함수를 융합할 때는 tl.where를 출력 타일의 원소별 조건식으로 사용할 수 있다.

python
c = tl.where(accumulator > 0.0, accumulator, 0.0)

tl.where는 두 값 분기를 모두 평가한다. 분기 중 하나가 안전하지 않은 메모리 접근을 일으킬 수 있다면 tl.where 대신 tl.load와 tl.store의 마스크를 사용한다. 자세한 내용은 tl.where API 문서를 참고한다.

axis 인수는 축소할 차원을 지정하며, 해당 차원은 결과에서 제거된다. 모양이 [B0, B1]인 2차원 텐서에서 axis=0으로 축소하면 [B1], axis=1로 축소하면 [B0]이 된다. 다음 예에서 axis=1은 각 행의 열을 더한다.

text
k = [
  [1, 2, 3, 4],
  [10, 20, 30, 40],
]
python
tl.sum(k, axis=1)
text
[1+2+3+4, 10+20+30+40] = [10, 100]

정확성 테스트와 성능 측정

성능을 재기 전에 정확성을 확인해야 한다. 대표적인 모양, 타일 크기로 나누어떨어지지 않는 모양, 커널이 지원한다고 내세우는 각 자료형과 메모리 배치를 시험한다. 행렬 곱셈은 정사각형만 검사하지 말고 M, N, K를 각각 바꿔 테스트한다. 신뢰할 수 있는 기준 구현으로는 보통 PyTorch를 사용한다.

python
def test_matmul(M, N, K):
    torch.manual_seed(0)
    a = torch.randn((M, K), device=DEVICE, dtype=torch.float16)
    b = torch.randn((K, N), device=DEVICE, dtype=torch.float16)
    actual = matmul(a, b)
    expected = torch.matmul(a, b)
    torch.testing.assert_close(
        actual,
        expected,
        atol=1e-2,
        rtol=1e-2,
    )
if __name__ == "__main__":
    for shape in [
        (128, 128, 128),  # regular tiles
        (1, 17, 33),      # small problem
        (127, 65, 129),   # partial tiles on every axis
    ]:
        test_matmul(*shape)

torch.testing.assert_close()는 절대 오차와 상대 오차를 함께 확인한다. atol은 허용하는 절대 차이를 지정해 0에 가까운 값에서 중요하고, rtol은 기대값에 비례해 허용 오차를 조정한다. Triton과 PyTorch가 연산 순서나 중간 정밀도를 다르게 사용할 수 있어 작은 수치 차이가 생길 수 있으므로 두 값을 함께 둔다.

커널 테스트를 통과한 뒤 같은 모양·자료형·장치에서 의미 있는 기준과 비교한다. GPU 작업은 비동기식이므로 일반 CPU 타이머로 커널 실행 시간을 재면 GPU 작업이 아니라 실행 요청을 대기열에 넣는 시간만 측정할 수 있다. triton.testing.do_bench는 준비 실행과 반복 GPU 측정을 수행해 실행 시간을 밀리초로 반환한다.

python
def benchmark_matmul(M, N, K):
    a = torch.randn((M, K), device=DEVICE, dtype=torch.float16)
    b = torch.randn((K, N), device=DEVICE, dtype=torch.float16)
    # Check correctness and trigger compilation before reporting performance.
    torch.testing.assert_close(
        matmul(a, b),
        torch.matmul(a, b),
        atol=1e-2,
        rtol=1e-2,
    )
    triton_ms = triton.testing.do_bench(
        lambda: matmul(a, b),
        return_mode="median",
    )
    torch_ms = triton.testing.do_bench(
        lambda: torch.matmul(a, b),
        return_mode="median",
    )
    # A matrix multiplication performs approximately 2 * M * N * K FLOPs.
    flops = 2 * M * N * K
    to_tflops = lambda ms: flops / (ms * 1e-3) / 1e12
    print(
        f"shape=({M}, {N}, {K}) | "
        f"Triton={to_tflops(triton_ms):6.1f} TFLOP/s | "
        f"PyTorch={to_tflops(torch_ms):6.1f} TFLOP/s"
    )
if __name__ == "__main__":
    for shape in [
        (256, 256, 256),
        (512, 512, 512),
        (1024, 1024, 1024),
        (1025, 769, 513),  # irregular dimensions
    ]:
        benchmark_matmul(*shape)

작은 연산은 지연 시간(ms 또는 µs)이 가장 이해하기 쉬운 지표다. 처리량(TFLOP/s)은 행렬 곱셈의 부동소수점 연산을 얼마나 빠르게 완료하는지 나타낸다. 여러 모양을 대상으로 표와 그래프를 만들 때는 triton.testing.Benchmark와 triton.testing.perf_report를 사용할 수 있다.

디버깅할 때 인터프리터 모드를 켜면 Triton 프로그램을 NumPy에 해당하는 연산으로 CPU에서 순차적으로 시뮬레이션한다. 중간 값을 Python print()나 pdb로 살펴볼 수 있다.

python
import os
os.environ["TRITON_INTERPRET"] = "1"

인터프리터 모드는 GPU 병렬 실행, 스케줄링, 성능을 재현하지 않으므로 벤치마크에 사용하면 안 된다. bfloat16 값에 대한 연산과 간접 메모리 접근 패턴 등 지원되지 않는 경우도 있다. 정확성 확인에 활용한 뒤에는 대상 GPU에서도 커널을 시험해야 한다. 자세한 내용은 인터프리터 모드 디버깅 안내를 참고한다.

Triton 프로그램은 컴파일된 커널의 한 실행 인스턴스로, 보통 타일 하나를 맡지만 여러 타일을 처리하기도 한다. 그리드는 각 축에서 실행할 프로그램 수를 정한다. 예를 들어 (1, 1, 1)은 프로그램 하나, (3, 2, 1)은 여섯 개를 실행한다. 후자의 경우 tl.program_id(0)은 0부터 2까지, tl.program_id(1)은 0부터 1까지다. 대부분의 비영속 커널은 프로그램 하나를 출력 타일 하나에 대응시키며, 영속 커널은 프로그램마다 여러 타일을 맡길 수 있다.

각 프로그램은 프로그램 ID를 이용해 자신이 출력에 기여할 타일을 찾고, 논리적 텐서 좌표를 메모리 포인터로 바꾼 뒤 데이터를 읽고 계산해 결과를 쓴다. 프로그램 ID, 블록 좌표, tl.arange()와 스트라이드에서 나온 표현식은 보통 인덱스·메모리 오프셋·포인터를 계산한다. 실제 텐서 값은 해당 위치에서 tl.load()로 읽으며, 데이터에 대한 연산은 그 뒤에 수행한다.

부록: 원문 예시

text
def matmul(a: torch.Tensor, b: torch.Tensor):     assert a.is_cuda and b.is_cuda     assert a.device == b.device     assert a.ndim == 2 and b.ndim == 2     assert a.shape[1] == b.shape[0]     assert a.dtype == b.dtype == torch.float16     M, K = a.shape     _, N = b.shape     c = torch.empty((M, N), device=a.device, dtype=a.dtype)     grid = lambda meta: (         triton.cdiv(M, meta["BM"]),         triton.cdiv(N, meta["BN"]),     )     _matmul_kernel[grid](         a, b, c,         M, N, K,         a.stride(0), a.stride(1),         b.stride(0), b.stride(1),         c.stride(0), c.stride(1),         BM=32, BN=32, BK=32,     )     return c