单线程版本#
逐元素相加：
Triton 实现#
在 Triton 中，向量加法内核通过将向量划分为多个块（blocks），并在每个 Grid 中的线程（threads）并行计算，实现高效的向量加法操作。每个线程负责加载两个向量中对应位置的元素，进行相加并存储结果。
核心步骤#
- 线程并行计算：每个 Grid 中的线程独立处理向量中的一部分元素。
- 加载元素：每个线程加载向量 A 和向量 B 中对应位置的元素。
- 元素相加：将加载的元素进行相加。
- 存储结果：将相加后的结果存储到输出向量中。
tl.constexpr 的使用#
tl.constexpr 用于声明编译时常量。这意味着使用这个修饰符的变量的值在编译时就已经确定，而不是在运行时。编译器可以基于这些常量值进行更 aggressive 的优化来提升内核的执行效率
@triton.jit
def kernel_vector_addition(a_ptr, b_ptr, out_ptr,
num_elems: tl.constexpr,
block_size: tl.constexpr):
# 内核代码
上述代码中，num_elems 和 block_size 被声明为编译时常量，使得 Triton 可以在编译阶段优化内核代码。
确定当前块与 Program ID#
每个线程块（block）在 Triton 中都有一个唯一的 Program ID，用于标识当前线程所在的块。通过 tl.program_id，我们可以确定当前线程所在的块，从而计算处理的数据偏移量。
pid = tl.program_id(axis=0)
block_start = pid * block_size
处理最后一个 Block#
由于向量长度可能不被块大小整除，最后一个块可能只有部分线程需要工作。通过掩码操作，可以确保只有有效的线程进行计算，避免无效的内存访问和计算。
掩码的作用#
Triton 提供了掩码（mask）操作，用于屏蔽那些不需要工作的线程（最后一个 Grid 中 NA 的线程）。
mask = thread_offsets < num_elems
a_pointers = tl.load(a_ptr + thread_offsets, mask=mask, other=0.0)
b_pointers = tl.load(b_ptr + thread_offsets, mask=mask, other=0.0)
ceil_div 函数的作用#
ceil_div 函数用于计算块的数量，确保即使向量长度不被块大小整除，也能覆盖所有元素。例如 vec_size=10，block_size=3，ceil_div(10, 3)=4，这样就能确保所有 10 个元素都被处理。
def ceil_div(x: int, y: int) -> int:
return (x + y - 1) // y
说白了，该函数的作用就是高效实现 “向上取整”。
数值精度验证#
在实现向量加法内核后，验证数值精度是确保内核正确性的关键步骤。通过与 PyTorch 的内置加法操作进行对比，可以确认 Triton 实现的准确性。
def verify_numerics() -> bool:
torch.manual_seed(2020) # seed both cpu and gpu
vec_size = 8192
a = torch.rand(vec_size, device='cuda')
b = torch.rand_like(a)
torch_res = a + b
triton_res = vector_addition(a, b)
fidelity_correct = torch.allclose(torch_res, triton_res)
print(f"{fidelity_correct=}")
return fidelity_correct
验证了解到我们的 Triton 实现与 PyTorch 原生的数值精度一致，可以进行后面的操作了。
下面是完整的 Kernel 实现：
@triton.jit
def kernel_vector_addition(a_ptr, b_ptr, out_ptr,
num_elems: tl.constexpr,
block_size: tl.constexpr,):
pid = tl.program_id(axis=0)
# tl.device_print("pid", pid)
block_start = pid * block_size # 0 * 2 = 0, 1 * 2 = 2,
thread_offsets = block_start + tl.arange(0, block_size)
mask = thread_offsets < num_elems
a_pointers = tl.load(a_ptr + thread_offsets, mask=mask)
b_pointers = tl.load(b_ptr + thread_offsets, mask=mask)
res = a_pointers + b_pointers
tl.store(out_ptr + thread_offsets, res, mask=mask)
def ceil_div(x: int,y: int) -> int:
return (x + y - 1) // y
def vector_addition(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
output_buffer = torch.empty_like(a)
assert a.is_cuda and b.is_cuda
num_elems = a.numel()
assert num_elems == b.numel() # todo - handel mismatched sizes
block_size = 1024
grid_size = ceil_div(num_elems, block_size)
grid = (grid_size,)
num_warps = 8
k2 = kernel_vector_addition[grid](a, b, output_buffer,
num_elems,
block_size,
num_warps=num_warps
)
return output_buffer
基准测试与性能调优#
为了评估 Triton 向量加法内核的性能，下面进行基准测试并探讨性能调优的方法。
Benchmark API 介绍#
Triton 提供了丰富的基准测试 API，允许用户测量内核的执行时间和吞吐量。以下代码是使用 triton.testing.perf_report 得到性能报告的一个示例：
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=['size'], # 用作图表x轴的参数名
x_vals=[2**i for i in range(10, 28)], # `x_name`的可能取值
x_log=True, # x轴使用对数刻度
line_arg='provider', # 图表中不同线条对应的参数名
line_vals=['triton', 'torch'], # `line_arg`的可能取值
line_names=["Triton", "Torch"], # 线条的标签名
styles=[('blue', '-'), ('green', '-')], # 线条颜色和样式
ylabel='GB/s', # y轴标签
plot_name='vector-add-performance', # 图表名称，也用作保存文件的文件名
args={}, # 不在`x_names`和`y_name`中的函数参数的值
)
)
def benchmark(size, provider):
x = torch.rand(size, device='cuda', dtype=torch.float32)
y = torch.rand(size, device='cuda', dtype=torch.float32)
quantiles = [0.5, 0.2, 0.8] # 设定分位数
# 根据 provider 选择不同的计算实现
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: vector_addition(x, y), quantiles=quantiles)
# 计算GB/s
def gbps(ms):
return 12 * size / ms * 1e-06
# 返回中位数、最大值和最小值对应的GB/s
return gbps(ms), gbps(max_ms), gbps(min_ms)
性能报告：
性能对比：
总的来说，vector add 只是一个相对简单的 kernel，相较于复杂的内核更难获得 Triton 实现带来的优势（大部分常用操作 PyTorch 已经通过 CUDA/cuBLAS 等优化到极质了）
调优参数：Num Warps & Block size#
调优内核性能的关键在于合理配置 Warp 数量和块大小。Warp 是 GPU 中的基本执行单元，合理的 Warp 数量和块大小能够充分利用 GPU 的并行计算能力，提升内核的执行效率。
block_size = 1024 # 决定每个线程块处理的元素数量,较大的块大小可以减少块的数量，但可能增加每个块的计算负担。
grid_size = ceil_div(num_elems, block_size)
grid = (grid_size,)
num_warps = 8 # 每个块中包含的 Warp 数量,合理配置 Warp 数量可以优化线程的调度和资源利用。
上节 (Softmax in OpenAI Triton)便给出了通过驱动程序来动态调整参数的方式：
# 计算 block_size，为大于等于 cols 的最小 2 的幂
block_size = triton.next_power_of_2(cols)
# 根据 block_size 动态调整 num_warps
num_warps = 4 # 每个 warp 有 32 个线程
if block_size > 2047:
num_warps = 8
if block_size > 4095:
num_warps = 16
参考资料#
感谢：
- 本文内容主要基于这位老师的视频教程系列 SOTA Deep Learning Tutorials - YouTube
- Triton’s documentation
- o1-preview 语言模型，画 svg 的能手