从模型计算到 GPU 执行:读懂算子、Kernel 与 GPU 性能

14795 字
37 分钟

从模型计算到 GPU 执行:读懂算子、Kernel 与 GPU 性能

发布于

你不要那样看着我的眼睛 我就怕你会看穿我的心

从模型计算到 GPU 执行:读懂算子、Kernel 与 GPU 性能

这篇文章想解决一个很具体的问题:模型中的一行数学公式,最后是怎样变成 GPU 上真正执行的程序的?

如果你已经看过 Transformer 和 Attention 公式,却经常把模型架构、算法、算子、Kernel 和 GPU 混在一起,那么可以从这篇文章开始重新建立一条清晰的认知链路。

开始之前:先给整篇文章画一张地图

在大模型推理中,我们通常从模型结构开始思考:这里有 Attention,那里有 MLP,还有归一化和残差连接。但 GPU 并不直接理解“这是一个 Transformer”或“这是一次 Attention”。一层层往下看,同一件计算会经过几个不同的描述层次:

模型架构 → 算法 → 算子 → GPU Kernel → GPU 硬件

这是一张帮助理解的教学地图,不是严格的软件调用栈。真实执行过程中,还会经过框架、编译器、运行时和驱动程序;而且算法、算子与 Kernel 之间也不是固定的一对一关系。

从模型架构到 GPU 硬件的五层分层图
从模型架构到 GPU 硬件的五层分层图

可以先用一句话记住每一层的职责:

  • 模型架构决定系统由哪些模块组成,例如 Attention、MLP、Norm 和残差连接。
  • 算法决定要做什么计算,以及如何组织这些计算,例如普通的 Dense Attention 或 FlashAttention。
  • 算子描述一个明确的张量计算功能,例如 MatMul、Softmax 和 RMSNorm。
  • GPU Kernel是一段在 GPU 上并行执行的程序,规定线程怎样加载数据、计算结果并写回数据。
  • GPU 硬件是真实存在的执行设备,包括 SM、CUDA Core、Tensor Core 和多级存储系统。

这五个概念之所以容易混淆,是因为它们描述的是同一件事情的不同层次。接下来用一次 LLM 推理和 Attention 公式,把它们逐层拆开。

第 1 章:先分清五个概念

1.1 从一次 LLM 推理开始

假设我们输入一句话,希望模型生成下一个 Token。一次简化的推理过程可以表示为:

文本
  ↓
Tokenize
  ↓
Embedding
  ↓
多个 Transformer Block
  ↓
LM Head
  ↓
下一个 Token 的概率分布

Tokenize 把文本转换为模型能够处理的 Token ID;Embedding 把这些离散的 ID 映射成向量;之后,向量会依次经过多个 Transformer Block,最后由输出层产生下一个 Token 的概率分布。

我们只展开其中一个 Transformer Block。为了便于理解,可以把它看成两类主要工作:

  1. Attention:让一个位置的表示读取其他位置的信息。
  2. MLP 或 FFN:对每个位置的表示进行逐位置的变换和加工。

实际模型还会包含归一化、残差连接等操作,不同模型在模块顺序、FFN 形式和 Attention 细节上也可能不同。这里的目的不是完整复述模型结构,而是观察:模型中的一个模块,最终会由许多具体的张量计算组成。

以一个 Linear 层为例,它的数学表达式可以写成:

Y = XW + b

这里要把几件事分开:

  • X 是输入张量;
  • W 和 b 是训练得到的模型参数;
  • XW + b 是要完成的数学计算;
  • 矩阵乘法和加法分别可以由 MatMul、Add 等算子描述;

其中有一个容易被忽略的区别:模型权重是数据,Kernel 是程序。 加载模型时,加载的是训练得到的参数;运行推理时,还需要有相应的计算程序来读取这些参数、完成计算并产生输出。

1.2 用 Attention 拆出算子

Attention 的核心公式是:

Attention(Q, K, V) = softmax(QKᵀ / √d) V

从数学上看,这是一条公式;从程序实现上看,它可以拆成几个更明确的计算步骤:

Q、K、V
  ↓
MatMul:S = QKᵀ
  ↓
Scale:S = S / √d
  ↓
Softmax:P = softmax(S)
  ↓
MatMul:O = PV
  ↓
输出 O

这里至少出现了三类算子:

  • MatMul:完成矩阵乘法;
  • Scale:对结果进行缩放;
  • Softmax:把分数转换为归一化的权重。

因此,算子可以先理解为:一个具有明确输入、输出和计算规则的功能单元。

例如,下面这些代码都可以对应到具体算子:

