项目地址：https://github.com/triton-lang/triton
windows版本地址：https://github.com/woct0rdho/triton-windows
官网手册：https://triton.hyper.ai/docs/getting-started/installation/
项目论文：https://www.eecs.harvard.edu/~htk/publication/2019-mapl-tillet-kung-cox.pdf
Triton 是一种用于编写高效自定义深度学习基元的语言和编译器。Triton 的目标是提供一个开源环境，以比 CUDA 更高的生产力编写快速代码，但也比其他现有的 DSL 具有更高的灵活性。其本质就是，提供了一个统一的python接口，让算法人员可以快速的编写各类cuda操作算子。同时提供了不少算子的计算优化策略（矩阵运算的分组实现），从而实现计算效率或者计算成本的降低。
总体来说，triton是提供了一个简易的cuda操作接口，基于算法人员对于各类算子的优化理解，可以快速使用triton进行实现，降低算子的计算量、或者访存成本，从而实现模型训练或者推理的提效。 在当前的开源生态下，可以基于各种使用triton实现的算子替换torch原生算子，从而实现提效。随着torch、tensorrt等推理框架的优化，各种底层算子的计算优化操作也会被加入。基于triton可以让我们先行一步。
1、快速安装
linux安装
pip install triton
windows安装
pip install -U triton-windows
如果要支持Blackwell ，请参考项目中的源码安装指南。
2、基本案例
如果在案例代码中执行出现了 Error #15: Initializing libiomp5md.dll, but found libiomp5md.dll already initialized 的报错
请在每一行代码最前面添加
import os
os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
关于下面的每一个案例，官网都提供了python代码、jupyter notebook代码的案例
2.1 向量相加
https://triton.hyper.ai/docs/getting-started/tutorials/vector-addition
基于官方公布的效果，可以发现基于triton实现的相加操作，性能与torch原始的tensor操作是差不多的
"""
Vector Addition
===============
In this tutorial, you will write a simple vector addition using Triton.
In doing so, you will learn about:
* The basic programming model of Triton.
* The `triton.jit` decorator, which is used to define Triton kernels.
* The best practices for validating and benchmarking your custom ops against native reference implementations.
"""
# %%
# Compute Kernel
# --------------
import torch
import triton
import triton.language as tl
DEVICE = triton.runtime.driver.active.get_active_torch_device()
@triton.jit
def add_kernel(x_ptr, # *Pointer* to first input vector.
y_ptr, # *Pointer* to second input vector.
output_ptr, # *Pointer* to output vector.
n_elements, # Size of the vector.
BLOCK_SIZE: tl.constexpr, # Number of elements each program should process.
# NOTE: `constexpr` so it can be used as a shape value.
):
# There are multiple 'programs' processing different data. We identify which program
# we are here:
pid = tl.program_id(axis=0) # We use a 1D launch grid so axis is 0.
# This program will process inputs that are offset from the initial data.
# For instance, if you had a vector of length 256 and block_size of 64, the programs
# would each access the elements [0:64, 64:128, 128:192, 192:256].
# Note that offsets is a list of pointers:
block_start = pid * BLOCK_SIZE
offsets = block_start + tl.arange(0, BLOCK_SIZE)
# Create a mask to guard memory operations against out-of-bounds accesses.
mask = offsets < n_elements
# Load x and y from DRAM, masking out any extra elements in case the input is not a
# multiple of the block size.
x = tl.load(x_ptr + offsets, mask=mask)
y = tl.load(y_ptr + offsets, mask=mask)
output = x + y
# Write x + y back to DRAM.
tl.store(output_ptr + offsets, output, mask=mask)
# %%
# Let's also declare a helper function to (1) allocate the `z` tensor
# and (2) enqueue the above kernel with appropriate grid/block sizes:
def add(x: torch.Tensor, y: torch.Tensor):
# We need to preallocate the output.
output = torch.empty_like(x)
assert x.device == DEVICE and y.device == DEVICE and output.device == DEVICE
n_elements = output.numel()
# The SPMD launch grid denotes the number of kernel instances that run in parallel.
# It is analogous to CUDA launch grids. It can be either Tuple[int], or Callable(metaparameters) -> Tuple[int].
# In this case, we use a 1D grid where the size is the number of blocks:
grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), )
# NOTE:
# - Each torch.tensor object is implicitly converted into a pointer to its first element.
# - `triton.jit`'ed functions can be indexed with a launch grid to obtain a callable GPU kernel.
# - Don't forget to pass meta-parameters as keywords arguments.
add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=1024)
# We return a handle to z but, since `torch.cuda.synchronize()` hasn't been called, the kernel is still
# running asynchronously at this point.
return output
# %%
# We can now use the above function to compute the element-wise sum of two `torch.tensor` objects and test its correctness:
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))}')
# %%
# Seems like we're good to go!
# %%
# Benchmark
# ---------
#
# We can now benchmark our custom op on vectors of increasing sizes to get a sense of how it does relative to PyTorch.
# To make things easier, Triton has a set of built-in utilities that allow us to concisely plot the performance of our custom ops.
# for different problem sizes.
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=['size'], # Argument names to use as an x-axis for the plot.
x_vals=[2**i for i in range(12, 28, 1)], # Different possible values for `x_name`.
x_log=True, # x axis is logarithmic.
line_arg='provider', # Argument name whose value corresponds to a different line in the plot.
line_vals=['triton', 'torch'], # Possible values for `line_arg`.
line_names=['Triton', 'Torch'], # Label name for the lines.
styles=[('blue', '-'), ('green', '-')], # Line styles.
ylabel='GB/s', # Label name for the y-axis.
plot_name='vector-add-performance', # Name for the plot. Used also as a file name for saving the plot.
args={}, # Values for function arguments not in `x_names` and `y_name`.
))
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)
# %%
# We can now run the decorated function above. Pass `print_data=True` to see the performance number, `show_plots=True` to plot them, and/or
# `save_path='/path/to/results/' to save them to disk along with raw CSV data:
benchmark.run(print_data=True, show_plots=True)
2.2 softmax操作融合
该操作在某些类别的矩阵上比 PyTorch 的原生操作快得多：即那些可以适应 GPU 静态随机存取存储器 (SRAM) 的行。这个操作涉及：
1、内核融合对于带宽受限操作的优势。
2、Triton 中缩减操作。
triton实现的算子性能如下所示，可以看到在特定维度的计算效率上是比torch要高。
原始的softmax技术步骤
import torch
import triton
import triton.language as tl
from triton.runtime import driver
def naive_softmax(x):
"""Compute row-wise softmax of X using native pytorch
使用原生 PyTorch 计算 X 的逐行 softmax
We subtract the maximum element in order to avoid overflows. Softmax is invariant to
this shift.
我们减去最大元素以避免溢出。Softmax 对于这种偏移是不变的。
"""
# read MN elements ; write M elements
# 读取 MN 个元素；写入 M 个元素
x_max = x.max(dim=1)[0]
# read MN + M elements ; write MN elements
# 读取 MN + M 个元素；写入 MN 个元素
z = x - x_max[:, None]
# read MN elements ; write MN elements
# 读取 MN 个元素；写入 MN 个元素
numerator = torch.exp(z)
# read MN elements ; write M elements
# 读取 MN 个元素；写入 M 个元素
denominator = numerator.sum(dim=1)
# read MN + M elements ; write MN elements
# 读取 MN + M 个元素；写入 MN 个元素
ret = numerator / denominator[:, None]
# in total: read 5MN + 2M elements ; wrote 3MN + 2M elements
# 总计：读取 5MN + 2M 个元素；写入 3MN + 2M 个元素
return ret
triton优化
softmax 内核工作原理如下：每个程序加载输入矩阵 X 的一组行，按程序数量跨步处理，对其进行归一化，并将结果写回输出 Y。
注意，Triton 的一个重要限制是每个块必须具有 2 的幂次数的元素，因此，如果我们要处理任意可能的输入形状，我们需要在内部「填充」每一行，并适当保护内存操作。
@triton.jit
def softmax_kernel(output_ptr, input_ptr, input_row_stride, output_row_stride, n_rows, n_cols, BLOCK_SIZE: tl.constexpr,
num_stages: tl.constexpr):
# starting row of the program
# 程序起始行
row_start = tl.program_id(0)
row_step = tl.num_programs(0)
for row_idx in tl.range(row_start, n_rows, row_step, num_stages=num_stages):
# The stride represents how much we need to increase the pointer to advance 1 row
# 步长表示我们需要对指针增加多少以推进 1 行
row_start_ptr = input_ptr + row_idx * input_row_stride
# The block size is the next power of two greater than n_cols, so we can fit each
# 块大小是大于 n_cols 的下一个二的幂，因此我们可以适配
# row in a single block
# 单个块中的行
col_offsets = tl.arange(0, BLOCK_SIZE)
input_ptrs = row_start_ptr + col_offsets
# Load the row into SRAM, using a mask since BLOCK_SIZE may be > than n_cols
# 将行加载到 SRAM 中，使用掩码，因为 BLOCK_SIZE 可能大于 n_cols
mask = col_offsets < n_cols
row = tl.load(input_ptrs, mask=mask, other=-float('inf'))
# Subtract maximum for numerical stability
# 为了数值稳定性而减去最大值
row_minus_max = row - tl.max(row, axis=0)
# Note that exponentiation in Triton is fast but approximate (i.e., think __expf in CUDA)
# 请注意，Triton 中的指数运算速度很快，但是是近似的（例如，类似于 CUDA 中的 __expf）。
numerator = tl.exp(row_minus_max)
denominator = tl.sum(numerator, axis=0)
softmax_output = numerator / denominator
# Write back output to DRAM
# 将输出写回 DRAM
output_row_start_ptr = output_ptr + row_idx * output_row_stride
output_ptrs = output_row_start_ptr + col_offsets
tl.store(output_ptrs, softmax_output, mask=mask)
这里的代码，并不完整，完整的代码参考：https://triton.hyper.ai/docs/getting-started/tutorials/fused-softmax
2.3 矩阵乘法
基于triton可以实现一个非常简短的高性能 FP16 矩阵乘法内核，其性能可以与 cuBLAS 或 rocBLAS 相媲美。
这里的关键在于矩阵分块计算(block倍数制约)、L2 缓存优化(降低内存访问成本)
与cuBLAS库的技术对比如下，可以发现性能基本不分上下
动机
矩阵乘法是现代大多数高性能计算系统的关键构建块。
矩阵乘法难以优化，因此其实现通常由硬件供应商自行完成，作为所谓「内核库」（例如 cuBLAS ）的一部分。
这些库通常是专有的，不能轻易定制以满足现代深度学习工作负载的需求（例如融合激活函数）。
triton实现的矩阵分块乘法案例如下所示
# Do in parallel
# 并行进行
for m in range(0, M, BLOCK_SIZE_M):
# Do in parallel
# 并行进行
for n in range(0, N, BLOCK_SIZE_N):
acc = zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=float32)
for k in range(0, K, BLOCK_SIZE_K):
a = A[m : m+BLOCK_SIZE_M, k : k+BLOCK_SIZE_K]
b = B[k : k+BLOCK_SIZE_K, n : n+BLOCK_SIZE_N]
acc += dot(a, b)
C[m : m+BLOCK_SIZE_M, n : n+BLOCK_SIZE_N] = acc
实际上，上述算法在 Triton 中实现起来相当简单。
主要困难在于计算内循环中必须读取 A 和 B 块的内存位置。为此，我们需要多维指针算术。
指针算术
因此，对于行主序的二维张量 X，X[i, j] 的内存位置由 &X[i, j] = X + istride_xi + jstride_xj 给出。
因此，A[m : m+BLOCK_SIZE_M, k:k+BLOCK_SIZE_K]和B[k : k+BLOCK_SIZE_K, n:n+BLOCK_SIZE_N] 的指针块可以用伪代码定义为：
&A[m : m+BLOCK_SIZE_M, k:k+BLOCK_SIZE_K] = a_ptr + (m : m+BLOCK_SIZE_M)[:, None]*A.stride(0) + (k : k+BLOCK_SIZE_K)[None, :]*A.stride(1);
&B[k : k+BLOCK_SIZE_K, n:n+BLOCK_SIZE_N] = b_ptr + (k : k+BLOCK_SIZE_K)[:, None]*B.stride(0) + (n : n+BLOCK_SIZE_N)[None, :]*B.stride(1);
当 M 不是 BLOCK_SIZE_M 的倍数或 N 不是 BLOCK_SIZE_N 的倍数时，我们需要额外的取模运算来应对，这种情况下我们可以用一些无用的值填充数据，这些值不会对结果有影响。对于 K 维度，我们将在后面使用掩码加载语义来处理。
L2 缓存优化
正如上面提到的，每个程序实例计算 C 的一个 [BLOCK_SIZE_M, BLOCK_SIZE_N] 块。
重点要记住这些块的计算顺序，因为它会影响我们程序的 L2 缓存命中率，而且，简单的行主序排序是行不通。需要进行计算顺序的调整
以下的矩阵乘法示例中，每个矩阵都是 9*9 个块。可以看到，如果按行主序计算输出，我们需要加载 90 个块到 SRAM 中来计算前 9 个输出块，但如果按组顺序计算，我们只需要加载 54 个块。同样是9次访问，基于grouped顺序，数据加载量大幅度降低
2.4 其他算子优化
一共有以下算子优化案例，这里只罗列了部分易懂的案例。
Dropout **
Dropout 是在 [SRIVASTAVA2014] 中引入的一种技术，用于改善低数据条件下深度神经网络的性能，通常用于正则化。它接受一个向量作为输入，并生成相同 shape 的输出向量。输出中的每个标量都有概率
p 被设为零**，否则直接从输入复制，最后进行（1-p）的缩放。
基于triton实现的Dropout具有以下优势:
1、更小的内存占用。
2、较少的数据移动。
3、简化了在多次调用内核函数时持久化随机性的管理。
layer normalization
层标准化 (LayerNorm) 算子最先在 BA2016 中提出，旨在提高序列模型（例如 Transformers）或小 batchsize 神经网络的性能。它以向量 x作为输入，并生成与输入 shape 相同的向量 y作为输出。 标准化是通过减去均值并除以 x 的标准差来实现的。 标准化后，会应用带有权重 w 和偏置 b 的可学习线性变换。
基于triton可以实现LayerNorm的前向传播与反向传播。根据官网描述，主要是对反向传播操作进行访存优化。
1、首先分析出 ，在同一批次中的所有行使用相同的权重 w 和偏差 b，它们的梯度需要累加
2、每个内核实例将某些行的部分 ∇ w 和 ∇ b累积到 GROUP_SIZE_M 个独立缓冲区之一中。这些缓冲区保存在 L2 缓存中
基于以上示例，对于M个数据，分为GROUP_SIZE_M 个组进行计算，使访存操作降低到M/GROUP_SIZE_M+1
3、总结
Triton 的核心理念是基于分块的编程范式可以促进神经网络的高性能计算核心的构建。CUDA 编写属于传统的 “单程序，多数据” GPU 执行模型，在线程的细粒度上进行编程，Triton 是在分块的细粒度上进行编程。例如，在矩阵乘法的情况下，CUDA和Triton有以下不同
扫一扫
841
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
