[Multi-Core & GPU Programming] Triton Introduction

"Triton Introduction"

[Why Triton?]

· PyTorch 연산을 여러 줄로 나누면 각 연산마다 Memory Load/Store가 반복되어 GPU Memory Bandwidth 낭비

· "torch.compile"을 사용하면 "Naive PyTorch"보다 개선되지만, 여전히 이상적인 성능에는 도달 불가

→ "torch.compile"이 자동 최적화를 해주긴 하지만, Memory 접근 Pattern과 연산 융합을 직접 제어하는 Triton Kernel이 더 유리

[CUDA Software Stack]

▣ Application → GPU Path

① Application

② CUDA Libraries(Ex. cuBLAS)

③ CUDA Runtime API(libcudart.so) : # include <cuda_runtime.h>, cudaMalloc, cudaMemcpy, Kernel Launch

④ CUDA Driver API(libcuda.so) : User Mode APIs, # include <cuda.h>, cuMemcpyDtoA, cuMemAlloc

⑤ NVIDIA GPU Driver(nvidia.ko) : Open-Source, Kernel Mode in OS, PTX → SASS in case of nvrtc

⑥ GPU

> CUDA Runtime API

· 비교적 사용하기 쉽고 일반 CUDA C Code에 가까움

> CUDA Device API

· 더 낮은 수준에서 Module Load, 함수 실행, Memory 제어를 직접 구현

[GPU Inside]

▣ H100(Datacenter GPU)

· 여러 개의 GPC/SM Block

· 큰 L2 Cache

· HBM Memory Interface

→ 하나의 거대한 Core가 아니라, 수많은 SM(Streaming Multi-Processor)와 Memory 계층으로 구성된 병렬 처리 장치

** Gaming GPU보다 더 많은 SM, 더 큰 HBM Memory, 더 넓은 Bandwidth, Server용 안정성 기능, AI/HPC 특화 기능

    → Gaming GPU도 Triton Kernel을 실행 가능

    → 성능 차이는 주로 SM 개수, Memory Bandwidth, Tensor Core 세대, VRAM 용량에서 발생

① CPU Memory에 있는 데이터를 GPU에서 계산하기 위해 DRAM에 있는 데이터를 먼저 GPU Global Memory로 복사

    → PyTorch : ".cuda( )" / CUDA : "cudaMemcpy"

② Global Memory의 데이터가 L2 Cache를 거쳐 SM 내부의 L1 Cache / Shared Memory에 Load

③ Register를 통해 CUDA 및 Tensor Core로 전달되어 연산에 활용

[Operation Fusion]

· 여러 개의 연산을 하나의 GPU Kernel로 합치는 최적화

· Softmax를 PyTorch로 작성하면 {max, subtract, exp, sum, divide}가 각각 따로 실행

    → 중간 결과를 계속 Global Memory에 Store하고 다시 Load

· Triton은 이런 연산들을 하나의 Kernel 안에 넣어서 Global Memory 접근↓

▣ Naive PyTorch Softmax

1 2 3 4 5 6 7 8

def naive_softmax(x):     x_max = x.max(dim = 1)[0]     z = x - x_max[:, None]     numerator = torch.exp(z)     denominator = numerator.sum(dim = 1)     ret = numerator / denominator[:, None]       return ret

cs

> x_max = x.max(dim = 1)[0]

· N x N 형태의 입력에서 각 Row마다 최댓값 찾기

① Global Memory에서 데이터를 Load

② 연산을 통해 각 Row의 최댓값 찾기

③ 최댓값을 다시 Memory에 Store

→ Load : N² / Store : N

> z = x - x_max[:, None]

· Softmax에서 Overflow를 막기 위한 안정화 단계

① 위에서 저장한 최댓값과 원본 데이터를 Memory에서 또다시 Load하여 뺄셈 연산

② 뺄셈 연산 결과를 Memory에 다시 Store

→ Load : N² + N / Store : N²

> numerator = torch.exp(z)

· z의 모든 원소에 Exponential 적용하여 Softmax 공식의 분자 생성(특정 값)

① 위에서 저장한 값을 다시 Memory로 가져와서 Exponential 적용

② Exponential 적용 결과를 Memory에 다시 Store

→ Load : N² / Store : N²

> denominator = numerator.sum(dim = 1)

· Softmax 공식의 분모 생성(전체 값의 합)

① Exponential 적용 결과를 Memory에서 다시 Load

② Row 단위로 모든 값을 더함

③ 덧셈 연산 결과를 Memory에 다시 Store

→ Load : N² / Store : N

> ret = numerator / denominator[:, None]

· Softmax 공식을 적용한 최종 확률 계산

① Global Memory에서 "numerator"와 "denominator"를 둘 다 Load

② 나눗셈 수행

③ 나눗셈 결과를 다시 Global Memory에 Store

→ Load : N² + N / Store : N²

∴ Total Load : 5N² + 2N / Total Store : 3N² + 2N → Total : 8N² + 4N

▣ Triton Softmax

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24

