介绍一下Triton。开源的神经网络GPU编程

我们将发布Triton 1.0，这是一种开源的类似Python的编程语言，使没有CUDA经验的研究人员能够编写高效的GPU代码--大多数情况下与专家所能产生的代码相当。Triton使其有可能以相对较少的努力达到硬件性能的峰值；例如，它可以用来编写FP16矩阵乘法内核，其性能与cuBLAS相当--这是许多GPU程序员无法做到的，其代码不到25行。我们的研究人员已经用它来编写内核，其效率比同等的Torch实现高出2倍，我们很高兴与社区合作，使GPU编程对每个人来说都更容易。

查看CodeRead文档

---

深度学习领域的新研究理念一般都是使用本地框架运算符的组合来实现的。虽然方便，但这种方法往往需要创建（和/或移动）许多临时张量，这可能会损害神经网络的规模性能。这些问题可以通过编写专门的GPU内核来缓解，但由于GPU编程的许多复杂性，这样做可能出乎意料地困难。 而且，尽管最近出现了各种系统 ，以使这一过程更容易，但我们发现它们要么过于冗长，缺乏灵活性，要么生成的代码明显比我们手工调整的基线慢。这促使我们扩展和改进了Triton ，这是一种最新的语言和编译器，其原始创建者现在在OpenAI工作。

GPU编程的挑战

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

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

一个GPU的基本架构。

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

	CUDA	Triton
内存凝聚	手动	自动
共享内存管理	手动	自动
调度（在SM内）	手动	自动
调度(跨SM)	手动	手动

CUDA与Triton的编译器优化。

编程模型

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

#triton-vector-add .code-toolbar { height: 100%; } #triton-vector-add pre { margin-top: 0; margin-bottom:0; height: 100%; }

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

Triton中的矢量添加。

虽然这可能对令人尴尬的并行（即，元素-wise）计算没有特别的帮助，但它可以大大简化更复杂的GPU程序的开发。

例如，考虑融合softmax内核的情况（如下），其中每个实例都对给定的输入张量$X的不同行进行归一化处理，在\\mathbb{R}^{M\\times N}$。这种并行化策略的标准CUDA实现在编写上是很有挑战性的，当线程同时减少$X$的同一行时，需要在线程之间进行明确的同步。在Triton中，这种复杂性大部分都消失了，每个内核实例都会加载感兴趣的行，并使用类似NumPy的基元对其进行顺序标准化。

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中融合的softmax。

请注意，Triton JIT将X和Y视为指针，而不是张量；我们认为保留对内存访问的低级控制对于解决更复杂的数据结构（如块状稀疏张量）很重要。

重要的是，softmax的这一特定实现在整个规范化过程中保持$X$的行在SRAM中，这在适用时最大限度地提高了数据的重复利用（~<32K列）。这与PyTorch的内部CUDA代码不同，后者对临时内存的使用使其更具通用性，但速度明显较慢（如下）。这里的底线不是说Triton本质上更好，而是说它简化了专门内核的开发，可以比通用库中的内核快很多。

M=4096的融合softmax的A100性能。

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]

使用Torch JIT的融合softmax。

矩阵乘法

能够为向元素运算和还原编写融合内核是很重要的，但鉴于神经网络中矩阵乘法任务的突出性，这还不够。事实证明，Triton对这些任务也非常有效，只用了大约25行Python代码就达到了峰值性能。另一方面，在CUDA中实现类似的东西将需要更多的努力，甚至可能会取得更低的性能。

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

Triton中的矩阵乘法。

手写矩阵乘法内核的一个重要优势是，它们可以根据需要定制，以适应其输入（如切片）和输出（如Leaky ReLU）的融合转化。如果没有像Triton这样的系统，对矩阵乘法内核的非实质性修改对于没有特殊GPU编程专长的开发者来说是遥不可及的。

在BLOCK$\_M$、BLOCK$\_N$、BLOCK$\_K$、GROUP$\_M$的适当调整下，V100张量核心的矩阵乘法性能。

高级别的系统结构

Triton的良好性能来自于以Triton-IR为中心的模块化系统架构，Triton-IR是一种基于LLVM的中间表示，其中多维值块是一等公民。

Python

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;
}

蟒蛇

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;
}

Triton的高层架构。

@triton.jit 装饰器的工作原理是：行走所提供的Python函数的抽象语法树（AST），以便使用常见的SSA构建算法来即时生成Triton-IR。 然后，生成的IR代码被我们的编译器后端简化、优化并自动并行化，然后被转换为高质量的LLVM-IR，最终被转换为PTX，以便在最近的NVIDIA GPU上执行。目前不支持CPU和AMD的GPU，但我们欢迎社区为解决这一限制而做出贡献。

编译器后端

