BBuf的CUDA笔记系列：OpenAI Triton从零开始入门笔记详解第一篇_AI热点日报

OpenAITriton是一种开源类PythonGPU编程语言，旨在简化CUDAkernel编写。它通过块级操作抽象并行性，自动化内存合并、共享内存同步和TensorCores调度等优化，使无CUDA经验的研究人员也能写出接近专家手写性能的代码，例如FP16矩阵乘法在25行内可媲美cuBLAS。

0x1. OpenAI Triton 介绍与解读

首先来看官方说明：https://openai.com/research/triton。通过官方文档，我们可以深入了解Triton的诞生背景、设计目标，并看到一些经典算法的实现示例。

官网页面标题为“Introducing Triton: Open-source GPU programming for neural networks”，中文可译为《介绍 Triton：面向神经网络的开源 GPU 编程语言》。下方关键语句强调了：Triton 1.0 是一种开源的类 Python 编程语言，即使没有 CUDA 经验的开发者也能高效编写 GPU 代码，且在大多数场景下性能可与专家手写代码媲美。这直接点明了 Triton 的核心目的：简化 CUDA kernel 的编写。接下来，我们逐段解读官方介绍的细节。为准确起见，我会摘录原文并附上翻译与理解。

这段内容表明：Triton 能让用户以极低成本编写出接近硬件峰值性能的 kernel。例如，用 Triton 实现一个 FP16 矩阵乘法核函数，只需不到 25 行代码，性能即可与 cuBLAS 持平。研究人员已基于 Triton 开发出多种高效实现，与功能相同的 Torch 实现相比，性能可提升一倍。后续段落进一步指出：使用 CUDA 将原始的 PyTorch 实现改写为自定义算子通常更高效，但难度较大，且现有工具无法覆盖所有情况——这正是 OpenAI Triton 诞生的原因。

这里讨论的是 GPU 编程的难点。现代 GPU 架构大致可分为三个主要部分：DRAM、SRAM 和 ALU。在优化 CUDA 代码时，必须综合考虑以下因素：

从 DRAM 的内存传输必须合并为大型事务，以充分利用现代内存接口的高带宽（即内存合并访问）。
数据需在重复使用前手动存储到 SRAM 中，并进行合理管理以最小化 bank conflict。
计算必须仔细划分和调度，不仅在流式多处理器（SM）之间，还要在 SM 内部，以促进指令级/线程级并行性，并充分利用专用 ALU（如 Tensor Cores）。

综合考虑所有这些因素，即便对拥有多年经验的资深 CUDA 程序员而言也极具挑战。Triton 的目标正是将这些优化完全自动化，让开发者能更专注于并行代码的高层逻辑。同时，Triton 力求广泛适用，因此不会自动在 SM 之间调度工作——一些重要的算法考量（如 tiling、跨 SM 同步）仍留给开发者自行决定。

下面是一个表格，对比了 CUDA 编译器与 Triton 的区别。

在现有的领域特定语言和 JIT 编译器中，Triton 与 Numba 最为相似：kernel 被定义为一个被装饰的函数，并通过不同的 program_id 在网格实例上并行启动。然而，如下面的代码片段所示，相似之处仅限于此。Triton 通过操作块（block）来暴露实例内部的并行性——这些块是大小为 2 的幂次的小数组——而不是使用单指令多线程（SIMT）执行模型。这样一来，Triton 有效抽象掉了所有与 CUDA 线程块内部并发相关的问题（如内存合并、共享内存同步/冲突、Tensor Cores 调度）。

注意，Triton 的 JIT 编译器将 X 和 Y 视为指针而非张量。官方认为，保留对内存访问的低级控制对于处理更复杂的数据结构（如块稀疏张量）非常重要。关键在于，这种特定的 softmax 实现在整个归一化过程中将 X 的行保留在 SRAM 中，在适用时（大约小于 32K 列）最大化数据重用。这与 PyTorch 内部的 CUDA 代码不同，PyTorch 使用临时内存以提高通用性，但速度明显更慢（如下所示）。这里并非 Triton 本质上更优，而是它简化了专用 kernel 的开发，这些内核可能比通用库中的内核快得多。

Torch（v1.9）JIT 编译器的较低性能，凸显了从高级张量操作序列自动生成 CUDA 代码的难度。

