Giới thiệu Triton: Lập trình GPU nguồn mở cho mạng nơ-ron

Các ý tưởng nghiên cứu mới trong lĩnh vực Học sâu thường được triển khai bằng cách kết hợp các toán tử khung gốc. Mặc dù tiện lợi, nhưng cách tiếp cận này thường yêu cầu tạo (và/hoặc di chuyển) nhiều tenxơ tạm thời, điều này có thể làm giảm hiệu suất của mạng nơ-ron ở quy mô lớn. Những vấn đề này có thể được giảm thiểu bằng cách viết các hạt nhân GPU chuyên dụng, nhưng việc thực hiện như vậy có thể cực kỳ khó khăn do nhiều phức tạp của lập trình GPU. Và mặc dù nhiều hệ thống đã xuất hiện gần đây để giúp quá trình này dễ dàng hơn, chúng tôi thấy chúng quá dài dòng, thiếu tính linh hoạt hoặc tạo mã chậm hơn đáng kể so với các đường cơ sở do chúng tôi điều chỉnh thủ công. Điều này đã khiến chúng tôi mở rộng và cải thiện Triton, một ngôn ngữ và trình biên dịch gần đây mà người tạo ra ban đầu hiện đang làm việc tại OpenAI.

Kiến trúc của GPU hiện đại có thể được chia thành ba thành phần chính—DRAM, SRAM và ALU—mỗi thành phần phải được xem xét khi tối ưu hóa mã CUDA:

+ Việc chuyển bộ nhớ từ DRAM phải được hợp nhất thành các giao dịch lớn để tận dụng được độ rộng bus lớn của các giao diện bộ nhớ hiện đại.

+ Dữ liệu phải được lưu trữ thủ công vào SRAM trước khi sử dụng lại và được quản lý sao cho giảm thiểu xung đột giữa các ngân hàng bộ nhớ dùng chung khi truy xuất.

+ Các phép tính phải được phân vùng và lên lịch cẩn thận, cả trên và trong Bộ xử lý đa luồng (SM), để thúc đẩy tính song song ở cấp độ lệnh/luồng và tận dụng ALU có mục đích đặc biệt (ví dụ: lõi tensor).

Kiến trúc cơ bản của GPU

Việc lý giải tất cả các yếu tố này có thể là một thách thức, ngay cả đối với các lập trình viên CUDA dày dạn kinh nghiệm với nhiều năm kinh nghiệm. Mục đích của Triton là tự động hóa hoàn toàn các tối ưu hóa này, để các nhà phát triển có thể tập trung tốt hơn vào logic cấp cao của mã song song của họ. Triton hướng đến mục tiêu có thể áp dụng rộng rãi và do đó không tự động lên lịch công việc trên các SM -- để lại một số cân nhắc thuật toán quan trọng (ví dụ như lát gạch, đồng bộ hóa giữa các SM) tùy theo quyết định của các nhà phát triển.

	KHÁC BIỆT	TRITON
Bộ nhớ hợp nhất	Thủ công	Tự động
Quản lý bộ nhớ chia sẻ	Thủ công	Tự động
Lên lịch (Trong SM)	Thủ công	Tự động
Lên lịch (Trên khắp SM)	Thủ công	Thủ công

Tối ưu hóa trình biên dịch trong CUDA so với Triton.

Mô hình lập trình

Trong số tất cả các Ngôn ngữ miền cụ thể và trình biên dịch JIT có sẵn, Triton có lẽ giống với Numba nhất: các hạt nhân được định nghĩa là các hàm Python được trang trí và được khởi chạy đồng thời với các program_id' khác nhau trên một lưới các cái gọi là các thể hiện . Tuy nhiên, như được hiển thị trong đoạn mã bên dưới, sự giống nhau dừng lại ở đó: Triton phơi bày tính song song trong thể hiện thông qua các hoạt động trên các khối — các mảng nhỏ có kích thước là lũy thừa của hai — thay vì mô hình thực thi SIMT (Single Instruction, Multiple Thread). Khi làm như vậy, Triton thực sự trừu tượng hóa tất cả các vấn đề liên quan đến tính đồng thời trong các khối luồng CUDA (ví dụ: hợp nhất bộ nhớ, đồng bộ hóa/xung đột bộ nhớ được chia sẻ, lập lịch lõi tensor).

BLOCK = 512
# This is a GPU kernel in Numba.
# Different instances of this
# function may run in parallel.
@jit
def add(X, Y, Z, N): # In Numba/CUDA, each kernel # instance itself uses an SIMT execution # model, where instructions are executed in # parallel for different values of threadIdx tid = threadIdx.x bid = blockIdx.x # scalar index idx = bid * BLOCK + tid if id < N: # There is no pointer in Numba. # Z,X,Y are dense tensors Z[idx] = X[idx] + Y[idx]
...
grid = (ceil_div(N, BLOCK),)
block = (BLOCK,)
add[grid, block](x, y, z, x.shape[0])

