从模型计算到 GPU 执行:读懂算子、Kernel 与 GPU 性能
发布于
你不要那样看着我的眼睛 我就怕你会看穿我的心
从模型计算到 GPU 执行:读懂算子、Kernel 与 GPU 性能
这篇文章想解决一个很具体的问题:模型中的一行数学公式,最后是怎样变成 GPU 上真正执行的程序的?
如果你已经看过 Transformer 和 Attention 公式,却经常把模型架构、算法、算子、Kernel 和 GPU 混在一起,那么可以从这篇文章开始重新建立一条清晰的认知链路。
开始之前:先给整篇文章画一张地图
在大模型推理中,我们通常从模型结构开始思考:这里有 Attention,那里有 MLP,还有归一化和残差连接。但 GPU 并不直接理解“这是一个 Transformer”或“这是一次 Attention”。一层层往下看,同一件计算会经过几个不同的描述层次:
模型架构 → 算法 → 算子 → GPU Kernel → GPU 硬件
这是一张帮助理解的教学地图,不是严格的软件调用栈。真实执行过程中,还会经过框架、编译器、运行时和驱动程序;而且算法、算子与 Kernel 之间也不是固定的一对一关系。
可以先用一句话记住每一层的职责:
- 模型架构决定系统由哪些模块组成,例如 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。为了便于理解,可以把它看成两类主要工作:
- Attention:让一个位置的表示读取其他位置的信息。
- 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 + b | Add | 逐元素相加 |
a @ b | MatMul | 矩阵乘法 |
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,减少中间结果的显存读写。
这也是理解 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 启动请求 |
| SM | GPU 硬件 | 调度并执行 GPU 线程 |
| CUDA Core | GPU 硬件 | 执行通用计算中的一部分 |
| Tensor Core | GPU 硬件 | 加速高吞吐矩阵运算 |
| HBM / GDDR | GPU 显存 | 保存权重和张量等数据 |
当 PyTorch 执行 torch.matmul(A, B) 时,用户看到的是框架层面的算子调用。框架会结合设备、数据类型、张量形状和布局等条件,选择合适的底层实现;在 NVIDIA GPU 上,某些矩阵乘法可能由 CUDA 生态中的高性能库来完成。最终,硬件真正执行的仍然是编译后的 GPU 指令,以及这些指令对应的数据读取、计算和写回。
第 2 章:GPU 的硬件结构
上一章的最后一层是“GPU 硬件”,但只列了几个名字。这一章把 GPU 打开来看:一块 GPU 由哪些部件组成,SM、CUDA Core、Tensor Core 和显存之间是什么关系。
先说结论:GPU 由很多个 SM 组成。每个 SM 都有自己的调度器、两类计算单元和片上存储;显存在芯片外,离计算单元最远。 下面按“整块芯片 → 一个 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 Core | 128 个,整卡 16,896 个 |
| 每个 SM 的 Tensor Core | 4 个(第四代),整卡 528 个 |
| 每个 SM 的寄存器文件 | 256 KB |
| 每个 SM 的 Shared Memory + L1 | 合计 256 KB |
| L2 Cache | 50 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 / L1 | SM 内的片上存储,同一个 Block 的线程可以在 Shared Memory 中共享数据 |
Warp 和 Block 这两个词这里先出现名字,下一章再解释。现在只需要知道:Warp 是一组一起被调度的线程,Block 是一组可以协作的线程。
这几类部件可以按职责分成三组:
- 调度:Warp Scheduler 决定下一步执行哪些线程的指令。
- 计算:CUDA Core 和 Tensor Core 真正完成运算。
- 存储:寄存器、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 Memory | SM 内部 | 同一个 Block 内的线程协作和复用数据 |
| L1 Cache | SM 内部或附近 | 由硬件自动缓存近期访问的数据 |
| L2 Cache | GPU 芯片上 | 所有 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(分支发散),它会降低执行效率。不过影响有多大,要看分支的长短和 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 的线程要共享 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 的“线程核心”。
对照这张表,有三个最容易混淆的地方:
- Thread 不等于 CUDA Core。 CUDA Core 是物理计算单元,Thread 是程序的逻辑执行单位。随着调度进行,一个 CUDA Core 在不同时刻可以执行不同线程的指令。
- Block 不等于 SM。 一个 Block 执行期间驻留在一个 SM 上,但一个 SM 可以同时驻留多个 Block,Block 的总数也通常远多于 SM。
- 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 个线程,它们都要用同一批数据。有两种做法:
做法 B 有机会减少访问片外显存的次数。不过 GPU 上还有 L1 和 L2 Cache,它们由硬件自动缓存近期读过的数据,所以做法 A 的重复读取也不一定每次都真的访问显存。Shared Memory 的价值在于,程序可以明确安排哪些数据搬进来、给哪些线程复用。缓存能不能命中由硬件决定,Shared Memory 里放什么由程序决定。
Shared Memory 也有代价:
- 容量有限,放不下太多数据;
- 线程之间要同步,确保数据搬完了再开始用;
所以用了 Shared Memory 不代表一定更快,用得不好反而会变慢。
4.2 Memory Coalescing:读同样多的数据,也可能快慢不同
数据搬运的代价不只取决于读了多少字节,还取决于怎样读。两个读取字节数相同的 Kernel,速度也可能差很多。
显存不是一个字节一个字节地送数据,而是按一段一段的连续区域搬运。当一个 Warp 的线程同时发出读取请求时,硬件会看这些地址落在哪些段里,尽量把它们合并成较少的内存事务。这种行为叫 Memory Coalescing(访存合并)。
对比一个 Warp 的两种读取方式:
- 连续访问时,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 次。
放大到 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 指算术强度。它的定义是:
三个容易混淆的量需要分清:
| 量 | 含义 | 回答的问题 |
|---|---|---|
| FLOPs | 浮点运算次数 | 一共算了多少? |
| FLOP/s | 每秒浮点运算次数 | 算得多快? |
| Bytes Transferred | 数据传输字节数 | 一共搬了多少数据? |
例如,1000 FLOPs 表示完成了 1000 次浮点运算;1000 FLOP/s 表示每秒完成 1000 次浮点运算。常见的乘加操作按一次乘法、一次加法计为 2 FLOPs。
用同样的计算量对比
假设两个 Kernel 完成同样的计算,但数据传输量不同:
| 指标 | Kernel A | Kernel B |
|---|---|---|
| 计算量 | 1000 FLOPs | 1000 FLOPs |
| 显存读写总量 | 1000 Bytes | 100 Bytes |
| AI | 1 FLOP/Byte | 10 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 取这两个上限中较小的一个:
- :模型给出的性能上限。
- :计算峰值。
- :显存带宽。
- :算术强度。
用一块假设的 GPU 算一遍
沿用原始素材中的假设:计算峰值为 20 TFLOP/s,显存带宽为 500 GB/s。这组数字用于演示公式。
| AI(FLOP/Byte) | 带宽决定的上限 BW × AI | 计算峰值 | Roofline 上限 | 理想模型中的限制 |
|---|---|---|---|---|
| 10 | 5 TFLOP/s | 20 TFLOP/s | 5 TFLOP/s | 带宽 |
| 20 | 10 TFLOP/s | 20 TFLOP/s | 10 TFLOP/s | 带宽 |
| 40 | 20 TFLOP/s | 20 TFLOP/s | 20 TFLOP/s | 两个上限相交 |
| 80 | 40 TFLOP/s | 20 TFLOP/s | 20 TFLOP/s | 计算 |
把这些数画出来,就得到一条先上升、再变平的“屋顶线”:
横轴是 AI,纵轴是计算性能。左侧斜线表示显存带宽支撑的计算量:AI 越高,相同带宽就能支撑越多计算。右侧水平线表示计算峰值:继续提高 AI,也无法突破计算单元本身的能力。
图中采用线性坐标,便于直接核对表中的数字;实际性能分析也常使用对数坐标,以便展示跨度更大的 AI 和性能范围。
Ridge Point:两个上限相交的位置
斜线与水平线相交的位置叫 Ridge Point,可以理解为屋顶的拐点。此时:
这里的 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,我们现在有了一条分析路径:先看要完成什么计算,再看工作怎样分配、数据怎样流动,最后用模型提出瓶颈假设,并用实际测量确认。