代码或数学表达式算子计算含义
a + bAdd逐元素相加
a @ bMatMul矩阵乘法
softmax(x)Softmax计算归一化概率
sum(x)ReduceSum对元素求和
relu(x)ReLU将负数置零

需要注意的是,算子不一定小到只有一次加法或乘法。上面的 Attention 既可以拆成多个基本算子,也可以在更高层被看成一个复合算子。

算子描述的是“要完成什么计算”,并不强制规定“底层必须怎样执行”。

1.3 算子与 Kernel 不是一对一关系

算子描述计算任务,但 GPU 需要执行的是具体程序。这个程序就是 Kernel。

还是以加法为例。对于:

C = A + B

如果 A 和 B 各有一百万个元素,可以让许多 GPU 线程并行处理不同位置:

Thread 0:C[0] = A[0] + B[0]
Thread 1:C[1] = A[1] + B[1]
Thread 2:C[2] = A[2] + B[2]
...

Add 是对计算任务的描述;add_kernel 则是一个规定“线程如何找到自己的位置、读取哪些数据、完成什么计算并写回结果”的 GPU 程序。现在不需要记住 CUDA 语法,只需要抓住这个区别:

算子回答“算什么”,Kernel 回答“让 GPU 怎样算”。

它们之间常见的关系可以浓缩成三种:

  • 简单的 Add 算子通常对应一个 Kernel;
  • 复杂算子可能分成多个阶段,由多个 Kernel 共同完成;
  • MatMul、缩放和 Softmax 等相邻步骤也可以融合进一个 Kernel,减少中间结果的显存读写。
算子与 Kernel 的三种对应关系
算子与 Kernel 的三种对应关系

这也是理解 FlashAttention 的一个重要入口:它不只是把公式换一种写法,还会重新组织计算和数据流,尽量让中间结果留在更靠近计算单元的位置,减少不必要的显存读写。这里关注的是 Kernel 和数据流的组织方式,而不是把某个高层算法名称简单等同于一个 Kernel。

1.4 从 Kernel 到硬件还隔着什么

Kernel 也不是直接从 Python 代码跳到 CUDA Core 上执行的。中间还存在编程工具、编译器、运行时和驱动程序等环节。可以用下面这条简化路径理解:

算子实现
  ↓
CUDA C++ / Triton Kernel
  ↓
编译器
  ↓
PTX / SASS 等底层指令表示
  ↓
Runtime / Driver 提交执行
  ↓
GPU 上的 SM 调度并执行线程

这些概念的职责不同:

概念所属层次作用
CUDA C++GPU 编程工具编写 GPU 程序
Triton编程语言与编译工具以更高层的方式编写 Kernel
PTX / SASS指令表示描述编译后的底层执行形式
Runtime / Driver软件运行环境管理显存并提交 Kernel 启动请求
SMGPU 硬件调度并执行 GPU 线程
CUDA CoreGPU 硬件执行通用计算中的一部分
Tensor CoreGPU 硬件加速高吞吐矩阵运算
HBM / GDDRGPU 显存保存权重和张量等数据

当 PyTorch 执行 torch.matmul(A, B) 时,用户看到的是框架层面的算子调用。框架会结合设备、数据类型、张量形状和布局等条件,选择合适的底层实现;在 NVIDIA GPU 上,某些矩阵乘法可能由 CUDA 生态中的高性能库来完成。最终,硬件真正执行的仍然是编译后的 GPU 指令,以及这些指令对应的数据读取、计算和写回。

第 2 章:GPU 的硬件结构

上一章的最后一层是“GPU 硬件”,但只列了几个名字。这一章把 GPU 打开来看:一块 GPU 由哪些部件组成,SM、CUDA Core、Tensor Core 和显存之间是什么关系。

先说结论:GPU 由很多个 SM 组成。每个 SM 都有自己的调度器、两类计算单元和片上存储;显存在芯片外,离计算单元最远。 下面按“整块芯片 → 一个 SM → 存储位置”的顺序展开。

GPU 的硬件结构:从整块芯片放大到一个 SM
GPU 的硬件结构:从整块芯片放大到一个 SM

2.1 先看整块芯片:很多个 SM 重复排列

SM 是 Streaming Multiprocessor 的缩写,中文通常叫流式多处理器。刚开始可以把它理解为 GPU 内部的一个计算集群:它有自己的执行单元、调度器和片上存储,可以调度和执行一批线程。

一块 GPU 上有多个 SM,它们结构相同,像格子一样重复排列。所有 SM 共用芯片上的 L2 Cache。芯片外面连着显存,通常是 HBM 或 GDDR,模型权重、输入输出张量和各种中间结果都存放在这里。

