在机器学习领域，即使是微小的优化，也能带来巨大的价值。例如，训练像GPT-4这样的大型模型，成本超过1亿美元，这意味着1%的效率提升就能节省超过一百万美元。提升机器学习模型效率的强大方法之一，是直接在GPU上编写部分组件。对于许多开发者而言，提到CUDA内核往往令人望而却步，因为它们以编写和调试的复杂性而闻名。
幸运的是，OpenAI于2021年发布了Triton，这是一种全新的语言和编译器，它抽象了CUDA的许多复杂性，让经验较少的开发者也能编写出高性能的内核。一个显著的例子是Unsloth，这项LLM训练服务宣称能将训练速度提高30倍，内存使用量减少60%，这都归功于将原本用PyTorch编写的层替换为Triton内核。
在本系列教程中，读者将学习GPU架构的基础知识，并掌握如何实现高性能的Triton内核！本系列所有代码均可在https://github.com/RPegoud/Triton-Kernels获取。
GPU架构基础
在本节中，文章将深入探讨（Nvidia）GPU最基础的架构知识，旨在帮助读者在本文结束时编写出第一个Triton内核。
从最小的软件执行单元开始，可以将GPU的执行单元层级结构描述如下：
- 线程 (Threads)：最小的工作单元，负责运行用户定义的内核代码。
- Warp (Warp)：最小的调度单元，始终由32个并行线程组成，每个线程都有独立的指令地址计数器和寄存器状态。Warp中的线程同时启动，但可以自由分支并独立执行。
- 线程块 (Thread Blocks)：Warp的组合，其中所有线程可以通过共享内存和同步屏障进行协作。线程块必须能够独立执行，并且可以以任何顺序并行或顺序执行。这种独立性使得线程块能够以任意顺序在任意数量的核心上调度，从而使GPU程序能够随着核心数量的增加而高效扩展。如有需要，可以在内核中的特定点同步块内的线程，例如同步内存访问。
- 流式多处理器 (Streaming Multiprocessor, SM)：负责并行执行多个Warp的单元，拥有共享内存和L1缓存（用于存储SM最近访问的全局内存行）。一个SM拥有专用的Warp调度器，它从准备运行的线程块中拉取Warp进行处理。
在硬件层面，最小的工作单元是CUDA核心，它是执行线程算术运算（或其部分运算）的物理算术逻辑单元（ALU）。
为总结本节内容，可以用一个比喻来形容：可以将CUDA核心视为单个工人，而一个Warp则是一个由32名工人组成的班组，同时接收到相同的指令。他们可能会以不同的方式执行任务（分支），也可能在不同的时间点完成任务（独立性）。一个线程块由多个班组组成，它们共享一个公共工作区（即拥有共享内存），工作区内所有班组的工人都可以在同一时间等待彼此一起吃午饭。一个流式多处理器（SM）可以被看作是一个工厂车间，有许多班组协同工作，并共享工具和存储空间。最终，GPU则是一座拥有多层车间的完整工厂。
Nvidia GPU架构的层级结构。虚线矩形表示内存块（作者绘制）
优化基础
在优化深度学习模型时，主要需要权衡三个核心组成部分：
- 计算 (Compute)：GPU执行浮点运算（FLOPS）所花费的时间。
- 内存 (Memory)：在GPU内部传输张量所花费的时间。
- 开销 (Overhead)：所有其他操作（如Python解释器、PyTorch调度等）。
牢记这些组成部分有助于找到解决瓶颈的正确方法。例如，如果大部分时间都花在内存传输上，那么仅仅增加计算能力（例如，使用更强大的GPU）是无济于事的。然而，理想情况下，大部分时间应该用于计算，更精确地说，是用于矩阵乘法，因为这是GPU专门为之优化的操作。
这意味着要最小化数据传输的成本，无论是从CPU到GPU（“数据传输成本”），从一个节点到另一个节点（“网络成本”），还是从CUDA全局内存（DRAM，廉价但速度慢）到CUDA共享内存（SRAM，昂贵但设备上最快的内存）。后者被称为带宽成本，也将是本文目前的重点。减少带宽成本的常见策略包括：
- 数据重用 (Reusing)：在多个步骤中重用加载到共享内存中的数据。一个典型的例子是平铺矩阵乘法（tiled matrix multiplication），这将在后续文章中深入探讨。
- 操作融合 (Fusing)：将多个操作融合到一个内核中（因为每次内核启动都意味着数据从DRAM移动到SRAM），例如，可以将矩阵乘法与激活函数融合。通常，操作融合可以带来巨大的性能提升，因为它避免了大量的全局内存读写，并且任意两个操作都存在融合的机会。
在这个例子中，执行一次矩阵乘法x@W，并将结果存储在中间变量a中。然后对a应用ReLU函数，并将结果存储在变量y中。这需要GPU从全局内存读取x和W，将结果写入a，再次从a读取，最后写入y。而操作融合则允许通过在一个内核中执行矩阵乘法和应用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 + tl.arange(0, BLOCK_SIZE)来获取该块剩余元素的地址，并将它们加载到内存中。
然而，需要记住的是，程序2只负责元素6，但其偏移量为[6, 7, 8]。为了避免任何索引错误，Triton允许定义一个掩码来识别有效的目标元素，在这里是mask = offsets < n_elements。
现在可以安全地加载X和Y并将它们相加，然后以类似的方式将结果写回全局内存中的输出变量Z。
下面，详细查看代码，这是Triton内核的实现：
import triton
import triton.language as tl
@triton.jit
def add_kernel(
x_ptr, # 指向x的第一个内存条目的指针
y_ptr, # 指向y的第一个内存条目的指针
output_ptr, # 指向输出结果第一个内存条目的指针
n_elements, # x和y的维度
BLOCK_SIZE: tl.constexpr, # 单个块的大小
):
# --- 计算偏移量和掩码 ---
pid = tl.program_id(axis=0) # 块索引
block_start = pid * BLOCK_SIZE # 当前块的起始索引
offsets = block_start + tl.arange(0, BLOCK_SIZE) # 索引范围
mask = offsets < n_elements # 掩盖越界元素
# --- 从全局内存加载变量 ---
x = tl.load(x_ptr + offsets, mask=mask)
y = tl.load(y_ptr + offsets, mask=mask)
# --- 执行操作 ---
output = x + y
# --- 将结果保存到全局内存 ---
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级别的封装器，它负责提供内存指针并定义内核网格（kernel grid）。通常，内核网格是一个一维、二维或三维的元组，其中包含分配给内核的沿每个轴的线程块数量。在前面的例子中，使用了包含3个线程块的一维网格：grid = (3,)。
为了处理不同大小的数组，默认情况下将网格设置为grid = (ceil(n_elements / BLOCK_SIZE),)。
def add(X: torch.Tensor, Y: torch.Tensor) -> torch.Tensor:
"""`add_kernel`的PyTorch封装器。"""
output = torch.zeros_like(x) # 为输出分配内存
n_elements = output.numel() # X和Y的维度
# cdiv = 向上取整的除法，计算需要使用的块数量
grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
# 调用内核将自动把`BLOCK_SIZE`存储到`meta`中
# 并更新`output`
add_kernel[grid](X, Y, output, n_elements, BLOCK_SIZE=1024)
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
本次Triton入门就到这里，在后续的文章中，将深入学习如何实现更复杂的内核，例如平铺矩阵乘法，并探索如何利用autograd将Triton内核集成到PyTorch模型中。
下次再见！ 
