GPU 基础原理

说明：本文翻译自 Damek Davis 的博客 Basic facts about GPUs，原文最后更新于 2025 年 6 月，版权归原作者所有。

我一直想对 GPU 的工作原理有更深入的了解。我在网上读了很多资料，其中以下文章对我帮助最大：

1. Making Deep Learning Go Brrrr From First Principles
2. What Shapes Do Matrix Multiplications Like?
3. How to Optimize a CUDA Matmul Kernel for cuBLAS-like Performance: a Worklog

本文汇集了从这些资源中学到的各种事实。

• Hacker News 讨论
• Twitter 讨论串

致谢：感谢 Alex McKinney 对独立线程调度相关内容的评论。

1. 计算与内存层级

GPU 的设计天然存在不平衡：它计算的速度远快于访问主内存的速度。例如，一块 NVIDIA A100 GPU 每秒可执行 19.5 万亿次 32 位浮点运算（TFLOPS），但其内存带宽只有约 1.5 万亿字节每秒（TB/s）。在读入一个 4 字节数字所需的时间里，GPU 本可以执行超过 50 次计算。

下面是 NVIDIA A100 GPU 计算与内存层级的示意图。我引用的 FLOP/s 与 TB/s 数据仅针对 A100。

+---------------------------------------------------------------------------------+
|                               Global Memory (VRAM)                              |
|                            (~40 GB, ~1.5 TB/s on A100)                          |
+----------------------------------------+----------------------------------------+
                                         | (Slow off-chip bus)
+----------------------------------------v----------------------------------------+
|                            Streaming Multiprocessor (SM)                        |
|                     (1 of 108 SMs on an A100, each ~(19.5/108) TFLOPS)          |
|                           (2048 threads, 64 warps, 32 blocks)                   |
| +-----------------------------------------------------------------------------+ |
| |                        Shared Memory (SRAM) / L1 Cache                        |
| |                    (~192 KB on-chip workbench, 19.5 TB/s)                     |
| +-----------------------------------------------------------------------------+ |
| |                        Register File (~256 KB, ? TB/s)                        |
| +-----------------------------------------------------------------------------+ |
| |                                                                             | |
| |                //-- A "Block" of threads runs on one SM --//                | |
| | +--------------------------+ +------------------------+                     | |
| | |      Warp 0 (32 thr)     | |      Warp 1 (32 thr)   | ... (up to 32 warps)| |
| | | +----------------------+ | +----------------------+ |                     | |
| | | | Thread 0 Registers   | | | Thread 32 Registers  | |                     | |
| | | | [reg0: float]        | | | [reg0: float]        | |                     | |
| | | | [reg1: float] ...    | | | [reg1: float] ...    | |                     | |
| | | +----------------------+ | +----------------------+ |                     | |
| | +--------------------------+ +------------------------+                     | |
| |                                                                             | |
+---------------------------------------------------------------------------------+

这张图展示了性能层级（性能层级示意）。全局内存（Global Memory，VRAM） 是容量大、速度慢的片外内存池，所有数据初始都存放在这里。流式多处理器（Streaming Multiprocessor，SM） 是 GPU 的计算单元。要进行计算，它必须通过慢速总线获取数据。为缓解这一瓶颈，每个 SM 都配有快速的片上共享内存（Shared Memory，SRAM），带宽达 19.5 TB/s（A100 的 SRAM 带宽为 19.5 TB/s）。程序员把它当作手工管理的缓存来使用。

线程（Thread） 是最小的执行单元。每个线程都有一组私有的寄存器（Register） 来保存即时计算所需的数值，访问速度超过 ?? TB/s（寄存器访问速度原文标注为 '?? TB/s'，未给出确切数值）。硬件把线程组织成每组 32 个的 Warp。本文使用简化的锁步执行（lockstep execution）模型来分析性能：一个 Warp 内的全部 32 个线程在同一时刻执行同一条指令（独立线程调度在 Ampere 及以后架构可用，此处为简化模型）。在 A100 上，一个 SM 最多支持 64 个 Warp。程序员把线程组织成线程块（Block）——一个保证运行在单个 SM 上的线程网格。块可以是一维、二维或三维的。为简单起见，本文聚焦于 BLOCK_DIM x BLOCK_DIM 的方形二维线程块，其总线程数不能超过 1024 的硬件上限。块内的所有线程共享对同一块片上共享内存的访问权。

2. 两种性能区间

我们分析的对象是一个 Kernel——由宿主机（CPU）启动、由大量 GPU 线程并行执行的函数。一个 Kernel 的性能要么受限于内存带宽，要么受限于计算吞吐量。这两个上限定义了性能区间。

一个操作是内存受限（Memory-Bound）的，如果它的运行时间由数据从全局内存传送到 SM 的速度决定。以逐元素加法 y = x + 1 为例，SM 对读入的每个元素只执行微不足道的 FLOP。SM 大部分时间处于空闲，等待数据到达。