这里提到，Triton 仅需约 25 行 Python 代码即可实现接近峰值的矩阵乘法。（后续有专门章节讲解代码原理）代码如下：

@triton.jit
def matmul(A, B, C, M, N, K,
           stride_am, stride_ak,
           stride_bk, stride_bn,
           stride_cm, stride_cn,
           **META):
    # extract metaparameters
    BLOCK_M, GROUP_M = META['BLOCK_M'], META['GROUP_M']
    BLOCK_N = META['BLOCK_N']
    BLOCK_K = META['BLOCK_K']
    # programs are grouped together to improve L2 hit rate
    _pid_m = tl.program_id(0)
    _pid_n = tl.program_id(1)
    pid_m = _pid_m // GROUP_M
    pid_n = (_pid_n * GROUP_M) + (_pid_m % GROUP_M)
    # rm (resp. rn) denotes a range of indices
    # for rows (resp. col) of C
    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    # rk denotes a range of indices for columns 
    # (resp. rows) of A (resp. B)
    rk = tl.arange(0, BLOCK_K)
    # the memory addresses of elements in the first block of
    # A and B can be computed using numpy-style broadcasting
    A = A + (rm[:, None] * stride_am + rk[None, :] * stride_ak)
    B = B + (rk[:, None] * stride_bk + rn[None, :] * stride_bn)
    # initialize and iteratively update accumulator
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for k in range(K, 0, -BLOCK_K):
        a = tl.load(A)
        b = tl.load(B)
        # block level matrix multiplication
        acc += tl.dot(a, b)
        # increment pointers so that the next blocks of A and B
        # are loaded during the next iteration
        A += BLOCK_K * stride_ak
        B += BLOCK_K * stride_bk
    # fuse leaky ReLU if desired
    # acc = tl.where(acc >= 0, acc, alpha * acc)
    # write back result
    C = C + (rm[:, None] * stride_cm + rn[None, :] * stride_cn)
    mask = (rm[:, None] < M) & (rn[None, :] < N)
    tl.store(C, acc, mask=mask)

手写矩阵乘法 kernel 的一个重要优势在于，它们可以根据需要定制，以融合输入（例如切分）和输出（例如 LeakyReLU）的变换。如果没有像 Triton 这样的系统，缺乏出色 GPU 编程专长的开发者将无法对矩阵乘法 kernel 进行定制修改。

这里说明，Triton 的良好性能源于一个以 Triton-IR 为中心的模块化系统架构。Triton-IR 是一种基于 LLVM 的中间表示，在该 IR 中，多维值块（MLIR 的概念）被视为一等公民。

@triton.jit 装饰器的工作原理是：遍历所提供 Python 函数的抽象语法树（AST），以便使用常见的 SSA 构建算法即时生成 Triton-IR。随后，编译器后端会对生成的 IR 代码进行简化、优化和自动并行化，再将其转换为高质量的 LLVM-IR——最终生成 PTX——以在近期的 NVIDIA GPU 上执行。目前不支持 CPU 和 AMD GPU，但官方欢迎社区贡献，旨在解决这一限制。

我们发现，通过 Triton-IR 使用块级程序表示，使编译器能够自动执行多种重要的程序优化。例如，可以通过观察计算密集型块级操作（如 tl.dot）的操作数，自动将数据暂存到共享内存中，并使用标准的活性分析技术进行分配和同步。

另一方面，如下所示，Triton 程序可以高效且自动地并行化，既可以通过（1）并发执行不同的 kernel 实例在 SM 之间并行，又可以通过（2）分析每个块级操作的迭代空间，并在不同的 SIMD 单元间适当分配，从而在 SM 内部并行。

0x2. 教程 1：向量加法解析

本教程将介绍 Triton 编程模型定义 kernel 的基本写法，并展示如何实现良好的基准测试。下面来看计算 kernel 的实现，注释已转换为中文：

import torch
import triton
import triton.language as tl