BLOCK = 512
# This is a GPU kernel in Triton.
# Different instances of this
# function may run in parallel.
@jit
def add(X, Y, Z, N): # In Triton, each kernel instance # executes block operations on a # single thread: there is no construct # analogous to threadIdx pid = program_id(0) # block of indices idx = pid * BLOCK + arange(BLOCK) mask = idx < N # Triton uses pointer arithmetics # rather than indexing operators x = load(X + idx, mask=mask) y = load(Y + idx, mask=mask) store(Z + idx, x + y, mask=mask)
...
grid = (ceil_div(N, BLOCK),)
# no thread-block
add[grid](x, y, z, x.shape[0])

Phép cộng vectơ trong Triton.

Mặc dù điều này có thể không hữu ích lắm đối với các phép tính song song phức tạp (tức là từng phần tử), nhưng nó có thể đơn giản hóa đáng kể quá trình phát triển các chương trình GPU phức tạp hơn.

Ví dụ, hãy xem xét trường hợp của một hạt nhân softmax hợp nhất (bên dưới) trong đó mỗi trường hợp chuẩn hóa một hàng khác nhau của tenxơ đầu vào đã cho X_∈R_M_×_N . Các triển khai CUDA chuẩn của chiến lược song song hóa này có thể khó viết, đòi hỏi phải đồng bộ hóa rõ ràng giữa các luồng khi chúng đồng thời giảm cùng một hàng của X . Hầu hết sự phức tạp này biến mất với Triton, trong đó mỗi trường hợp hạt nhân tải hàng quan tâm và chuẩn hóa nó theo trình tự bằng cách sử dụng các nguyên hàm giống NumPy.

con trăn

123456789101112131415161718192021222324252627282930313233343536

1

import triton

2

import triton.language as tl

3

4

@triton.jit

5

def softmax(Y, stride_ym, stride_yn, X, stride_xm, stride_xn, M, N):

6

# row index

7

m = tl.program_id(0)

8

# col indices

9

# this specific kernel only works for matrices that

10

# have less than BLOCK_SIZE columns

11

BLOCK_SIZE = 1024

12

n = tl.arange(0, BLOCK_SIZE)

13

# the memory address of all the elements

14

# that we want to load can be computed as follows

15

X = X + m * stride_xm + n * stride_xn

16

# load input data; pad out-of-bounds elements with 0

17

x = tl.load(X, mask=n < N, other=-float('inf'))

18

# compute numerically-stable softmax

19

z = x - tl.max(x, axis=0)

20

num = tl.exp(z)

21

denom = tl.sum(num, axis=0)

22

y = num / denom

23

# write back to Y

24

Y = Y + m * stride_ym + n * stride_yn

25

tl.store(Y, y, mask=n < N)

26

27

import torch

28

# Allocate input/output tensors

29

X = torch.normal(0, 1, size=(583, 931), device='cuda')

30

Y = torch.empty_like(X)

31

# SPMD launch grid

32

grid = (X.shape[0], )

33

# enqueue GPU kernel

34

softmax[grid](Y, Y.stride(0), Y.stride(1),

35

X, X.stride(0), X.stride(1),

36

X.shape[0] , X.shape[1])

Lưu ý rằng Triton JIT xử lý X và Y như con trỏ chứ không phải tenxơ; chúng tôi cảm thấy việc duy trì quyền kiểm soát cấp thấp đối với quyền truy cập bộ nhớ là quan trọng để giải quyết các cấu trúc dữ liệu phức tạp hơn (ví dụ: tenxơ thưa khối).

Quan trọng là, việc triển khai softmax cụ thể này giữ các hàng của X trong SRAM trong suốt quá trình chuẩn hóa, giúp tối đa hóa việc sử dụng lại dữ liệu khi có thể (~<32K cột). Điều này khác với mã CUDA nội bộ của PyTorch, sử dụng bộ nhớ tạm thời khiến nó tổng quát hơn nhưng chậm hơn đáng kể (bên dưới). Điểm mấu chốt ở đây không phải là Triton vốn tốt hơn, mà là nó đơn giản hóa quá trình phát triển các hạt nhân chuyên biệt có thể nhanh hơn nhiều so với những hạt nhân có trong các thư viện mục đích chung.

5,00010,000N05001000 GB/sTritonTorch (native)Torch (jit)

Hiệu suất A100 của softmax hợp nhất cho M=4096.