一个操作是计算受限（Compute-Bound）的，如果它的运行时间由 SM 的算术速度决定。大型矩阵乘法是典型例子。一旦数据被加载进 SM，就会执行海量计算。在 SM 忙碌时，内存总线处于空闲。

算术强度（Arithmetic Intensity，AI） 是判定所处区间的正式指标，即计算量与内存流量之比：

$$\text{Arithmetic Intensity}=\frac{\text{Total FLOPs}}{\text{Total Bytes Accessed}}$$

在 Roofline 模型中，Total Bytes Accessed 特指全局内存（HBM）与片上 SM 之间传输的数据量。这是因为该模型以最慢的片外内存总线这一主要瓶颈来评估 Kernel 的性能。片上流量（例如从共享内存到寄存器）不计入该计算。

Roofline 模型把 Kernel 的可达性能（以每秒 FLOP 计）绘制为其算术强度的函数。图中的两条"屋顶线"是 GPU 的硬性物理极限。

  ^ Performance (TFLOPS)
  |
  | Memory-Bound Region ¦ Compute-Bound Region
  |                     ¦
  |                    /¦----------------------  <-- Peak Compute (~19.5 TFLOPS)
  |                   / ¦
  |                  /  ¦
  | Peak Global     /<--¦------ Inefficient Compute Roof (e.g., using scalar ops, transcendental functions)
  | Mem BW (~1.5   /    ¦
  | TB/s)         /     ¦
  |              /      ¦
  +---------------------¦---------------------------> Arithmetic Intensity (FLOPs/Byte)
                        ^
                        ¦
                  Hardware Ridge Point (~13)

Kernel 的性能按如下方式确定：

• 内存受限时，SM 因等待数据而停滞。运行时间就是搬运这些数据所需的时间：$\text{Runtime}=\text{Bytes\_Accessed}/\text{Memory\_Bandwidth}$。因此 Kernel 的性能为 $\text{Performance}=\text{Total\_FLOPs}/\text{Runtime}=\text{AI}\times \text{Memory\_Bandwidth}$。在对数-对数图上，这是对角线。
• 计算受限时，SM 被充分利用。性能受限于其峰值算术吞吐量：$\text{Performance}=\text{Peak\_Compute\_FLOPs}$。这是水平线。

Kernel 的实际性能是这两个值中的较小者。脊点（Ridge Point） 是两条性能上限相交处的算术强度。对 A100 而言，这是 $19.5\text{ TFLOPS}/1.5\text{ TB/s}\approx 13\text{ FLOPs/Byte}$。一个 Kernel 必须超过这一算术强度才能进入计算受限区间。AI 低于 13 的 Kernel 工作在内存受限区间；AI 高于 13 的 Kernel 工作在计算受限区间。优化的目标是提高 AI，把 Kernel 的工作点向右推移，使其性能不断上升，直到触及计算屋顶。

19.5 TFLOPS 的"峰值计算"屋顶是一个理想值，只有使用诸如 Tensor Core 矩阵乘法等高度优化的指令且功率上限足够高时才能达到。一个操作可以处于计算受限区间，但性能仍远低于该峰值。例如，一个 AI 很高、却以标量算术或复杂的超越函数（sin、exp）为主的 Kernel，将受限于这些特定慢指令的吞吐量。这为该 Kernel 形成了一个更低的有效"屋顶"，如图所示。提高 AI 是必要的，但不充分；FLOP 本身也必须是高效的。

提高 AI 的首要策略是：数据一旦被加载进 SM 的快速片上内存，就最大化对它的复用。下面是一个简化模型：线程直接从全局内存把数据读入自己的私有寄存器。该分析计算的是所需的最小数据传输量；实际内存流量取决于访问模式，我们稍后会讨论。

考虑计算 C = A@B，其中所有矩阵都是 $N\times N$，元素为 4 字节浮点数。

策略 1：一个线程计算一个元素 C[i,j]

• FLOP：为计算 C[i,j]，线程执行 N 次乘加运算，即 $2N$ 个 FLOP。
• 访问字节数：线程必须读取 A 的第 i 行（N 个浮点数）和 B 的第 j 列（N 个浮点数），共 $2N$ 个浮点数，即 $8N$ 字节。
• 算术强度：$(2N\text{ FLOP})/(8N\text{ Bytes})=0.25\text{ FLOPs/Byte}$。

这个 AI 很低，Kernel 将是内存受限的。

策略 2：一个线程计算 C 的一个 $2\times 2$ 分块。要计算一个 $2\times 2$ 分块（C[i,j]、C[i,j+1]、C[i+1,j]、C[i+1,j+1]），线程必须完成全部四个元素的计算。

• FLOP：$4\text{ elements}\times 2N\text{ FLOPs/element}=8N$ 个 FLOP。
• 访问字节数：线程必须读取 A 的两行（A[i,:]、A[i+1,:]）和 B 的两列（B[:,j]、B[:,j+1]），共 $2N+2N=4N$ 个浮点数，即 $16N$ 字节。
• 算术强度：$(8N\text{ FLOP})/(16N\text{ Bytes})=0.5\text{ FLOPs/Byte}$。