@triton.jit
def add_kernel(x_ptr,  # *指针*，指向第一个输入向量。
               y_ptr,  # *指针*，指向第二个输入向量。
               output_ptr,  # *指针*，指向输出向量。
               n_elements,  # 向量的大小。
               BLOCK_SIZE: tl.constexpr,  # 每个程序应处理的元素数量。
               # 注意：`constexpr`这样可以被用作形状值。
               ):
    # 这里有多个“程序”处理不同的数据。我们在这里识别我们是哪一个程序：
    pid = tl.program_id(axis=0)  # 我们使用一维启动网格，所以轴是0。
    # 该程序将处理从初始数据偏移的输入。
    # 例如，如果你有一个长度为256的向量和块大小为64，那么程序
    # 将分别访问元素[0:64, 64:128, 128:192, 192:256]。
    # 注意偏移量是一个指针列表：
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    # 创建一个掩码以防止内存操作越界访问。
    mask = offsets < n_elements
    # 从DRAM加载x和y，屏蔽任何额外的元素以防输入不是块大小的倍数。
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    output = x + y
    # 将x + y写回DRAM。
    tl.store(output_ptr + offsets, output, mask=mask)

还声明了一个辅助函数来（1）分配 z 张量，（2）使用适当的网格/块大小排队上述 kernel：

def add(x: torch.Tensor, y: torch.Tensor):
    # 我们需要预分配输出。
    output = torch.empty_like(x)
    assert x.is_cuda and y.is_cuda and output.is_cuda
    n_elements = output.numel()
    # SPMD启动网格表示并行运行的kernel实例的数量。
    # 它类似于CUDA启动网格。它可以是Tuple[int]，也可以是Callable(metaparameters) -> Tuple[int]。
    # 在这种情况下，我们使用一个1D网格，其大小是块的数量：
    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), )
    # 注意：
    #  - 每个torch.tensor对象都隐式地转换为指向其第一个元素的指针。
    #  - 使用`triton.jit`装饰的函数可以用一个启动网格索引来获得可调用的GPU内核。
    #  - 不要忘记将元参数作为关键字参数传递。
    add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=1024)
    # 我们返回一个指向z的句柄，但是因为`torch.cuda.synchronize()`还没有被调用，所以这时kernel仍然
    # 在异步运行。
    return output

现在可以用上面定义的函数来计算两个 torch.tensor 对象的逐元素求和，并验证正确性：

torch.manual_seed(0)
size = 98432
x = torch.rand(size, device='cuda')
y = torch.rand(size, device='cuda')
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))}')

输出：

tensor([1.3713, 1.3076, 0.4940,  ..., 0.6724, 1.2141, 0.9733], device='cuda:0')
tensor([1.3713, 1.3076, 0.4940,  ..., 0.6724, 1.2141, 0.9733], device='cuda:0')
The maximum difference between torch and triton is 0.0

我们可以对不同大小的向量进行自定义操作的性能基准测试，观察其与 PyTorch 的对比表现。Triton 提供了一套内置工具，能够简洁地绘制自定义操作在不同问题规模下的性能图表。

