Developing Triton Kernels on AMD GPUs — ROCm Blogs

2024 年 4 月 15 日，作者： Clint Greene.

简介

OpenAI 开发了一种强大的、专注于 GPU 的编程语言和编译器，名为 Triton，它可以无缝地在 AMD GPU 上运行。Triton 的目标是使 AI 工程师和科学家能够用最少的专业知识编写高性能的 GPU 代码。Triton 内核之所以性能优越，是因为其块状程序表示，使其能够被编译成高度优化的二进制代码。Triton 还利用 Python 来进行内核开发，使其既熟悉又容易上手。并且，内核只需在定义前简单声明 triton.jit python 装饰器即可轻松编译。

在这篇博客中，我们将深入探讨如何开发 GELU（高斯误差线性单元）内核，并将其性能与 PyTorch 的类似实现进行基准测试。

先决条件

要运行本博文中的内容，你需要以下条件：

• 一块 AMD GPU: 请参考兼容 GPU 列表

• Linux: 请参考支持的 Linux 发行版

• ROCm 5.7+: 请参考安装说明

开始

让我们先安装需要用到的库。如果你已经安装了 ROCm 6.0 和最新版本的 PyTorch，可以跳过这一步。然而，如果你在运行任何命令时遇到问题，我们建议使用 nightly wheels 进行更新。这也会安装与 ROCm 的 PyTorch 兼容的 Triton 版本。

pip install matplotlib pandas -q
pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/rocm6.0/ -q

并导入它们

import torch
import triton
import triton.language as tl

现在必要的库已经安装和导入，我们可以开始开发一个使用_tanh_近似 GELU 的 Triton 内核。但让我们首先讨论什么是 GELU 以及为什么我们需要为它开发一个内核。

GELU

在神经网络中，GELU（Gaussian Error Linear Unit）是最广泛使用的激活函数之一，因为它能够引入网络的非线性特性，其梯度在训练过程中平滑且稳定，具有计算效率且在深度学习应用中表现良好。GELU通常用于Transformer模型如GPT和BERT中，并且已经在各种计算机视觉任务中的卷积神经网络（CNN）中取得了优异的性能。GELU接受一个任意实值的向量_x_作为输入，并输出一个相同大小的向量。通过tanh函数近似的GELU定义如下：

GELU(x)≈x²[1+tanh(√(2/π)⋅(x+0.044715x³))]

我们可以很容易地在Triton中从头编写这个代码。

# GELU计算的常数
k = (2.0 / torch.pi) ** 0.5

@triton.jit  # Triton即时编译的装饰器
def gelu_kernel(
    output_buffer,
    input_buffer,
    input_row_stride,
    output_row_stride,
    n_cols,
    block_size: tl.constexpr,
):
    """
    使用Triton应用GELU的函数。

    参数
    ----------
    output_buffer (pointer): 存储输出张量数据的缓冲区指针。
    input_buffer (pointer): 存储输入张量数据的缓冲区指针。
    input_row_stride (int): 输入张量中行之间的步幅（跳过的元素数量）。
    output_row_stride (int): 输出张量中行之间的步幅（跳过的元素数量）。
    n_cols (int): 输入和输出张量中的列数。
    block_size (tl.constexpr int): 在编译时已知的用于高效处理张量的块大小。

    返回值：无，输出张量缓冲区就地修改。
    """
    # 根据程序ID确定起始行指针
    row_idx = tl.program_id(0)
    row_start_ptr = input_buffer + row_idx * input_row_stride

    # 为块处理生成一系列列偏移
    col_offsets = tl.arange(0, block_size)
    input_ptrs = row_start_ptr + col_offsets

    # 加载当前行的数据，处理潜在的越界元素
    row = tl.load(input_ptrs, mask=col_offsets < n_cols, other=-float('inf'))

    # 使用tanh近似值应用GELU激活函数
    t = (row + 0.044715 * row * row * row) * k
    numerator = 2
    denominator = (1 + tl.exp(-2 * t))
    gelu_output = 0.5 * row * (1 + (numerator / denominator) - 1)

    # 计算输出指针并存储计算出的GELU值
    output_row_start_ptr = output_buffer + row_idx * output_row_stride
    output_ptrs = output_row_start_ptr + col_offsets
    tl.store(output_ptrs, gelu_output, mask=col_offsets < n_cols)