这些 AI 值远低于 A100 约 13 FLOPs/Byte 的脊点。仅靠寄存器这一简单模型不足以让矩阵乘法进入计算受限区间（纯寄存器模型不足以让矩阵乘法达到计算受限）。实现高 AI 的关键在于：块内的线程相互协作，把 A 和 B 的更大分块加载进共享的片上 SRAM。1024 个线程的块通过协作处理这份共享数据，可以达到大于 13 的 AI。我们将在共享内存一节详述其机制。

3. 第三种区间：开销

性能还可能受限于宿主机侧的开销。这是 CPU（宿主机）为 GPU 准备工作所花费的时间，例如 Python 解释器或某个框架的调度系统中。

如果 GPU Kernel 太小或太多，应用就是开销受限（Overhead-Bound）的。GPU 快速执行完每个小任务，然后空闲等待 CPU 发出下一条命令。运行时间由 CPU 无法足够快地"投喂"GPU 所主导。

现代框架使用异步执行来缓解这一问题。宿主机可以把一串命令排队给 GPU，而不必等待每条命令完成。只要单个 GPU 操作足够大，宿主机就能"超前运行"，启动一个 Kernel 的开销被前一个 Kernel 的执行所掩盖。

本文余下部分假设 Kernel 足够大，开销不是主要限制因素，而专注于内存与计算（本文假设 Kernel 足够大，开销不是主要限制因素）。

4. 提升性能的两大基本策略：融合与分块

只要 Kernel 足够大、使启动开销可以忽略，性能就由 GPU 的两大物理极限决定：内存带宽与计算吞吐量。因此，提升 Kernel 的性能意味着把它在 Roofline 模型上的工作点向右上方推移。实现这一点有两大基本策略。

• 对于一串各自内存受限的操作，策略是把它们融合（Fuse）进单个 Kernel，以消除中间内存流量。
• 对于单个具有高潜在算术强度的复杂操作（如矩阵乘法），策略是使用分块（Tiling）来最大化 SM 快速内存中的数据复用。

我们逐一讨论这两种策略。

4.1 算子融合

像 y = relu(x + 1) 这样的简单操作链很常见。每个操作（add、relu）的算术强度都很低，是内存受限的。把它们作为独立的、串行的 GPU Kernel 来执行是低效的。优化这些序列的首要策略就是算子融合（Operator Fusion）。

问题出在中间的流量。考虑 y = relu(x + 1) 未融合的执行：

1. Kernel 1（add）：从全局内存读入整个张量 x，计算 tmp = x + 1，再把整个中间张量 tmp 写回全局内存。
2. Kernel 2（relu）：从全局内存读入整个张量 tmp，计算 y = relu(tmp)，再把最终张量 y 写回全局内存。

这种做法很浪费。它带来两次独立的 Kernel 启动开销，并迫使中间张量 tmp 往返慢速全局内存。

融合把这几步合并成一个更高效的 GPU Kernel。像 Triton 或 torch.compile 的 Inductor 后端这样的 JIT 编译器可以自动完成这一变换。

在融合的 Kernel 中：

1. 单个线程从全局内存把一个元素读入自己的私有寄存器。
2. 它执行所有计算，即 tmp = x + 1，然后 y = relu(tmp)，完全在那些快速寄存器中完成。
3. 它只把最终结果 y 写回全局内存。

# 未融合（概念性）
def unfused_add_relu(x):
    tmp = torch.add(x, 1) # 从 HBM 读取 x，把 tmp 写入 HBM
    y = torch.relu(tmp)   # 从 HBM 读取 tmp，把 y 写入 HBM
    return y

# 融合（概念性）
@torch.compile
def fused_add_relu(x):
    # 编译器把下面这些融合进一个 Kernel。
    # x+1 的中间结果永远不会触及 HBM。
    return torch.relu(x + 1)

中间张量 tmp 变成临时量，永远不会在全局内存中物化。这使内存流量减半（一次 x 的读取、一次 y 的写入），并消除了第二个 Kernel 的启动开销。

4.2 分块：面向计算受限 Kernel 的策略

我们对 C=A@B 的纯寄存器模型得到算术强度只有 0.25 FLOPs/Byte，远低于 A100 约 13 的脊点。原因是单个线程读取 $2N$ 个值只为执行 $2N$ 个 FLOP；数据只用一次就被丢弃。要提高数据复用、进入计算受限区间，块内的线程必须协作，把输入矩阵的大分块加载进 SM 快速片上共享内存。

这种协作的逻辑基于对矩阵乘积的分解。单个元素 C[i,j] 的计算是对 k 维的求和：$C[i,j]=\sum _{k}A[i,k]B[k,j]$。这个和可以被拆分为对分块的偏和之和。对方形分块，内层 k 维被拆成大小为 BLOCK_DIM 的块，与外层维度匹配。公式变为：