GPU 擅长大规模并行计算,原因就在这种结构上。CPU 也能并行执行任务,但 GPU 拥有更多适合并行计算的执行资源:芯片上有多个 SM,每个 SM 又能同时驻留和调度很多线程。一个 Kernel 启动后,它的工作会被分给许多 SM 同时处理。

用一款真实 GPU 建立数量感

只说“多个 SM”还比较抽象。以 NVIDIA H100 SXM5 为例:

项目H100 SXM5
SM 数量132 个
每个 SM 的 FP32 CUDA Core128 个,整卡 16,896 个
每个 SM 的 Tensor Core4 个(第四代),整卡 528 个
每个 SM 的寄存器文件256 KB
每个 SM 的 Shared Memory + L1合计 256 KB
L2 Cache50 MB
显存80 GB HBM3

不同架构差异很大。上一代 A100(SXM4)有 108 个 SM,每个 SM 有 64 个 FP32 CUDA Core。这些数字只用来建立数量感,后文的内容不依赖具体型号。

2.2 放大一个 SM:里面有什么

把其中一个 SM 放大,可以看到几类部件:

部件作用
Warp Scheduler挑选已经准备好的 Warp,发出它的下一条指令
CUDA Core执行通用的标量运算,例如索引、地址计算和部分逐元素运算
Tensor Core执行高吞吐的矩阵乘加,适合 Linear、MatMul
寄存器文件分给每个线程私用,离计算单元最近
Shared Memory / L1SM 内的片上存储,同一个 Block 的线程可以在 Shared Memory 中共享数据

Warp 和 Block 这两个词这里先出现名字,下一章再解释。现在只需要知道:Warp 是一组一起被调度的线程,Block 是一组可以协作的线程。

这几类部件可以按职责分成三组:

  1. 调度:Warp Scheduler 决定下一步执行哪些线程的指令。
  2. 计算:CUDA Core 和 Tensor Core 真正完成运算。
  3. 存储:寄存器、Shared Memory 和 L1 把数据放在离计算单元很近的地方。

寄存器和 Shared Memory 的容量都有限,而且由 SM 上驻留的所有线程分着用。一个线程占用的寄存器越多,同一个 SM 能同时驻留的线程就越少。第 3 章讲 Block 怎样分配到 SM 时,还会用到这一点。

另外,在部分架构中,L1 和 Shared Memory 共用同一块物理存储。二者的区别在管理方式:Shared Memory 由程序决定放什么,L1 Cache 由硬件自动管理。

2.3 CUDA Core 与 Tensor Core 的分工

一个 SM 里有两类计算单元,它们擅长的事情不同:

硬件主要作用在 LLM 中的典型场景
CUDA Core通用的标量运算索引、地址计算、部分逐元素运算
Tensor Core高吞吐的矩阵乘加Linear、MatMul、部分 Attention 计算

用 Attention 中的 S = QKᵀ 举例。这一步的主体是矩阵乘法,如果矩阵形状、数据类型和执行方式合适,GPU 可以用 Tensor Core 加速它。但同一个 Kernel 里还有很多普通工作,例如计算每个线程该读哪个地址、判断有没有越界、对结果做缩放,这些通常由 CUDA Core 完成。

需要注意的是,调用 MatMul 不等于一定用满了 Tensor Core。数据类型、矩阵形状、内存布局和底层 Kernel 的实现都会影响结果。另一方面,如果一个 Kernel 的主要时间花在大量不规则的数据读取上,瓶颈就不在矩阵计算,Tensor Core 再强也帮不上忙。这也是为什么只了解算法还不够,还要了解硬件。

2.4 存储放在哪里:由近到远

把上面提到的存储按离计算单元的远近排开,可以得到一条简化的层次:

存储物理位置主要特点
寄存器SM 内部很快,容量有限,通常由单个线程使用
Shared MemorySM 内部同一个 Block 内的线程协作和复用数据
L1 CacheSM 内部或附近由硬件自动缓存近期访问的数据
L2 CacheGPU 芯片上所有 SM 共享
HBM / GDDR 显存GPU 芯片外容量最大,访问成本最高

一般来说,离计算单元越近的存储越快、容量越小;越远的越慢、容量越大。模型权重和张量平时放在容量最大的显存里,计算单元要用它们,就得把数据一路搬到 SM 内部。

这一节只回答“数据放在哪”。数据从显存搬到计算单元要付出多少代价,以及怎样减少搬运,是第 4 章的内容。

2.5 哪些是硬件,哪些不是

读 GPU 资料时,软件概念和硬件部件经常混在一起讲。本章讲的都是芯片上真实存在的东西:

  • 硬件:SM、Warp Scheduler、CUDA Core、Tensor Core、寄存器、Shared Memory、L1 / L2 Cache、HBM / GDDR 显存。

