逐个学习-Triton-One-内核-向量加法

逐个学习 Triton One 内核：向量加法

原文：towardsdatascience.com/learning-triton-one-kernel-at-a-time-vector-addition/

在亿参数模型的时代，一点优化就能走得很远。像GPT4这样的模型训练成本超过一亿美元，这使得1%的效率提升****价值超过一百万美元。一种强大的优化机器学习模型效率的方法是将它们的一些组件直接在 GPU 上编写。如果你像我一样，简单的提及 CUDA 内核就足以让人脊背发凉，因为它们是出了名的难以编写和调试。

幸运的是，OpenAI在 2021 年发布了Triton，这是一种新的语言和编译器，它抽象掉了 CUDA 的许多复杂性，并允许经验较少的从业者编写高性能的内核。一个显著的例子是Unsloth，一个 LLM 训练服务，承诺30 倍的培训速度和60%的内存使用减少，这一切都归功于用 Triton 内核替换了用 PyTorch 编写的层。

在本教程系列中，我们将学习 GPU 架构的基础知识以及如何实现高性能的 Triton 内核！本系列中展示的所有代码都将可在github.com/RPegoud/Triton-Kernels找到。

GPU 架构基础

在本节中，我们将介绍 (Nvidia) GPU 的基础知识，以便我们开始，并在本文末尾编写我们的第一个 Triton 内核。

从最小的软件单元开始，我们可以描述执行单元的层次结构如下：

线程：最小的工作单元，它们运行用户定义的内核代码。
Warps：最小的调度单元，它们总是由 32 个并行线程组成，每个线程都有自己的指令地址计数器和寄存器状态。warp 中的线程同时开始，但可以自由分支和独立执行。
线程块：warp 的组，其中所有线程都可以通过共享内存协作并同步屏障。要求线程块可以独立且按任意顺序执行，以并行或顺序方式。这种独立性允许线程块可以在任何数量的核心上按任意顺序调度，这样 GPU 程序就可以随着核心数量的增加而高效扩展。如果需要，我们可以在内核的特定点同步块内的线程，例如同步内存访问。
流式多处理器 (SM)：负责并行执行多个 warps的单元，它拥有共享内存和一个 L1 缓存（存储 SM 访问过的最新全局内存行）。一个 SM 拥有一个专门的warp 调度器，从准备运行的线程块中提取 warps。

在硬件方面，最小的作业单元是CUDA 核心，物理的算术逻辑单元（ALU），它为线程（或其部分）执行算术运算。

用类比来总结本节，我们可以将CUDA 核心看作个体工人，而warp是32 个工人组成的队伍，他们同时接收到相同的指令。他们可能以相同或不同的方式执行这项任务（分支），并且可能在不同的时间点完成它（独立性）。线程块由几个共享相同工作空间的队伍组成（即有共享内存），工作空间中所有队伍的工人可以同时等待午餐。流式多处理器是一个有许多队伍共同工作并共享工具和存储的工厂车间。最后，GPU是一个整个工厂，有许多楼层。

英伟达 GPU 架构的层次结构。虚线矩形表示内存块（由作者绘制）

优化基础

当优化深度学习模型时，我们正在处理三个主要组件：

计算: GPU 执行浮点运算（FLOPS）所花费的时间。
内存: 在 GPU 内部传输张量所花费的时间。
开销: 所有其他操作（Python 解释器、PyTorch 调度等）。

考虑这些组件有助于找到解决瓶颈的正确方法。例如，如果大部分时间都花在内存传输上，那么增加计算（例如使用更强大的 GPU）是没有帮助的。理想情况下，大部分时间应该花在计算上，更确切地说是在矩阵乘法上，这是 GPU 优化的精确操作。

这意味着最小化数据传输的成本，无论是从 CPU 到 GPU（“数据传输成本”），从一个节点到另一个节点（“网络成本”），还是从 CUDA 全局内存（DRAM，便宜但速度慢）到 CUDA 共享内存（SRAM，昂贵但设备上最快的内存）。后者被称为带宽成本，目前将是我们的主要关注点。降低带宽成本的常见策略包括：

重用在共享内存中加载的数据以进行多个步骤。一个典型的例子是分块矩阵乘法，我们将在未来的文章中介绍。
融合单个内核中的多个操作（因为每次内核启动都意味着将数据从 DRAM 移动到 SRAM），例如我们可以将矩阵乘法与激活函数融合。通常，算子融合可以提供巨大的性能提升，因为它防止了大量的全局内存读写，并且任何两个操作都提供了一个融合的机会。

在没有算子融合的情况下，矩阵乘法后跟 ReLU 激活。（由作者绘制）

在这个例子中，我们执行矩阵乘法x@W并将结果存储在中间变量a中。然后我们对a应用relu并将结果存储在变量y中。这需要 GPU 从全局内存中的x和W读取，将结果写入a，再次从a读取，最后写入y。相反，操作融合允许我们在单个内核中执行矩阵乘法和应用 ReLU，从而将全局内存的读写次数减半。

融合矩阵乘法和 ReLU 激活。（作者制作）

Triton

现在，我们将编写我们的第一个 Triton 内核，一个简单的向量加法。首先，让我们了解一下这个操作是如何在 GPU 上分解和执行的。

考虑想要对两个向量X和Y的条目求和，每个向量有 7 个元素（n_elements=7）。