@triton.testing.perf_report(
    triton.testing.Benchmark(
        x_names=['size'],  # 用作绘图x轴的参数名。
        x_vals=[2**i for i in range(12, 28, 1)],  # `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]
    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: 12 * size / ms * 1e-6
    return gbps(ms), gbps(max_ms), gbps(min_ms)

gbps = lambda ms: 12 * size / ms * 1e-6 这里的 12 表示数据读写的比特数，因为有 x、y 和 z 三个张量，所以是 3×4=12 bit。现在可以执行上述装饰函数了。传递 print_data=True 参数查看性能数据，传递 show_plots=True 参数绘制图表，和/或传递 save_path='/path/to/results/' 参数将它们连同原始 CSV 数据一起保存到磁盘：

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

可以看到，对于逐元素任务，Triton 的性能几乎与 PyTorch 持平，但 Triton 的编写方式要简单得多。

0x3. 教程 2：融合 Softmax 解读

在本教程中，我们将编写一个融合的 softmax 操作，对于特定类型的矩阵（即行大小可放入 GPU SRAM 的矩阵）而言，它比 PyTorch 的原生操作快得多。

通过本教程，我们将学到：

kernel 融合对于带宽受限操作的好处。
Triton 中的 reduce 操作符。

动机

自定义 GPU kernel 用于逐元素加法在教育上很有价值，但在实际应用中可能作用有限。来看一个简单的（数值稳定的）softmax 操作的情况：

import torch
import triton
import triton.language as tl

@torch.jit.script
def naive_softmax(x):
    """使用原生pytorch计算X的逐行softmax
    我们减去最大元素是为了避免溢出。Softmax对这种偏移是不变的。
    """
    # 读取 MN 个元素；写入 M 个元素
    x_max = x.max(dim=1)[0]
    # 读取 MN + M 个元素；写入 MN 个元素
    z = x - x_max[:, None]
    # 读取 MN 个元素；写入 MN 个元素
    numerator = torch.exp(z)
    # 读取 MN 个元素；写入 M 个元素
    denominator = numerator.sum(dim=1)
    # 读取 MN + M 个元素；写入 MN 个元素
    ret = numerator / denominator[:, None]
    # 总计：读取 5MN + 2M 个元素；写入 3MN + 2M 个元素
    return ret

计算 Kernel

我们的 softmax kernel 工作方式如下：每个程序加载输入矩阵 X 的一行，对其进行归一化处理，然后将结果写回输出 Y 中。需要注意的是，Triton 的一个重要限制是每个块必须包含 2 的幂次方个元素，因此如果想处理任何可能的输入形状，我们需要在内部对每行进行“填充”并对内存访问操作进行保护（即防止越界）：

@triton.jit
def softmax_kernel(output_ptr, input_ptr, input_row_stride, output_row_stride, n_cols, BLOCK_SIZE: tl.constexpr):
    # softmax的各行是独立的，所以我们在这些行上进行并行处理
    row_idx = tl.program_id(0)
    # 步长代表我们需要增加多少指针来前进1行
    row_start_ptr = input_ptr + row_idx * input_row_stride
    # 块大小是大于n_cols的下一个2的幂次，因此我们可以将每一行放入单个块中
    col_offsets = tl.arange(0, BLOCK_SIZE)
    input_ptrs = row_start_ptr + col_offsets
    # 将行加载到SRAM中，使用掩码因为BLOCK_SIZE可能大于n_cols
    row = tl.load(input_ptrs, mask=col_offsets < n_cols, other=-float('inf'))
    # 减去最大值以实现数值稳定性
    row_minus_max = row - tl.max(row, axis=0)
    # 注意在Triton中指数运算快但是近似的（即，类似于CUDA中的__expf）
    numerator = tl.exp(row_minus_max)
    denominator = tl.sum(numerator, axis=0)
    softmax_output = numerator / denominator
    # 将输出写回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=col_offsets < n_cols)

接下来创建一个辅助函数，该函数为任何给定的输入张量排队执行 kernel 并设置启动参数。

def softmax(x):
    n_rows, n_cols = x.shape
    # 块大小是大于`x`中列数的最小2的幂
    BLOCK_SIZE = triton.next_power_of_2(n_cols)
    # 另一个技巧：要求编译器通过增加每行分布的warp数（`num_warps`）来使用更多线程。
    # 下一个教程中会看到如何以更自然的方式自动调整这个值，这样就不必自己想出手动启发式方法。
    num_warps = 4
    if BLOCK_SIZE >= 2048:
        num_warps = 8
    if BLOCK_SIZE >= 4096:
        num_warps = 16
    # 分配输出
    y = torch.empty_like(x)
    # 排队执行内核。一维启动网格很简单：我们有每行一个内核实例
    # 输入矩阵
    softmax_kernel[(n_rows, )](
        y,
        x,
        x.stride(0),
        y.stride(0),
        n_cols,
        num_warps=num_warps,
        BLOCK_SIZE=BLOCK_SIZE,
    )
    return y

这里验证一下 Triton 实现的融合 softmax 和 PyTorch 的 naive 实现是否等价——显然它们等价。

基准测试

这里将矩阵的行数固定为 4096 进行基准测试。

@triton.testing.perf_report(
    triton.testing.Benchmark(
        x_names=['N'],  # 用作绘图x轴的参数名
        x_vals=[128 * i for i in range(2, 100)],  # `x_name`的不同可能值
        line_arg='provider',  # 其值对应于图中不同线条的参数名
        line_vals=[
            'triton',
            'torch-native',
            'torch-jit',
        ],  # `line_arg`的可能值
        line_names=[
            "Triton",
            "Torch (原生)",
            "Torch (jit)",
        ],  # 线条的标签名称
        styles=[('blue', '-'), ('green', '-'), ('green', '--')],  # 线条样式
        ylabel="GB/s",  # y轴的标签名称
        plot_name="softmax-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 == 'torch-native':
        ms, min_ms, max_ms = triton.testing.do_bench(lambda: torch.softmax(x, axis=-1), quantiles=quantiles)
    if provider == 'triton':
        ms, min_ms, max_ms = triton.testing.do_bench(lambda: softmax(x), quantiles=quantiles)
    if provider == 'torch-jit':
        ms, min_ms, max_ms = triton.testing.do_bench(lambda: naive_softmax(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)

这里提到，虽然 Triton 实现的 softmax 性能更好且易于理解和维护，但 PyTorch 的 torch.softmax 则更加通用。

0x4. 教程 3：矩阵乘法解读

首先，教程指出，这里要编写一个块级别的矩阵乘法。这涉及多维指针操作、程序重排以更好地命中 L2 缓存，以及自动调优。

动机

矩阵乘法是大多数现代高性能计算系统的关键构建块。它们众所周知地难以优化，因此其实现通常由硬件供应商以所谓的“内核库”（例如 cuBLAS）的形式提供。不幸的是，这些库通常是专有的，无法轻易定制以适应现代深度学习工作负载的需求（例如融合激活函数）。在本教程中，你将学习如何使用 Triton 自己实现高效的矩阵乘法，这种方法易于定制和扩展。

大致来说，我们将要编写的内核将实现以下块级算法来乘以一个 (M, K) 矩阵和一个 (K, N) 矩阵：

# 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

其中，双重嵌套 for 循环的每次迭代都由一个专用的 Triton program 实例执行。

计算 Kernel

上述算法在 Triton 中相当容易实现。主要的难点在于内循环中计算必须读取 A 和 B 块的内存位置。为此，我们需要多维指针运算。

指针运算

对于一个 2D Tensor X，X[i, j] 的内存位置为 &X[i, j] = X + i*stride_xi + j*stride_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);

这意味着 A 和 B 块的指针可以在 Triton 中初始化，例如 k=0 如下代码所示。注意，需要额外的模运算来处理 M 不是 BLOCK_SIZE_M 的倍数或 N 不是 BLOCK_SIZE_N 的倍数的情况，此时可以用一些无用的值填充数据，这些值不会对结果产生影响。对于 K 维度，稍后将使用掩码加载语义来处理。

offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
offs_k = tl.arange(0, BLOCK_SIZE_K)
a_ptrs = a_ptr + (offs_am[:, None]*stride_am + offs_k [None, :]*stride_ak)
b_ptrs = b_ptr + (offs_k [:, None]*stride_bk + offs_bn[None, :]*stride_bn)

然后在内循环中按如下方式更新：

a_ptrs += BLOCK_SIZE_K * stride_ak;
b_ptrs += BLOCK_SIZE_K * stride_bk;

如上所述，每个 program 实例计算一个 [BLOCK_SIZE_M, BLOCK_SIZE_N] 大小的 C 矩阵块。这些块的计算顺序很重要，因为它会影响程序的 L2 缓存命中率，不幸的是，一个简单的行优先顺序是不够的。

pid = triton.program_id(0);
grid_m = (M + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M;
grid_n = (N + BLOCK_SIZE_N - 1) // BLOCK_SIZE_N;
pid_m = pid / grid_n;
pid_n = pid % grid_n;

L2 缓存优化

每个程序实例计算一个 [BLOCK_SIZE_M, BLOCK_SIZE_N] 大小的 C 矩阵块。重要的是要记住，这些块的计算顺序很重要，因为它会影响我们程序的 L2 缓存命中率，一个简单的行主序排序是不够的。

一个可能的解决方案是以一种促进数据重用的顺序启动块。这可以通过在切换到下一列之前将块在 GROUP_M 行的 super group 中分组来实现：

# 程序ID
pid = tl.program_id(axis=0)
# 沿M轴的程序ID数量
num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
# 沿N轴的程序ID数量
num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
# 组中的程序数量
num_pid_in_group = GROUP_SIZE_M * num_pid_n
# 该程序所在组的ID
group_id = pid // num_pid_in_group
# 组中第一个程序的行ID
first_pid_m = group_id * GROUP_SIZE_M
# 如果`num_pid_m`不能被`GROUP_SIZE_M`整除，最后一个组更小
group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
# *在组内*，程序按列主序排列
# 程序在*启动网格*中的行ID
pid_m = first_pid_m + (pid % group_size_m)
# 程序在*启动网格*中的列ID
pid_n = (pid % num_pid_in_group) // group_size_m

例如，在下面的矩阵乘法中，每个矩阵由 9 个块乘以 9 个块组成，我们可以看到，如果按行主序计算输出，需要将 90 个块加载到 SRAM 中以计算前 9 个输出块，但按 grouped ordering 进行计算，只需要加载 54 个块。

在实际应用中，这可以在某些硬件架构上提高矩阵乘法内核的性能超过 10%（例如，在 A100 上从 220 提升到 245 TFLOPS）。

L2 缓存优化原理补充讲解

上面的 group ordering 的访问代码比较难理解，这里更详细地解析一下。

# 程序ID
pid = tl.program_id(axis=0)
# 沿M轴的程序ID数量
num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
# 沿N轴的程序ID数量
num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)

这里的 num_pid_m 和 num_pid_n 就是分别求要在 M 和 N 方向循环多少次。

上面图中的黑色数字可以理解为 program id，program id 增加的方向其实就代表了遍历的 ordering。对于 row major 来说就是在行方向上顺序遍历，而对于 group ordering 来说就是按照一个 BLOCK_SIZE_M×BLOCK_SIZE_N 这么大的一个小组来遍历。下面这段代码就是完成 group ordering 的遍历：

num_pid_in_group = GROUP_SIZE_M * num_pid_n
group_id = pid // num_pid_in_group
first_pid_m = group_id * GROUP_SIZE_M
group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
pid_m = first_pid_m + (pid % group_size_m)
pid_n = (pid % num_pid_in_group) // group_size_m

以上面图来看，num_pid_m=3，num_pid_n=3，num_pid_in_group=group_id * GROUP_SIZE_M=9×3=27，也就是下面的红色框里面的 program 个数。从名字也可以看出来这个红色框划分的区域也是一个 group。

group_id 表示当前的这次“循环”是在第几个红色框里。以 program 0 为例，group_id = pid // num_pid_in_group = 0 // 27 = 0。而 first_pid_m 代表当前 group 中的第一个黄色 program 在全局的 M 维度上是第几个 program，这里 first_pid_m = 0。group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) 是为了处理最后一个 group 可能占不满数据（存在 padding）的情况，做一个截断。

pid_m = first_pid_m + (pid % group_size_m)
pid_n = (pid % num_pid_in_group) // group_size_m

这两行代码计算当前的 program 处理的黄色小块坐标（[pid_m, pid_n]）。pid_m 这行是在行方向上移动，pid_n 这行则保证在红色框里面一定是一列一列来访问的。

作为对比，在 Row-major 的方法中，访问方式应该是这样的：

pid_m = pid // num_pid_n
pid_n = pid % num_pid_n

计算最后的结果

有了上面的铺垫，我们就可以计算最终的结果了。下面的代码展示了完整的 Triton 矩阵乘法 kernel 实现。

# 使用`triton.jit`装饰的函数可以通过`triton.autotune`装饰器进行自动调优，该装饰器包括：
#   - 一系列定义不同配置的`triton.Config`对象，
#       这些配置涉及元参数（例如`BLOCK_SIZE_M`）和编译选项（例如`num_warps`）的不同设置
#   - 一个自动调优*关键字*，其值的变化将触发对所有
#       提供的配置的评估
@triton.autotune(
    configs=[
        # 每个Config定义了一组特定的配置参数和编译选项
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=3,
                      num_warps=8),
        triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4,
                      num_warps=4),
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4,
                      num_warps=4),
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4,
                      num_warps=4),
        triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4,
                      num_warps=4),
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4,
                      num_warps=4),
        triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=5,
                      num_warps=2),
        triton.Config({'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=5,
                      num_warps=2),
    ],
    key=['M', 'N', 'K'], # 自动调优关键字
)
@triton.jit
def matmul_kernel(
        # 指向矩阵的指针
        a_ptr, b_ptr, c_ptr,
        # 矩阵维度
        M, N, K,
        # 步长变量表示在特定维度上移动1个元素时指针增加的量。
        # 例如`stride_am`是将`a_ptr`增加多少以获取下一行的元素（A有M行）。
        stride_am, stride_ak,  # A矩阵的步长
        stride_bk, stride_bn,  # B矩阵的步长
        stride_cm, stride_cn,# C矩阵的步长
        # 元参数
        BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,
        GROUP_SIZE_M: tl.constexpr,
        ACTIVATION: tl.constexpr  # 激活函数
):
    """用于计算矩阵乘法C = A x B的内核。
    A的形状为(M, K)，B的形状为(K, N)，C的形状为(M, N)。
    """
    # -----------------------------------------------------------
    # 将程序ID `pid`映射到它应该计算的C矩阵的块。
    # 这是以grouped ordering完成的，以促进L2数据重用。
    # 详细解释看上一节
    pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
    num_pid_in_group = GROUP_SIZE_M * num_pid_n
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * GROUP_SIZE_M
    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
    pid_m = first_pid_m + (pid % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m
    # ----------------------------------------------------------
    # 为A和B的第一个块创建指针。
    # 我们将在K方向移动时推进这个指针并累加
    # `a_ptrs`是[BLOCK_SIZE_M, BLOCK_SIZE_K]块的指针
    # `b_ptrs`是[BLOCK_SIZE_K, BLOCK_SIZE_N]块的指针
    # 有关详细信息，请参阅上方“指针算术”部分
    offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
    offs_k = tl.arange(0, BLOCK_SIZE_K)
    a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
    b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
    # -----------------------------------------------------------
    # 迭代以计算C矩阵的一个块。
    # 我们将累加到一个`[BLOCK_SIZE_M, BLOCK_SIZE_N]`块
    # 的fp32值以获得更高的精度。
    # `accumulator`在循环后会转换回fp16。
    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
        # Load the next block of A and B, generate a mask by checking the K dimension.
        # If it is out of bounds, set it to 0.
        a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0)
        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
        # We accumulate along the K dimension.
        accumulator += tl.dot(a, b)
        # Advance the ptrs to the next K block.
        a_ptrs += BLOCK_SIZE_K * stride_ak
        b_ptrs += BLOCK_SIZE_K * stride_bk
    # 当累加器仍然是FP32时，可以融合任意激活函数
    if ACTIVATION == "leaky_relu":
        accumulator = leaky_relu(accumulator)
    c = accumulator.to(tl.float16)
    # -----------------------------------------------------------
    # 使用掩码将输出矩阵C的块写回。
    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
    tl.store(c_ptrs, c, mask=c_mask)

