GPU 内存层级
GPU 内存被组织为一种层级结构。小容量的片上存储紧邻计算单元,提供低延迟;而高带宽内存(HBM)则在片外提供大得多的容量。大多数 GPU 内核优化技术归根结底就是搬运更少的数据,或在层级中更快的层级上复用数据。
寄存器
寄存器是 GPU 上最快的存储。每个线程都拥有一组逻辑上私有的寄存器,即其他线程无法访问它们。然而在物理上,这些寄存器来自 SM 上的一个大寄存器文件(例如 H100 上每个 SM 256 KB),它被所有驻留线程共享并分割。
访问寄存器比访问共享内存或 HBM 便宜得多,因为数据已经在片上,并且计算单元可以直接使用。
寄存器是一种有限资源。当寄存器分配成为瓶颈时,每个线程更高的寄存器使用量会减少驻留线程和 warp 的数量。这是占用率权衡的主要驱动因素。
共享内存(SMEM)与 L1 缓存
现代 GPU 在每个 SM 内提供两种快速的片上内存:共享内存(SMEM)和 L1 缓存。它们通常使用相同的物理 SRAM,但作用不同。
L1 缓存由硬件管理。当线程从全局内存(HBM)读取时,GPU 可能会自动把数据存到 L1 中。如果再次访问相同的数据,就可以从 L1 更快地获得(缓存命中),而不必回到 HBM。这个过程对程序员是透明的:你不需要显式地把数据加载到 L1,也无法控制什么会留在那里。因此,最好把 L1 视为一种尽力而为的优化,当内存访问模式存在时间局部性或空间局部性时,它能提升性能。
相比之下,共享内存由程序员管理。它是一个容量较小、由程序员显式分配的内存空间,由线程块中的所有线程共享。线程可以直接读写共享内存,并使用同步机制(例如 __syncthreads())来协调访问。这使得共享内存成为线程间协作的可预测、可控的工作空间。
共享内存的主要目的是减少代价高昂的 HBM 访问。一种常见模式是先把数据从全局内存加载到共享内存一次,然后在多个线程之间多次复用。由于共享内存在片上且比 HBM 快得多,这可以显著提升性能。这种模式出现在许多高性能内核中,包括矩阵乘法和注意力。
共享内存被组织成存储体(bank,通常 32 个)。当 warp 中的多个线程同时访问同一存储体中不同的地址时,就会发生存储体冲突(bank conflict),这些访问会被串行化(读取相同地址的线程由广播服务,不会冲突)。避免存储体冲突是内核调优中常见的微优化。
下面是对比:
| 项目 | 共享内存 | L1 缓存 |
|---|---|---|
| 管理方 | 程序员 | 硬件 |
| 作用范围 | 线程块 | 每个 SM(该 SM 上的所有线程块) |
| 持久性 | 线程块生命周期 | 由硬件驱逐 |
| 存储体冲突 | 可能发生 | 不会发生(由硬件处理) |
| 最适合 | 可规划的数据复用 | 不规则或不可预测的访问 |
一个常见问题:既然已经有了寄存器,为什么还需要共享内存(SMEM)和 L1 缓存?
寄存器是 GPU 上最快的存储,但光靠它们还不够。
- 寄存器是私有的:一个线程无法直接读取分配给另一个线程的寄存器。warp shuffle 指令可以在 warp 内交换寄存器值,而共享内存在整个线程块内支持复用。L1 还能防止重复访问到达 HBM。
- 寄存器是有限的:每个线程只能获得少量寄存器。像完整的权重张量或大分块这样的数据放不下,因此必须来自内存。L1 帮助自动缓存它们,共享内存则让你显式地暂存和复用。
- 线程之间没有协调手段:线程无法通过寄存器通信。共享内存提供了用于协作和同步的共享工作空间。
L2 缓存
L2 缓存是容量最大的片上内存,由 GPU 上的所有 SM 共享,是前往 HBM 之前的最后一个快速中转站。在 H100 上,L2 缓存约为 50 MB;在 A100 上,约为 40 MB。
L2 缓存由硬件管理。你不会显式地把数据加载进去,而是由它自动缓存最近的全局内存访问。对于在跨线程块之间存在一定程度数据复用的工作负载,L2 可以显著减少 HBM 流量。
L1 缓存和共享内存在单个 SM 内工作,而 L2 的存在是为了捕捉跨 SM、跨线程块的复用。那些放不进片上内存、或会被多个线程块访问的数据,仍可以通过 L2 复用,而不必反复从 HBM 获取。
在 LLM 推理中,L2 缓存可以用于:
- 分组查询(GQA)或多查询(MQA)注意力中被多个查询头共享的 KV 缓存条目
- 在批内多个元素之间复用的权重分块
- 小型查找表或元数据
然而,当工作集非常大时(这在 LLM 推理中很常见),L2 命中率会下降,HBM 带宽成为瓶颈。
HBM(高带宽内存)
HBM 是 GPU 的主内存(通常称为显存或全局内存)。它存储模型权重、KV 缓存、激活以及所有其他大型数据结构。现代数据中心 GPU 使用 HBM2e、HBM3 或 HBM3e:
| GPU | HBM 容量 | HBM 峰值带宽 |
|---|---|---|
| NVIDIA A100 SXM | 80 GB | 2.0 TB/s |
| NVIDIA H100 SXM | 80 GB | 3.35 TB/s |
| NVIDIA H200 SXM | 141 GB | 4.8 TB/s |
对于 LLM 推理:
- HBM 容量决定了单个 GPU 上能保留多少模型状态和 KV 缓存。
- HBM 带宽通常决定解码期间的 token 吞吐量。
与片上内存相比,HBM 容量大但速度慢。一次 HBM 访问需要数百个周期。这就是内核优化专注于最小化 HBM 与计算单元之间流量的原因。
内存金字塔
下面的数值提供了近似 Hopper 级的示例。确切的容量、延迟和带宽因 GPU、访问模式和测量方法而异。
| 层级 | 容量(H100) | 带宽 | 作用范围 | 延迟 | 管理方 |
|---|---|---|---|---|---|
| 寄存器 | 每个 SM 256 KB | 最高 | 每线程 | 约 1 个周期 | 编译器 |
| 共享内存 / L1 | 256 KB 组合池;SMEM 最多 228 KB | 有效约 20 TB/s | 每线程块(SMEM)/ 每个 SM(L1) | 约 20-30 个周期 | 程序员(SMEM)/ 硬件(L1) |
| L2 缓存 | 50 MB | 约 12 TB/s | 所有 SM | 约 200 个周期 | 硬件 |
| HBM | 80 GB | 3.35 TB/s | 全局 | 约 400+ 个周期 | 程序员 / 运行时 |
许多内核优化技术,包括分块、融合和数据布局调整,通过把常用数据保存在寄存器或共享内存中,而不是反复从 HBM 读取,来减少金字塔底层的流量。