Triton 실전 — 벡터 덧셈 커널 만들기

Triton 실전 — 벡터 덧셈 커널 만들기

처음으로 Triton 커널을 만들려면 가장 단순한 벡터 덧셈부터 시작해요. 이 튜토리얼을 통해 Triton의 기본 프로그래밍 모델, @triton.jit 데코레이터, 그리고 커스텀 연산을 네이티브 구현과 검증·벤치마크하는 요령을 익힐 수 있어요.

출처: https://triton-lang.org/main/getting-started/tutorials/01-vector-add.html

커널 정의

import torch
import triton
import triton.language as tl

DEVICE = triton.runtime.driver.active.get_active_torch_device()


@triton.jit
def add_kernel(x_ptr,        # *Pointer* to first input vector.
               y_ptr,        # *Pointer* to second input vector.
               output_ptr,   # *Pointer* to output vector.
               n_elements,   # Size of the vector.
               BLOCK_SIZE: tl.constexpr,  # Number of elements each program should process.
               ):
    pid = tl.program_id(axis=0)   # We use a 1D launch grid so axis is 0.
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    output = x + y
    tl.store(output_ptr + offsets, output, mask=mask)

mask로 경계 밖 접근을 막고, 입력이 블록 크기의 배수가 아닐 때 남는 요소를 안전하게 처리해요.

래퍼 함수

def add(x: torch.Tensor, y: torch.Tensor):
    output = torch.empty_like(x)
    assert x.device == DEVICE and y.device == DEVICE and output.device == DEVICE
    n_elements = output.numel()
    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), )
    add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=1024)
    return output

grid는 병렬로 실행할 커널 인스턴스 수를 정의하는 SPMD 런치 그리드예요. CUDA 런치 그리드와 유사해요.

검증

torch.manual_seed(0)
size = 98432
x = torch.rand(size, device=DEVICE)
y = torch.rand(size, device=DEVICE)
output_torch = x + y
output_triton = add(x, y)
print(f'The maximum difference between torch and triton is '
      f'{torch.max(torch.abs(output_torch - output_triton))}')

이 경우 The maximum difference between torch and triton is 0.0이 출력돼요.

벤치마크

@triton.testing.perf_report로 문제 크기별 성능을 비교할 수 있어요.

@triton.testing.perf_report(
    triton.testing.Benchmark(
        x_names=['size'],
        x_vals=[2**i for i in range(12, 28, 1)],
        x_log=True,
        line_arg='provider',
        line_vals=['triton', 'torch'],
        line_names=['Triton', 'Torch'],
        styles=[('blue', '-'), ('green', '-')],
        ylabel='GB/s',
        plot_name='vector-add-performance',
        args={},
    ))
def benchmark(size, provider):
    ...

benchmark.run(print_data=True, show_plots=True)로 결과를 보고, save_path로 CSV를 저장할 수 있어요.

더 알아보기