我们将指示 GPU 以每次 3 个元素的分块处理这个问题（BLOCK_SIZE=3）。因此，为了覆盖输入向量的所有 7 个元素，GPU 将启动 3 个并行的“程序”，即我们的内核的独立实例，每个实例都有一个唯一的程序 ID，pid：

程序 0 分配了元素0, 1, 2。
程序 1 分配了元素3, 4, 5。
程序 2 分配了元素6。

然后，这些程序将结果写回到全局内存中存储的向量Z。

一个重要的细节是，内核不会接收整个向量X，而是接收指向第一个元素内存地址的指针，X[0]。为了访问X的实际值，我们需要手动从全局内存中加载它们。

我们可以通过程序 ID 访问每个块的数据：block_start = pid * BLOCK_SIZE。从那里，我们可以通过计算offsets = block_start + range(0, BLOCK_SIZE)来获取该块的剩余元素地址，并将它们加载到内存中。

然而，请记住，程序 2 只分配了元素 6，但其偏移量是[6, 7, 8]。为了避免任何索引错误，Triton 允许我们定义一个掩码来识别有效的目标元素，这里mask = offsets < n_elements。

我们现在可以安全地加载X和Y并将它们相加，然后将结果以类似的方式写回到全局内存中的输出变量Z。

每块向量索引。X、Y 和 Z 的切片被发送到独立的线程块，每个块由一个唯一的 ID 索引。（图像由作者提供）

让我们更仔细地看看代码，这里是 Triton 内核：

import triton
import triton.language as tl

@triton.jit
def add_kernel(
	x_ptr, # pointer to the first memory entry of x
	y_ptr, # pointer to the first memory entry of y
	output_ptr, # pointer to the first memory entry of the output
	n_elements, # dimension of x and y
	BLOCK_SIZE: tl.constexpr, # size of a single block
):
	# --- Compute offsets and mask ---
	pid = tl.program_id(axis=0) # block index
	block_start = pid * BLOCK_SIZE # start index for current block
	offsets = block_start + tl.arange(0, BLOCK_SIZE) # index range
	mask = offsets < n_elements # mask out-of-bound elements

	# --- Load variables from global memory ---
	x = tl.load(x_ptr + offsets, mask=mask)
	y = tl.load(y_ptr + offsets, mask=mask)

	# --- Operation ---
	output = x + y	

	# --- Save results to global memory ---
	tl.store(pointer=output_ptr + offsets, value=output, mask=mask)

让我们分解一些 Triton 特定的语法：

首先，一个 Triton 内核总是由<a href="http://twitter.com/triton" target="_blank" rel="noreferrer noopener">@triton</a>.jit装饰。
第二，一些参数需要声明为静态的，这意味着它们在计算时是已知的。这是对BLOCK_SIZE的要求，通过添加tl.constexpr类型注解来实现。此外，请注意，我们没有注解其他变量，因为它们不是合适的 Python 变量。
我们使用 tl.program_id 来访问当前块的 ID，tl.arange 的行为类似于 Numpy 的 np.arange。
通过调用 tl.load 和 tl.store 并传入指针数组来实现变量的加载和存储。请注意，这里没有 return 语句，这个角色被委托给了 tl.store。

要使用我们的内核，我们现在需要编写一个 PyTorch 级别的包装器，它提供内存指针并定义一个 内核网格。通常，内核网格是一个包含沿着每个轴分配给内核的 线程块数量的 1D、2D 或 3D 元组。在我们之前的例子中，我们使用了一个包含 3 个线程块的 1D 网格：grid = (3, )。

为了处理不同的数组大小，我们默认使用 grid = (ceil(n_elements / BLOCK_SIZE), )。

def add(X: torch.Tensor, Y: torch.Tensor) -> torch.Tensor:
	"""PyTorch wrapper for `add_kernel`."""
	output = torch.zeros_like(x) # allocate memory for the output
	n_elements = output.numel()  # dimension of X and Y

	# cdiv = ceil div, computes the number of blocks to use
	grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
	# calling the kernel will automatically store `BLOCK_SIZE` in `meta`
	# and update `output`
	add_kernelgrid

	return output

关于包装器的最后两点说明：

你可能已经注意到 grid 被定义为 lambda 函数。这允许 Triton 在启动时计算要启动的线程块数量。因此，我们根据存储在 meta 中的块大小来计算网格大小，meta 是一个包含编译时常数的字典，这些常数被暴露给内核。

当调用内核时，output 的值将就地修改，因此我们不需要重新分配 output = add_kernel[…]。

我们可以通过验证内核是否正常工作来结束这个教程：

x, y = torch.randn((2, 2048), device="cuda")

print(add(x, y))
>> tensor([ 1.8022, 0.6780, 2.8261, ..., 1.5445, 0.2563, -0.1846], device='cuda:0')

abs_difference = torch.abs((x + y) - add(x, y))
print(f"Max absolute difference: {torch.max(abs_difference)}")
>> Max absolute difference: 0.0

这就是本次介绍的结束，在接下来的文章中，我们将学习如何实现更有趣的内核，例如分块矩阵乘法，并了解如何使用 autograd 将 Triton 内核集成到 PyTorch 模型中。

下次见！👋

参考资料

[训练 GPT4 的成本]( https://en.wikipedia.org/wiki/GPT-4#:~:text=Sam Altman stated that the ,51)
Unsloth 内核
Triton 教程：向量加法
从第一性原理开始构建深度学习

posted @ 2026-03-29 09:50 布客飞龙III 阅读(38) 评论(0) 收藏举报

刷新页面返回顶部

龙哥盟