Triton：用于神经网络的开源 GPU 编程

我们正在发布 Triton 1.0，这是一种类似于 Python 的开源编程语言，它使没有 CUDA 经验的研究人员能够编写高效的 GPU 代码——大部分时间与专家能够编写的代码相当。 Triton 使得以相对较少的努力达到峰值硬件性能成为可能；例如，它可以用不到 25 行代码编写与 cuBLAS 性能相匹配的 FP16 矩阵乘法内核——这是许多 GPU 程序员无法做到的。我们的研究人员已经使用它来生成效率比同等 Torch 实现高 2 倍的内核，我们很高兴与社区合作，让每个人都更容易使用 GPU 编程。在深度学习领域，一般都是使用原生框架算子的组合来实现的。虽然方便，但这种方法通常需要创建（和/或移动）许多临时张量，这可能会大规模损害神经网络的性能。这些问题可以通过编写专门的 GPU 内核来缓解，但由于 GPU 编程的许多复杂性，这样做可能会非常困难。而且，尽管最近出现了各种系统来简化此过程，但我们发现它们要么过于冗长，要么缺乏灵活性，要么生成的代码明显慢于我们手动调整的基线。这促使我们扩展和改进 Triton，这是一种最新的语言和编译器，其最初的创建者现在在 OpenAI 工作。现代 GPU 的架构大致可以分为三个主要部分——DRAM、SRAM 和 ALU——在优化 CUDA 代码时必须考虑到每个部分： 必须将来自 DRAM 的内存传输合并为大型事务，以利用现代内存的大总线宽度接口。数据必须在重新使用之前手动存储到 SRAM 中，并进行管理以最大限度地减少检索时的共享内存库冲突。计算必须在流式多处理器 (SM) 之间和内部仔细分区和调度，以促进指令/线程级并行性并利用专用 ALU（例如，张量核心）。对所有这些因素进行推理可能具有挑战性，即使对于具有多年经验的 CUDA 程序员也是如此。 Triton 的目的是将这些优化完全自动化，让开发人员可以更好地专注于并行代码的高级逻辑。 Triton 旨在广泛适用，因此不会自动安排跨 SM 的工作——将一些重要的算法考虑因素（例如平铺、SM 间同步）留给开发人员自行决定。

