Triton Internals

Triton 사용자라면 누구나 안다 — 같은 커널 소스가 launch 설정만으로 perf 가 한 자리수 차이가 난다. 왜? 이 강의는 그 “왜” 를 컴파일러 단계로 분해한다. Kapil 의 두 편짜리 블로그 시리즈를 영상 형태로 옮긴 것.

강의가 답하려는 두 줄 —

1. @triton.jit 데코레이터 안에서 무엇이 일어나는가 — 첫 호출, 두 번째 호출, cache hit 이 각각 어떤 작업을 한다.
2. BLOCK_SIZE, num_warps, num_stages 가 어느 단계에서 어떤 결정으로 바뀌는가 — 사용자가 주는 hint 가 컴파일러의 어디에 흘러가는가.

강의의 frame 은 명시적으로 — “CUDA 의 nvcc lowering 은 cuc → PTX → SASS 의 3 단계로 잘 알려져 있다. Triton 은 그 위에 세 단계가 더 얹혀 있을 뿐이다”. 그 다섯 단계가 강의 전체를 끌고 간다.

강의 끝에 손에 잡혀야 하는 건 — 5 단계의 이름과 입출력, 각 단계에서 일어나는 핵심 결정(layout, pipeline, allocation), cache 디렉터리의 구조, 그리고 perf 디버깅을 어디부터 시작해야 하는지 의 한 페이지 mental model.

강의의 본론이 시작되는 자리. Kapil 이 같은 vector_add 커널을 가지고 5 단계를 한 단계씩 따라간다. 각 단계가 정확히 무엇을 하고, 무엇을 결정하고, 무엇을 다음 단계로 넘기는가.

# L0 — 우리가 짜는 코드
@triton.jit
def add_kernel(x_ptr, y_ptr, out_ptr, n,
                BLOCK: tl.constexpr):
    pid    = tl.program_id(0)
    offs   = pid * BLOCK + tl.arange(0, BLOCK)
    mask   = offs < n
    x      = tl.load(x_ptr + offs, mask=mask)
    y      = tl.load(y_ptr + offs, mask=mask)
    tl.store(out_ptr + offs, x + y, mask=mask)

// L1 — Triton IR (간략)
tt.func @add_kernel(%x: !tt.ptr<f32>, ...) {
  %pid    = tt.get_program_id x : i32
  %offs   = tt.make_range {start=0, end=1024}
  %addrs  = tt.addptr %x_ptr, %offs
  %x      = tt.load %addrs, %mask
  %y      = tt.load %y_addrs, %mask
  %sum    = arith.addf %x, %y : tensor<1024xf32>
  tt.store %out_addrs, %sum, %mask
  tt.return
}

Triton IR (L1) 에서는 텐서가 그냥 tensor<1024xf32>. 어느 thread 가 어느 element 를 들고 있는지 안 정해져 있다. L2 의 TritonGPU IR 에서 그 mapping 이 박힌다 — 그게 layout. perf 의 결정자가 거의 다 여기서 정해진다.

L2 의 가장 흥미로운 pass — software pipelining. matmul 의 K-loop 처럼 load → compute → load → compute 가 반복되는 패턴에서, load 와 compute 를 시간적으로 겹치는 변환. num_stages 가 이 변환의 stage 수를 결정.

변환 전 (num_stages=1):
  for k in K_BLOCKS:
    a = load(A[k])           ┐
    b = load(B[k])           ├ memory stall — load latency 다 노출
    acc += dot(a, b)         ┘

변환 후 (num_stages=2 — pipelined):
  a0 = load(A[0])            ┐
  b0 = load(B[0])            │ prologue
  for k in K_BLOCKS[1:]:
    a_next = load(A[k])      │ next 의 load 가
    b_next = load(B[k])      │ 현재 iter 의 dot 와 겹침
    acc += dot(a, b)         │
    a = a_next; b = b_next   │
  acc += dot(a, b)           │ epilogue

L2 까지는 “tile 의 element 가 thread 의 어디에 있다” 라는 abstract level. L3 의 LLVM IR 에서 그 element 가 구체적인 virtual register에 박힌다. 그 다음 LLVM 의 register allocator 가 physical register 에 mapping. 이 자리에서 spill 이 발생할 수 있다.

// L3 — LLVM NVVM (간략)
define void @add_kernel(...) {
  %tid    = call i32 @llvm.nvvm.read.ptx.sreg.tid.x()
  %bid    = call i32 @llvm.nvvm.read.ptx.sreg.ctaid.x()
  %off    = ... // pid * BLOCK + tid

  // 8 element vector load — coalesced
  %vec_x  = load <8 x float>, <8 x float>* %addr
  %vec_y  = load <8 x float>, <8 x float>* %addr
  %vec_s  = fadd <8 x float> %vec_x, %vec_y

  store <8 x float> %vec_s, <8 x float>* %out
  ret void
}

강의의 가장 실용적인 부분 — Kapil 이 터미널에서 ~/.triton/cache 디렉터리를 열어 안의 파일들을 한 줄씩 보여준다. 각 specialization 마다 한 디렉터리, 그 안에 모든 IR 단계의 텍스트가 떨어져 있다.

