翻译自 ：https://huggingface.co/spaces/nanotron/ultrascale-playbook

作者：nanotron

校正：pprp

GPU 深度挖掘——融合、线程化、混合

截至目前，我们的讨论主要集中在模型操作的high-level组织结构上。我们已经在不同加速器上关注了计算，同时考虑到一般内存限制和计算单元的高级调度。

但这忽略了我们可以在更低层次上通过仔细理解我们的模型操作如何在每个GPU上调度和执行来做的所有优化。

本节将深入介绍GPU架构的更多细节，特别是NVIDIA的GPU架构，但通常的想法可以在类似的加速器单元上重复使用。

在覆盖Flash-Attention革命如何有效调度GPU工作负载之前，我们将简要解释GPU的组织方式，并最终解释如何在GPU上有效使用各种精度。

GPU入门

通常，GPU 具有非常层次化的组织结构。在本指南中，我们将保持讨论在支撑我们后续展示所需的概念层面。

(1)在计算方面 ，GPU由一组称为 流多处理器 Streaming Multiprocessors （ SM ）的计算单元组成并控制。每个SM包含并控制一组流处理器，也称为 核心 Cores 。例如，Nvidia H100 GPU具有132个SM，每个SM有128个核心，总共有16,896个核心（有关张量核心的详细信息，请参见 张量核心文档 [1]），每个核心可以同时处理多个线程 Thread。