Hiệu suất thấp hơn của JIT Torch (v1.9) làm nổi bật khó khăn trong việc tạo mã CUDA tự động từ chuỗi các hoạt động tenxơ cấp cao.

con trăn

1234567

1

@torch.jit.script

2

def softmax(x):

3

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

4

z = x - x_max[:, None]

5

numerator = torch.exp(x)

6

denominator = numerator.sum(dim=1)

7

return numerator / denominator[:, None]

Kết hợp softmax với Torch JIT.

Phép nhân ma trận

Có thể viết các hạt nhân hợp nhất cho các phép toán và phép giảm từng phần tử là quan trọng, nhưng không đủ vì các tác vụ nhân ma trận trong mạng nơ-ron rất nổi bật. Hóa ra, Triton cũng hoạt động rất tốt cho những tác vụ đó, đạt hiệu suất cao nhất chỉ với ~25 dòng mã Python. Mặt khác, việc triển khai một cái gì đó tương tự trong CUDA sẽ tốn nhiều công sức hơn (mở trong cửa sổ mới) và thậm chí có thể đạt được hiệu suất thấp hơn.

con trăn

1234567891011121314151617181920212223242526272829303132333435363738394041

1

@triton.jit

2

def matmul(A, B, C, M, N, K, stride_am, stride_ak,

3

stride_bk, stride_bn, stride_cm, stride_cn,

4

**META):

5

# extract metaparameters

6

BLOCK_M, GROUP_M = META['BLOCK_M'], META['GROUP_M']

7

BLOCK_N = META['BLOCK_N']

8

BLOCK_K = META['BLOCK_K']

9

# programs are grouped together to improve L2 hit rate

10

_pid_m = tl.program_id(0)

11

_pid_n = tl.program_id(1)

12

pid_m = _pid_m // GROUP_M

13

pid_n = (_pid_n * GROUP_M) + (_pid_m % GROUP_M)

14

# rm (resp. rn) denotes a range of indices

15

# for rows (resp. col) of C

16

rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)

17

rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)

18

# rk denotes a range of indices for columns

19

# (resp. rows) of A (resp. B)

20

rk = tl.arange(0, BLOCK_K)

21

# the memory addresses of elements in the first block of

22

# A and B can be computed using numpy-style broadcasting

23

A = A + (rm[:, None] * stride_am + rk[None, :] * stride_ak)

24

B = B + (rk [:, None] * stride_bk + rn[None, :] * stride_bn)

25

# initialize and iteratively update accumulator

26

acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)

27

for k in range(K, 0, -BLOCK_K):

28

a = tl.load(A)

29

b = tl.load(B)

30

# block level matrix multiplication

31

acc += tl.dot(a, b)

32

# increment pointers so that the next blocks of A and B

33

# are loaded during the next iteration

34

A += BLOCK_K * stride_ak

35

B += BLOCK_K * stride_bk

36

# fuse leaky ReLU if desired

37

# acc = tl.where(acc >= 0, acc, alpha * acc)

38

# write back result

39

C = C + (rm[:, None] * stride_cm + rn[None, :] * stride_cn)

40

mask = (rm[:, None] < M) & (rn[None, :] < N)

41

tl.store(C, acc, mask=mask)

Phép nhân ma trận trong Triton.

Một lợi thế quan trọng của các hạt nhân nhân ma trận viết tay là chúng có thể được tùy chỉnh theo ý muốn để phù hợp với các phép biến đổi hợp nhất của đầu vào (ví dụ: cắt lát) và đầu ra (ví dụ: Leaky ReLU). Nếu không có một hệ thống như Triton, các sửa đổi không tầm thường của các hạt nhân nhân ma trận sẽ nằm ngoài tầm với của các nhà phát triển không có chuyên môn lập trình GPU đặc biệt.

1,0002,0003,0004,000M = N = K020406080100 TFLOPscuBLASTritonTriton + LeakyReLUcuBLAS +torch.nn.LeakyReLU

Hiệu suất lõi tenxơ V100 của phép nhân ma trận với các giá trị được điều chỉnh phù hợp cho BLOCKTôiTôi, KHỐINN, KHỐIKK, NHÓMTôiTôi.

Kiến trúc hệ thống cấp cao

Hiệu suất tốt của Triton xuất phát từ kiến trúc hệ thống mô-đun tập trung xung quanh Triton-IR, một biểu diễn trung gian dựa trên LLVM trong đó các khối giá trị đa chiều là công dân hạng nhất.

Trăn

Triton-IR

LLVM-IR

PTX

@jit
def add(X, Y, Z, N): pid = program_id(0) idx= pid * 512 + arange(512) mask = idx < N x = load(X + idx, mask=mask) y = load(Y + idx, mask=mask) store(Z + idx, x + y, mask=mask)