# 我们可以通过将其作为`ACTIVATION`元参数提供给`_matmul`来融合`leaky_relu`。
@triton.jit
def leaky_relu(x):
    x = x + 1
    return tl.where(x >= 0, x, 0.01 * x)

现在可以创建一个方便的封装函数，它只需要两个输入张量，并且会：（1）检查任何形状约束；（2）分配输出；（3）启动上述 kernel。

def matmul(a, b, activation=""):
    # Check constraints.
    assert a.shape[1] == b.shape[0], "Incompatible dimensions"
    assert a.is_contiguous(), "Matrix A must be contiguous"
    assert b.is_contiguous(), "Matrix B must be contiguous"
    M, K = a.shape
    K, N = b.shape
    # Allocates output.
    c = torch.empty((M, N), device=a.device, dtype=a.dtype)
    # 1D launch kernel where each block gets its own program.
    grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(N, META['BLOCK_SIZE_N']), )
    matmul_kernel[grid](
        a, b, c,
        M, N, K,
        a.stride(0), a.stride(1),
        b.stride(0), b.stride(1),
        c.stride(0), c.stride(1),
        ACTIVATION=activation
    )
    return c

计算过程的补充说明

上面的《L2 Cache优化原理补充讲解》一节明确了 kernel 的 group ordering 的访问方式以及实现。现在来看对于当前的 program 实例具体是怎么计算的。以计算 C 中的第一个 Block 的 (0,0) 为例，它需要从 A 和 B 分别加载 9 个黄色的小块数据相乘并累加，最后得到 C 中的 (0,0) 位置结果。如下图所示：