我们发现，通过Triton-IR使用阻塞的程序表示法，我们的编译器可以自动执行各种重要的程序优化。例如，通过查看计算密集型块级操作（如tl.dot ）的操作数，可以将数据自动存储到共享内存中，并使用标准的有效性分析技术进行分配/同步。

Triton编译器通过分析计算密集型操作中使用的块变量的实时范围来分配共享内存。

另一方面，Triton程序可以有效地自动并行化：(1)通过并发执行不同的内核实例在SM上并行，(2)在SM内通过分析每个块级操作的迭代空间并在不同的SIMD单元上充分划分，如下所示。

逐个元素

S1 float A[4,4] = ...
S2 float B[4,4] = ...
S3 float C[4,4] = A + B

FP16矩阵乘法

S1 half A[4,2] = ...
S2 half B[2,2] = ...
S3 float C[4,2] = dot(A,B)

1. 定义一个由三个语句组成的Triton程序P S1,S2 。S3

矢量化

张扬的

2. 迭代空间的S3

SM

3. 将S3 映射到流式多处理器（SM）上

GPU

4. 将P映射到GPU上

逐元素

S1 float A[4,4] = ...
S2 float B[4,4] = ...
S3 float C[4,4] = A + B

FP16矩阵多乘法

S1 half A[4,2] = ...
S2 half B[2,2] = ...
S3 float C[4,2] = dot(A,B)

矢量化

张扬的

SM

GPU

1. 由三个语句组成的Triton程序P的定义S1,S2 。S3

2. 迭代空间的S3

3. 将S3 映射到流式多处理器（SM）上

4. 将P映射到GPU上

Triton中的自动并行化。每个块级操作都定义了一个封锁的迭代空间，该空间被自动并行化以利用流式多处理器（SM）上的可用资源。

贡献

我们打算让Triton成为一个社区驱动的项目。欢迎在GitHub上分叉我们的存储库。

如果你有兴趣加入我们的团队，从事Triton和GPU内核的工作，我们正在招聘!

/* prism.scss重写 */* .token.property, .token.tag, .token.boolean, .token.number, .token.constant, .token.symbol, .token.delete { color:#FF5828; /*暗红色->橙色 */ } .token.altrule, .token.attr-value, .token.keyword { color:#21B5C2; /*中蓝->茶色 */ } import {Runtime, Inspector, Library} from "unpkg.com/@observable…"; import notebook_a100 from "api.observablehq.com/d/adcb73899…"; import notebook_v100 from "api.observablehq.com/d/1f5fc5b6e…"; const customWidth = function (selector) { return (new Library).Generators.observe(function(change) { var width = change(document.querySelector(selector).clientWidth); function resized() { var w = document.querySelector(selector) .clientWidth; if (w !== width) change(width = w); } window.addEventListener("resize", resized); return function() { window.removeEventListener("resize", resized); }; }; const a100_renders = { "图表":"#a100", }; new Runtime(Object.assign(new Library, {width: customWidth("#a100")}).module(notebook_a100, name => { const selector = a100_renders[name]; if (selector) { // key exists return new Inspector(document.querySelector(selector)); } else { return true; }); const v100_renders = { "chart":"#v100", }; new Runtime(Object.assign(new Library, {width: customWidth("#v100")}).module(notebook_v100, name => { const selector = v100_renders[name]; if (selector) { // key exists return new Inspector(document.querySelector(selector)); } else { return true; }) 。

---

鸣谢

大岩（香港科技大学），DeepSpeed（微软），Anthropic

---

参考文献

1. Gray, S. (2017).SGEMM演练。URLgithub.com/NervanaSyst…

2. Kerr, A. (2020).开发CUDA内核，将NVIDIA A100上的Tensor Cores推至绝对极限。URLdeveloper.nvidia.com/gtc/2020/vi….

3. Yan, D., Wang, W., & Chu, X. (2020, May).揭开张量核心的神秘面纱，优化半精度矩阵乘法。In2020 IEEE International Parallel and Distributed Processing Symposium (IPDPS).IEEE.

4. NVIDIA CUTLASS

5. Apache TVM

6. Tillet, P., Kung, H. T., & Cox, D. (2019, June).Triton：用于分层神经网络计算的中间语言和编译器。在第三届ACM SIGPLAN机器学习和编程语言国际研讨会论文集（第10-19页）。

7. Lin, Y. & Grover, V. (2018).使用CUDA Warp-Level Primitives.URLdeveloper.nvidia.com/blog/using-…](developer.nvidia.com/blog/using-…)

8. Braun, M., Buchwald, S., Hack, S., Leißa, R., Mallon, C., & Zwinkau, A. (2013, March).简单而高效地构建静态单一赋值形式。InInternational Conference on Compiler Construction(pp. 102-122).Springer, Berlin, Heidelberg.