Shared Memory 是真实存在的硬件存储,只是由程序决定往里放什么、怎样使用,并不是软件模拟出来的内存。

程序里写的 Thread、Block、Grid 不在这张硬件图上。它们是软件对工作的分组方式,芯片上并没有一个固定对应某个 Thread 的“线程核心”,也没有一个专门对应某个 Block 的部件。

第 3 章:GPU 怎样把一个 Kernel 跑起来

上一章看的是 GPU 里有什么硬件。这一章换到程序这一侧:Kernel 里写的成千上万个线程,是怎样被分组的,又是怎样落到 SM 上执行的。

先说结论:程序把工作分成 Thread、Warp、Block 和 Grid 几层;GPU 把 Block 分给 SM,SM 再以 Warp 为单位调度指令,交给计算单元执行。 下面从最小的 Thread 开始,一层层往上看。

3.1 Thread:一个线程负责一小块工作

还是用最简单的数组加法:

C = A + B

假设 A 和 B 各有 100 万个元素。最直接的做法,是让一个线程负责一个元素:

Thread 0:     C[0] = A[0] + B[0]
Thread 1:     C[1] = A[1] + B[1]
Thread 2:     C[2] = A[2] + B[2]
...
Thread 999999:C[999999] = A[999999] + B[999999]

第 1 章的 add_kernel 做的就是这件事。每个线程执行同一段程序,只是根据自己的编号去找自己负责的位置。这里的线程是程序定义的逻辑执行单位,一个线程也可以负责一组元素,不一定只算一个。

有一点要先说清:线程总数不等于同一时刻真正在执行的线程数。 100 万个线程不会同时跑起来,GPU 会把它们分批调度执行。具体怎么分批,就要看下面几层分组。

3.2 Warp:32 个线程一起执行

在 NVIDIA GPU 上,通常每 32 个线程组成一个 Warp。**同一个 Warp 里的线程一起执行同一条指令,但各自处理自己的数据。**例如 32 个线程同时执行“加法”这条指令,线程 0 算 C[0],线程 1 算 C[1],以此类推。

这种执行方式叫 SIMT(Single Instruction, Multiple Threads),意思是一条指令、多个线程。可以这样理解:

  • 多个线程执行相同的程序逻辑;
  • 每个线程处理自己的数据;
  • Warp 是 GPU 调度和执行线程的重要单位。

上一章 SM 里的 Warp Scheduler,调度的正是 Warp:它挑出一个已经准备好的 Warp,发出这个 Warp 的下一条指令,再由 CUDA Core 或 Tensor Core 等执行单元完成运算。

分支发散:同一个 Warp 走了不同的路

既然一个 Warp 里的线程一起执行同一条指令,那如果程序里有 if,而线程们的判断结果不一样,会怎样?

if (x[i] > 0) {
    分支 A
} else {
    分支 B
}

假设线程 0–15 走分支 A,线程 16–31 走分支 B。两条分支没法在同一步里一起执行,只能先后进行:先让走分支 A 的线程执行,另一半线程等着;再让走分支 B 的线程执行,前一半线程等着。原本一步能干完的事,变成了两步,而且每一步都有一半线程闲着。

Warp Divergence:同一个 Warp 的线程走向不同分支
Warp Divergence:同一个 Warp 的线程走向不同分支

这种现象叫 Warp Divergence(分支发散),它会降低执行效率。不过影响有多大,要看分支的长短和 Kernel 的具体写法,并不是只要写了 if 就一定会很慢。另外,只有同一个 Warp 内部的线程走向不同分支时才会这样。

3.3 Thread Block 与 Grid:再往上分组

Warp 之上还有两层分组。

Thread Block(线程块) 是一组可以协作的线程。例如,一个 Block 可以包含 128 个线程,也就是 4 个 Warp:

Block 0
├── Warp 0:Thread 0~31
├── Warp 1:Thread 32~63
├── Warp 2:Thread 64~95
└── Warp 3:Thread 96~127

Block 的作用不只是分组。同一个 Block 里的线程可以协作完成工作,例如:

  • 共享 SM 里 Shared Memory 中的数据;
  • 通过同步操作协调执行进度;
  • 分工完成矩阵中的一小块计算。

最后一点在第 4 章讲 Tiling 时会用到。

Grid 是一次 Kernel 启动产生的全部 Block。如果数组加法要用 100 万个线程,每个 Block 放 128 个线程,那么这次启动的 Grid 就包含大约 7813 个 Block(1,000,000 ÷ 128 向上取整)。

把几层放在一起,嵌套关系是:

一次 Kernel 启动
└── Grid:全部 Block
    └── Block:一组可以协作的线程
        └── Warp:32 个线程,一起执行同一条指令
            └── Thread:负责一个或一组元素