下面的代码先把 program 实例当前要处理的 A 和 B 的第一个 Block 加载上来：

# ----------------------------------------------------------
# 为A和B的第一个块创建指针。
# 我们将在K方向移动时推进这个指针并累加
# `a_ptrs`是[BLOCK_SIZE_M, BLOCK_SIZE_K]块的指针
# `b_ptrs`是[BLOCK_SIZE_K, BLOCK_SIZE_N]块的指针
# 有关详细信息，请参阅上方“指针算术”部分
offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
offs_k = tl.arange(0, BLOCK_SIZE_K)
a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)

这里的 a_ptr 是整个 A 矩阵第一个元素的地址，offs_am 和 offs_bn 表示当前的 program id 在 M 维度和 K 维度的坐标，这个坐标是一个 list，用 tl.arange(0, BLOCK_SIZE_K) 来获取。

得到 M 维度和 K 维度的坐标后，让它们各自和 M 维度、K 维度的 stride 相乘，然后和 a_ptr 相加，就可以得到 A 矩阵 9 个 block 中第一个 block 里每个元素的地址。b_ptr 也是同理。

最后一部分就是累加了，在 K 维度上进行累加，每次计算输出的一个块。

# 迭代以计算C矩阵的一个块。
# 我们将累加到一个`[BLOCK_SIZE_M, BLOCK_SIZE_N]`块
# 的fp32值以获得更高的精度。
# `accumulator`在循环后会转换回fp16。
accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
    # Load the next block of A and B, generate a mask by checking the K dimension.
    # If it is out of bounds, set it to 0.
    a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0)
    b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
    # We accumulate along the K dimension.
    accumulator += tl.dot(a, b)
    # Advance the ptrs to the next K block.
    a_ptrs += BLOCK_SIZE_K * stride_ak
    b_ptrs += BLOCK_SIZE_K * stride_bk