# 직접 열어보는 명령 — 강의 데모 그대로
$ ls ~/.triton/cache/
8a3b...c4d/    # add_kernel BLOCK=1024 · fp32
9912...abc/    # add_kernel BLOCK=1024 · fp16

$ cd 8a3b...c4d/
$ ls
add_kernel.ttir          add_kernel.ptx
add_kernel.ttgir         add_kernel.cubin
add_kernel.llir          add_kernel.json

# shared memory 와 register 사용량 한 줄에 — JSON 안
$ jq '.shared, .num_warps, .num_ctas' add_kernel.json
24576
4
1

# SASS 까지 보고 싶으면 — cuobjdump
$ cuobjdump --dump-sass add_kernel.cubin | head -30

Triton 의 두 가지 검사 도구 — 커널 실행 검사(interpret 모드)와 컴파일러 단계 검사(MLIR dump). 강의에서 Kapil 이 두 가지 모두 데모.

@triton.jit(interpret=True)
def add_kernel(...):
    pid    = tl.program_id(0)
    offs   = pid * BLOCK + tl.arange(0, BLOCK)
    breakpoint()                    # pdb 가 동작
    x = tl.load(x_ptr + offs)
    # pdb 안에서 print(x.tensor) 가능

# 환경변수로도 가능 — 커널 변경 없이
# TRITON_INTERPRET=1 python script.py

# 모든 pass 의 in/out IR 을 stdout 에
$ MLIR_ENABLE_DUMP=1 python script.py 2> trace.log

# 특정 pass 만
$ MLIR_ENABLE_DUMP=1 \
  TRITON_REPRODUCER_PATH=/tmp/repro.mlir \
  python script.py

# 한 IR 파일을 input 으로 직접 컴파일
$ triton-opt --convert-triton-to-tritongpu in.ttir

Triton 3.x (강의 시점에 베타) 가 가져오는 큰 변화들. 아키텍처가 새 명령을 내놓을 때마다 컴파일러가 따라가야 한다 — H100 의 wgmma, TMA, async barrier 가 그 핵심.

손에 새기기 — 실습 시퀀스

1. vector_add 한 번 실행 — 강의 repo 의 vector_add.py. ~/.triton/cache 안에 새 디렉터리가 생기는 것 확인.
2. 5 단계 IR 직접 열기 — 그 디렉터리 안의 .ttir, .ttgir, .llir, .ptx 를 열어 같은 vector_add 가 어떻게 변하는지 본다.
3. BLOCK_SIZE sweep + cache — BLOCK_SIZE 를 다섯 가지로. 새 디렉터리 다섯 개. 각 디렉터리의 .ttgir 의 layout 차이 확인.
4. num_warps 의 효과 — num_warps=1, 4, 8 로 같은 커널. .ttgir 에서 BlockedLayout 의 warpsPerCTA 가 어떻게 박히는지.
5. num_stages 의 효과 — matmul 커널로. num_stages=1 vs 3 의 .llir 비교 — pipelined load 의 prologue/epilogue 가 보이는지.
6. spill 발생시키기 — 매우 큰 BLOCK 으로 register 를 일부러 압박. .json 의 spill_count > 0 인지 확인. SASS 에서 LDL(local memory load) 명령이 보이는지.
7. MLIR pass 추적 — MLIR_ENABLE_DUMP=1 으로 한 컴파일 전체 dump. layout conversion pass 가 무엇을 변환하는지 본다.
8. SASS 까지 내려가기 — cuobjdump --dump-sass. 자기 GPU(A100/H100/RTX) 별 명령 차이 확인. wgmma 명령이 보이는 자리.

• autotune 과 cache 의 관계 — autotune 의 각 config 가 모두 별도 cache 디렉터리. cache 디렉터리 수가 수백 개로 폭발할 수 있는데 어떻게 관리하는가.
• Triton 3.x 의 cluster mode — H100 의 cluster (여러 SM 이 묶임)을 Triton 이 어떻게 표현하는가. 새 dialect 인가, 기존 BlockedLayout 의 확장인가.
• AMD MI300 의 mfma 와 Triton — 같은 GEMM 커널 코드가 NVIDIA mma 와 AMD mfma 둘 다로 lower 되는가? 코드 변경이 정말 0 인가.
• Intel XPU backend 의 현재 상태 — Triton 3.x 의 Intel backend 가 stable 인지. SYCL/SPIR-V 와 어떻게 합쳐지는가 (L026 과의 연결).
• spill 의 자동 검출 — .json 의 spill_count 외에 — 컴파일러가 사용자에게 “여기서 register 가 모자랐다” 를 직접 알려주는 길이 있는가.
• flash attention v3 의 IR — fa3 의 .ttgir 가 fa2 와 비교해서 어떤 새 layout/dialect 를 쓰는가. wgmma 와 TMA 의 흔적.
• Triton 의 backward 자동 미분 — Triton 자체에는 autograd 가 없다. tl.dot 의 backward 를 어떻게 수동으로 짜는가의 표준 패턴.