在所有可用的领域特定语言和 JIT 编译器中，Triton 可能与 Numba 最相似：内核被定义为修饰的 Python 函数，并在所谓的实例网格上与不同的程序 ID 同时启动。但是，如下面的代码片段所示，相似之处仅止于此：Triton 通过对块（维数为 2 的幂的小数组）的操作来公开实例内并行性，而不是单指令多线程 (SIMT) 执行模型。这样做时，Triton 有效地抽象出了与 CUDA 线程块内的并发相关的所有问题（例如，内存合并、共享内存同步/冲突、张量核心调度）。 BLOCK = 512# 这是 Numba 中的 GPU 内核。# this# 函数的不同实例可以并行运行。@jitdef add(X, Y, Z, N): # 在 Numba/CUDA 中，每个内核 # 实例本身使用一个SIMT 执行 # 模型，其中指令针对不同的 threadIdx 值并行执行 # tid = threadIdx.x bid = blockIdx.x # 标量索引 idx = bid * BLOCK + tid if id < N: # Numba 中没有指针。 # Z,X,Y 是密集张量 Z[idx] = X[idx] + Y[idx]...grid = (ceil_div(N, BLOCK),)block = (BLOCK,)add[grid, block]( x, y, z, x.shape[0]) BLOCK = 512# 这是 Triton 中的 GPU 内核。# this# 函数的不同实例可以并行运行。@jitdef add(X, Y, Z, N): # 在 Triton 中，每个内核实例 # 在单个线程上执行块操作：没有构造 # 类似于 threadIdx pid = program_id(0) # 索引块 idx = pid * BLOCK + arange(BLOCK) mask = idx < N # Triton 使用指针算术 # 而不是索引运算符 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-blockadd[grid](x, y, z, x.shape[0]) 虽然这对于令人尴尬的并行（即元素方式）可能不是特别有用) 计算，它可以大大简化更复杂的 GPU 程序的开发。例如，考虑融合 softmax 内核（如下）的情况，其中每个实例标准化给定输入张量 $X \in \mathbb{R}^{M \times N}$ 的不同行。这种并行化策略的标准 CUDA 实现可能难以编写，需要线程之间的显式同步，因为它们同时减少 $X$ 的同一行。大多数这种复杂性在 Triton 中消失了，其中每个内核实例加载感兴趣的行并使用类似 NumPy 的原语按顺序对其进行规范化。 import tritonimport triton.language as [email protected] softmax(Y, stride_ym, stride_yn, X, stride_xm, stride_xn, M, N): # row index m = tl.program_id(0) # col indices # 这个特定的内核才有效对于 # 少于 BLOCK_SIZE 列的矩阵 BLOCK_SIZE = 1024 n = tl.arange(0, BLOCK_SIZE) # 我们想要加载的所有元素的内存地址 # 可以计算如下 X = X + m * stride_xm + n * stride_xn # 加载输入数据；用 0 x = tl.load(X, mask=n < N, other=-float('inf')) 填充越界元素 # 计算数值稳定的 softmax z = x - tl.max(x, axis =0) num = tl.exp(z) denom = tl.sum(num,axis=0) y = num / denom # 写回 YY = Y + m * stride_ym + n * stride_yn tl.store(Y, y , mask=n < N)import torch# 分配输入/输出 tensorsX = torch.normal(0, 1, size=(583, 931), device='cuda')Y = torch.empty_like(X)# SPMD 启动 gridgrid = (X.shape[0], )# 入队 GPU kernelsoftmax[grid](Y, Y.stride(0), Y.stride(1), X, X.stride(0), X.stride(1), X.shape[0] , X.shape[1]) 请注意，Triton JIT 将 X 和 Y 视为指针而不是张量；我们觉得保留对内存访问的低级控制对于处理更复杂的数据结构（例如，块稀疏张量）很重要。