3.4 Block 怎样被分配到 SM

Kernel 启动后,GPU 会把 Grid 里的 Block 分配到有空闲资源的 SM 上执行。

Block 怎样被分配到 SM
Block 怎样被分配到 SM

这个分配过程有三个要点。

第一,一个 Block 在执行期间只驻留在一个 SM 上。 同一个 Block 的线程要共享 Shared Memory、互相同步,所以它们必须待在同一个 SM 里。

第二,一个 SM 可以同时驻留多个 Block。 能放几个,取决于每个 Block 占用多少资源:线程数、寄存器和 Shared Memory。上一章提到,寄存器和 Shared Memory 的容量有限,由驻留在 SM 上的所有线程分着用。一个 Block 占得越多,同一个 SM 能同时放下的 Block 就越少。

第三,Block 的数量通常远多于 SM。 放不下的 Block 先排队,等某个 SM 上的 Block 执行完、腾出资源,再分配上去。前面数组加法的例子有大约 7813 个 Block,它们会被分成很多批,陆续交给各个 SM。这也说明了 3.1 节的那句话:线程总数不等于同一时刻正在执行的线程数。

所以,Block 和 SM 不是一一对应的。 Block 是程序组织工作的方式,SM 是执行工作的硬件。

3.5 软件概念与硬件概念对照

到这里,可以把第 2 章和第 3 章出现的概念放在一起分类:

概念属于什么解释
Kernel软件GPU 执行的程序
Thread软件抽象程序定义的逻辑执行线程
Thread Block软件抽象一组可以协作的线程
Grid软件抽象一次 Kernel 启动产生的全部 Block
Warp软硬件结合NVIDIA GPU 上 32 个线程组成的执行分组
SM硬件GPU 内的计算与调度单元
CUDA Core硬件执行标量运算的计算单元
Tensor Core硬件加速矩阵乘加的计算单元
寄存器、Shared Memory硬件SM 内的片上存储
L1 / L2 Cache、HBM / GDDR硬件缓存和显存

Warp 最特殊。它既是编程模型里的线程分组,又直接对应 NVIDIA GPU 的硬件调度方式,所以归为“软硬件结合”。

说 Thread 是软件抽象,并不是说线程不需要硬件就能执行。意思是,芯片上没有一个固定的、专门对应某个 Thread 的“线程核心”。

对照这张表,有三个最容易混淆的地方:

  1. Thread 不等于 CUDA Core。 CUDA Core 是物理计算单元,Thread 是程序的逻辑执行单位。随着调度进行,一个 CUDA Core 在不同时刻可以执行不同线程的指令。
  2. Block 不等于 SM。 一个 Block 执行期间驻留在一个 SM 上,但一个 SM 可以同时驻留多个 Block,Block 的总数也通常远多于 SM。
  3. Warp 不是一块独立的芯片。 硬件上没有一个叫 Warp 的部件。Warp 是 32 个线程组成的执行分组,由 SM 里的 Warp Scheduler 调度。

本章小结

把一次 Kernel 启动从头到尾串起来:

程序启动 Kernel
  ↓
产生一个 Grid,里面有很多 Block
  ↓
GPU 把 Block 分配到有空闲资源的 SM,放不下的先排队
  ↓
每个 Block 里的线程按 32 个一组,组成 Warp
  ↓
SM 里的 Warp Scheduler 挑出准备好的 Warp,发出下一条指令
  ↓
CUDA Core、Tensor Core 等执行单元完成运算

软件负责把工作分组,硬件负责把这些分组分配给 SM,再以 Warp 为单位交给计算单元执行。

到这里,我们知道了计算是怎样被组织和调度的。但一个 Kernel 跑得快不快,还取决于另一件事:数据能不能及时送到计算单元。下一章讨论这个问题。

第 4 章:GPU 性能首先是数据搬运问题

前两章讲了 GPU 有哪些硬件,以及线程怎样被分配到这些硬件上。这一章回答一个更实际的问题:计算单元已经很强了,为什么 Kernel 还是会慢?

先说结论:很多时候,计算单元在等数据。 数据平时放在芯片外的显存里,计算单元在 SM 里,中间的搬运常常比计算本身更花时间。所以优化 Kernel,不只是看算了多少次,还要看数据怎样搬、搬了几次、能不能重复使用。

4.1 计算单元在等数据

假设 GPU 正在执行矩阵乘法:

C = A × B

Tensor Core 每秒能完成海量的乘加运算,但它必须先拿到 A 和 B 的数据才能开始算。回到第 2 章的存储位置:A 和 B 放在芯片外的显存里,计算单元在 SM 内部,数据要从显存一路搬进 SM,才能参与计算。

