最近在看 MIT 6.S191 Lecture 9 的时候，反复听到一句话：Data Movement Matters。
之前写过 Arithmetic Intensity 和 Compute/Memory Bound，也刚写完 FlashAttention 的分块和 Online Softmax。这些文章里"HBM"和"片上存储"出现了很多次，但我发现自己其实一直没有认真想过一个问题：HBM 到底是什么？它全称叫 High Bandwidth Memory——高带宽内存——既然带宽已经很高了，为什么还要费这么大劲减少对它的访问？
这篇把 GPU 存储层级理一遍，算是给前面那些文章补一个前置背景。
GPU 的数据存在哪#
做 CPU 性能优化的人对存储层级不陌生：Register 最快但最小，L1 / L2 / L3 Cache 逐级变大变慢，最后是 DRAM。程序要跑得快，就得让热点数据尽量待在靠近计算单元的地方。
GPU 也有类似的层级结构，但名字和用法都有区别。
HBM#
HBM 是 High Bandwidth Memory 的缩写。物理上，它是多层 DRAM 芯片堆叠在一起，通过硅中介层（interposer）连接到 GPU 计算 Die 的旁边。
关键点：HBM 不在 GPU 的计算 Die 上。 它是独立的存储芯片，靠近 GPU 但不是片上资源。当说"GPU 显存 80 GB"的时候，指的就是 HBM 的容量。模型权重、输入数据、中间 tensor——这些占大头的数据通常都存在 HBM 里。
HBM 的带宽确实很高。A100 SXM 的 HBM2e 带宽大约 2 TB/s，比 CPU 端的 DDR4/DDR5 快了一个数量级。但它毕竟是片外存储——和计算 Die 内部的 SRAM 相比，访问延迟和每次传输的能耗都要高得多。
Shared Memory#
Shared Memory 是 GPU 每个 SM（Streaming Multiprocessor）内部的片上 SRAM。
它和 HBM 有两个本质区别。第一，它在计算 Die 上，物理距离极短，带宽可以到 19 TB/s 量级，远高于 HBM。第二，它由程序显式管理——CUDA 程序员自己决定把什么数据搬进来、什么时候用、用完怎么处理。这一点和 CPU 的 Cache（硬件自动管理）思路不同。
代价是容量极小。A100 上每个 SM 的 Shared Memory 最大可配置到 228 KB。和 HBM 的 80 GB 比，差了五个数量级。
Registers#
Register 是离运算单元最近的存储，每个线程私有，保存计算过程中的临时值。速度最快，但资源也最有限（A100 每个 SM 有 256 KB Register File）。
L1 / L2 Cache#
GPU 也有硬件管理的 Cache 层级。L1 和 Shared Memory 在物理上共享同一块 SRAM（可配置比例）；L2 是全局共享的，A100 上大约 40 MB。它们由硬件自动管理，程序不直接控制。这里不展开 L1 和 Shared Memory 的具体硬件组织。
全局视图#
把这些层级放在一起：
容量越大的层级离计算单元越远，访问越慢。这是 GPU 和 CPU 共享的基本矛盾：存储容量和访问速度不可兼得。
需要说明：上图是简化的逻辑关系。真实的数据通路不一定是严格的逐级串行——比如 L2 Cache 可能命中后直接返回数据，不需要真正访问 HBM。所以"每次读 Global Memory 都一定访问 HBM"是不准确的说法。
为什么数据搬运可能比计算贵#
现在可以回答开头的问题了：HBM 带宽确实高，但 GPU 的计算吞吐更高。
A100 的 BF16 Tensor Core 峰值算力约 312 TFLOPS，HBM 带宽约 2 TB/s。算一下：每从 HBM 搬 1 byte 数据，GPU 能做大约 156 次浮点运算。如果一个操作每搬 1 byte 只做两三次运算（比如 LayerNorm），那计算单元大部分时间在等数据——这就是 memory bound。
用一个极端例子。假设某块数据需要参与 100 次计算：
方案 A 每次计算都从 HBM 重新加载，搬运 100 次。方案 B 先把数据搬到 Shared Memory，然后在片上反复使用，HBM 只搬运 1 次。
方案 B 不一定在所有场景都更快——取决于数据大小是否放得进 Shared Memory、计算模式是否允许复用等具体条件。但当数据确实能被多次复用时，减少 HBM 搬运可以显著提高计算单元的有效利用率。
这就是 GPU 做 GEMM 时使用 Tiling 的核心动机：把矩阵切成小块搬到片上，让每个元素参与尽可能多的计算再丢弃。之前那篇文章详细讨论了 Arithmetic Intensity、Roofline Model，以及为什么矩阵乘法天然适合这种复用模式。
关于 Compute Bound 和 Memory Bound 的严格定义、常见 ML 操作的分类、以及 Accelerator Intensity 随硬件代际增长导致 memory wall 加剧的趋势，这里就不重复了，直接参考那篇。
FlashAttention 把这些思想用在了哪#
标准 Attention 的计算过程中会生成一个 \(S \times S\) 的 Score Matrix（\(S\) 是 Sequence Length）。朴素实现需要把这个矩阵完整写入 HBM，再读回来做 Softmax，再写回去，再读回来和 V 相乘——中间矩阵在 HBM 里反复搬运。
FlashAttention 的核心思想是：这个中间矩阵不需要完整存下来。 通过分块处理 Q、K、V，配合 Online Softmax 在线维护归一化统计量，每个 Block 的 Score 在片上计算、片上消费、片上丢弃，避免了完整 \(S \times S\) 矩阵对 HBM 的写入和读取。
这里面有两件事需要分开理解。Online Softmax 是算法层面的等价重构——它让 Softmax 从全局操作变成可分块的在线操作，打破了"必须看到整行 Score 才能归一化"的依赖。Tiling、Kernel Fusion 和片上数据复用则是调度层面的优化——不改变数学，但改变数据的搬运模式。FlashAttention 同时做了这两件事。
完整的推导——分块策略的选择、Online Softmax 的数值例子、Value 累加器如何在线更新、和 Halide Algorithm/Schedule 分离的类比——都在这篇里。
值得强调的是：FlashAttention 并没有减少 Attention 的 \(O(S^2 d)\) 计算量。 它减少的是昂贵的中间结果存储和 HBM 数据搬运。这正是 IO-Aware Algorithm 的含义——算法设计不仅考虑计算复杂度，也主动适应硬件的存储层级。
写在最后#
回过头看，这篇其实是在补一块拼图。之前讨论 Arithmetic Intensity 的时候，“HBM 带宽"是一个现成的参数；讨论 FlashAttention 的时候，“片上存储"是一个理所当然的概念。但把 GPU 存储层级真正摊开来看，才能把这些讨论串起来。
HBM 是片外的高带宽 DRAM，Shared Memory 是片上的程序可控 SRAM。 两者在容量、速度和使用方式上有本质区别。理解这个区别，是理解后续所有 IO-Aware 优化的前提。
参考资料#
- 从 CS336 到 ML Infra：理解 Arithmetic Intensity
- FlashAttention 的核心思想：为什么 Attention Matrix 根本不需要存下来？
- 从参数量到 Peak Memory：模型显存到底应该怎么算？
- NVIDIA. CUDA C++ Programming Guide — Memory Hierarchy.
- NVIDIA. A100 Tensor Core GPU Architecture Whitepaper.
- Dao et al. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS, 2022.