这行代码 a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0) 考虑到 K 可能不能被 BLOCK_SIZE_K 整除，到每一行最后一个 block 的时候，实际大小不足 BLOCK_SIZE_K，所以需要把超出的那部分元素 mask 掉。

最后这部分代码把当前的算子和 LeakyReLU 激活函数进行融合：

# 当累加器仍然是FP32时，可以融合任意激活函数
if ACTIVATION == "leaky_relu":
    accumulator = leaky_relu(accumulator)
c = accumulator.to(tl.float16)

单元测试

基准测试

这里使用一个方阵来对比 Triton 实现的 matmul kernel 和 cuBLAS 的 matmul kernel 的性能。

@triton.testing.perf_report(
    triton.testing.Benchmark(
        x_names=['M', 'N', 'K'],  # 用作图表x轴的参数名
        x_vals=[128 * i for i in range(2, 33)],  # `x_name`的不同可能值
        line_arg='provider',  # 其值对应于图表中不同线条的参数名
        # `line_arg`的可能值
        line_vals=['cublas', 'triton'],
        # 线条的标签名称
        line_names=["cuBLAS", "Triton"],
        # 线条样式
        styles=[('green', '-'), ('blue', '-')],
        ylabel="TFLOPS",  # y轴的标签名称
        plot_name="matmul-performance",  # 图表的名称，也用作保存图表的文件名。
        args={},  # 其他参数
    ))
def benchmark(M, N, K, provider):
    # 初始化张量
    a = torch.randn((M, K), device='cuda', dtype=torch.float16)
    b = torch.randn((K, N), device='cuda', dtype=torch.float16)
    quantiles = [0.5, 0.2, 0.8]  # 分位数
    # 如果提供者是cublas
    if provider == 'cublas':
        ms, min_ms, max_ms = triton.testing.do_bench(lambda: torch.matmul(a, b), quantiles=quantiles)
    # 如果提供者是triton
    if provider == 'triton':
        ms, min_ms, max_ms = triton.testing.do_bench(lambda: matmul(a, b), quantiles=quantiles)
    # 性能计算函数
    perf = lambda ms: 2 * M * N * K * 1e-12 / (ms * 1e-3)
    return perf(ms), perf(max_ms), perf(min_ms)

# 运行基准测试，展示图表和打印数据
benchmark.run(show_plots=True, print_data=True)

可以看到，基于 Triton 实现的矩阵乘 kernel 性能大体可以和高度优化的 cuBLAS 持平。