重要的是，softmax 的这种特殊实现在整个规范化过程中将 $X$ 的行保留在 SRAM 中，这在适用时最大化了数据重用（~<32K 列）。这与 PyTorch 的内部 CUDA 代码不同，后者使用临时内存使其更通用但速度明显更慢（如下）。这里的底线不是 Triton 本质上更好，而是它简化了专用内核的开发，这些内核比通用库中的内核快得多。 Torch (v1.9) JIT 的较低性能凸显了从高级张量操作序列自动生成 CUDA 代码的难度。 @torch.jit.scriptdef softmax(x): x_max = x.max(dim=1)[0] z = x - x_max[:, None] numerator = torch.exp(x) denominator = numerator.sum(dim= 1) return numerator / denominator[:, None] 能够为元素运算和归约编写融合内核很重要，但考虑到神经网络中矩阵乘法任务的重要性，这还不够。事实证明，Triton 也非常适合这些应用，只需约 25 行 Python 代码即可实现最佳性能。另一方面，在 CUDA 中实现类似的东西会花费更多的精力，甚至可能会降低性能。 @triton.jitdef matmul(A, B, C, M, N, K, stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn, **META): # 提取元参数 BLOCK_M, GROUP_M = META['BLOCK_M'], META ['GROUP_M'] BLOCK_N = META['BLOCK_N'] BLOCK_K = META['BLOCK_K'] # 程序组合在一起以提高 L2 命中率 _pid_m = tl.program_id(0) _pid_n = tl.program_id(1) pid_m = _pid_m // GROUP_M pid_n =（_pid_n * GROUP_M）+（_pid_m％GROUP_M）＃RM（分别地，RN）表示用于行的范围的索引＃（相应的栏）C RM = pid_m * BLOCK_M + tl.arange的（0， BLOCK_M) rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) # rk 表示 A (resp. B) 的列 # (resp. rows) 的索引范围 (resp. B) rk = tl.arange(0, BLOCK_K) # the # A 和 B 的第一个块中元素的内存地址可以使用 numpy 风格的广播计算 A = A + (rm[:, None] * stride_am + rk[None, :] * stride_ak) B = B + (rk [:, None] * stride_bk + rn[None, :] * stride_bn) # 初始化并迭代更新累加器 acc = tl.zeros((BLOCK_M, BLOCK _N), dtype=tl.float32) for k in range(K, 0, -BLOCK_K): a = tl.load(A) b = tl.load(B) # 块级矩阵乘法 acc += tl.dot( a, b) # 递增指针，以便在下一次迭代期间加载 A 和 B 的下一个块 # A += BLOCK_K * stride_ak B += BLOCK_K * stride_bk # 如果需要，融合泄漏的 ReLU # acc = tl.where(acc > = 0, acc, alpha * acc) # 写回结果 C = C + (rm[:, None] * stride_cm + rn[None, :] * stride_cn) mask = (rm[:, None] < M) & ( rn[None, :] < N) tl.store(C, acc, mask=mask) 手写矩阵乘法内核的一个重要优点是它们可以根据需要进行定制以适应其输入的融合变换（例如，切片）和输出（例如，Leaky ReLU）。如果没有像 Triton 这样的系统，对于没有特殊 GPU 编程专业知识的开发人员来说，矩阵乘法内核的非平凡修改将是遥不可及的。 Triton 的良好性能来自于以 Triton-IR 为中心的模块化系统架构，这是一种基于 LLVM 的中间表示，其中多维值块是一等公民。

@triton.jit 装饰器通过遍历提供的 Python 函数的抽象语法树 (AST) 来工作，以便使用常见的 SSA 构造算法即时生成 Triton-IR。生成的 IR 代码随后由我们的编译器后端进行简化、优化和自动并行化，然后再转换为高质量的 LLVM-IR，最终转换为 PTX，以便在最新的 NVIDIA GPU 上执行。目前不支持 CPU 和 AMD GPU，但我们欢迎旨在解决此限制的社区贡献。我们发现通过 Triton-IR 使用阻塞的程序表示允许我们的编译器自动执行各种重要的程序优化。例如，数据可以通过查看计算密集型块级操作（例如，tl.dot）的操作数自动存储到共享内存中，并使用标准活性分析技术进行分配/同步。另一方面，Triton 程序可以（1）通过同时执行不同的内核实例跨 SM 进行高效和自动并行化，以及（2）通过分析每个块级操作的迭代空间并在不同的 SIMD 中对其进行充分分区来在 SM 内进行并行化单位，如下图。我们打算让 Triton 成为一个社区驱动的项目。随意在 GitHub 上 fork 我们的存储库！如果您有兴趣加入我们的团队并研究 Triton 和 GPU 内核，我们正在招聘！ Yan, D.、Wang, W. 和 Chu, X.（2020 年 5 月）。揭秘张量核心以优化半精度矩阵乘法。在 2020 年 IEEE 国际并行和分布式处理研讨会 (IPDPS)。 IEEE。 Tillet, P.、Kung, HT 和 Cox, D.（2019 年 6 月）。 Triton：一种用于平铺神经网络计算的中间语言和编译器。在第三届 ACM SIGPLAN 机器学习和编程语言国际研讨会论文集（第 10-19 页）中。

Braun, M.、Buchwald, S.、Hack, S.、Leißa, R.、Mallon, C. 和 Zwinkau, A.（2013 年 3 月）。简单高效的静态单赋值形式构造。在编译器构建国际会议上（第 102-122 页）。斯普林格，柏林，海德堡。