在Triton中，还需要创建一个辅助函数，该函数使用所需的参数排队(kernel)。如果您需要有关参数和Triton中GPU编程的更多信息，请参阅 official docs。现在让我们开发我们的排队函数。

def gelu(x: torch.Tensor):
    n_rows, n_cols = x.shape

    # 确定用于并行处理元素的有效块大小
    block_size = triton.next_power_of_2(n_cols)

    # 根据块大小调整warp（线程组）的数量以获得最佳性能
    # 对于较小的块大小使用4个warp
    num_warps = 4
    
    # 对于较大的块，将warp 增加到8个
    if block_size >= 2048:
        num_warps = 8
    
    # 对于更大的块，将warp 增加到16个
    elif block_size >= 4096:
        num_warps = 16

    y = torch.empty_like(x)

    gelu_kernel[(n_rows,)](# Launch with a one-element tuple for grid size (number of rows)
      y,  # 输出张量
      x,  # 输入张量
      x.stride(0),  # 输入张量中行之间的步幅
      y.stride(0),  # 输出张量中行之间的步幅
      n_cols,  # 张量中的列数
      num_warps=num_warps,  # 并行执行的warp数量
      block_size=block_size,  # 高效处理的块大小
    )

    return y

在开发了Triton内核后，进行单元测试总是一个好主意。这将确保我们正确地编码了gelu近似，并且我们的内核能够处理不规则数量的行和列。

torch.manual_seed(0)
x = torch.randn(4097, 311, device='cuda')
y_triton = gelu(x)
y_torch = torch.nn.functional.gelu(x, approximate='tanh')
assert torch.allclose(y_triton, y_torch, rtol=1e-04,atol=1e-04), (y_triton, y_torch)

单元测试确认我们的内核正常工作。现在我们可以对内核进行基准测试并评估其性能。

基准测试

我们使用Triton的基准测试工具对Triton内核在不同大小的张量上的性能进行了基准测试，并将其与PyTorch内部的`gelu`函数的性能进行比较。

@triton.testing.perf_report(
    triton.testing.Benchmark(
        x_names=['N'],  # 用作x轴的参数名
        x_vals=[128 * i for i in range(2, 100)],  # d`x_name`的不同可能值
        line_arg='provider',  # 参数名，其值对应于图中不同的线
        line_vals=[
            'triton',
            'pytorch',
        ],  # `line_arg`的可能值
        line_names=[
            "Triton",
            "PyTorch",
        ],  # 线条标签名
        styles=[('blue', '-'), ('green', '-'), ('green', '--')],  # 线条样式
        ylabel="GB/s",  # y轴标签名
        plot_name="gelu-performance",  # 图的名称。也用作保存图像文件的文件名
        args={'M': 4096},  #函数参数`x_names`和`y_name`之外的参数值
    ))
def benchmark(M, N, provider):
    x = torch.randn(M, N, device='cuda', dtype=torch.float32)
    quantiles = [0.5, 0.2, 0.8]
    if provider == 'pytorch':
        ms, min_ms, max_ms = triton.testing.do_bench(lambda: torch.nn.functional.gelu(x), quantiles=quantiles)
    if provider == 'triton':
        ms, min_ms, max_ms = triton.testing.do_bench(lambda: gelu(x), quantiles=quantiles)
    gbps = lambda ms: 2 * x.nelement() * x.element_size() * 1e-9 / (ms * 1e-3)
    return gbps(ms), gbps(max_ms), gbps(min_ms)

benchmark.run(show_plots=True, print_data=True)

现在我们运行基准测试，看看性能表现。

benchmark.run(print_data=True, show_plots=True)

我们可以看到，我们的GELU内核（蓝色）在带宽上比PyTorch的GELU实现（绿色）领先多达10%。性能提升归功于在片上SRAM中手动融合操作，而不是从HBM DRAM多次读写。

总结

在这篇博客中，我们演示了如何开发一个Triton内核并对其性能进行基准测试。有关理解GPU编程更多详细信息和开发更高级融合内核的其他示例，请访问Triton的官方教程.