如果计算单元处理数据的速度,比显存提供数据的速度还快,就会出现这种情况:

计算单元有能力继续算,但数据还没送过来。

这是 GPU 性能瓶颈的一个重要来源。第 2 章那套由近到远的多级存储,目的就是让常用的数据尽量留在离计算单元近的地方,少跑几趟显存。

Shared Memory:由程序决定的复用

多级存储里,Shared Memory 最值得关注,因为它由程序直接控制。

假设一个 Block 里有 128 个线程,它们都要用同一批数据。有两种做法:

独立读取显存与通过 Shared Memory 在片上复用的对比
独立读取显存与通过 Shared Memory 在片上复用的对比

做法 B 有机会减少访问片外显存的次数。不过 GPU 上还有 L1 和 L2 Cache,它们由硬件自动缓存近期读过的数据,所以做法 A 的重复读取也不一定每次都真的访问显存。Shared Memory 的价值在于,程序可以明确安排哪些数据搬进来、给哪些线程复用。缓存能不能命中由硬件决定,Shared Memory 里放什么由程序决定。

Shared Memory 也有代价:

  • 容量有限,放不下太多数据;
  • 线程之间要同步,确保数据搬完了再开始用;

所以用了 Shared Memory 不代表一定更快,用得不好反而会变慢。

4.2 Memory Coalescing:读同样多的数据,也可能快慢不同

数据搬运的代价不只取决于读了多少字节,还取决于怎样读。两个读取字节数相同的 Kernel,速度也可能差很多。

显存不是一个字节一个字节地送数据,而是按一段一段的连续区域搬运。当一个 Warp 的线程同时发出读取请求时,硬件会看这些地址落在哪些段里,尽量把它们合并成较少的内存事务。这种行为叫 Memory Coalescing(访存合并)。

对比一个 Warp 的两种读取方式:

Memory Coalescing:连续访问与分散访问
Memory Coalescing:连续访问与分散访问
  • 连续访问时,32 个线程要的数据挤在相邻的地址里,硬件通常能把它们合并成很少的几次内存事务。
  • 分散访问时,每个线程要的数据落在不同的段里,硬件就要发出更多次内存事务,而每次搬来的一整段里,大部分数据又用不上,有效带宽就降下来了。

这也是 Memory Coalescing 被归为“软硬件结合”的原因:合并是硬件的行为,但能不能合并,取决于程序让线程按什么顺序访问数据。

4.3 Tiling:把数据搬到片上重复使用

前面说 Shared Memory 能让数据被重复使用。Tiling(分块计算)就是组织这种复用的典型方法:把大的计算任务切成小块,一次把一小块输入搬到片上,让它服务多个输出。

先看一个最小的例子

用一个 2×2 的矩阵乘法:

A = [1 2]    B = [5 6]
    [3 4]        [7 8]

每个输出元素都是 A 的一行和 B 的一列做点积:

C₀₀ = 1×5 + 2×7 = 19
C₀₁ = 1×6 + 2×8 = 22
C₁₀ = 3×5 + 4×7 = 43
C₁₁ = 3×6 + 4×8 = 50

仔细看会发现,A 的第 0 行 [1 2] 在算 C₀₀ 和 C₀₁ 时都用到了;B 的第 0 列 [5 7] 在算 C₀₀ 和 C₁₀ 时也都用到了。每个输入都被两个输出用到。

朴素做法是每算一个输出,就去显存把它需要的 4 个数读一遍。4 个输出一共读 16 次,每个输入都被读了 2 次。

Tiling 的做法是先把 A 和 B 这 8 个数一次性搬进 Shared Memory,然后 4 个输出都从 Shared Memory 里取数。显存只需要读 8 次。

Tiling:一块数据只搬一次,供多个输出使用
Tiling:一块数据只搬一次,供多个输出使用

放大到 4×4

同样的道理放到 4×4 矩阵乘法上,差距更明显。16 个输出元素,每个输出要读 A 的一行和 B 的一列,各 4 个数。朴素做法的读取次数是:

16 × (4 + 4) = 128 次

如果 A 和 B 的全部 32 个元素都能在片上被充分复用,理论上每个数只需加载一次:

16 + 16 = 32 次

这是理想化的逻辑读取次数,忽略了 L1/L2 缓存、写回结果和管理开销,并不代表实际显存流量一定减少 75%。真实的大矩阵也没法整个放进 Shared Memory,所以实际做法是把矩阵切成很多块,一块一块地搬、一块一块地算。块切得越大,复用越多,但 Shared Memory 和寄存器的容量、同步的开销都会限制块的大小,需要权衡。

Tiling 减少的是什么

