Transformer 中的 QKV 投影、Attention 输出投影和 FFN 都包含大量矩阵乘法。数学上，它们可以统一写成 GEMM；工程上，性能却主要取决于矩阵怎样分块、数据怎样搬运，以及一个 SM 能同时容纳多少 block 和 warp。
本文从矩阵乘法公式出发，依次解释 tile、HBM、共享内存、寄存器、SM、block、thread、warp、算术强度、occupancy、尾块和浮点误差之间的关系。
1. GEMM 是什么？
GEMM 是 General Matrix Multiplication，即通用矩阵乘法。简化形式是：
$C=AB$
更一般的形式是：
$C\leftarrow\alpha AB+\beta C$
假设：
$A\in\mathbb{R}^{M\times K},\qquad B\in\mathbb{R}^{K\times N}$
那么：
$C\in\mathbb{R}^{M\times N}$
shape 规则是：
1
(M, K) @ (K, N) -> (M, N)
矩阵 C 有 MN 个输出元素，每个元素都是 A 的一行与 B 的一列做点积：
$C_{i,j}=\sum_{k=1}^{K}A_{i,k}B_{k,j}$
每个输出元素大约需要 K 次乘法和 K 次加法，所以 GEMM 的计算量约为：
$\operatorname{FLOPs}\approx2MKN$
分块不会减少这个数量。它优化的是数据搬运和硬件利用率。
2. 为什么矩阵乘法可以分块？
将 A、B 按相容的行列切成块：
$A=\begin{bmatrix}A_{11}&A_{12}\\A_{21}&A_{22}\end{bmatrix},\qquad B=\begin{bmatrix}B_{11}&B_{12}\\B_{21}&B_{22}\end{bmatrix}$
输出左上角的块是：
$C_{11}=A_{11}B_{11}+A_{12}B_{21}$
输出右上角的块是：
$C_{12}=A_{11}B_{12}+A_{12}B_{22}$
也就是说，每个输出块都可以通过若干局部矩阵乘法，沿归约维 K 不断累加得到。
假设一个 block 负责输出 tile：
$C_{tile}\in\mathbb{R}^{B_M\times B_N}$
每一轮 K 循环读取：
$A_{tile}\in\mathbb{R}^{B_M\times B_K},\qquad B_{tile}\in\mathbb{R}^{B_K\times B_N}$
并执行：
$C_{tile}\mathrel{+}=A_{tile}B_{tile}$
伪代码如下：
1
2
3
4
5
6
7
8
C_tile = 0
for k in range(0, K, BK):
A_tile = A[m:m+BM, k:k+BK]
B_tile = B[k:k+BK, n:n+BN]
C_tile += A_tile @ B_tile
C[m:m+BM, n:n+BN] = C_tile
分块改变了计算顺序，却没有改变目标公式。
3. 分块为什么会更快？
GPU 的不同存储层级在容量和访问代价上差异很大。可以先用一个简化模型理解：
1
2
3
4
5
6
7
HBM / Global Memory：容量大，访问代价高
↓
L2 Cache
↓
Shared Memory / L1：每个 SM 上的片上存储
↓
Registers：每个线程直接使用的片上状态
朴素 GEMM 如果每次乘加都重新从 HBM 读取 A[i,k] 和 B[k,j]，同一元素会被反复搬运。
实际上：
A[i,k]可以被同一输出行的多个列复用；B[k,j]可以被同一输出列的多个行复用；C[i,j]的中间累加值可以长时间保留在寄存器中。
分块 Kernel 通常会：
- 从 HBM 搬入一块
A和一块B； - 将它们放进共享内存；
- 让同一 block 的多个线程重复读取这些数据；
- 在寄存器或 Tensor Core 累加器中保存局部输出；
- 完成整个
K循环后，将输出 tile 写回 HBM。
因此，高性能 GEMM 的核心通常不是“少做乘法”，而是：
让一次昂贵的数据搬运支撑尽可能多的计算。
4. 什么是 SM？
SM 是 Streaming Multiprocessor，通常翻译为流式多处理器，是 NVIDIA GPU 中执行线程块、调度 warp 和管理片上资源的主要计算单元。其他厂商有相似概念，例如 AMD GPU 中的 Compute Unit，但具体组织方式并不完全相同。
一块 GPU 包含多个 SM。每个 SM 内部通常包括：
- warp scheduler；
- CUDA Core 等标量/向量执行单元；
- Tensor Core；
- register file；
- shared memory / L1；
- load/store 等执行单元。
可以把 SM 理解成一个能够独立接收 block，并在多个 warp 之间切换执行的“小型并行处理器”。
不同 GPU 架构中，SM 的寄存器数量、共享内存容量、warp 槽位和执行单元数量都可能不同。因此本文后面的容量数字只用于说明资源约束，不代表某一款具体 GPU。
5. Grid、Block 和 Thread 是什么关系？
CUDA Kernel 启动时，需要指定 grid 中有多少个 block，以及每个 block 有多少个线程。
例如：
1
kernel<<<100, 256>>>(...);
表示：
1
2
3
Grid 中有 100 个 block
每个 block 有 256 个逻辑线程
总逻辑线程数为 100 × 256 = 25600
线程数量是 Kernel 启动配置分配给每个 block 的。在一次普通 Kernel 启动中，所有 block 的 blockDim 相同。
一个 block 通常会被完整调度到某一个 SM 上执行，不会把同一个 block 的线程拆到多个 SM。这样，同一个 block 内的线程才能：
- 访问同一块 shared memory；
- 通过
__syncthreads()同步； - 合作计算同一个输出 tile。
不同 block 的 shared memory 相互隔离。一个 block 开始驻留后，其共享内存和线程寄存器通常会一直保留到该 block 执行结束。
需要注意：
256 个 CUDA 线程是 256 个逻辑执行上下文，不表示这个 block 永久独占 256 个物理计算核心。
6. 什么是 Warp？
在 NVIDIA GPU 中，一个 warp 包含 32 个线程。SM 的调度器通常以 warp 为基本线程组发射指令。
一个 256 线程的 block 会被划分为：
$256/32=8\text{ 个 warp}$
1
2
3
4
5
Warp 0：thread 0～31
Warp 1：thread 32～63
Warp 2：thread 64～95
...
Warp 7：thread 224～255
同一个 warp 中的线程执行同一条指令，但处理各自的数据。遇到条件分支时，如果一部分线程走 if，另一部分线程走 else，硬件可能需要分别执行两条路径，这就是 warp divergence。
warp scheduler 会从已经驻留且准备好的 warp 中选择下一条指令执行。如果一个 warp 正在等待内存或依赖，SM 可以切换到其他 ready warp。
这不是像 CPU 线程那样保存、恢复大量上下文，因为多个 resident warp 的寄存器状态已经同时保存在 SM 上。
7. 什么是 Resident Block、Resident Warp 和 Occupancy？
GPU 程序可以启动成千上万个 block，但同一时刻只有一部分 block 真正驻留在 SM 上。已经进入 SM、资源已经分配并且尚未执行结束的 block，称为 resident block。
这些 block 包含的 warp 就是 resident warp。
Occupancy 通常表示：
$\text{Occupancy}=\dfrac{\text{SM 上的 active/resident warps}}{\text{该 SM 支持的最大 active warps}}$
较多的 resident warp 有助于隐藏延迟。例如：
1
2
3
4
5
6
7
Warp 0 等待内存
↓
调度 Warp 1
↓
Warp 1 等待依赖
↓
调度 Warp 2
如果 resident warp 太少，就可能出现所有 warp 都在等待、执行单元暂时空闲的情况。
Occupancy 不是越高性能就一定越好，但过低通常会削弱延迟隐藏能力。
8. 为什么 tile 太大会减少驻留 block 和 warp？
每个 SM 的共享内存和寄存器总量固定，而 resident block 的资源需要同时放在 SM 上。
8.1 共享内存按 block 分配
假设一个 SM 有 64 KB 共享内存。
如果每个 block 使用 16 KB：
$\left\lfloor64/16\right\rfloor=4\text{ blocks}$
如果每个 block 使用 32 KB：
$\left\lfloor64/32\right\rfloor=2\text{ blocks}$
如果每个 block 使用 48 KB：
$\left\lfloor64/48\right\rfloor=1\text{ block}$
剩余 16 KB 因为放不下第二个完整 block，不能用来驻留半个 block。
8.2 寄存器按线程分配
假设一个 block 有 256 个线程，每个线程使用 32 个寄存器，那么一个 block 需要：
$256\times32=8192\text{ registers}$
如果 tile 变大，每个线程需要保存更多输出累加器，寄存器使用量增加到 128 个，那么一个 block 需要：
$256\times128=32768\text{ registers}$
在寄存器总量固定的情况下，同一个 SM 能放下的 block 数就会减少。
一个 SM 的 resident block 上限大致同时受到以下条件约束：
$N_{blocks}\le\min(N_{smem},N_{regs},\text{其他硬件上限})$
其中共享内存给出的上限是：
$N_{smem}=\left\lfloor\dfrac{S_{SM}}{S_{block}}\right\rfloor$
寄存器给出的上限是：
$N_{regs}=\left\lfloor\dfrac{R_{SM}}{R_{thread}T_{block}}\right\rfloor$
其他硬件上限还包括每个 SM 的最大线程数、warp 数和 block 数。
哪种资源最先耗尽，哪种资源就是当前 Kernel 的驻留限制。
由于一个 256 线程 block 包含 8 个 warp：
1
2
3
驻留 4 blocks -> 32 resident warps
驻留 2 blocks -> 16 resident warps
驻留 1 block -> 8 resident warps
所以完整的因果关系是：
1
2
3
4
5
6
7
8
9
10
11
tile 增大
↓
A/B tile 和 C 累加器变大
↓
每 block 的 shared memory、每线程的 registers 增加
↓
每个 SM 能完整容纳的 block 减少
↓
resident warp 减少
↓
隐藏内存和指令延迟的能力可能下降
如果每个线程需要的寄存器超过硬件和编译器能够有效提供的范围，还可能发生 register spilling，把部分线程局部状态放到 local memory。这里的 local memory 位于设备内存体系，访问代价远高于寄存器。
9. 算术强度：每搬一个字节能做多少计算？
算术强度定义为：
$\text{Arithmetic Intensity}=\dfrac{\text{FLOPs}}{\text{Bytes transferred from target memory}}$
理想化地假设 FP16 GEMM 的 A、B 只从 HBM 读取一次，C 只写一次，数据量约为：
$2(MK+KN+MN)\text{ bytes}$
于是：
$AI\approx\dfrac{2MKN}{2(MK+KN+MN)}$
- 算术强度较低：更可能受内存带宽限制；
- 算术强度较高：数据复用充分，更可能受计算吞吐限制。
实际情况还会受到 cache、对齐、读改写、融合操作和数据类型的影响，所以这只是帮助形成直觉的粗略估算。
10. Tile 为什么不是越大越好？
更大的 tile 通常意味着：
- 同一块
A、B数据被更多输出元素复用； - HBM 访问相对减少；
- 算术强度可能提高。
但同时也意味着：
- shared memory 占用增加；
- register accumulator 增加；
- block 同步和边界处理更复杂；
- resident block / warp 可能减少；
- 某些 shape 下会产生更多无效计算。
因此：
1
2
tile 太小：数据复用不足，频繁访问 HBM
tile 太大：片上资源压力过高，驻留并行度下降
最优 tile 是数据复用、occupancy、指令效率、硬件结构和矩阵 shape 之间的折中。
11. 尾块与 Padding
如果 M、N、K 不是 tile size 的整数倍，最后一块就可能不完整。
例如：
1
2
3
4
矩阵长度 = 35
tile size = 16
35 = 16 + 16 + 3
最后 3 个元素构成尾块。常见处理方式包括：
- 边界判断；
- mask；
- padding 到对齐尺寸；
- 规则部分走快路径，尾块走通用路径。
Padding 会增加无效存储和计算，边界判断则会引入控制开销。哪种方式更快，需要结合实际 shape 和硬件测试。
12. 为什么不同分块结果可能略有不同？
实数加法满足结合律，但浮点加法不严格满足：
$(a+b)+c\ne a+(b+c)$
不同 tile、线程归约树和 Tensor Core 路径会改变累加顺序，因此优化前后的结果可能无法逐位一致。
正确性测试应该根据数据类型和归约长度使用合适的绝对、相对容差：
$|\hat{x}-x|\le\text{atol}+\text{rtol}\cdot|x|$
除了观察最大误差，还应关注误差是否随 K 的长度系统性放大。
13. 把完整工程逻辑串起来
分块 GEMM 的执行过程可以总结为：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
大矩阵切成输出 tile
↓
Grid 中的 block 分别负责不同 tile
↓
block 被调度到 SM
↓
block 内线程组成多个 warp
↓
A/B tile 从 HBM 搬到 shared memory
↓
线程/warp 复用数据并执行乘加
↓
C tile 在 registers 中沿 K 累加
↓
完成后将输出写回 HBM
性能优化需要同时考虑两类目标：
1
2
数据复用：减少 HBM 搬运
并行调度：保留足够 resident warp 隐藏延迟
二者通过 tile size、每 block 线程数、shared memory 用量和每线程 register 用量联系起来。这也是后续理解 CUDA GEMM Kernel、Tensor Core、Triton、FlashAttention 和算子自动调优的基础。