$$C[i,j]=\sum _{t=0}^{\text{NUM\_K\_TILES}-1}(\sum _{k=t\cdot \text{BLOCK\_DIM}}^{(t+1)\cdot \text{BLOCK\_DIM}-1}A[i,k]B[k,j])$$

分块算法每次迭代计算外层和中的一项（一个偏乘积）。一个线程块通过沿 k 维迭代、加载 A 和 B 的分块、在片上计算它们的乘积并累加结果，来计算一个输出块 C_tile。这通过一个加载（Load）、同步（Synchronize）、计算（Compute）的三阶段模式实现。

# 概念性算法：一个线程块计算一个输出块 C_tile。
# C_tile 对应例如 C[block_row_start:end, block_col_start:end]。

# 块中的每个线程在寄存器里持有 C_tile 的一部分。初始化为零。
thread_private_C_accumulator = zeros(...)

# 沿 k 维遍历 A 和 B 的分块。
# 每次迭代计算上面求和中一个偏乘积。
for k_tile_idx in range(NUM_K_TILES):
    # 阶段 1：加载
    # 块中所有线程协作，把一个 A 分块和一个 B 分块
    # 从慢速全局内存加载到快速共享内存。
    A_tile = load_A_tile_from_global_mem(k_tile_idx)
    B_tile = load_B_tile_from_global_mem(k_tile_idx)

    # 阶段 2：同步
    # 在任何线程开始计算之前，等待所有线程完成加载。
    # 这保证 A_tile 和 B_tile 已被完整填充。
    __syncthreads()

    # 阶段 3：计算
    # 每个线程计算自己在片上矩阵乘法中的那一份。
    # A_tile 和 B_tile 中的数据从共享内存中被大量复用。
    thread_private_C_accumulator += on_chip_matmul_piece(A_tile, B_tile)

    # 在加载下一个分块之前，等待所有线程完成计算。
    __syncthreads()

# 循环结束后，把最终累加结果写入全局内存。
write_C_tile_to_global_mem(thread_private_C_accumulator)

下面我们来考察加载、同步、计算三阶段模式的机制。

4.3 合并加载：从 HBM 到 SRAM

第一阶段把 A 和 B 的分块从慢速全局内存（HBM）加载进快速片上共享内存（SRAM）。目标是以尽可能高的内存带宽完成这次传输。这需要合并内存访问（Coalesced Memory Access）。当一个 Warp 的全部 32 个线程在一次事务中访问 HBM 中一段连续的 128 字节块时，内存访问就是合并的。

为此，Kernel 把线程索引映射到内存地址。对于一个加载同样大小数据分块的 BLOCK_DIM x BLOCK_DIM 线程块，一种常见映射是让线程 (tx, ty) 负责把 A[global_row + ty, global_k + tx] 加载进共享内存的 A_tile[ty, tx]。在这个例子中，BLOCK_DIM 为 32。

考虑一个 ty 固定、tx 从 0 到 31 变化的 Warp。

• 线程 (0, ty) 读取 A[global_row + ty, global_k + 0]
• 线程 (1, ty) 读取 A[global_row + ty, global_k + 1]
• …
• 线程 (31, ty) 读取 A[global_row + ty, global_k + 31]