编者注：计算分层概念： SM → Core → Thread 实际编程概念：SM → Grid → Block → Thread Warps(线程束）→ Thread

(2)内存方面也高度分层 ，具有多层缓存和内存： 寄存器 Registers 是最小的单位，在执行过程中是私有的， 共享内存 Shared Memory 和 L1 Cache 在单个SM上运行的线程之间共享，更高层次是所有SM共享的 L2缓存 Cache ，最后是 全局内存 Global Memory ，这是GPU上最大的内存（例如H100的80GB），但访问和查询速度也是最慢的。

编者注：内存分层：Global Mem → L2 Cache → L1 Cache → Shared Mem

GPU的目标是 通过利用计算/内存的这种分层组织，尽可能并行地在GPU核心上运行尽可能多的工作负载 。

在GPU核心上运行的代码片段称为 内核 Kernel 。它可以在高级别上用 CUDA或Triton 等语言编写，然后编译为NVIDIA GPU使用的低级汇编 Parallel Thread Execution（PTX） 。

要运行内核，你还需要一个特定的代码部分，称为 主机代码 Host Code ，它在CPU/主机上执行，并负责准备数据分配和加载数据和代码。

ounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(line// Host code                void vecAdd(float* h_A, float *h_B, float *h_c, int n) {    // Allocate vectors in device memory    int size = n * sizeof(float);    float *d_A, *d_B, *d_C;    cudaMalloc(&d_A, size);    cudaMalloc(&d_B, size);    cudaMalloc(&d_C, size);
    // Copy vectors from host memory to device memory    cudaMemcpy(d_A, h_A, size, cudaMemcpyHostToDevice);    cudaMemcpy(d_B, h_B, size, cudaMemcpyHostToDevice);
    // Invoke kernel    int threadsPerBlock = 256;    int blocksPerGrid =            (N + threadsPerBlock - 1) / threadsPerBlock;    VecAdd<<>>(d_A, d_B, d_C, N);
    // Copy result from device memory to host memory    // h_C contains the result in host memory    cudaMemcpy(h_C, d_C, size, cudaMemcpyDeviceToHost);
    // Free device memory    cudaFree(d_A);    cudaFree(d_B);




    
    cudaFree(d_C);}

内核通常按如下方式调度：

• 线程被分组成大小为32的 线程束（warps） 。线程束中的所有线程被同步以同时执行指令，但在数据的不同部分上。
• 线程束 被分组成更大的 块（blocks） ，大小更灵活（例如大小为256），每个块仍然分配给单个SM。一个SM可以并行运行多个块，但是根据资源情况，并非所有块都会立即分配执行，有些可能会等待资源。

从这些细节中最重要的是记住，有各种大小和分配约束（各种内存的大小，每个线程束和块中的线程数），需要考虑使用GPU架构的最有效方式。

大多数情况下，你不需要这么精确，幸运的是，你可以重用社区其他成员准备的内核和代码。但无论如何，我们希望为你提供有关如何开始使用内核的入门指南！

How to improve performance with Kernels ?

如果你想添加一个缺少优化过的内核的新操作或加快现有的 PyTorch 函数，从头编写内核可能看起来是最直接的方法。然而，从头创建高性能的 CUDA 内核需要丰富的经验和陡峭的学习曲线。通常，更好的入门方法是利用 torch.compile ，它通过捕获你的操作并在 triton 中生成低级、高性能内核来动态优化 PyTorch 代码。

假设你想编写一个名为指数线性单元 ELU 的激活函数的内核：

你可以从一个简单的 PyTorch 实现开始，然后只需在顶部添加 @torch.compile 装饰器即可：

ounter(lineounter(lineounter(line@torch.compiledef elu(x, alpha=




    
1.0):    return torch.where(x < 0, alpha * (torch.exp(x) - 1), x)

编译版和非编译版之间的区别非常明显，尤其是在只添加了一个装饰器的情况下。这种显著的不同在下图中得到了体现（N 表示列数）：

然而，如果这种性能提升不足，你可以考虑实现 Triton 内核。作为一个起点，你可以查看由@torch.compile 生成的 triton 内核。为此，你只需设置环境变量 TORCH_LOGS 为 "output_code" ：

ounter(lineexport TORCH_LOGS="output_code"

一旦运行带有 @torch.compile 装饰器的 Python 脚本，它将生成并输出相应的 Triton 内核，在这种情况下，是：

ounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(line@triton.jitdef triton_(in_ptr0, out_ptr0, xnumel, XBLOCK : tl.constexpr):    xnumel = 100000000    xoffset = tl.program_id(0) * XBLOCK    xindex = xoffset + tl.arange(0, XBLOCK)[:]    xmask = xindex < xnumel    x0 = xindex    tmp0 = tl.load(in_ptr0 + (x0), xmask)    tmp1 = 0.0    tmp2 = tmp0 < tmp1    tmp3 = tl_math.exp(tmp0)    tmp4 = 1.0    tmp5 = tmp3 - tmp4    tmp6 = tl.where(tmp2, tmp5, tmp0)    tl.store(out_ptr0 + (x0), tmp6, xmask)

为了提高可读性，我们可以修改变量名，添加注释，并进行轻微调整（或者请求LLM为我们完成），如下所示：

ounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(line@triton.jitdef elu_kernel(input_ptr, output_ptr, num_elements, BLOCK_SIZE: tl.constexpr):    # Calculate the starting index for this block    block_start = tl.program_id(0) * BLOCK_SIZE    # Create an array of indices for this block    block_indices = block_start + tl.arange(0, BLOCK_SIZE)[:]    # Create a mask to ensure only valid indices are processed    valid_mask = block_indices < num_elements    # Load input values from the input pointer based on valid indices    input_values = tl.load(input_ptr + block_indices, valid_mask)    # Define the ELU parameters    zero_value = 0.0  # Threshold for ELU activation    negative_mask = input_values < zero_value    exp_values = tl.math.exp(input_values)    # Define the ELU output shift    one_value = 1.0    shifted_exp_values = exp_values - one_value
    output_values = tl.where(negative_mask, shifted_exp_values, input_values)
    # Store the computed output values back to the output pointer    tl.store(output_ptr + block_indices, output_values, valid_mask)

此处， tl.program_id(0) 提供一个唯一的 Block ID，我们用它来确定该块将处理哪个数据部分。使用此 Block ID， block_start 计算每个块的起始索引，而 block_indices 指定该部分内的索引范围。 valid_mask 确保仅处理 num_elements 内的索引，安全地使用 tl.load 加载数据。然后应用 ELU 函数，根据数值是否为负修改值，并将结果使用 tl.store 写回内存。

当使用 triton.testing.Benchmark 对生成的内核进行基准测试时，其性能如下：

这个独立的内核在较小规模下甚至表现出比 @torch.compile 更优的性能，但这可能仅仅是 torch.compile 的编译时间影响所致。无论如何，与其从零开始，不如从这些生成的内核出发，并将精力集中在优化其性能上，从而节省大量时间。

即使在 Triton 中，有时也无法完全达到设备的峰值性能，因为该语言在处理共享内存和流多处理器（SMs）内的调度等低级细节方面存在限制。 Triton 的能力仅限于块及其在 SMs 之间的调度 。为了获得更深入的控制，你需要直接在 CUDA 中实现内核，在那里你将能够访问所有底层低级细节。

CUDA 方面，可以采用各种技术来提高内核的效率。这里仅介绍其中几个： 优化内存访问模式以降低延迟 、 使用共享内存存储频繁访问的数据 以及 管理线程工作负载以最小化空闲时间。

在深入 CUDA 示例之前，总结一下看到的工具，这些工具使我们能够编写内核代码以在 GPU 上执行指令：

1. PyTorch：简单但速度较慢
2. torch.compile：简单且速度快，但灵活性不足
3. Triton：更难，但更快、更灵活
4. CUDA：最难，但最快、最灵活（如果掌握得当）

下面讨论 CUDA 中最常见的优化技术之一： 优化内存访问 。GPU 的全局内存（在前面的图表中是最大的内存）相比缓存来说，延迟较高，带宽较低，这通常是大多数应用程序的主要瓶颈。 高效地访问全局内存的数据 可以极大地提高性能。

内存合并

为了有效利用全局内存的带宽，理解其架构至关重要。在CUDA设备中，全局内存是使用DRAM实现的。

内存归约（Memory coalescing） 利用 DRAM 在访问内存地址时 以突发或连续内存位置范围的形式提供数据 的特点。每次访问 DRAM 位置时，包括请求的位置在内的连续位置序列由 DRAM 芯片中的多个传感器并行读取。一旦读取，这些数据可以快速传输到处理器。在 CUDA 中， 归约 coalescing 利用这种突发行为，通过确保 warp 中的线程（32 个执行相同指令的线程，SIMD）访问连续的内存位置，以最大化内存访问效率。

例如，如果线程 0 访问位置 M，线程 1 访问 M + 1，线程 2 访问 M + 2，依此类推，GPU 硬件将这些请求 归约或合并 为一个大型、高效的 DRAM 突发访问请求，而不是单独处理每个访问。

以矩阵乘法为例。一个简单直接的实现方式是，每个线程计算输出矩阵的一个元素，如下：

ounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(line
__global__ void matmul_naive(int M, int N, int K, const float *A, const float *B, float *C) {    const uint x = blockIdx.x * blockDim.x + threadIdx.x;    const uint y = blockIdx.y * blockDim.y + threadIdx.y;
    if (x < M && y < N) {        float tmp = 0.0;        for (int i = 0; i < K; ++i) {            tmp += A[x * K + i] * B[i * N + y];        }        C[x * N + y] = tmp;    }}

这是一篇精彩博客文章[2]中内核的优秀可视化：

然而，当使用类似 ncu 的工具对内核进行性能分析时，可以看到问题，包括 低内存吞吐量 和 未归约的内存访问 。

原因在于，在此内核中，同一块中的两个线程（线程 ID 为 (0, 0) 和 (1, 0)，最终将位于同一 warp 中）将同时从矩阵 B 的同一列加载，但矩阵 A 的不同行。由于 矩阵元素按行主序存储 （意味着行元素位于连续的内存地址中，如图所示），线程 (0, 0) 将在第一次迭代 中加载 ，而线程 (1, 0) 将加载 。这些元素在内存中并不相邻，这种错位将在每次迭代中存在，从而防止内存访问归约。

为了提高我们内核的性能，我们可以改变坐标 x 和 y 的计算方式，如下所示：

ounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(line
const int x = blockIdx.x * BLOCKSIZE + (threadIdx.x / BLOCKSIZE);const int y = blockIdx.y * BLOCKSIZE + (threadIdx.x % BLOCKSIZE);
if (x < M && y < N) {	float tmp = 0.0;	for (int i = 0; i < K; ++i) {	    tmp += A[x * K + i] * B[i * N + y];		}	 C[x * N + y] = tmp;}

而不是使用二维块，我们切换到一维块，并重新定义确定 x 和 y 值的方法。在这种新方法中，同一 warp（具有接近的 threadIdx.x 值）内的线程将共享相同的 x 值，但具有不同的 y 值。这意味着它们将加载矩阵 A 的同一行，但矩阵 B 的不同列。因此，可以合并行主序矩阵的内存访问。

当我们对新的内核进行性能分析时，注意到关于未归约内存访问的警告已经消失，GPU 的内存吞吐量大约提高了 10 倍。

内核的执行时间降低了 10 倍！惊人。

现在让我们介绍另一种在文献中经常提到的技术： 分块Tiling 。

分块处理（Tiling）

分块处理 是一种利用 共享内存 Shared Memory 优化内存访问模式的技术。正如我们前面提到的，共享内存是一种小而快速的存储，块内的所有线程都可以访问它。这使得 数据可以被多个线程重复使用 ，从而减 少了从较慢的全局内存中重复加载数据的需求 。

以矩阵乘法为例，块中的每个线程可能需要从两个矩阵（如 A 和 B）中获取元素。如果每个线程 独立地从全局内存加载所需的行和列 ，就会出现大量冗余加载，因为块中的多个线程会访问重叠的数据。相反，我们可以使用分块处理 Tiling，将 A 和 B 的一个块（或 Tile）一次性加载到共享内存中，让该块中的所有线程重复使用相同的共享数据。

在分块处理的方法中，每次迭代时，块内的所有线程协同工作，将两个 Tile（一个来自矩阵 A，另一个来自矩阵 B）加载到共享内存中。具体来说，线程加载矩阵 A 的一个Tile（大小为 BLOCK_SIZE_M × BLOCK_SIZE_K ）以及矩阵 B 的一个Tile（大小为 BLOCK_SIZE_K × BLOCK_SIZE_N ）。一旦这些Tile存入共享内存，线程就可以在这些Tile上执行矩阵乘法，从而实现高效计算，因为所有必要的数据都可以被快速访问。Tile乘法的结果存储在一个累积矩阵中，该矩阵保存中间结果。在每次迭代后，当前Tile乘法的结果都会累加到该矩阵中，直到两个矩阵的所有Tile都被处理完毕。

让我们来看看实现中的关键部分：

ounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(lineounter(line
// Set pointers to the starting elementsA += blockRow * TILE_SIZE * K; // Start at row = blockRow, column = 0B += blockCol * TILE_SIZE; // Start at row = 0, column = blockColC += blockRow * TILE_SIZE * N + blockCol * TILE_SIZE; // Start at row = blockRow, column = blockColfloat sum = 0.0;// The outer loop moves through tiles of A (across columns) and B (down rows)for (int tileIdx = 0; tileIdx < K; tileIdx += TILE_SIZE) {sharedA[localRow * TILE_SIZE + localCol] = A[localRow * K + localCol];sharedB[localRow * TILE_SIZE + localCol] = B[localRow * N + localCol];
// Ensure all threads in the block have completed data loading__syncthreads();
// Shift pointers to the next tileA += TILE_SIZE;B += TILE_SIZE * N;
// Compute the partial dot product for this tilefor (int i = 0; i < TILE_SIZE; ++i) {    sum += sharedA[localRow * TILE_SIZE + i] * sharedB[i * TILE_SIZE + localCol];}// Synchronize again to prevent any thread from loading new data// into shared memory before others have completed their calculations__syncthreads();}C[localRow * N + localCol] = sum;

每个线程首先从 矩阵 A 和 矩阵 B 中加载一个元素到共享内存。在这种情况下，实现合并内存访问（coalesced memory access）非常直观：通过将 threadIdx.x 作为 局部列索引（localCol） ，同一个 warp 中的线程可以访问相邻的矩阵元素。块内所有线程完成数据加载后（通过调用 __syncthreads() 确保同步），它们就会计算这两个Tile的点积。当所有Tile遍历完成—— 矩阵 A 在水平方向移动， 矩阵 B 在垂直方向移动——最终计算出的结果存入 矩阵 C 的对应位置。

当我们使用 ncu 对这个内核进行基准测试时，我们发现内存吞吐量增加到了 410 Gb/s，内核执行时间减少了约 43%，实现了约 6.6 TFLOPs 的性能。

线程粗化（Thread Coarsening）

分块处理技术显著提高了我们内核的性能。但是，当分析量化每个状态中花费的周期的warp状态时，我们观察到以下情况：

这些神秘状态名称的含义可以在 NVidia的性能指南 [3]中找到，在 Warp Stall Reasons 部分可以阅读到：

*"smsp__pcsamp_warps_issue_stalled_mio_throttle : 等待MIO（内存输入/输出）指令队列不再满的Warp被停顿。在MIO管道（包括特殊数学指令、动态分支以及共享内存指令）极端利用的情况下，此停顿原因较高。当由共享内存访问引起时，尝试使用更少但更宽的加载可以减少管道压力。*

所以看起来Warp正在等待共享内存访问返回！为了解决这个问题，我们可以应用一种称为 线程粗化 Thread Coarsening 的技术，该技术涉及将 多个线程合并为一个粗化线程 。这将显著减少共享内存访问，因为每个粗化线程可以处理多个输出元素。

在写入或改进 自定义内核 时，一个最重要的考虑因素： 最小化控制分歧 Minimizing Control Divergence 。

最小化控制分歧

流多处理器（SM）被设计为使用 单指令多数据（SIMD） 模型执行 warp 中的所有线程。这意味着在任何给定时刻， 一个指令同时为warp中的所有线程获取和执行 。当执行warp时，其中的线程在数据的不同段上操作，但遵循相同的指令，因此得名 单指令多数据 。SIMD的主要优势在于其效率；负责指令获取和调度的控制硬件在多个执行单元之间共享。 这种设计最小化了与控制功能相关的硬件开销，使得更大比例的硬件专注于提高算术吞吐量 。

当warp内的线程采取不同的执行路径时，就会发生控制分歧 。例如，如果条件语句（如 if 语句）导致一些线程执行一个代码块，而其他线程执行另一个代码块，那么warp必须串行执行这些执行，导致空闲线程等待其他线程完成。为了最小化控制分歧，我们需要设计内核，**确保warp内的线程遵循相同的执行路径。这可以通过重构代码以减少分支、**使用确保所有线程遵循类似执行路径的数据结构，或使用预测等技术来实现。

编者注：简单理解为不要有if等判断语句

我们已经介绍了写入自定义内核和改进GPU操作性能和内存占用的一些主要考虑因素。但在转向实际示例之前，还有一个重要的概念需要讨论：“融合内核 Fused Kernel”。

融合内核（Fused Kernels）

之前提到GPU和CPU操作可以异步进行。特别是，CPU上的 Host Code 主机代码可以以 非阻塞方式调度 GPU的工作负载。

非阻塞对于重叠通信和计算非常有用——可以扩展到更一般的想法，即 尽量避免来回在主机和GPU内核命令之间切换 。

这个想法在 Horace He [4] 的这些图中得到了很好的诠释：

如何避免这种来回？最好的办法是尽可能让我们的 GPU 实现自主。这通过将尽可能多的连续计算操作打包在一个单独的内核中来实现，这个内核被称为“ 融合内核 Fused Kernel ”。

融合内核对于 独立于彼此在各个输入Tokens上执行的一系列点状操作 特别高效且易于编写。在这种情况下，在将计算值移动到 SM 内存并启动新内核之前，没有必要将计算值返回到全局内存。在完成计算序列之前，将所有值保留在本地要高效得多。

Transformer 模型中有许多地方可以应用这种“融合”方法：每次我们遇到一系列逐点point-wise操作，例如在层归一化计算中。

现在我们已经掌握了欣赏内核工程的真正杰作所必需的所有理解： Flash Attention

Flash Attention 1-3

Flash attention是由 Tri Dao [5] 引入，并提出通过编写自定义CUDA内核来优化注意力计算，使其更快且更内存高效。Flash Attention的核心思想是充分利用GPU的各种内存，避免过度依赖最慢的内存之一：GPU的全局内存。

编者注： 在Flash attention中，HBM - 高带宽内存 High band Memory 就是GPU全局内存。

注意机制的基本实现涉及在内存和worker之间进行大量传输。它要求在HBM中实现S和P矩阵，这意味着结果需要发送到HBM，然后再次发送到SRAM进行下一步计算：

由于HBM的带宽较低，这在注意力计算中引入了严重的瓶颈。关键元素是将S矩阵计算成可以适应SM较小共享内存的小块。但可以做得更好，不仅仅是分块计算S矩阵，而是 完全避免存储庞大的S矩阵 ， 仅保留计算Softmax归一化因子所需的统计信息 。这样，可以直接在SRAM中一次性计算部分 O ，而无需在中间结果之间来回传输数据。这不仅利用了共享内存，还消除了由于 存储注意力矩阵 （在长上下文长度下是模型中最大的激活矩阵之一）而导致的内存瓶颈。

Flash Attention 的理念解决了模型训练中的众多瓶颈，因此迅速成为所有Transformer模型执行注意力计算的默认方法：

• 通过避免存储S矩阵， 降低了注意力计算的内存负担
• 消除了大部分 注意力计算的平方复杂度(S²)所带来的影响

因此，自Transformer架构发明后不久发展出的所有线性注意力变体和次二次近似注意力方法大多被搁置，取而代之的是这种精准且快速的Flash Attention实现和机制。

在Flash Attention 1发布之后，同一实验室相继推出了两个改进版本：Flash Attention 2 和 3。与Flash Attention 1相比， Flash Attention 2 和 3 的改进更多体现在对GPU的底层优化 ，而不是对注意力机制本身的改动。具体来说：

• 减少非矩阵乘法（matmul）操作的数量
• 精细划分计算任务至warp和线程块 （适用于Flash Attention 2）
• 在最新的Hopper架构（H100）上 优化FP8和Tensor Core的支持 （适用于Flash Attention 3）

Flash Attention 是一个典型案例，展示了当深入考虑当前GPU加速器的内存/计算设计时，所能带来的突破性改进。

到目前为止，我们讨论的算子融合技术要求对模型代码进行改动，并为特定操作编写自定义内核，以加速训练。

在计算操作的底层优化的最后部分，我们将探索一系列与模型代码无关的方法，这些方法适用于任何模型，并且已经成为业界标准： 混合精度训练（Mixed Precision Training）！

混合精度训练（Mixed Precision Training）

在本书的多个章节中，我们讨论了低精度数值格式及其对存储激活值、参数和优化器状态的内存需求的影响。现在，我们将深入了解这些格式的细节，并更好地理解它们的权衡、优势和局限性。

顾名思义，混合精度训练涉及在训练过程中混合使用不同的数值精度。PyTorch张量的默认数值精度是 单精度浮点格式，即FP32 （float32），这意味着每个存储的数值占用32位（4字节）。这些位被分为三个部分：

• 符号位（Sign）：第一个比特决定数值是正数还是负数
• 尾数（Mantissa）：决定数值的有效数字
• 指数（Exponent）：控制数值的数量级

浮点数的基本原理可以通过科学计数法轻松理解，例如−5.734× ，其中首先是符号位，然后是尾数和指数。这样可以在广泛的数值范围内以自适应精度表示数值。虽然float32是默认格式，但PyTorch还支持多种浮点格式：