[Multi-Core & GPU Programming] HW5 - Triton Puzzle
"Triton Puzzle"
[Objective]
High-Level의 PyTorch 연산을 효율적인 Low-Level Triton Kernel Code로 변환하는 방법을 이해하기
[Implementation]
▣ Vector Add
> Triton Kernel
· PyTorch로 작성하면 전체 동작은 단순히 "output = x + y"
→ Triton에서는 해당 연산을 GPU Program 단위로 직접 나누어 처리
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 | import torch import triton import triton.language as tl DEVICE = 'cuda' @triton.jit def add_kernel(x_ptr, y_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr, ): pid = tl.program_id(axis=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) | cs |
1 : PyTorch Tensor를 만들고 비교하기 위해 import
2 : Triton Kernel을 실행하기 위해 import
3 : Triton Kernel 내부에서 쓰는 문법을 위해 import
5 : GPU에서 실행
7~8 : 일반 Python 함수가 아니라 GPU에서 여러 개의 Program으로 나누어 병렬 실행될 Triton Kernel 선언
· x_ptr : 1번째 입력 Vector(x)의 시작 주소
· y_ptr : 2번째 입력 Vector(y)의 시작 주소
· output_ptr : 결과를 저장할 Vector(output)의 시작 주소
· n_elements : 전체 Vector의 원소 개수
· BLOCK_SIZE : 하나의 Triton Program이 처리할 원소 개수(tl.constexpr : Compile Time 상수)
→ Triton이 Kernel을 Compile할 때 BLOCK_SIZE 값을 알고 있어야 Tensor Shape을 만드는데 사용 가능
9 : GPU에서 이 Kernel이 여러 Program으로 동시에 실행되며, "pid"는 현재 실행 중인 Program 번호(axis = 0 : 1차원 Grid를 사용)
· 전체 원소 = 256개, BLOCK_SIZE = 64 → Program = 4개 필요
10 : 현재 Program이 처리할 1번째 원소 번호를 계산(= Block 시작 위치 계산)
11 : Offset을 계산하여 현재 Program이 처리할 실제 Index를 계산
12 : 범위를 벗어난 Index를 무시
· n_elements(10) > BLOCK_SIZE(4) → Program = 3개 (3번째 Program = [8, 9, 10, 11] ▶ [True, True, False, False])
13 : "x" Vector에서 현재 Program이 담당하는 원소들을 Read (PyTorch : x = x_tensor[offsets])
· Ex. x = [10, 20, 30, 40, 50, 60, 70, 80], offsets = [4, 5, 6, 7] ▶ x = [50, 60, 70, 80]
14 : "y" Vector에서 현재 Program이 담당하는 원소들을 Read (∴ 현재 Program은 같은 Index의 x, y를 Load)
15 : 원소별 덧셈 연산(output[i] = x[i] + y[i])
16 : 계산한 결과를 "output" Tensor에 저장(PyTorch : output_tensor[offsets] = output)
- Flow
① 현재 Program 번호 확인
② 현재 Program이 담당할 시작 Index 계산
③ 현재 Program이 담당할 Index 목록 생성
④ 유효한 Index인지 확인
⑤ GPU Memory에서 x, y Load
⑥ 원소별 덧셈 연산 후 결과 저장
> Wrapper 함수
· GPU Kernel인 "add_kernel"을 Python에서 실행하기 쉽게 감싼 Wrapper 함수
→ 사용자는 복잡한 Triton 실행 문법을 몰라도 그냥 "z = add(x, y)"로 작성하여 내부적으로 GPU에서 "x + y"를 실행
1 2 3 4 5 6 | def add(x: torch.Tensor, y: torch.Tensor): output = torch.empty_like(x) 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 | cs |
1 : 입력으로 PyTorch Tensor 2개(x, y)를 받아 "output = x + y"를 Triton Kernel로 계산
2 : "x"와 같은 Shape, 같은 Data Type, 같은 Device를 가진 빈 Tensor 생성
3 : "numel( )"을 사용하여 Tensor 안의 전체 원소 개수를 반환
4 : Triton Program을 몇 개 실행할지 결정(cdiv : Ceiling Division = 올림 나눗셈)
· meta : Kernel 실행 시 넘겨지는 Meta-Parameter 정보를 보유
· Triton의 Launch Grid는 Tuple 형태여야 함
5 : GPU Kernel을 실행(Python과 달리 Triton Kernel은 Launch Grid가 필요하므로 [grid] 추가)
· BLOCK_SIZE = 1024 : 하나의 Triton Program이 1024개의 원소를 처리
6 : Triton Kernel이 결과를 "output" Tensor에 저장했으므로, 해당 Tensor를 반환
** GPU 연산은 기본적으로 비동기 → CPU는 Kernel 실행을 GPU에 요청하고 바로 다음 줄로 넘어감
∴ 정확한 시간 측정이나 결과 사용 전에 필요하면 "torch.cuda.synchronize( )"를 호출해서 GPU 작업이 끝날 때까지 대기
- Flow
① 입력 Tensor 2개 받기
② 결과 저장 공간 만들기
③ 전체 원소 수 계산
④ 몇 개의 GPU Program을 띄울지 계산
⑤ Triton Kernel 실행
⑥ 결과 Tensor 반환
** PyTorch : "def add_torch(x, y): return x + y" → 단순
> PyTorch vs Triton Kernel
· PyTorch로 계산한 Vector Add 결과와 Triton Kernel로 계산한 Vector Add 결과가 같은지 확인
1 2 3 4 5 6 7 8 9 10 | 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(output_torch) print(output_triton) print(f'The maximum difference between torch and triton is ' f'{torch.max(torch.abs(output_torch - output_triton))}') | cs |
1 : Random Seed를 고정하여 Random 값을 항상 똑같이 생성
2 : Vector 크기 설정
3~4 : 0 이상 1 미만의 랜덤 Float 값으로 이루어진 Vector 생성 (DEVICE = 'cuda'이므로 GPU Memory에 생성)
5 : PyTorch가 제공하는 기본 덧셈으로 결과 계산
6 : Triton 방식으로 덧셈 계산
7 : PyTorch로 계산한 결과 출력
8 : Triton Kernel로 계산한 결과 출력
9~10 : 두 계산 결과를 원소별로 차이를 계산한 후, 모든 원소의 차이 중 가장 큰 값을 구함 (0 = Triton 방식이 PyTorch 방식과 동일)
- Flow
① Random 값 고정
② Vector 길이 설정
③ GPU에 Random Vector 2개 생성
④ PyTorch로 정답 계산
⑤ Triton Kernel로 결과 계산
⑥ 두 결과 출력
⑦ 두 결과의 최대 차이 확인
> Benchmark
· PyTorch의 Vector Add와 Triton Kernel로 계산한 Vector Add의 성능을 비교
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.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): x = torch.rand(size, device=DEVICE, dtype=torch.float32) y = torch.rand(size, device=DEVICE, dtype=torch.float32) quantiles = [0.5, 0.2, 0.8] if provider == 'torch': ms, min_ms, max_ms = triton.testing.do_bench(lambda: x + y, quantiles=quantiles) if provider == 'triton': ms, min_ms, max_ms = triton.testing.do_bench(lambda: add(x, y), quantiles=quantiles) gbps = lambda ms: 3 * x.numel() * x.element_size() * 1e-9 / (ms * 1e-3) return gbps(ms), gbps(max_ms), gbps(min_ms) benchmark.run(print_data=True, show_plots=True) | cs |
1~2 : "benchmark" 함수를 여러 조건에서 반복 실행하고, 결과를 표 또는 그래프로 정리해주는 Triton 기능
3 : x축으로 사용할 변수 = "size" → Vector 크기를 바꿔가면서 성능 측정
4 : 측정할 Vector 크기 목록 (212, 213, …, 227)
5 : x축을 Log Scale로 표시
6 : 그래프에서 "provider"를 기준으로 여러 선을 구분
7 : "provider"가 가질 수 있는 값
8 : 그래프에 표시될 이름
9 : 그래프 선 스타일(Triton : 파란색 실선 / PyTorch : 초록색 실선)
10 : y축은 실행 시간이 아닌 Memory Bandwidth(∵ Vector Add는 계산량보다 Memory Read/Write 비용이 더 중요)
11 : 생성되는 그래프 이름
14 : 특정 "size"와 특정 "provider"에 대해 성능 측정
15~16 : 길이가 "size"인 Float32 Tensor 2개를 GPU에 생성
17 : Benchmark를 여러 번 실행한 후 실행 시간의 분위수를 가져옴
18~19 : PyTorch의 기본 덧셈 성능을 측정
20~21 : Triton의 덧셈 성능 측정
22 : 성능을 GB/s로 변환
· Vector Add는 각 원소마다 Memory 접근 3번 수행 (x, y Read & output Write)
· x.numel( ) : 원소 개수
· x.element_size( ) : 원소 하나의 Byte 크기
23 : 대표 성능 / 느린 실행 시간 기준 성능 / 빠른 실행 시간 기준 성능 반환
** 실행 시간↑ → GB/s↓
24 : Benchmark 실행
> Result
→ Vector Add = 매우 단순한 연산이라 계산량보다 메모리 Read/Write 비용↑
· 작은 Size : GPU는 Kernel Launch - Scheduler 동작 - Warp 배치 - Memory Request 과정 먼저 발생
→ 느림(실제 계산 시간 << GPU 준비 시간)
· 중간 Size : GPU가 더 많은 SM을 활용 → 빨라짐
· 큰 Size : GPU Memory Bandwidth 한계 도달 → 포화
▣ Constant Add
1 2 3 4 5 6 7 8 9 10 11 12 | def add_spec(x: Float32[Tensor, "32"]) -> Float32[Tensor, "32"]: return x + 10. @triton.jit def add_kernel(x_ptr, z_ptr, N0, B0: tl.constexpr): range = tl.arange(0, B0) x = tl.load(x_ptr + range) x = x + 10 tl.store(z_ptr + range, x) return test(add_kernel, add_spec, nelem={"N0": 32}, viz=False) | cs |
· x 값에 상수 10을 더하는 연산(전체 원소 수 = 한 Program이 처리하는 원소 수)
· 한 Program이 처리하는 원소(B) = Test 함수에서 32로 고정
4~5 : 일반 Python 함수가 아니라 GPU에서 여러 개의 Program으로 나누어 병렬 실행될 Triton Kernel 선언
→ 원소 수가 32로 Program 1개만 실행
· x_ptr : 입력 Vector의 시작 주소
· z_ptr : 출력 Vector의 시작 주소
· N0 : 전체 원소 수
· B0 : 한 Program이 처리하는 원소 수
6 : Index 생성
7 : 데이터 Load(PyTorch : x = x_tensor[:32])
8 : 상수 10을 더하는 연산
9 : 계산 결과를 출력 Vector에 저장(PyTorch : z[:] = x)
10 : Triton Kernel은 값 반환 X (∵ 이미 "tl.store" GPU Memory에 저장)
▣ Constant Add Block
1 2 3 4 5 6 7 8 9 10 11 12 13 14 | def add2_spec(x: Float32[Tensor, "200"]) -> Float32[Tensor, "200"]: return x + 10. @triton.jit def add_mask2_kernel(x_ptr, z_ptr, N0, B0: tl.constexpr): pid = tl.program_id(0) offsets = pid * B0 + tl.arange(0, B0) mask = offsets < N0 x = tl.load(x_ptr + offsets, mask = mask) z = x + 10 tl.store(z_ptr + offsets, z, mask = mask) return test(add_mask2_kernel, add2_spec, nelem={"N0": 200}) | cs |
· x 값에 상수 10을 더하는 연산(전체 원소 수 ≠ 한 Program이 처리하는 원소 수)
· 한 Program이 처리하는 원소(B) = Test 함수에서 32로 고정
· Program 0 : 0~31 / Program 1 : 32~63 / … / Program 6 : 192~223
→ N = 200개인데, 한 Program은 32개 원소만 처리 (∴ 여러 Program 필요)
→ 전체 원소는 0~199이므로 Program 6에서 200~223 존재 X (∴ Making 필요)
4~5 : 일반 Python 함수가 아니라 GPU에서 여러 개의 Program으로 나누어 병렬 실행될 Triton Kernel 선언
· x_ptr : 입력 Vector의 시작 주소
· z_ptr : 출력 Vector의 시작 주소
· N0 : 전체 원소 수
· B0 : 한 Program이 처리하는 원소 수
6 : Program 번호
7 : 현재 Program이 맡을 Index 생성
8 : 유효한 Index만 True로 표시
9 : 유효한 위치의 "x"만 Load
10 : 각 원소에 상수 10을 더하는 연산
11 : 유효한 위치에만 결과 저장
12 : Triton Kernel은 값 반환 X (∵ 이미 "tl.store" GPU Memory에 저장)
▣ Outer Vector Add
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 | def add_vec_spec(x: Float32[Tensor, "32"], y: Float32[Tensor, "32"]) -> Float32[Tensor, "32 32"]: return x[None, :] + y[:, None] @triton.jit def add_vec_kernel(x_ptr, y_ptr, z_ptr, N0, N1, B0: tl.constexpr, B1: tl.constexpr): i = tl.arange(0, B0) j = tl.arange(0, B1) x = tl.load(x_ptr + i) y = tl.load(y_ptr + j) z = x[None, :] + y[:, None] offsets = j[:, None] * N0 + i[None, :] tl.store(z_ptr + offsets, z) return test(add_vec_kernel, add_vec_spec, nelem={"N0": 32, "N1": 32}) | cs |
· 크기가 다른 배열을 자동으로 늘려서 연산하는 Broadcasting 덧셈 연산(전체 원소 수 = 한 Program이 처리하는 원소 수)
· 한 Program이 처리하는 원소(B) = Test 함수에서 32로 고정
Ex. x = [1, 2, 3] / y = [10, 20, 30]
→ [[11, 12, 13],
[21, 22, 33],
[31, 32, 33]]
2 : "x"를 행 방향으로 생성 / "y"를 열 방향으로 생성
4~5 : 일반 Python 함수가 아니라 GPU에서 여러 개의 Program으로 나누어 병렬 실행될 Triton Kernel 선언
→ 원소 수가 32로 Program 1개만 실행
· x_ptr : 입력 Vector의 시작 주소
· y_ptr : 입력 Vector의 시작 주소
· z_ptr : 출력 Vector의 시작 주소
· N0, N1 : 전체 원소 수
· B0, B1 : 한 Program이 처리하는 원소 수
6 : Column Index 생성
7 : Row Index 생성
8~9 : 데이터 Load
10 : Broadcasting
11 : 2차원 결과를 1차원 Memory 주소로 변환
12 : 계산한 결과를 "z_ptr"에 저장
13 : Triton Kernel은 값 반환 X (∵ 이미 "tl.store" GPU Memory에 저장)
▣ Outer Vector Add Block
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 | def add_vec_block_spec(x: Float32[Tensor, "100"], y: Float32[Tensor, "90"]) -> Float32[Tensor, "90 100"]: return x[None, :] + y[:, None] @triton.jit def add_vec_block_kernel(x_ptr, y_ptr, z_ptr, N0, N1, B0: tl.constexpr, B1: tl.constexpr): pid_0 = tl.program_id(0) pid_1 = tl.program_id(1) i = pid_0 * B0 + tl.arange(0, B0) j = pid_1 * B1 + tl.arange(0, B1) mask_i = i < N0 mask_j = j < N1 x = tl.load(x_ptr + i, mask = mask_i, other = 0) y = tl.load(y_ptr + j, mask = mask_j, other = 0) z = x[None, :] + y[:, None] offsets = j[:, None] * N0 + i[None, :] mask = mask_j[:, None] & mask_i[None, :] tl.store(z_ptr + offsets, z, mask = mask) return test(add_vec_block_kernel, add_vec_block_spec, nelem={"N0": 100, "N1": 90}) | cs |
· 크기가 다른 배열을 자동으로 늘려서 연산하는 Broadcasting 덧셈 연산(전체 원소 수 ≠ 한 Program이 처리하는 원소 수)
· 한 Program이 처리하는 원소(B) = Test 함수에서 32로 고정
→ z = 90x100인데, 한 Program은 32x32 Block만 처리 (∴ 여러 Program 필요)
2 : "x"를 행 방향으로 생성 / "y"를 열 방향으로 생성
4~5 : 일반 Python 함수가 아니라 GPU에서 여러 개의 Program으로 나누어 병렬 실행될 Triton Kernel 선언
· x_ptr : 입력 Vector의 시작 주소
· y_ptr : 입력 Vector의 시작 주소
· z_ptr : 출력 Vector의 시작 주소
· N0, N1 : 전체 원소 수
· B0, B1 : 한 Program이 처리하는 원소 수
6~7 : 현재 Program의 2차원 위치
8 : Column Index 생성
9 : Row Index 생성
10 : 100~127 범위 Masking = False
11 : 96~127 범위 Masking = False
12~13 : 유효한 데이터 Load
14 : Broadcasting
15 : 2차원 결과를 1차원 Memory 주소로 변환
16 : 2차원 Mask 생성(j < N1 & i < N0)
17 : 유효한 위치에만 결과 저장
18 : Triton Kernel은 값 반환 X (∵ 이미 "tl.store" GPU Memory에 저장)
** 범위를 벗어난 마지막 Block을 처리하기 위해 2차원 Masking
** 결과 Matrix의 Shape = (90, 100) ▶ 가로 방향 Program 수 : ceil(100 / 32) = 4 / 세로 방향 Program 수 : ceil(90 / 32) = 3
→ 전체 Grid = (4, 3) / 총 Program 수 = 12개
▣ Fused Outer Multiplication
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 | def mul_relu_block_spec(x: Float32[Tensor, "100"], y: Float32[Tensor, "90"]) -> Float32[Tensor, "90 100"]: return torch.relu(x[None, :] * y[:, None]) @triton.jit def mul_relu_block_kernel(x_ptr, y_ptr, z_ptr, N0, N1, B0: tl.constexpr, B1: tl.constexpr): pid_0 = tl.program_id(0) pid_1 = tl.program_id(1) i = pid_0 * B0 + tl.arange(0, B0) j = pid_1 * B1 + tl.arange(0, B1) mask_i = i < N0 mask_j = j < N1 x = tl.load(x_ptr + i, mask=mask_i, other=0.0) y = tl.load(y_ptr + j, mask=mask_j, other=0.0) z = x[None, :] * y[:, None] z = tl.maximum(z, 0.0) offsets = j[:, None] * N0 + i[None, :] mask = mask_j[:, None] & mask_i[None, :] tl.store(z_ptr + offsets, z, mask=mask) return test(mul_relu_block_kernel, mul_relu_block_spec, nelem={"N0": 100, "N1": 90}) | cs |
· 크기가 다른 배열을 자동으로 늘려서 연산하는 Broadcasting 곱셈 후 ReLU 적용(전체 원소 수 ≠ 한 Program이 처리하는 원소 수)
· 한 Program이 처리하는 원소(B) = Test 함수에서 32로 고정
→ z = 90x100인데, 한 Program은 32x32 Block만 처리 (∴ 여러 Program 필요)
Ex. x = [1, 2, 3] / y = [10, 20, 30]
→ [[10, 20, 30],
[20, 40, 60],
[30, 60, 90]]
2 : "x"를 행 방향으로 생성 / "y"를 열 방향으로 생성
4~5 : 일반 Python 함수가 아니라 GPU에서 여러 개의 Program으로 나누어 병렬 실행될 Triton Kernel 선언
· x_ptr : 입력 Vector의 시작 주소
· y_ptr : 입력 Vector의 시작 주소
· z_ptr : 출력 Vector의 시작 주소
· N0, N1 : 전체 원소 수
· B0, B1 : 한 Program이 처리하는 원소 수
6~7 : 현재 Program의 2차원 위치
8 : Column Index 생성
9 : Row Index 생성
10 : 100~127 범위 Masking = False
11 : 96~127 범위 Masking = False
12~13 : 유효한 데이터 Load
14 : Broadcasting으로 2D 곱셈 결과 생성
15 : ReLU 적용(PyTorch : torch.relu(z))
16 : 2차원 결과를 1차원 Memory 주소로 변환
17 : 2차원 Mask 생성(j < N1 & i < N0)
18 : 유효한 위치에만 결과 저장
19 : Triton Kernel은 값 반환 X (∵ 이미 "tl.store" GPU Memory에 저장)
** 범위를 벗어난 마지막 Block을 처리하기 위해 2차원 Masking
** 결과 Matrix의 Shape = (90, 100) ▶ 가로 방향 Program 수 : ceil(100 / 32) = 4 / 세로 방향 Program 수 : ceil(90 / 32) = 3
→ 전체 Grid = (4, 3) / 총 Program 수 = 12개
▣ Long Sum
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 | def sum_spec(x: Float32[Tensor, "4 200"]) -> Float32[Tensor, "4"]: return x.sum(1) @triton.jit def sum_kernel(x_ptr, z_ptr, N0, N1, T, B0: tl.constexpr, B1: tl.constexpr): pid_0 = tl.program_id(0) row = pid_0 * B0 + tl.arange(0, B0) col = tl.arange(0, B1) acc = tl.zeros((B0,), dtype=tl.float32) for t in range(0, T, B1): cols = t + col mask = (row[:, None] < N0) & (cols[None, :] < T) x = tl.load( x_ptr + row[:, None] * T + cols[None, :], mask=mask, other=0.0 ) acc += tl.sum(x, axis=1) tl.store(z_ptr + row, acc, mask=row < N0) return test(sum_kernel, sum_spec, B={"B0": 1, "B1": 32}, nelem={"N0": 4, "N1": 32, "T": 200}) | cs |
· 2차원 Tensor에서 Column 방향을 모두 더해서 Row마다 하나의 값만 남기는 Sum Reduction 연산
· 한 Program은 Row 1(=B0)개를 담당하고, Column은 한 번에 32(=B1)개씩 나누어 처리
· x.shape = (4, 200) → z.shape = (4, ) : 각 Row마다 200개 값을 모두 더하는 연산
→ Program 0 : z[0] = x[0, 0] + x[0, 1] + … + x[0, 199]
→ Program 1 : z[1] = x[1, 0] + x[1, 1] + … + x[1, 199]
→ Program 2 : z[2] = x[2, 0] + x[2, 1] + … + x[2, 199]
→ Program 3 : z[3] = x[3, 0] + x[3, 1] + … + x[3, 199]
4~5 : 일반 Python 함수가 아니라 GPU에서 여러 개의 Program으로 나누어 병렬 실행될 Triton Kernel 선언
· x_ptr : 입력 Vector의 시작 주소
· z_ptr : 출력 Vector의 시작 주소
· N0 : Row 개수
· N1 : Grid 계산용 보조 차원
· T : 전체 Column 개수
· B0 : 한 Program이 처리하는 Row 개수
· B1 : 한 번에 읽는 Column 개수
6 : 현재 Program 번호
7 : Row Index 생성
8 : Column Index 생성
9 : Sum 결과를 누적할 변수(B0 = 1 → Shape = (1, ) → Row 하나의 합계를 저장)
11 : Column 200개를 32개씩 나누어 Read
12 : 현재 Loop에서 읽을 Column Index
13 : 유효한 위치만 읽기 위한 2D Mask(row < N0 & cols < T)
14~18 : 2차원 Tensor "x[row, col]"를 Read(범위 밖 Column = 0)
19 : 현재 32개 Column 값을 더해서 누적(x.shape = (B0, B1) → z.shape(B0, ) ∴ Column 방향으로 합침 = 현재 Row의 부분합을 누적)
21 : 각 Row의 최종 합계를 "z[row]"에 저장
22 : Triton Kernel은 값 반환 X (∵ 이미 "tl.store" GPU Memory에 저장)
** Column 200개를 한 번에 다 읽지 않고, 32(=B1)개씩 나누어 읽으면서 "acc"에 누적
▣ Long Sum 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 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 | def softmax_spec(x: Float32[Tensor, "4 200"]) -> Float32[Tensor, "4 200"]: x_max = x.max(1, keepdim=True)[0] x = x - x_max x_exp = x.exp() return x_exp / x_exp.sum(1, keepdim=True) @triton.jit def softmax_kernel(x_ptr, z_ptr, N0, N1, T, B0: tl.constexpr, B1: tl.constexpr): pid_0 = tl.program_id(0) log2_e = 1.44269504 row = pid_0 * B0 + tl.arange(0, B0) col = tl.arange(0, B1) row_mask = row < N0 # 1단계: row마다 max 찾기 x_max = tl.full((B0,), -float("inf"), dtype=tl.float32) for t in range(0, T, B1): cols = t + col mask = row_mask[:, None] & (cols[None, :] < T) x = tl.load(x_ptr + row[:, None] * T + cols[None, :], mask=mask, other=-float("inf")) x_max = tl.maximum(x_max, tl.max(x, axis=1)) # 2단계: exp(x - max)의 합 구하기 denom = tl.zeros((B0,), dtype=tl.float32) for t in range(0, T, B1): cols = t + col mask = row_mask[:, None] & (cols[None, :] < T) x = tl.load(x_ptr + row[:, None] * T + cols[None, :], mask=mask, other=-float("inf")) x_shifted = x - x_max[:, None] x_exp = tl.exp2(x_shifted * log2_e) denom += tl.sum(x_exp, axis=1) # 3단계: softmax 값 저장 for t in range(0, T, B1): cols = t + col mask = row_mask[:, None] & (cols[None, :] < T) x = tl.load(x_ptr + row[:, None] * T + cols[None, :], mask=mask, other=-float("inf")) x_shifted = x - x_max[:, None] x_exp = tl.exp2(x_shifted * log2_e) z = x_exp / denom[:, None] tl.store(z_ptr + row[:, None] * T + cols[None, :], z, mask=mask) return test(softmax_kernel, softmax_spec, B={"B0": 1, "B1":32}, nelem={"N0": 4, "N1": 32, "T": 200}) | cs |
· x.shape = (4, 200) → z.shape = (4, 200) : 각 Row마다 Softmax 계산
· 한 Program은 Row 1(=B0)개를 담당하고, Column은 한 번에 32(=B1)개씩 나누어 처리
→ Program 0 : Row 0 Softmax
→ Program 1 : Row 1 Softmax
→ Program 2 : Row 2 Softmax
→ Program 3 : Row 3 Softmax
7~8 : 일반 Python 함수가 아니라 GPU에서 여러 개의 Program으로 나누어 병렬 실행될 Triton Kernel 선언
· x_ptr : 입력 Vector의 시작 주소
· z_ptr : 출력 Vector의 시작 주소
· N0 : Row 개수
· N1 : Grid 계산용 보조 차원
· T : 전체 Column 개수
· B0 : 한 Program이 처리하는 Row 개수
· B1 : 한 번에 읽는 Column 개수
9 : 현재 Program 번호
11 : 현재 Program이 담당할 Row Index 생성
12 : 한 번에 읽을 Column Block Index 생성
13 : 유효한 Row인지 확인 (row = 0~3 & N0 = 4 → 모두 True)
> Row Max 구하기
16 : 각 Row의 최댓값을 저장할 변수(처음에는 아주 작은 값인 "- ∞"로 시작 / B0 = 1 → Shape = (1, ))
18 : Column 200개를 32개씩 나눠서 Load
19 : 이번 반복에서 실제로 Load 할 Column Index
20 : 유효한 위치만 Load하기 위한 2차원 Mask(row < N0 & col < T)
21 : 현재 Row의 Column Block을 Load(범위 밖 Column = - ∞)
22 : 현재 Load 한 32개 값 중 최댓값을 구하고, 기존 "x_max"와 비교해서 갱신
→ 이 반복이 끝나면 해당 Row 전체 200개 값의 최댓값이 "x_max"에 저장
> exp(x - max)의 합 구하기
25 : Softmax의 분모를 저장할 변수
27 : Column 200개를 32개씩 나눠서 Load
28 : 이번 반복에서 읽을 Column Index
29 : 유효한 위치만 Load하기 위한 2차원 Mask(row < N0 & col < T)
30 : 현재 Row의 Column Block을 Load(범위 밖 Column = - ∞)
31 : Softmax Overflow를 막기 위해 각 값에서 Row 최댓값을 빼는 연산 진행
32 : "exp(x_shifted)" 계산(Triton에서 "exp"대신 "exp2"를 사용한 형태)
33 : 현재 32개 Column의 "exp" 값을 더해서 분모에 누적
> exp(x - max) / sum{exp(x - max)} 저장
36 : Column 200개를 32개씩 나눠서 Load
37 : 이번 반복에서 저장할 Column Index
38 : 유효한 위치에 저장하기 위한 Mask
39 : 현재 Column Block 값을 다시 Load
40 : 다시 Row 최댓값을 빼는 연산 진행
41 : 분자 "exp(x - max)" 계산
42 : Softmax 값 계산
43 : 계산한 Softmax 결과를 출력 Tensor "z"에 저장(범위 밖 Column은 저장 X)
44 : Triton Kernel은 값 반환 X (∵ 이미 "tl.store" GPU Memory에 저장)
[Reference]
· HW5 - Triton Puzzle (MGP) - Yongjun Park
· 11_Triton Indroduction_mgp_2026 (MGP) - Yongjun Park
댓글
댓글 쓰기