假设按行优先存储，这些线程访问 32 个连续的 4 字节浮点数，即一段连续的 128 字节。这是一次完美的合并读取。整个 $32\times 32$ 分块通过 32 次这样的合并读取完成加载，块中每个 Warp 一次。

   Thread Block (32x32)          Global Memory (HBM)
                                 (One row of A's tile)
   +--------------------+
   | Warp 0 (ty=0)      | ----> [A_ij, A_i,j+1, ..., A_i,j+31]  (128 bytes)
   | (tx = 0..31)       |       (One coalesced memory transaction)
   +--------------------+
   | Warp 1 (ty=1)      | ----> [A_i+1,j, ..., A_i+1,j+31] (128 bytes)
   +--------------------+
   | ...                |
   +--------------------+
   | Warp 31 (ty=31)    | ----> [A_i+31,j, ..., A_i+31,j+31] (128 bytes)
   +--------------------+

这种加载还可以通过向量化访问（Vectorized Access）进一步提效。对合并读取而言，物理内存事务无论如何都会从 HBM 取回完整的 128 字节。区别在于 SM 如何请求这些数据。

使用标量加载时，Warp 必须发出 32 条独立的 32 位加载指令。使用向量化加载时，它只需发出 8 条更宽的 128 位加载指令。这更高效，因为 SM 每个时钟周期的指令发射槽有限。用 8 条宽指令请求数据，比用 32 条窄指令消耗更少的硬件资源。这确保内存控制器始终被全宽请求的连续流所占用，通过减少 SM 侧瓶颈来提高实际利用的内存带宽。

向量化访问通过在设备代码中转换指针（例如从 float* 到 float4*）来启用，向编译器承诺内存与向量大小对齐。

这些向量化加载的效率依赖于内存对齐（Memory Alignment）。一条 float4 指令加载一个 16 字节向量。对 4 字节浮点矩阵，该向量包含 4 个元素。只有当内存地址是 16 的倍数时，硬件才能高效执行这条指令。这意味着矩阵的内层维度 K（列数）必须是 4 的倍数。如果 K 不是 4 的倍数，行就会与 16 字节内存段错位。

考虑一个 4 字节浮点矩阵和 16 字节内存段的内存系统。

• 对齐（K=8，是 4 的倍数）：

  Memory: |<--- 16B --->|<--- 16B --->|
          [Seg 0       ][Seg 1       ]
  Row 0:  [e0 e1 e2 e3 | e4 e5 e6 e7]  (A float4 load for e0-e3 is aligned)
  Row 1:  [e0 e1 e2 e3 | e4 e5 e6 e7]  (A float4 load for e0-e3 is aligned)

• 未对齐（K=7）：

  Memory: |<--- 16B --->|<--- 16B --->|<--- 16B --->|
          [Seg 0       ][Seg 1       ][Seg 2       ]
  Row 0:  [e0 e1 e2 e3 e4 e5 e6]
  Row 1:                      [e0 e1 e2 e3 e4 e5 e6] (A float4 load for Row 1's e0-e3 spans Seg 0 and Seg 1)

这种错位迫使硬件发出更复杂、更慢的加载操作，降低内存带宽（引自 Andrej Karpathy 的相关讨论）。

重要：这种按行策略为矩阵 A 提供了合并访问。但对矩阵 B，所需的访问模式是相悖的。

1. HBM 要求：为保持合并，B 分块必须按行从 HBM 读取。
2. 计算要求：矩阵乘法本身需要访问 B 分块的列。

直接从行优先矩阵加载列是一种非合并的、带步长的访问，会把 HBM 事务串行化。因此，解决方案是用合并的行读取来加载 B 分块，然后在写入共享内存时重排数据。这种重排的结构由共享内存物理上的 Bank 架构所决定。

4.4 同步

__syncthreads() 调用起到屏障作用。在块内所有线程都到达该点之前，任何线程都不会继续执行。这保证了在计算阶段开始之前，A_tile 和 B_tile 已完整加载进共享内存（__syncthreads() 作为屏障）。

4.5 片上硬件：Bank 与 Warp

共享内存是位于流式多处理器（SM）上的物理资源。当一个线程块被调度到 SM 上运行时，它会被分配该 SM 总共享内存的一部分供其独占使用。

共享内存在物理上被划分为 32 个等大的独立存储模块，称为 Bank。这些 Bank 可以并行地响应内存请求。这个数字并非随意设定，而是与 Warp 大小 匹配。回顾一下，一个 Warp 由 32 个以锁步方式执行指令的线程组成，是内存访问的基本单元。32 个 Bank 被设计成在一个时钟周期内并行服务单个 Warp 的 32 个内存请求——前提是这些请求落在不同的 Bank 上。

以 4 字节字为单位的地址被交错分配到各个 Bank：

bank 0:  [word 0, word 32, word 64, ...]
bank 1:  [word 1, word 33, word 65, ...]
...
bank 31: [word 31, word 63, word 95, ...]

给定字地址所属的 Bank 由 $\text{bank\_id}=\text{address}mod32$ 决定。

4.6 Bank 冲突问题

要获得共享内存的全部带宽，一个 Warp 的 32 个线程必须访问落在 32 个不同 Bank 上的字。当多个线程访问映射到同一 Bank 的不同地址时，就发生 Bank 冲突（Bank Conflict）。硬件通过串行化这些请求来解决冲突，从而降低带宽。而 广播（Broadcast）——所有线程读取同一个地址——是一种快速、无冲突的操作。

这对矩阵乘法造成了一个问题。考虑共享内存中按行优先布局的一个 BLOCK_DIM x BLOCK_DIM 分块，其中 BLOCK_DIM=32。tile[row, col] 的地址是 row * 32 + col。

• 行访问（A_tile）：一个 Warp 访问 A_tile[fixed_row, t]，其中 t = 0..31。地址为 fixed_row * 32 + t。每个线程 t 的 Bank 是 $(fixed\mathrm{\_}row\times 32+t)mod32=tmod32$。由于每个线程的 t 唯一，线程访问 32 个不同的 Bank。这是一次无冲突、满带宽的访问。
• 列访问（B_tile）：一个 Warp 访问 B_tile[t, fixed_col]，其中 t = 0..31。地址为 t * 32 + fixed_col。每个线程 t 的 Bank 是 $(t\times 32+fixed\mathrm{\_}col)mod32=fixed\mathrm{\_}colmod32$。全部 32 个线程都指向同一个 Bank。这造成 32 路 Bank 冲突，使内存访问串行化。

解决方案是在共享内存中按转置布局存储 B_tile。

# 加载阶段线程 (tx, ty) 的动作
# A 直接加载，B 边加载边转置
A_tile[ty, tx] = A_global[global_row + ty, global_k + tx]
B_tile[tx, ty] = B_global[global_k + ty, global_j + tx]  # 索引被交换

这种"边加载边转置"（Load-and-Transpose）的操作改变了片上计算。偏乘积某个元素的计算不再是对 A_tile 的一行与 B_tile 的一列做点积。相反，使用转置后的片上 B_tile，公式变为：

$$C_{\text{partial}}[i,j]=\sum _k{A}_{\text{tile}}[i,k]\cdot B_{\text{tile}}[j,k]$$

在这种形式下，一个为固定 i 计算不同 j 值的 Warp，将访问 A_tile 的一行和片上 B_tile 的一行。两者都是无冲突的访问模式。这一招同时解决了 HBM 合并要求和 SRAM Bank 冲突问题。

   Load-and-Transpose Operation (Thread tx, ty)
   Reads row-wise from HBM, writes column-wise to SRAM

   Global Memory (HBM)                Shared Memory (SRAM)
   +-------------------------+        +-----------------------+
   | B[k_base+ty, j_base+tx] | -----> |      B_tile[tx, ty]   |
   +-------------------------+        +-----------------------+

   Result: HBM reads are coalesced, SRAM reads are conflict-free.

4.7 片上计算阶段：提高算术强度

数据就位共享内存后，线程块开始计算。目标是最大化对该快速片上内存的数据复用。我们分析两种组织片上计算的策略。

策略 1：一个线程计算一个输出

最简单的方法是把一个输出元素映射到一个线程。一个 BLOCK_DIM x BLOCK_DIM 线程块计算一个 TILE_DIM x TILE_DIM 数据分块，其中 BLOCK_DIM 与 TILE_DIM 相等。这一策略在概念上类似于 Boehm 文章中的 Kernel 3，后者引入了共享内存缓存（对应 Boehm 的 Kernel 3，引入共享内存缓存）。每块 1024 线程的硬件上限把 BLOCK_DIM 限制为最多 32。线程 (tx, ty) 负责单个输出元素 C_partial[ty, tx]。

# 单个线程 (tx, ty) 的动作，其中 BLOCK_DIM = TILE_DIM
c_accumulator = 0.0
for k in range(TILE_DIM):
    c_accumulator += A_tile[ty, k] * B_tile[tx, k]

这一策略的算术强度为 TILE_DIM / 4。

• 总 FLOP：线程块执行 $2\times {\text{TILE\_DIM}}^3$ 个 FLOP。
• 访问总字节数（HBM）：线程块加载两个数据分块，共 $8\times {\text{TILE\_DIM}}^2$ 字节。
• 算术强度（AI）：$(2\times {\text{TILE\_DIM}}^3)/(8\times {\text{TILE\_DIM}}^2)=\text{TILE\_DIM}/4$ FLOPs/Byte。

当 TILE_DIM 被限制为 32 时，最大 AI 为 $32/4=8$。这不足以跨越 A100 约 13 的脊点。Kernel 仍是内存受限的。

策略 2：一个线程计算多个输出

要提高 AI，必须在不增加线程数的前提下增大 TILE_DIM。这就需要把数据分块大小与线程块大小解耦，给每个线程分配更多工作。这一策略对应 Boehm 文章中 Kernel 5 的目标。

一个 $16\times 16$ 线程块（BLOCK_DIM = 16，256 个线程）可以计算一个 $64\times 64$ 数据分块（TILE_DIM = 64）。每个线程现在计算输出的一个 $4\times 4$ 子块。这要求 TILE_DIM=64 不超出共享内存容量（TILE_DIM=64 需不超出共享内存容量）。

# 一个线程计算一个 4x4 输出子块
# TILE_DIM = 64, BLOCK_DIM = 16
c_regs = [[0.0] * 4 for _ in range(4)]
a_regs = [0.0] * 4
b_regs = [0.0] * 4

for k in range(TILE_DIM):
    # 把 A_tile 和 B_tile 的一个条带加载进寄存器
    for i in range(4): a_regs[i] = A_tile[thread_row*4 + i, k]
    for j in range(4): b_regs[j] = B_tile[thread_col*4 + j, k]

    # 在寄存器中计算外积，累加进 c_regs
    for i in range(4):
        for j in range(4):
            c_regs[i][j] += a_regs[i] * b_regs[j]

AI 的计算仍是 TILE_DIM / 4。当 TILE_DIM = 64 时，AI 为 $64/4=16$ FLOPs/Byte。这超过了 A100 的脊点。Kernel 现在是计算受限的。

计算受限 Kernel 的运行时间受限于 SM 的算术吞吐量。这并不能保证高绝对性能。一个 Kernel 可能是计算受限的，但仍然很慢——如果它的 FLOP 低效（例如使用标量 FP32 数学而非 Tensor Core（Tensor Core 是专用矩阵乘法硬件） 等专用硬件），或者 GPU 因功率限制而低于峰值时钟运行。

上面代码中的内层循环还可以进一步优化。一个线程从 A_tile 加载四个独立的 float 值到 a_regs。它可以改为发出一条指令加载一个 16 字节的 float4 向量。这种来自共享内存的向量化加载减少了片上数据移动发出的指令数，提高了计算阶段的效率。这对应 Boehm 文章 Kernel 6 中使用的片上向量化改进。

最后一个考虑：分块的量化

如果矩阵维度不是分块大小的整数倍，Kernel 会启动执行无效计算的额外线程块。

要用 TILE_M x TILE_N 分块覆盖一个 $M\times N$ 矩阵，GPU 会启动一个 $\lceil M/\text{TILE\_M}\rceil \times \lceil N/\text{TILE\_N}\rceil$ 的线程块网格。用 $32\times 32$ 分块切分一个 $65\times 65$ 矩阵，需要 $\lceil 65/32\rceil \times \lceil 65/32\rceil =3\times 3$ 的块网格。Kernel 的逻辑是固定的：每个块都被编程为执行一个完整 $32\times 32$ 分块的算术。

      Columns 0-31      Columns 32-63     Columns 64-95
    +-----------------+-----------------+-----------------+
R 0 |                 |                 |                 |
o-31|   Block 0,0     |   Block 0,1     |   Block 0,2     |
w   | (Full work)     | (Full work)     | (Wasted work)   |
s   |                 |                 |                 |
    +-----------------+-----------------+-----------------+
R 32|                 |                 |                 |
o-63|   Block 1,0     |   Block 1,1     |   Block 1,2     |
w   | (Full work)     | (Full work)     | (Wasted work)   |
s   |                 |                 |                 |
    +-----------------+-----------------+-----------------+
R 64|                 |                 |                 |
o-95|   Block 2,0     |   Block 2,1     |   Block 2,2     |
w   | (Wasted work)   | (Wasted work)   | (Wasted work)   |
s   |                 |                 |                 |
    +-----------------+-----------------+-----------------+

NVIDIA 官方文档指出："虽然库保证任何分块都不会执行非法内存访问，但所有分块都会执行同样多的数学运算。"我的理解是（欢迎指正）：边界块之所以做无效工作，是因为 Kernel 显式地对数据做了填充。被分配去加载矩阵边界之外元素的线程，会因一个保护条件而被阻止这样做。相反，它们向片上共享内存分块中的对应位置写入零。算术循环不会缩短。Kernel 的逻辑在整个分块上保持一致。一个 Warp 内的所有线程都执行相同的乘加指令。数据对应填充零的线程仍然执行该指令，只是做了一次无用计算，例如 C += A * 0。硬件资源被使用了，但工作被丢弃。

5. 更多性能考量

我们已经让 Kernel 进入计算受限状态。它的性能现在受限于片上算术的速度。然而，通过管理硬件的其他方面，Kernel 仍可以变得更快。以下是三个这样的考量。还有其他方面，但我水平还不够，暂时写不了。其他的可参见 Boehm 文章。

5.1 Occupancy 与延迟隐藏

当一个 Warp 执行长延迟指令（例如从全局内存读取）时，它会停滞（Stall）。在数据到达之前它无法执行下一条指令，而数据到达可能需要几百个时钟周期。在这段时间里，如果停滞的 Warp 是唯一可用的工作，SM 的计算单元就会空闲。

SM 通过执行其他工作来隐藏这种延迟。它可以同时容纳多个线程块，形成一个驻留 Warp 池。当一个 Warp 停滞时，SM 的硬件调度器会立即从这个池中切换到另一个就绪可运行的 Warp。这一机制称为延迟隐藏（Latency Hiding）。

+-------------------------------------------------------------------+
| Streaming Multiprocessor (SM)                                     |
|                                                                   |
|  [Block A]              [Block B]                                 |
|   - Warp A1 (Ready)      - Warp B1 (Ready)                        |
|   - Warp A2 (Stalled -> waiting on HBM)                           |
|        |                  |                                       |
|        +------------------v------------------+                    |
|           [ Pool of Ready-to-Run Warps ]                          |
|           [ A1, B1 ]                                              |
|                           |                                       |
|                   +-------v-------+                               |
|                   | SM Scheduler  | --> [Execute instructions]    |
|                   +---------------+                               |
|                                                                   |
+-------------------------------------------------------------------+

Occupancy（占用率） 是 SM 上活跃 Warp 数与它能支持的最大 Warp 数之比。高占用率给调度器提供更大的 Warp 池可供选择。这增加了它在任意周期找到就绪 Warp 来执行的可能性，从而使计算单元保持活跃。

这引出了每块资源使用量与可驻留块数量之间的权衡。两个极端可以示意如下：

+------------------------------------+ +----------------------------------------------+
| SM with High AI, Low Occupancy     | | SM with Low AI, High Occupancy               |
|                                    | |                                              |
| +--------------------------------+ | | +----------+ +-----------+     +-----------+ |
| | Block 0 (uses 64KB SMEM)       | | | | Block 0  | | Block 1   | ... | Block N   | |
| | TILE_DIM=128 -> High AI        | | | | (8KB SMEM) | (8KB SMEM)|     | (8KB SMEM)| |
| +--------------------------------+ | | +----------+ +-----------+     +-----------+ |
|                                    | |                                              |
| --> Low # of resident blocks.      | | --> High # of resident blocks.               |
| --> Small pool of warps for        | | --> Large pool of warps for                  |
|     latency hiding.                | |     latency hiding.                          |
+------------------------------------+ +----------------------------------------------+

我们调节 Kernel 的资源使用，在高 AI 的好处与足够 Occupancy 的必要性之间取得平衡。调节的主要杠杆是线程块维度（BLOCK_DIM）、每块分配的共享内存量（由 TILE_DIM 决定）以及每线程使用的寄存器数（调节的主要杠杆）。

5.2 避免线程发散

当 Warp 内的线程对条件分支（if-else）的判断结果不一致时，就会发生线程发散（Thread Divergence）（线程发散导致串行执行）。发生发散时，硬件通过串行执行不同的代码路径来解决。首先，走 if 路径的线程执行它，其余线程处于非活跃状态。然后，else 路径的角色互换。

# 一个 32 线程的 Warp 遇到一个 `if` 语句：
if (thread_id < 16)
    # Path A
else
    # Path B

Execution Timeline:

Time ->
+------------------------------------------------------------------+
| Warp Execution                                                   |
|                                                                  |
|  Cycle 1: Path A is executed.                                    |
|   - Threads 0-15:  Active, execute Path A code.                  |
|   - Threads 16-31: Inactive, masked off.                         |
|                                                                  |
|  Cycle 2: Path B is executed.                                    |
|   - Threads 0-15:  Inactive, masked off.                         |
|   - Threads 16-31: Active, execute Path B code.                  |
|                                                                  |
| Result: Two cycles are required instead of one.                  |
|         Effective throughput is halved.                          |
+------------------------------------------------------------------+

这种串行化使发散代码的执行时间翻倍，Warp 的有效吞吐量减半。我们通过编写无分支代码来避免这一代价——在性能关键部分使用 min、max 等原语替代 if-else。

5.3 量化

量化降低张量元素的精度，例如从 FP32 降到 FP16 或 BF16。这有两方面效果。首先，它减少了存储每个元素所需的内存，例如减少 2 倍。因此，每秒能从全局内存传输到共享内存的元素数量翻倍。这把 AI 提高了 2 倍。

其次，像 A100 这样的 GPU 对低精度元素可以运行得更快。例如在 A100 上，某些 FP16 操作可达 312 TFLOPS，而 FP32 操作被限制在 19.5 TFLOPS。因此理论上可以把计算加速 16 倍。

因此，量化可以把我们在 Roofline 图上向右上方推移。

相关链接

• 让深度学习跑得飞快：从第一性原理理解性能
• 矩阵乘法背景指南
• GPU 与 NPU 硬件组成详解
• 大模型推理优化技术概览

脚注

（性能层级示意）: 原文脚注：性能层级示意。 （A100 的 SRAM 带宽为 19.5 TB/s）: 原文脚注：SRAM 带宽 19.5 TB/s。 （寄存器访问速度原文标注为 '?? TB/s'，未给出确切数值）: 原文脚注：寄存器访问速度原文标注为"?? TB/s"（未给出确切数值）。 （独立线程调度在 Ampere 及以后架构可用，此处为简化模型）: 原文脚注：独立线程调度（independent thread scheduling）在 Ampere 及以后架构可用，此处为简化模型。 （纯寄存器模型不足以让矩阵乘法达到计算受限）: 原文脚注：纯寄存器模型不足以达到计算受限。 （本文假设 Kernel 足够大，开销不是主要限制因素）: 原文脚注：本文假设 Kernel 足够大，开销不是主要限制。 （引自 Andrej Karpathy 的相关讨论）: 原文脚注：引自 Andrej Karpathy 的相关讨论。 （__syncthreads() 作为屏障）: 原文脚注：__syncthreads() 作为屏障。 （对应 Boehm 的 Kernel 3，引入共享内存缓存）: 原文脚注：对应 Boehm 的 Kernel 3，引入共享内存缓存。 （TILE_DIM=64 需不超出共享内存容量）: 原文脚注：TILE_DIM=64 需不超出共享内存容量。 （Tensor Core 是专用矩阵乘法硬件）: 原文脚注：Tensor Core 是专用矩阵乘法硬件。 （调节的主要杠杆）: 原文脚注：调节的三个主要杠杆。 （线程发散导致串行执行）: 原文脚注：线程发散导致串行执行。

参考

1. Basic facts about GPUs — Damek Davis
2. Making Deep Learning Go Brrrr From First Principles — Horace He
3. What Shapes Do Matrix Multiplications Like? — Thonking
4. How to Optimize a CUDA Matmul Kernel for cuBLAS-like Performance: a Worklog — Simon Boehm
5. Matrix Multiplication Background User's Guide — NVIDIA