@triton.jit def softmax_kernel(output_ptr, input_ptr,                    input_row_stride, output_row_stride,                    n_rows, n_cols,                    BLOCK_SIZE: tl.constexpr, num_stages: tl.constexpr);     row_start = tl.program_id(0)     row_step = tl.num_programs(0)       for row_idx in tl.range(row_start, n_rows, row_step, num_stages = num_stages):         row_start_ptr = input_ptr + row_idx * input_row_stride         col_offsets = tl.arange(0, BLOCK_SIZE)         input_ptrs = row_start_ptr + col_offsets         mask = col_offsets < n_cols           row = tl.load(input_ptrs, mask = mask, other = -float('inf'))         row_minus_max = row - tl.max(row, axis = 0)         numerator = tl.exp(row_minus_max)         denominator = tl.sum(numerator, axis = 0)         softmax_output = numerator / denominator           output_row_start_ptr = output_ptr + row_idx * output_row_stride         output_ptrs = output_row_start_ptr + col_offsets           tl.stor(output_ptrs, softmax-output, mask = mask)

cs

> Factor

· output_ptr, input_ptr : Input/Output Buffer

· input_row_stride, output_row_stride : Input/Output Buffer Strides

· n_rows, n_cols : Input Shape

· BLOCKS_SIZE, num_stages : Kernel Config → "tl.constexpr"로 지정하여 실행 중 바뀌는 값이 아니라 Compile 시점에 고정

> Block Info

· tl.program_id(0) : 현재 실행 중인 Triton Program/Block의 ID

· tl.num_programs(0) : 전체 Program/Block의 개수

→ 여러 Row를 여러 Triton Program이 나누서 처리할 때, 각 Program은 자기 "program_id"에 해당하는 Row부터 시작

> Index & Mask 생성 for Load

· tl.arange(0, BLOCK_SIZE) : 한 Row 안에서 접근할 Column Index 생성

· mask : "n_cols"가 "BLOCK_SIZE"보다 작을 수 있으므로 유효한 Column만 Load 하도록 Mask

→ Triton에서는 Vector 단위로 여러 원소를 한 번에 다루지만, 범위를 벗어나는 접근은 Mask로 안전하게 차단

> Softmax Logic

** row, row_minus_max, numerator, denominator가 Global Memory에 중간 저장 X

→ 대부분 Register 또는 On-Chip Memory에서 처리

→ Load : N² / Compute : 5

∴ Softmax의 여러 단계가 하나의 Triton Kernel 안에서 Fusion

> Index 생성 for Store

· 출력 위치도 Input과 마찬가지로 Row 시작 주소와 Column Offset으로 계산(Mask를 사용해 유효한 Column에만 저장)

→ Store : N²

∴ Total : 2N² (Naive PyTorch : 8N² + 4N → 약 4배 빠름)

→ Softmax 같은 연산은 대부분 Memory-Bound이므로 계산량보다 Load/Store 횟수가 성능을 좌우

[Triton Language & Features]

· triton-lang.org : Triton 공식 문서 페이지

· Triton에는 약 90개 이상의 Operation 존재

· Triton Kernel 안에서는 Python 문법처럼 보이지만, 실제로는 GPU용 연산으로 Compile 되는 Triton Language Operation을 사용

> @triton.jit

· Triton Kernel 함수 위에 붙이면 해당 함수는 일반 Python 함수가 아니라 GPU Kernel로 JIT Compile되는 함수

→ "tl.load", "tl.store", "tl.arange" 같은 Triton 연산을 사용하고, 실행 시점에 GPU Code로 Compile되어 실행

> @triton.autotune

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18

@triton.autotune(     configs = [         triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'NUM_SM': 84}),         triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'NUM_SM': 128}),         triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32, 'NUM_SM': 84}),         triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32, 'NUM_SM': 128}),     ],     key = ['group_size'], )   @triton.jit def grouped_matmul_kernel(     ...     ...     BLOCK_SIZE_M: tl.constexpr,     BLOCK_SIZE_N: tl.constexpr,     BLOCK_SIZE_K: tl.constexpr, ) Colored by Color Scripter

cs

· HW 구조에 맞게 Block Size나 SM 개수 등의 여러 버전의 Kernel Config를 자동으로 시험해보고 가장 빠른 설정을 찾는 기능

· key : Autotune을 다시 수행할 기준("x_size" 값이 바뀔 때마다 여러 config를 다시 평가)

· tl.constexpr : Compile-Time Constant → Runtime이 아니라, Kernel이 Compile될 때 고정되어야 하는 값

> @triton.benchmark

· Triton에서 Kernel 성능을 측정하며, 그래프를 만들거나 Latency/Throughtput을 비교할 때 유용

· 여러 입력 크기나 Provider(torch, torch.compile, tirton 등)를 비교할 때 사용

▣ Useful EnvVars

· Triton 개발과 Debugging에 유용한 환경 변수

· TRITON_PRINT_AUTOTUNING = 1 : 최적화 소요 시간과 결과를 Terminal에 출력

· TRITON_INTERPRET = 1 : 코드가 GPU가 아닌 일반 Python Interpreter 환경에서 실행되어 Breakpoint를 걸고 Debugging 가능

[Useful Resources]

· https://github.com/linkedin/Liger-Kernel/tree/main

    → LLM 학습/추론에서 자주 쓰이는 연산들을 Triton으로 최적화한 Project

· https://docs.pytorch.org/tutorials/recipes/torch_compile_user_defined_triton_kernel_tutorial.html

    → 사용자가 직접 작성한 Triton Kernel을 PyTorch Workflow에 통합하는 방법 제시

[Reference]

· 11_Triton_Introduction_mgp_2026 (MGP) - Yongjun Park