这里最容易误解的一点是:Tiling 没有减少计算量。 矩阵乘法要做的乘法和加法一次都没少,减少的是为了完成这些运算而重复搬运的数据。

还要再分清一次软件和硬件:Tiling 是软件的计算组织策略,Shared Memory 是真实的硬件存储。 第 3 章讲过,同一个 Block 的线程可以共享 Shared Memory、互相同步,这正是 Tiling 能够实现的前提:一个 Block 负责一个输出块,块里的线程一起把输入搬进 Shared Memory,同步之后再各自计算自己负责的输出元素。

回头看第 1 章提到的 FlashAttention,它的一个关键思路也是这样:把 Attention 的计算分块组织,让中间结果尽量留在片上,减少在显存和计算单元之间来回搬运的数据。

本章小结

这一章的三个知识点,都在回答同一个问题:怎样让数据更顺利地送到计算单元。

方法解决什么问题谁来负责
Shared Memory 复用同一批数据被多个线程反复读取程序决定放什么,硬件提供存储
Memory Coalescing地址分散导致内存事务变多硬件负责合并,程序决定访问顺序
Tiling输入数据被不同输出重复搬运程序组织分块,依赖 Shared Memory

一个 Kernel 的性能,不只取决于它做了多少次运算,还取决于数据怎样从显存来到计算单元,以及搬来的数据能被重复使用多少次。

那么,怎样判断一个 Kernel 到底是卡在数据搬运上,还是卡在计算上?下一章用 Arithmetic Intensity 和 Roofline 模型来回答。

第 5 章:用 Arithmetic Intensity 和 Roofline 找瓶颈

上一章讲了两种减少搬运代价的办法:让线程集中访问相邻地址,以及把数据搬到片上反复使用。接下来,我们需要一个尺度,把“搬了多少数据”和“做了多少计算”联系起来。

这一章分两步:先用 Arithmetic Intensity(算术强度) 衡量每搬运一个字节能做多少计算,再用 Roofline(屋顶线模型) 看显存带宽和计算峰值中,哪一个给出的性能上限更低。

5.1 Arithmetic Intensity:每搬运 1 Byte,做多少计算

Arithmetic Intensity 简称 AI,这里的 AI 指算术强度。它的定义是:

AI=浮点运算次数数据传输字节数=FLOPsBytes Transferred单位:FLOP/Byte\mathrm{AI} = \frac{\text{浮点运算次数}}{\text{数据传输字节数}} = \frac{\mathrm{FLOPs}}{\text{Bytes Transferred}} \qquad \text{单位:}\mathrm{FLOP/Byte}

三个容易混淆的量需要分清:

量含义回答的问题
FLOPs浮点运算次数一共算了多少?
FLOP/s每秒浮点运算次数算得多快?
Bytes Transferred数据传输字节数一共搬了多少数据?

例如,1000 FLOPs 表示完成了 1000 次浮点运算;1000 FLOP/s 表示每秒完成 1000 次浮点运算。常见的乘加操作按一次乘法、一次加法计为 2 FLOPs。

用同样的计算量对比

假设两个 Kernel 完成同样的计算,但数据传输量不同:

指标Kernel AKernel B
计算量1000 FLOPs1000 FLOPs
显存读写总量1000 Bytes100 Bytes
AI1 FLOP/Byte10 FLOP/Byte

Kernel B 每搬运一个字节,可以完成 10 次浮点运算。因此,同样的显存带宽能够为它支撑更多计算。

这也解释了 Tiling 的作用:如果计算量保持不变,而实际显存传输量减少,AI 就会提高。 搬来的数据被多个输出使用,每个字节能支持的计算也就更多。

5.2 Roofline:性能受到两个上限约束

知道 AI 之后,还要把它和 GPU 的能力放在一起看。这里有两个上限:

  • 计算峰值:计算单元每秒最多能完成多少次浮点运算。
  • 带宽决定的上限:显存每秒送来的数据,最多能够支撑多少次浮点运算。

第二个上限为什么是“带宽 × AI”?看单位就很清楚:

显存带宽 × 算术强度
(Byte/s) × (FLOP/Byte) = FLOP/s

假设显存每秒能传输 500 GB 数据,而一个 Kernel 每搬运 1 Byte 只能做 1 次运算,那么显存最多能支撑每秒 5000 亿次运算,也就是 0.5 TFLOP/s。即使计算单元能达到 20 TFLOP/s,数据供应也支撑不起这个速度。

反过来,如果数据复用得足够多,显存带宽能够支撑的计算量超过了计算单元本身的能力,性能就会碰到计算峰值这个上限。

因此,Roofline 取这两个上限中较小的一个:

