Triton：openai开源GPU编程神器

• • 序言
  • GPU编程的困难
  • 编程模型
  • 矩阵乘法
  • 高级系统架构
  • 编译器后端

序言

我们将发布Triton 1.0，这是一种开源的类似Python的编程语言，使没有CUDA经验的研究人员能够编写高效的GPU代码–大多数情况下与专家所能产生的代码相当。Triton使其有可能以相对较少的努力达到硬件性能的峰值；例如，它可以用来编写FP16矩阵乘法内核，其性能与cuBLAS相当–这是许多GPU程序员在25行代码以下无法做到的。我们的研究人员已经用它编写了比同等的Torch实现效率高2倍的内核，我们很高兴能与社区合作，使GPU编程对每个人来说都更容易。
深度学习领域的新的研究思路一般都是使用本地框架运算符的组合来实现的。虽然方便，但这种方法往往需要创建（移动）许多临时张量，这可能会降低神经网络的规模性能。这些问题可以通过编写专门的GPU内核来缓解，但由于GPU编程的许多复杂性，这样做可能会出乎意料地困难。而且，尽管最近出现了各种系统来简化这一过程，但我们发现它们要么过于啰嗦，要么缺乏灵活性，要么生成的代码明显比我们手工调整的基线慢。这促使我们对Triton进行了扩展和改进，这是一种最新的语言和编译器，其最初的开发者现在在OpenAI工作。
Triton代码
Triton文档

GPU编程的困难

现代GPU的架构可以大致分为三个主要部分–DRAM、SRAM和ALU–在优化CUDA代码时必须考虑到每一个部分。

• 来自DRAM的内存传输必须凝聚成大型事务，以利用现代内存接口的大总线宽度。
• 在重新使用之前，数据必须被手动存储到SRAM中，并进行管理，以便在检索时尽量减少共享内存库的冲突。
• 计算必须在流式多处理器(SM)之间和内部仔细划分和安排，以促进指令/线程级并行，并利用特殊用途的ALU（如张量核心）

GPU的基本架构如下：

对所有这些因素进行推理是具有挑战性的，即使是对具有多年经验的经验丰富的CUDA程序员也是如此。Triton的目的是将这些优化完全自动化，以便开发人员能够更好地专注于他们的并行代码的高层逻辑。Triton的目标是广泛适用，因此不会自动安排跨SM的工作–将一些重要的算法考虑（如tiling，SM间的同步）留给开发者决定。
CUDA和Triton中编译器优化的比较如下：

编程模型

在所有可用的特定领域语言和JIT编译器中，Triton也许与Numba最相似：内核被定义为装饰过的Python函数，并在所谓实例的网格上以不同的program_id并发启动。然而，正如下面的代码片段所示，相似之处仅此而已。Triton通过在块上的操作暴露了实例内的并行性，而不是单指令多线程（SIMT）7的执行模型，块的尺寸是2的幂。这样做，Triton有效地抽象了所有与CUDA线程块内并发有关的问题（例如，内存凝聚、共享内存同步/冲突、张量核心调度）。

下面是Numba中的向量加法：

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])

下面是Triton中的向量加法：

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])

虽然这可能对令人尴尬的并行（即逐个元素）计算没有特别帮助，但它可以大大简化更复杂的GPU程序的开发。
例如，考虑融合softmax内核的情况，其中每个实例都对给定的输入张量X的不同行进行归一化处理，在 X ∈ R M × N X \in \mathbb{R}^{M\times N} X∈RM×N中。这种并行化策略的标准CUDA实现在编写时可能具有挑战性，需要线程之间明确的同步，因为它们同时减少X的同一行。这种复杂性在Triton中大部分都消失了，每个内核实例加载感兴趣的行，并使用类似NumPy的基元对其进行顺序标准化。
在triton中编写fused softmax代码如下：

import triton
import triton.language as tl

@triton.jit
def softmax(Y, stride_ym, stride_yn, X, stride_xm, stride_xn, M, N):
    # row index
    m = tl.program_id(0)
    # col indices
    # this specific kernel only works for matrices that 
    # have less than BLOCK_SIZE columns
    BLOCK_SIZE = 1024
    n = tl.arange(0, BLOCK_SIZE)
    # the memory address of all the elements
    # that we want to load can be computed as follows
    X = X + m * stride_xm + n * stride_xn
    # load input data; pad out-of-bounds elements with 0 
    x = tl.load(X, mask=n < N, other=-float('inf'))
    # compute numerically-stable softmax
    z = x - tl.max(x, axis=0)
    num = tl.exp(z)
    denom = tl.sum(num, axis=0)
    y = num / denom
    # write back to Y
    Y = Y + m * stride_ym + n * stride_yn
    tl.store(Y, y, mask=n < N)

import torch
# Allocate input/output tensors
X = torch.normal(0, 1, size=(583, 931), device='cuda')
Y = torch.empty_like(X)
# SPMD launch grid
grid = (X.shape[0], )
# enqueue GPU kernel
softmax[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本质上更好，而是说它简化了专门内核的开发，可以比通用库中的内核快很多。
M=4096时，A100处理fused Softmax的性能如下：

Torch（v1.9）JIT的较低性能突出了从高级张量操作序列中自动生成CUDA代码的难度。

@torch.jit.script
def 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中矩阵乘法代码：

@triton.jit
def matmul(A, B, C, M, N, K, stride_am, stride_ak, 
            stride_bk, stride_bn, stride_cm, stride_cn,
            **META):
    # extract metaparameters
    BLOCK_M, GROUP_M = META['BLOCK_M'], META['GROUP_M']
    BLOCK_N = META['BLOCK_N']
    BLOCK_K = META['BLOCK_K']
    # programs are grouped together to improve L2 hit rate
    _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 (resp. rn) denotes a range of indices
    # for rows (resp. col) of C
    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    # rk denotes a range of indices for columns 
    # (resp. rows) of A (resp. B)
    rk = tl.arange(0, BLOCK_K)
    # the memory addresses of elements in the first block of
    # A and B can be computed using numpy-style broadcasting
    A = A + (rm[:, None] * stride_am + rk[None, :] * stride_ak)
    B = B + (rk [:, None] * stride_bk  + rn[None, :] * stride_bn)
    # initialize and iteratively update accumulator
    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)
        # block level matrix multiplication
        acc += tl.dot(a, b)
        # increment pointers so that the next blocks of A and B
        # are loaded during the next iteration
        A += BLOCK_K * stride_ak
        B += BLOCK_K * stride_bk
    # fuse leaky ReLU if desired
    # acc = tl.where(acc >= 0, acc, alpha * acc)
    # write back result
    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为中心的模块化系统架构，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)在SM内部通过分析每个块级操作的迭代空间并在不同的SIMD单元之间充分划分，如下所示。

参考链接: https://openai.com/blog/triton/.