A Practitioner's Guide to Triton

import triton
import triton.language as tl

@triton.jit
def add_kernel(x_ptr, y_ptr, out_ptr, n,
               BLOCK_SIZE: tl.constexpr):
    # 1) program_id — 나는 몇 번째 program?
    pid    = tl.program_id(0)
    # 2) arange — 내가 다룰 tile 의 element index
    offs   = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    # 3) mask — 마지막 tile 의 tail 처리
    mask   = offs < n
    # 4) load — HBM 에서 SRAM 으로
    x      = tl.load(x_ptr + offs, mask=mask)
    y      = tl.load(y_ptr + offs, mask=mask)
    # 5) compute — tile 단위 numpy-like
    out    = x + y
    # 6) store — SRAM 에서 HBM 으로
    tl.store(out_ptr + offs, out, mask=mask)

각 op 의 의미.

• program_id(0) — 1D grid 위의 program index. 0, 1, 2, … N-1 까지 N 개 program 이 launch 되고 각자 자기 pid 만 본다. CUDA 의 blockIdx 에 가장 가까움.
• tl.arange(0, BLOCK_SIZE) — 0, 1, …, BLOCK_SIZE-1 의 vector 를 만든다. 이게 Triton 의 가장 핵심 — “tile 의 모든 element 에 대한 인덱스 vector”.
• tl.load(ptr + offs, mask=mask) — pointer + offset vector 위치에서 BLOCK_SIZE 개 element 를 한 번에 read. mask 가 false 인 자리는 default 값 (0).
• tl.store(ptr + offs, val, mask=mask) — 반대로 BLOCK_SIZE 개 write.

호출하는 자리.

n = x.numel()
grid = (triton.cdiv(n, BLOCK_SIZE),)
add_kernel[grid](x, y, out, n, BLOCK_SIZE=1024)

import os
os.environ['TRITON_INTERPRET'] = '1'

@triton.jit
def add_kernel(x_ptr, y_ptr, out_ptr, n,
               BLOCK_SIZE: tl.constexpr):
    pid    = tl.program_id(0)
    offs   = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask   = offs < n
    x      = tl.load(x_ptr + offs, mask=mask)
    breakpoint()                      # ← CPU 에서 멈춤
    y      = tl.load(y_ptr + offs, mask=mask)
    tl.store(out_ptr + offs, x+y, mask=mask)

interpret 모드의 한계.

• 실행 속도: 매우 느림. 큰 N 으로 돌리면 안 됨.
• 일부 op 가 시뮬레이션이 정확하지 않을 수 있음 — async copy, atomic 같은 부분.
• numerical 결과가 GPU 와 약간 다른 경우가 드물게 있음 (rounding 모드).

Umer 의 helper — triton_util.py 의 breakpoint_if 와 print_if. 특정 program_id 에서만 breakpoint. “pid_0=0, pid_1=1 일 때만 멈춰라”.

def breakpoint_if(conds, pid_0=[0], pid_1=[0],
                  pid_2=[0]):
    from IPython.core.debugger import set_trace
    if test_pid_conds(conds, pid_0, pid_1, pid_2):
        set_trace()

# autotune — 사용자가 후보 set 을 선언
@triton.autotune(
    configs=[
        triton.Config({'BLOCK_M':64, 'BLOCK_N':64,
                       'BLOCK_K':32},
                      num_warps=4, num_stages=3),
        triton.Config({'BLOCK_M':128, 'BLOCK_N':128,
                       'BLOCK_K':32},
                      num_warps=8, num_stages=3),
        # … 보통 5-30 개 candidate
    ],
    key=['M', 'N', 'K'],   # 입력 모양에 따라 cache
)
@triton.jit
def matmul_kernel(...): ...

autotune 의 동작.

1. 처음 호출 시 모든 config 를 직접 실행해 측정.
2. 가장 빠른 config 를 key 별로 cache.
3. 이후 같은 key 의 호출은 cached config 사용.

함정 두 가지.

• 첫 호출의 wall time 이 매우 길다 — 30 config × 5ms = 150 ms 시 측정에 들어간다. test 환경에서 이상하게 느려보임.
• shape 가 매번 다르면 cache 가 안 됨 — dynamic shape 모델에서는 autotune 의 효과가 거의 없거나 매번 재측정. key 를 잘 정의해서 비슷한 shape 를 묶기.

matmul 의 inner loop (K 축 reduction) 한 iteration 의 일.

1. K tile load (HBM → SRAM)
2. MMA 누산 (SRAM → register)
3. 다음 iteration 준비

num_stages = 3 이면 — iteration n 의 load 가 iteration n−1 의 MMA 와, iteration n−2 의 다음 K tile pre-load 와 같은 시간에 도는 식. HBM latency 를 숨기는 핵심 메커니즘.

Triton 으로 멈출 신호

• NCU 가 “tail effect / occupancy / launch shape” 만 hint 로 줌. 이건 BLOCK_SIZE/num_warps 조정으로 잡힘.
• memory-bound op (elementwise, normalize, cross-entropy 등). HBM 트래픽이 거의 peak. compute 의 디테일이 무의미.
• fast iteration 이 중요 — 코드 변경 후 다시 컴파일이 1초.

직접 CUDA 로 내려가는 신호

• NCU 가 “shared memory bank conflict”, “uncoalesced reads”, “register spill” hint. Triton 의 자동 layout 으로는 못 잡는 자리.
• vendor library (cuBLAS, cuDNN) 의 80-90% 까지 가야 하는 GEMM/conv. CUTLASS template 위로.
• warp-level primitive 가 필요 (warp shuffle, ballot 등). Triton 은 표현 못 함.
• 같은 cycle 에 여러 thread 가 같은 SRAM bank 를 치는 자리를 손으로 swizzle.