Pmax⁡=min⁡(Pcompute,  BW×AI)P_{\max} = \min\left(P_{\mathrm{compute}},\; \mathrm{BW} \times \mathrm{AI}\right)
  • Pmax⁡P_{\max}:模型给出的性能上限。
  • PcomputeP_{\mathrm{compute}}:计算峰值。
  • BW\mathrm{BW}:显存带宽。
  • AI\mathrm{AI}:算术强度。

用一块假设的 GPU 算一遍

沿用原始素材中的假设:计算峰值为 20 TFLOP/s,显存带宽为 500 GB/s。这组数字用于演示公式。

AI(FLOP/Byte)带宽决定的上限 BW × AI计算峰值Roofline 上限理想模型中的限制
105 TFLOP/s20 TFLOP/s5 TFLOP/s带宽
2010 TFLOP/s20 TFLOP/s10 TFLOP/s带宽
4020 TFLOP/s20 TFLOP/s20 TFLOP/s两个上限相交
8040 TFLOP/s20 TFLOP/s20 TFLOP/s计算

把这些数画出来,就得到一条先上升、再变平的“屋顶线”:

Roofline:带宽与计算峰值共同决定性能上限
Roofline:带宽与计算峰值共同决定性能上限

横轴是 AI,纵轴是计算性能。左侧斜线表示显存带宽支撑的计算量:AI 越高,相同带宽就能支撑越多计算。右侧水平线表示计算峰值:继续提高 AI,也无法突破计算单元本身的能力。

图中采用线性坐标,便于直接核对表中的数字;实际性能分析也常使用对数坐标,以便展示跨度更大的 AI 和性能范围。

Ridge Point:两个上限相交的位置

斜线与水平线相交的位置叫 Ridge Point,可以理解为屋顶的拐点。此时:

BW×AIcritical=PcomputeAIcritical=PcomputeBW=20×1012500×109=40  FLOP/Byte\begin{aligned} \mathrm{BW} \times \mathrm{AI}_{\mathrm{critical}} &= P_{\mathrm{compute}} \\ \mathrm{AI}_{\mathrm{critical}} &= \frac{P_{\mathrm{compute}}}{\mathrm{BW}} \\ &= \frac{20 \times 10^{12}}{500 \times 10^{9}} \\ &= 40\;\mathrm{FLOP/Byte} \end{aligned}

这里的 T 和 G 分别按 10¹² 和 10⁹ 换算。对这块假设的 GPU:

  • AI 小于 40:带宽给出的上限更低,位于理想的 Memory-bound(带宽受限)区域。
  • AI 大于 40:计算峰值给出的上限更低,位于理想的 Compute-bound(计算受限)区域。

这个拐点随硬件能力而变化。同一个 Kernel 换到另一块 GPU 上,即使 AI 不变,也可能落到不同的区域。

5.3 屋顶是上限,实际 Kernel 可能离它很远

Roofline 最容易被误用的地方,是把“模型中的上限”当成“实际能跑到的速度”。

例如 AI 为 10 时,刚才算出的上限是 5 TFLOP/s。这表示在上述假设下,带宽最多能支撑这么多计算,并不保证 Kernel 能达到 5 TFLOP/s。

实际性能还可能受到这些因素影响:

因素为什么可能达不到屋顶
访存模式与延迟地址分散、并发请求不足时,可能无法用满显存带宽
并行度与资源占用Block 太少,或寄存器、Shared Memory 占用过多,会限制能同时推进的工作
分支与同步部分线程等待,执行资源没有持续做有效计算
Kernel 启动开销计算任务很小时,启动等固定开销可能占较大比例
计算单元利用率数据类型、矩阵形状或实现方式不合适,可能达不到选定的计算峰值

因此,落在图的左侧,只说明带宽项给出的理论上限更低;要确认运行时是否真正耗尽了显存带宽,还需要测量。同样,落在右侧也不代表计算单元已经满负荷运转。

本章小结

Arithmetic Intensity 把计算和搬运联系起来:每搬运 1 Byte,能完成多少浮点运算。Roofline 再把这个比例放到具体硬件上,取“计算峰值”和“带宽 × AI”中较低的那个,给出理想性能上限。

回到文章开头的问题:同样一条矩阵乘法公式,为什么有的实现快,有的实现慢?公式规定了计算结果,Kernel 则决定线程怎样分工、数据怎样搬运和复用,以及硬件资源能否被充分使用。

从模型架构、算子和 Kernel,到 SM、线程分组和存储,再到 AI 与 Roofline,我们现在有了一条分析路径:先看要完成什么计算,再看工作怎样分配、数据怎样流动,最后用模型提出瓶颈假设,并用实际测量确认。

Last updated on