def void add(i32* X .aligned(16) , i32* Y .aligned(16) , i32* Z .aligned(16) , i32 N .multipleof(2) )
{
entry: %0 = get_program_id[0] i32; %1 = mul i32 %0, 512; %3 = make_range[0 : 512] i32<512>; %4 = splat i32<512> %1; %6 = add i32<512> %4, %3; %9 = splat i32<512> N; %11 = icmp_slt i1<512> %6, %9; %14 = splat i32*<512> X; %16 = getelementptr i32*<512> %14, %6; %19 = broadcast i1<512> %11; %21 = splat i32<512> undef; %22 = masked_load i32<512> %16, %19, %21; %26 = splat i32*<512> Y; %28 = getelementptr i32*<512> %26, %6; %31 = broadcast i1<512> %11; %33 = splat i32<512> undef; %34 = masked_load i32<512> %28, %31, %33; %38 = splat i32*<512> Z; %40 = getelementptr i32*<512> %38, %6; %43 = add i32<512> %22, %34; %46 = broadcast i32<512> %43; %48 = broadcast i1<512> %11; masked_store void %40, %46, %48; ret void;
}

.visible .entry add( .param .u64 add_param_0, .param .u64 add_param_1, .param .u64 add_param_2, .param .u32 add_param_3
)
.maxntid 128, 1, 1
{ .reg .pred %p<4>; .reg .b32 %r<18>; .reg .b64 %rd<8>; ld.param.u64 %rd4, [add_param_0]; ld.param.u64 %rd5, [add_param_1]; mov.u32 %r13, %tid.x; ld.param.u32 %r14, [add_param_3]; shl.b32 %r15, %r13, 2; mov.u32 %r16, %ctaid.x; mad.lo.s32 %r17, %r16, 512, %r15; setp.ge.s32 %p3, %r17, %r14; setp.lt.s32 %p1, %r17, %r14; mul.wide.s32 %rd7, %r17, 4; add.s64 %rd2, %rd4, %rd7; @%p1 ld.global.cg.v4.b32 {%r5,%r6,%r7,%r8}, [ %rd2 + 0]; add.s64 %rd3, %rd5, %rd7; @%p1 ld.global.cg.v4.b32 {%r9,%r10,%r11,%r12}, [ %rd3 + 0]; @%p3 bra LBB0_2; ld.param.u64 %rd6, [add_param_2]; add.s64 %rd1, %rd6, %rd7; add.s32 %r1, %r5, %r9; add.s32 %r2, %r6, %r10; add.s32 %r3, %r7, %r11; add.s32 %r4, %r8, %r12; st.global.v4.u32 [%rd1], {%r1, %r2, %r3, %r4};
LBB0_2: ret;
}

Kiến trúc cao cấp của Triton.

Bộ @triton.jit trang trí hoạt động bằng cách duyệt Cây cú pháp trừu tượng (AST) của hàm Python được cung cấp để tạo Triton-IR ngay lập tức bằng thuật toán xây dựng SSA phổ biến. Mã IR kết quả sau đó được đơn giản hóa, tối ưu hóa và tự động song song hóa bởi trình biên dịch phụ trợ của chúng tôi, trước khi được chuyển đổi thành LLVM-IR chất lượng cao—và cuối cùng là PTX—để thực thi trên các GPU NVIDIA gần đây. CPU và GPU AMD hiện không được hỗ trợ, nhưng chúng tôi hoan nghênh các đóng góp của cộng đồng nhằm giải quyết hạn chế này.

Trình biên dịch phụ trợ

Chúng tôi đã phát hiện ra rằng việc sử dụng các biểu diễn chương trình bị chặn thông qua Triton-IR cho phép trình biên dịch của chúng tôi tự động thực hiện nhiều tối ưu hóa chương trình quan trọng. Ví dụ, dữ liệu có thể được tự động lưu trữ vào bộ nhớ dùng chung bằng cách xem các toán hạng của các hoạt động cấp khối tốn nhiều tính toán (ví dụ: tl.dot)—và được phân bổ/đồng bộ hóa bằng các kỹ thuật phân tích độ sống tiêu chuẩn.

Mặt khác, các chương trình Triton có thể được song song hóa hiệu quả và tự động cả (1) trên các SM bằng cách thực thi các phiên bản kernel khác nhau đồng thời và (2) trong các SM bằng cách phân tích không gian lặp của từng hoạt động cấp khối và phân vùng nó một cách thích hợp trên các đơn vị SIMD khác nhau, như được hiển thị bên dưới.

Tự động song song hóa trong Triton. Mỗi hoạt động cấp khối xác định một không gian lặp bị chặn được tự động song song hóa để sử dụng các tài nguyên có sẵn trên Bộ xử lý đa luồng (SM).