What Is a GPU Doing With My Tensors?

很多 LLM researcher 都能从模型这一侧解释一次 forward pass。Hidden state 经过 attention 和 MLP，linear layer 改写每个 token 的表示，最后一层产生 vocabulary logits。GPU 却经常停留在模型之外：我们知道一块卡有多少显存、峰值 FLOPs 是多少，也知道训练用了多少块卡，却未必能说清这些 tensor 真正放在哪里，计算又怎样发生。

平时这层模糊没有太大影响。Framework 会把模型 operation 交给 GPU，也会替我们管理大部分 memory。等到 performance 出现问题，HBM、memory bandwidth、Tensor Core、kernel fusion、NVLink、all-reduce、compute-bound 和 communication-bound 才会一起冒出来。每个词都解释了一部分现象，缺少的仍然是一张把它们连接起来的地图。

这篇 blog 想补上这张地图。从 weight、activation 和 KV Cache 在一块 GPU 里的位置出发，我们会一路追踪 tensor operation 如何落到实际硬件上，inference 与 training 如何改变数据的生命周期，以及模型跨越多块 GPU 后哪些 tensor 必须随之移动。HBM、Tensor Core、kernel fusion 和 all-reduce 最终都会落进同一个判断框架：capacity、bandwidth、compute 与 communication。

一块 GPU 里有什么

把许多硬件细节暂时收起来，一块用于运行 LLM 的 GPU 有两个主要部分：大量并行的计算资源，以及持续向这些资源供应数据的 memory system。计算资源决定同一时间可以完成多少工作，memory system 决定当前需要的 value 能否及时到达。GPU 的许多设计都在协调这两部分。

模型的数据还要先进入 GPU。Checkpoint 平时以文件形式保存在 SSD 上，加载时经过 CPU RAM，再被复制到 GPU memory。模型开始运行以后，weight、activation 和 cache 会在 GPU 内部的不同位置之间移动。Figure 1 把这两个尺度放在同一幅图里：外面的仓库负责长期保存和转运，GPU 内部的厨房负责为许多并行工位持续供货。

GPU 上容量最大的常用 memory 是 HBM，全称 high-bandwidth memory。我们平时说一块 GPU 有 80 GB 或 192 GB 显存，通常讨论的就是这一层可以容纳多少数据。推理时，HBM 通常保存 model weight、各层暂时产生的 activation，以及 attention 为过去 token 保留的 KV Cache。训练还要保存 parameter gradient 和 optimizer 使用的历史 state。Checkpoint 主要记录 weight，所以文件大小只覆盖了运行时 memory 的一部分。

HBM 容量很大，计算单元旁边的 storage 则小得多。GPU 的计算资源分布在许多 Streaming Multiprocessor 中，简称 SM。每个 SM 都拥有自己的执行单元、register，以及由 shared memory 和 L1 cache 使用的片上资源；一层更大的 L2 cache 由整块 GPU 上的 SM 共享。现代 SM 中的 Tensor Core 可以高速执行适合 Transformer 的矩阵乘加。

Framework 提交给 GPU 的一项具体工作叫作 kernel。一个 kernel 可以完成 matrix multiplication、normalization、elementwise operation，也可以把多个步骤融合在一起。GPU 会把 kernel 拆成大量并行工作，再分配给可用的 SM。SM 承接这些并行工作，memory hierarchy 持续向它供应数据。

这套层级有一个非常重要的方向感。越靠近计算单元，memory 越适合快速访问，容量也越有限。HBM 可以保存完整的大 tensor；L2、shared memory、L1 和 register 逐步把范围收缩到当前正在处理的 working set，也就是这一小段计算眼下需要的数据。GPU 运行时，大 tensor 中当前需要的部分会被带到计算单元附近，使用过的数据则尽量留在近处继续复用。实际路径由 kernel 和 hardware cache 共同决定，数据未必依次停留在图中的每一层。

Inference 和 training 会在同一套硬件上产生不同的显存占用。Inference 主要保留 weight 和每个请求的 KV Cache，layer 中间产生的 activation 大多只需短暂存在。Training 的 forward 需要保存 backward 之后还会用到的 activation，backward 产生 gradient，optimizer 还可能为每个 parameter 保存额外 state。因此，同一个模型能顺利 inference，也可能无法用相同的 GPU configuration 训练。

把这些组件放回厨房，大致是下面这组对应关系：

一块 GPU 就是在不断协调 storage、movement 和 compute。Model state 从 SSD 和 CPU RAM 进入 HBM，kernel 再把当前 working set 带到许多 SM 附近完成计算。所需 tensor 超过 HBM capacity 会触发 out of memory（OOM）；大量时间花在等待数据时，operation 会呈现 memory-bound；大量时间花在执行算术时，它会呈现 compute-bound。模型跨越多块 GPU 以后，还要判断多少数据需要跨设备传输，以及计算是否在等待这部分 communication。

数据怎样抵达计算单元

模型运行时，完整的 weight 和中间结果可能占据几十甚至几百 GB，计算单元在某一刻只会用到其中很小的一部分。把所有数据都放在远处的大 memory 里，计算单元会频繁等待；把所有数据都放到计算单元旁边，芯片又没有足够空间。

GPU 因此准备了几种容量和位置不同的 memory：远处的层级容量更大，靠近计算的层级更小，也更适合快速访问。这种安排叫作 memory hierarchy。当前这一步所需的那一小部分数据叫作 working set；一份数据在被移走前多用几次，叫作 reuse。Memory hierarchy 要同时做到两件事：让足够大的数据有地方住，让正在使用的数据尽量靠近计算。

HBM：模型运行时的主要住处

在 PyTorch 里看到 tensor 位于 cuda:0，通常表示它已经占用了第一块 GPU 的 device memory，也就是 GPU 自己可以直接访问的大容量 memory。数据中心 GPU 的 device memory 通常由 high-bandwidth memory（HBM） 提供。Model weight 会长期放在这里，计算一层网络时产生的中间结果，也就是 activation，也会在这里出现。Inference 还要保存 attention 为过去 token 算出的 key 和 value，这部分数据叫作 KV Cache。Training 则要保存 parameter gradient 和 optimizer 维护的历史 state。靠近计算的其他 memory 容量小得多，无法承担这些数据的长期存放。

理解 HBM 需要分清两个指标。Memory capacity 表示同一时间可以放下多少 bytes；memory bandwidth 表示一秒内最多可以在 HBM 和 GPU 芯片之间传送多少 bytes。Capacity 不够会触发 out of memory（OOM）；bandwidth 不够时，数据虽然放得下，计算单元仍会等待它抵达。NVIDIA H100 SXM 的官方规格是 80 GB HBM 和 3.35 TB/s 峰值 memory bandwidth (NVIDIA, 2026)。即使按峰值连续读取，扫过 80 GB 数据也要大约 24 ms。

HBM 的高 bandwidth 来自它的物理设计。制造商把多层存储芯片垂直堆叠，再用大量连接把几组堆叠同时接到 GPU 附近 (SK Group, 2024)。可以把传送数据的接口想成道路。提升单条车道的速度有物理极限，HBM 选择铺开很多并行车道，让更多数据同时出发。这里改善的是整体吞吐量，一次访问本身依然需要时间。

Training forward 为 backward 留下的中间结果称为 saved activation；kernel 或 communication 临时申请的执行空间称为 workspace。把这些数据与长期存在的 model state 放在一起，两份粗略的清单就能说明 LLM 为什么同时需要很大的 capacity 和很高的 bandwidth：

\[\text{training HBM} \approx \text{weights} + \text{saved activations} + \text{gradients} + \text{optimizer states} + \text{temporary workspace}.\] \[\text{inference HBM} \approx \text{weights} + \text{KV caches of active requests} + \text{temporary workspace}.\]

Training 要同时容纳 weight、为 backward 留下的 activation、gradient 和 optimizer state。这些数据还会在 forward、backward 与 parameter update 中持续读写。模型状态分布到更多 GPU 时，系统同时增加了更多 HBM。B200 每块 GPU 提供 180 GB HBM 和最高 8 TB/s bandwidth，八卡 server 合计拥有 1.44 TB；GB200 NVL72 的 72 块 Blackwell GPU 总计配置 13.4 TB HBM (NVIDIA, 2026; NVIDIA, 2026)。大型 training cluster 对 HBM 的需求因此十分集中。

Inference 的长期状态少一些，使用时间却可能更长。每部署一份完整的模型副本，也就是一个 model replica，都要让 weight 常驻 HBM；每个活跃请求还会保留随 context 增长的 KV Cache (Luo & Shen, 2026)。生成每个新 token 时，GPU 又要读取 weight 和已有 cache。Long context、agent workload 和 o1 一类 long-reasoning model 会延长这项工作，让同一份 weight 被读取更多轮，也让 request state 留得更久 (Gandhi et al., 2026; Micron Technology, 2026)。Training 和 inference 因而都会同时消耗 HBM capacity 与 bandwidth，只是 tensor 的生命周期和访问方式不同。

这也解释了为什么我们突然会在 GPU 之外频繁听到 SK hynix。HBM 要经过 memory die 制造、堆叠、连接、封装和测试，扩产同时受晶圆与先进封装能力约束。Micron 在 2025 年 12 月已经完成全部 2026 年 HBM supply 的价格和数量协议；SK hynix 随后也表示客户需求超过供给能力，并继续增加投资 (Micron Technology, 2025; SK hynix, 2026)。

AI 对 HBM 的争夺，也催生了 SK hynix 极其醒目的股价与利润增长。2021 年 8 月 10 日到 2026 年 8 月 10 日，SK hynix 股价从 ₩112,500 涨到 ₩1.42 million，接近原来的 13 倍。财报给出了同样清晰的变化：2023 到 2025 年，annual revenue 从 ₩32.8 trillion 增长到 ₩97.1 trillion，operating result 从 ₩7.7 trillion loss 转为 ₩47.2 trillion profit。公司引用的 Counterpoint Research 数据显示，SK hynix 在 2025 年第二季度占 HBM shipment 的 62%，第三季度占 HBM revenue 的 57% (SK hynix, 2026)。

所以 HBM 的价值来自一组很具体的压力：每块 accelerator 要更多 capacity 和 bandwidth，部署规模还在继续扩大，供给却无法瞬间跟上。回到一块 GPU 内部，数据进入芯片以后，计算仍然不希望每次都走回 HBM。更小的 L2 cache 开始接手下一段路程。

L2 cache：避免反复回到 HBM

即使 HBM 已经很快，从它取数据仍然比使用芯片内部的数据昂贵。一次计算往往还会在很短的时间内重复使用同一份 value。如果每次使用都重新访问 HBM，宝贵的 bandwidth 会花在重复搬运上。GPU 因此在芯片内部放置了一层更小的 L2 cache。Cache 是一块自动保留近期数据的临时区域，之后又需要同一份数据时，GPU 有机会直接使用这份近处的副本。

L2 由整块 GPU 上的所有 SM 共享。某个 SM 从 HBM 读取的数据如果仍留在 L2，另一个 SM 随后访问它时就可能省下一次 HBM 传输。这种“需要的数据已经在 cache 中”的情况叫作 cache hit。保留哪些数据通常由 hardware 自动决定，PyTorch 也没有 tensor.to("l2") 这样的操作。

L2 的容量远小于 HBM。H100 配有 80 GB HBM，L2 cache 只有 50 MB (Andersch et al., 2022)。它装不下完整模型，只能保留模型当前活跃的一小部分。比如一个 kernel 正在处理大 tensor 中的若干小块，这些小块或刚刚产生的中间结果可能在 L2 中短暂停留。每减少一次重复的 HBM 读取，就多留出一些 bandwidth 给真正的新数据。这正是 reuse 的价值。

Shared memory 和 L1：一个 SM 身边的工作区

L2 服务整块 GPU，继续靠近计算以后，每个 SM 还需要自己的局部工作区。Kernel 会让许多 thread 分别处理一小份工作，一组需要协作的 thread 组成一个 thread block。同一个 block 的 thread 可以共享 SM 身边的数据。

SM 身边有 L1 cache 和 shared memory。L1 与 L2 类似，由 hardware 自动保留近期数据，只是服务范围更靠近单个 SM。Shared memory 则是一块由 kernel 明确安排的临时空间，同一 thread block 中的 thread 都可以访问。在现代 NVIDIA GPU 中，L1 和 shared memory 可以共享一部分片上物理资源，但程序使用它们的方式不同 (NVIDIA, 2026)。

假设一个 kernel 要处理远大于 shared memory 的 tensor。完整 tensor 继续留在 HBM，kernel 会把这项工作切成许多小块，每次让一个 block 处理其中一块。这样一块为了局部计算而选出的数据叫作 tile。Tile 到达 SM 附近以后，block 中的 thread 可以从 shared memory 反复读取其中的数据，各自完成一部分计算，然后再换下一块。Linear layer 会复用 input tile 和 weight tile，attention、normalization 以及其他 operation 也会按照自己的数据依赖选择 working set。Tile 多大、哪些 value 值得保留，都会影响一份数据能复用多少次，以及需要访问多少次 HBM。

Shared memory 的生命周期跟随当前 block。Block 执行结束以后，这片空间就可以交给后续工作。它的容量也会反过来限制并行度：如果一个 block 占用了大量 shared memory，同一个 SM 能同时接纳的 block 就会减少。把数据放近计算通常能减少搬运，同时会消耗有限的片上空间，kernel 需要在两者之间取舍。

Register：一次计算眼下使用的 value

Register 位于这条路径最靠近执行的位置。每个 thread 都有自己的一组 register，用来保存当前指令正在使用的少量 value，例如输入数字、数组位置，或尚未完成的累加结果。把 kernel code 转换成 GPU 指令的 compiler 会分配这些 register，程序员通常不会像操作 tensor 那样直接管理它们。

以向量 dot product 中的一小段计算为例，完整向量可能在 HBM，当前 tile 可能经过 L2 或 shared memory，某个 thread 正在维护的 partial sum 则可以放在 register 中。Thread 逐个完成乘法和加法时，这个 partial sum 会被反复更新；小段计算结束后，结果才写回更大的 memory。同样的分工也存在于 matrix multiplication、normalization 和 elementwise operation 中。

Register 的容量同样有限。一个 thread 需要的 register 越多，一个 SM 同时容纳的 thread 往往越少。如果 compiler 找不到足够的 register，一部分临时 value 会落到更远、访问更慢的 memory，这种情况叫作 register spilling。片上空间省下了数据搬运，也会反过来限制 GPU 同时展开多少工作。

判断一项计算怎样使用 memory hierarchy，可以问三个问题：必须同时保存多少 bytes，必须从 HBM 搬运多少 bytes，同一份数据在近处能够使用多少次。完整的运行时数据主要留在 HBM；L2 自动保留整块 GPU 最近可能复用的数据；L1 和 shared memory 服务一个 SM 当前处理的 tile；register 保存每个 thread 眼下使用的 value。真实路径不保证经过每一层，cache 也可能没有留住下一次需要的数据。

GPU 怎样执行一个 Tensor Operation

Weight 和 activation 进入 GPU 以后，一行 linear、softmax 或 elementwise addition 还要变成芯片里同时进行的大量工作。从模型代码开始，operation 描述 tensor 应该得到什么数学结果；kernel 则是 GPU 实际运行的程序，规定工作如何拆分、数据从哪里读取、执行哪些指令以及结果写到哪里。一个 operation 可能调用几个 kernel，compiler 也可能把几个连续 operation 融合进同一个 kernel。CPU 上的 framework 发起一次 kernel launch，GPU 随后执行选中的 kernel。

理解这次 launch，需要同时保留三张地图：

一次 launch 创建的全部工作称为 grid。Grid 被切成许多可以独立调度的 thread block，每个 block 又包含许多 thread。一个 thread 表示最小的一份逻辑任务。

GPU 会把整个 block 放到一个有空闲资源的 Streaming Multiprocessor（SM） 上执行。一个 block 在执行期间只驻留在一个 SM 上，一个 SM 则可以同时容纳多个 block。

Block 到达 SM 后，其中连续的 32 个 thread 会被编成一组，称为 warp。Warp 是 SM 实际推进指令的基本单位，其中每个位置称为 lane，对应一个 logical thread。SM 中的 scheduler 选择已经准备好的 warp，再把它的指令送到负责 memory access 或 arithmetic 的硬件，例如 Tensor Core。这些硬件统称为 execution unit。Warp 因而连接了 logical thread 与 physical hardware。Figure 3 把这组关系画在同一个 scheduling snapshot 里。

一个 linear layer 在 GPU 上发生了什么

以 LLM 中常见的 linear layer 为例。把 batch 与 sequence position 合并成 $N$ 行，输入为 $X\in\mathbb{R}^{N\times d_{\text{in}}}$，weight 为 $W\in\mathbb{R}^{d_{\text{out}}\times d_{\text{in}}}$，输出为

\[Y = XW^\top, \qquad Y\in\mathbb{R}^{N\times d_{\text{out}}}.\]

下面用一个 $6\times6$ 的小例子展示这项工作为什么能被切开。三个 matrix 都被划成 $3\times3$ 个 tile。动画从 $Y$ 左上角的 $2\times2$ output tile 开始，它只依赖 $X$ 中对应的一排 tile 与 $W^\top$ 中对应的一列 tile。沿 inner dimension 前进时，三对 input tile 依次相乘，各自产生一个 $2\times2$ partial product，再逐步累加到同一个 output tile。完成以后，同样的过程继续扫过 $Y$ 中其余八个 tile。

真实 linear layer 的 feature dimension 可能达到数千甚至数万，同样的过程会展开成许多 output tile，每块内部再连续累加许多 partial result。能否分块取决于 operation 的 mathematical dependency；kernel 找出可以独立推进的工作，再把它们映射成 block、warp 和 thread。这是 GPU 展开大规模并行计算的一条核心思路。

Framework 会为这次 matrix multiplication 选择一个 kernel。Launch 产生许多 block，每个 block 负责输出 $Y$ 的一个 tile，再被调度到某个有空闲资源的 SM。SM 以 warp 为单位推进其中的 thread，并把适合的矩阵指令交给 Tensor Core。Grid 和 block 描述工作如何拆分，SM 和 Tensor Core 则实际承接这些工作。

同一时间，完整的 $X$、$W$ 和 $Y$ 仍在 HBM。计算一个 output tile 时，kernel 分批取来对应的 input tile 和 weight tile；近期访问可能命中 L2，需要在当前 block 内复用的小块进入 shared memory，更小的 input value 被送入 Tensor Core，尚未完成的 partial sum 则尽量留在 register。遍历完 $d_{\text{in}}$ 后，output tile 才写回 HBM。Grid 决定 output 如何切块，SM 和 Tensor Core 执行这些 tile，memory hierarchy 则供应 input 并保存 partial sum。

这种安排让一次搬运支持更多计算。同一块 $X$ 可以与多块 $W$ 组合，同一块 $W$ 也可以服务 $X$ 中的多行，partial sum 则在结果完成前留在 execution unit 附近。高性能 kernel 会根据 shape 和 dtype 选择更细的 tile size，也会尽量重叠数据搬运和计算 (NVIDIA, 2026)。

一个 kernel 怎样铺开上百万个 thread

Elementwise addition 能给这套层级一个具体的数量感。假设两个长度为 1,048,576 的 vector 满足 $z_i=x_i+y_i$。每个输出位置都能独立计算，kernel 可以创建 1,048,576 个 logical thread，每个 thread 根据自己的编号找到一个 $i$。如果一个 block 有 256 个 thread，这次 launch 会得到 4,096 个 block，每个 block 又包含 8 个 warp。GPU 按照各个 SM 当时的可用资源动态安排这些 block，无需让 4,096 个 block 同时驻留在芯片上。

GPU 怎样隐藏 memory latency

某个 warp 发出 memory request 后，数据需要一段时间才能回来。Warp scheduler 可以在等待期间推进另一个已经准备好的 warp，让 execution unit 继续工作。用其他计算覆盖 memory latency 的效果称为 latency hiding，这也是 GPU 需要大量并行任务的一个重要原因。

两件事会让这种并行打折。遇到 data-dependent branch 时，同一 warp 中的 lane 可能需要分批走不同路径，这叫作 branch divergence。Register 按 thread 分配，shared memory 按 block 分配；单个 block 占用过多资源时，同一 SM 能保留的 warp 会减少。活跃 warp 相对 hardware 上限的比例称为 occupancy。它能提示 scheduler 是否有足够工作可以切换，数值更高并不自动等于运行更快 (NVIDIA, 2026)。

Scheduler 选中一个 warp 后，指令会进入对应的 execution unit。Load/store unit 处理 memory access，普通 arithmetic unit 处理 addition、multiplication 和 address calculation，Tensor Core 处理适合它的 matrix multiply-accumulate。Elementwise、normalization、attention 和 matrix multiplication 都沿用同一套调度骨架，同时依赖不同的 execution unit、数据复用方式和 HBM traffic。

Kernel 边界也会改变数据移动。假设模型先做 elementwise addition，再做 activation，两个独立 kernel 通常要把中间 tensor 写回 HBM 并重新读入；compiler 如果把它们融合进一个 kernel，中间 value 就可能留在 register 中，同时减少一次 launch 和一轮 HBM traffic。这种合并称为 kernel fusion (Turnansky, 2026)。

工作如何拆分和执行，决定了一项 operation 的运行方式；tensor 的 lifetime 则决定这些数据会在什么时候叠加成 HBM peak。

模型运行时，HBM 里装了什么

一个 byte 在 memory hierarchy 中的位置决定访问成本，一个 tensor 的 lifetime 则决定它从什么时候开始占用 HBM，又要活到什么时候。Peak HBM 正是不同 tensor lifetime 重叠后的结果。

Checkpoint 加载后去了哪里

Checkpoint 是保存在 SSD 上的一组持久化文件。它通常包含 model parameter，也可能包含运行时需要的固定数据、某些低精度 weight 附带的 scale，以及 training checkpoint 中的 optimizer state 和当前 step。加载时，framework 根据这些文件创建具有确定 shape、dtype 和 layout 的 runtime tensor。一个 tensor 占用的原始数据量可以从最基本的公式估算：

\[\text{tensor bytes}=\text{number of elements}\times\text{bytes per element}.\]

一个 70 亿 parameter 的模型，如果全部 weight 都以 BF16 保存在 HBM 中，仅 weight 就约占 14 GB。再加上 memory allocator 的对齐与预留、metadata、temporary buffer 和 KV Cache，实际显存占用会更高。Checkpoint 还可能经过切分或压缩，runtime dtype 也可能改变，文件大小因此无法直接代表显存占用。

不同 runtime tensor 的寿命差异很大。Weight 通常陪伴整个进程；一份 temporary activation 完成最后一次使用后，空间便可以复用；KV Cache 跟随一个 inference request；optimizer state 则贯穿整个 training run。Training 中还有一些 activation 必须等到对应的 backward 完成以后才能释放。

另外还有容易被忽略的 workspace。某些 kernel 和 library 会申请临时空间，用于 reduction、sorting、矩阵计算或 communication。这些 bytes 不属于模型语义，却会进入实际 peak memory。Framework 的 memory allocator 负责申请和复用显存，还可能暂时保留已经释放的 memory block。因此，nvidia-smi 显示的进程占用与当前有效 tensor 的总量未必完全一致。

这些 tensor 要在 HBM 里活多久

把这些 state 按生命周期排开，显存峰值的来源会清楚很多：

这也解释了为什么“能 inference”离“能 training”还很远。继续用 70 亿 parameter 做粗略估算：BF16 weight 约 14 GB；BF16 gradient 再增加约 14 GB；Adam 的两份 FP32 moment 约占 56 GB。三项已经达到约 84 GB。某些 training 配置还保留一份用于更新的 FP32 parameter copy，通常称为 master weight，再增加约 28 GB。Saved activation、workspace 和 allocator overhead 此时仍未计入，具体数字取决于 runtime dtype、低精度表示以及 model state 是否分布在多块 GPU 上。

显存峰值由这些生命周期的重叠决定。Inference 的 weight 长期存在，KV Cache 随活跃 request 增长，temporary activation 在 layer 之间不断产生和复用；training 还让 saved activation、gradient 和 optimizer state 在同一个 step 中相遇。同一份模型在处理 prompt 和逐 token generation 时，也会呈现两种不同的 performance profile。

Inference 时 GPU 在做什么

Serving process 启动后，weight 会留在 HBM 里供所有 request 共享。一次 decoder-only LLM request 接着经过两个阶段：prefill 处理已经给出的 prompt，decode 反复生成新的 token。它们经过同一组 layer，使用同一份 weight，performance profile 却很不一样。关键差别在于每轮有多少 token row 可以一起计算，以及哪些历史 state 必须继续读取。

Prefill：一次处理已知的 prompt

假设 prompt 长度为 $T$。进入第一层时，GPU 已经拥有 $T$ 个 token embedding。Causal attention 规定第 $t$ 个位置只能读取不晚于自己的 position，但全部 prompt token 此时都已经已知，所以同一 layer 中各个 position 的 Q、K、V projection 可以一起计算。Causal mask 遮住不允许的 future pair，无需让 prompt 从左到右逐 token 运行。Layer 之间仍有依赖，第二层必须等待第一层产生完整 hidden state。

Prefill 经过每个 decoder layer 时，会为所有 prompt position 计算 K 和 V，并写入该 layer 的 KV Cache。最后一层产生各个 position 的 logits，其中最后一行给出第一个 generated token 的 distribution。

这一阶段有很多 input row。同一块 linear-layer weight tile 到达计算单元附近后，可以服务许多 prompt position，matrix operation 因而具有较好的 data reuse。Attention 还要比较大量 query-key pair。对于长度为 $T$ 的 prompt，每个 attention head 在逻辑上都有一个 $T\times T$ 的 score matrix；prompt 变长时，计算量与直接实现产生的 intermediate 都会快速增长。

Decode：新 token 逐步出现

Prefill 选出第一个 generated token 后，这个 token 成为下一次 model call 的输入。进入某个 layer 时，它只带来一行新的 hidden state。该 layer 为这行 state 计算新的 Q、K、V；query 读取此前所有可见 position 的 K/V，新产生的 K/V 则追加到 cache，供后续 token 使用。经过完整 layer stack 后，最后一行 logits 再选出下一个 token。

下一轮必须等待这次 token selection，因为未来 token 的 embedding 取决于刚刚选出的 token ID。这项依赖让 decode 沿生成长度串行推进。每一步内部依然有许多 parallel kernel，串行约束发生在相邻 generated token 之间。

Decode 的 input matrix 对单个 request 通常只有一行。每层仍要读取大量 weight，一块 weight tile 服务的 row 却很少；attention 还要读取持续增长的 KV Cache。计算单元可能很快做完当前 arithmetic，随后等待 HBM 提供下一批数据。这解释了小 batch decode 为什么经常受到 memory bandwidth 限制。短 context 下，weight traffic 可能占主导；context 变长以后，historical K/V 的读取会越来越重要。

Serving system 会把多个活跃 request 放进同一轮 decode。若一轮有 $B$ 个 request，每个 linear layer 的 weight 就能同时服务 $B$ 行 input，提高 weight reuse 和总体 token throughput。Prompt 与生成长度各不相同，完成的 request 需要及时离开，新 request 也要尽快加入。每生成一轮 token 就重新调整 batch 的方式称为 continuous batching。更大的 batch 通常提高 throughput，同时占用更多 KV Cache，也可能增加单个 request 等待调度的时间。

KV Cache 的动态增长还会带来 memory allocation 问题。为每个 request 预留最大 context 会浪费空间，要求每份 cache 在物理上连续则容易让剩余空间碎成许多难以利用的小段，这种情况叫作 fragmentation。PagedAttention 把一个 request 的逻辑 KV position 映射到固定大小、可以分散存放的显存 page。Request 增长时按需申请新 page，完成后再归还；同一 prefix 的 page 也可以在满足条件时共享。PagedAttention 改善的是 KV Cache 的分配和复用，每个 token 在给定 attention architecture 下需要保存的 K/V 内容依然存在 (Kwon et al., 2023)。

FlashAttention：Softmax 为什么是关键？

前面的 linear layer 可以从一开始就按 tile 计算。每个 output tile 有自己的 accumulator，input tile 到达以后产生 partial product，累加完成便写回 HBM。Attention 中的两次 matrix multiplication 也能这样切，真正打断这条流水线的是中间的 softmax。对一个 attention head，scaled dot-product attention 可以写成

\[S=\frac{QK^\top}{\sqrt d}, \qquad P=\operatorname{softmax}(S), \qquad O=PV.\]

第一步用 Q 和 K 产生 score matrix $S$，softmax 把 $S$ 的每一行变成 probability，最后再用这些 probability 对 V 做 weighted sum。问题在于，一行中的每个 probability 都依赖这一行的全部 score：

\[p_i=\frac{e^{s_i}}{\sum_j e^{s_j}}.\]

假设 kernel 只读到了前四个 key，它可以算出四个 score，却还不能确定对应的四个 probability。后面的 tile 可能出现一个更大的 score，也会增加 denominator，前面所有 probability 都要随之改变。当前 score tile 因而无法像 linear layer 的 partial product 一样，算完便直接进入最终结果。

最直接的实现会让三个 operation 分开执行。$QK^\top$ 先产生完整的 $S$ 并写入 HBM；softmax 再读回 $S$，产生完整的 $P$ 并写回 HBM；最后一个 matrix multiplication 读取 $P$ 和 V，得到 O。Softmax 在两次 matrix multiplication 之间形成了一道 materialization boundary。即使 $S$ 和 $P$ 只会短暂存在，它们仍然占用 HBM，也要完整地写入和读出。

这笔 temporary data 很快就会变大。Sequence length 为 $T$ 时，每个 attention head 在逻辑上都有 $T\times T$ 个 score。$T=8192$ 时，一张 matrix 包含约 6710 万个 element；即使每个 element 只占 2 bytes，也需要 128 MiB。$S$ 和 $P$ 合计达到 256 MiB，而最终 output O 的 shape 只有 $T\times d_v$。FlashAttention 处理的核心问题，正是怎样跨过 softmax，同时避免把这两张大型 intermediate matrix 放进 HBM。

答案来自 softmax 的一个简单性质：整行 score 同时减去同一个常数 $c$，probability 不会改变，因为 numerator 和 denominator 会同时多出同一个因子 $e^{-c}$，相除时正好抵消：

\[\frac{e^{s_i-c}}{\sum_j e^{s_j-c}} = \frac{e^{-c}e^{s_i}}{e^{-c}\sum_j e^{s_j}} = \frac{e^{s_i}}{\sum_j e^{s_j}}.\]

实际计算通常选择 $c=\max_j s_j$。此时最大的 exponent 是 $e^{\max_j s_j-\max_j s_j}=e^0=1$，其余 exponent 都不超过 1，可以避免很大的 exponent 造成数值溢出。对于已经读过的一组 score，kernel 只需要留下三个量：

\[m=\max_i s_i, \qquad \ell=\sum_i e^{s_i-m}, \qquad u=\sum_i e^{s_i-m}v_i.\]

$m$ 是目前见过的最大 score，$\ell$ 是当前 softmax denominator，$u$ 是尚未除以 denominator 的 weighted value sum。处理完整行之后，attention output 正好是 $u/\ell$。这三个量已经包含后续计算需要的全部信息，之前每一个 score 的具体数值可以丢弃。

用一个具体例子来看。某个 query 对八个 key 的 score 被切成两个 tile。第一个 tile 的 score 是 $[2,1,-1,0]$，对应 value 是 $[1,2,3,4]$。为了让数字容易追踪，这里暂时让每个 value 只有一个 dimension。第一个 tile 的 maximum 是 2，因此

\[\begin{aligned} \ell_1&=e^{2-2}+e^{1-2}+e^{-1-2}+e^{0-2}\approx1.55,\\ u_1&=e^{2-2}\cdot1+e^{1-2}\cdot2+e^{-1-2}\cdot3+e^{0-2}\cdot4\approx2.43. \end{aligned}\]

此时，第一个 tile 的四个 score 已经可以丢掉。$(m_1,\ell_1,u_1)=(2,1.55,2.43)$ 足以代表它对最终 output 的全部贡献。

第二个 tile 的 score 是 $[3,-2,1,2]$，对应 value 是 $[5,6,7,8]$。同样计算后，它可以压缩成 $(m_2,\ell_2,u_2)=(3,1.51,8.93)$。新的 maximum 是 3，因此第一块保存的两个 sum 需要从“以 2 为基准”换算成“以 3 为基准”。换算只需乘上 $e^{2-3}$：

\[\begin{aligned} \ell&=e^{2-3}\ell_1+\ell_2\approx2.08,\\ u&=e^{2-3}u_1+u_2\approx9.82,\\ O&=u/\ell\approx4.72. \end{aligned}\]

如果一次性保留八个 score，对它们做完整 softmax，再加权八个 value，结果同样是 $4.72$。原因也很直接：$\ell$ 保存了所有 exponent 的总量，$u$ 保存了同一批 exponent 对 V 的 weighted sum，$m$ 让来自不同 tile 的两组数可以换到同一个基准下相加。真实 attention 中的 value 是 vector，$u$ 也成为 vector，每个 dimension 共享同一个 $m$ 和 $\ell$。这种逐块更新 softmax 的方法称为 online softmax (Dao et al., 2022)。

有了这个性质，attention 就能形成一条连续的 tiled pipeline。Kernel 把一个 Q tile 留在 shared memory 或 register 附近，依次从 HBM 读取 K/V tile。每到一块 K/V，Tensor Core 先计算当前 score tile；online softmax 随即更新 $(m,\ell,u)$；当前 score 和 probability 用完便释放。下一块 K/V 重复同样的过程。遍历完整段 sequence 后，kernel 才计算 $u/\ell$，把最终 output tile 写回 HBM。完整的 $S$ 和 $P$ 从未在 HBM 中 materialize。

FlashAttention 保持相同的 dense attention semantics，也仍然要计算相同数量的 query-key pair。它改变的是这些 operation 的执行顺序：临时产生的 score 和 probability 尽量停留在计算单元附近，完成消费以后立即消失。Prefill 同时处理大量 query row，这样可以显著减少 $T\times T$ intermediate 带来的 temporary capacity 和 HBM traffic。Training forward 也获得同样的收益；backward 可以从 Q、K、V 重新计算当前 tile 的 score 和 probability，用一部分 recomputation 换取更少的 saved intermediate。

Decode 中仍然可以使用相同的 tiled attention 与 online softmax，此时的 shape 改变了优化重点。每个 request 在一轮里通常只有一个新 query，score 是一条长度等于 context length 的 row，prefill 中完整的 $T\times T$ intermediate 已经不再出现。主要成本逐渐转向读取不断增长的 KV Cache。Flash-Decoding 会把 K/V sequence 切成多个 chunk，让不同 block 并行计算各自的 $(m,\ell,u)$，最后用一次 reduction 合并这些 summary。它沿 context dimension 展开更多并行工作，必须读取的 K/V bytes 仍然存在 (Dao et al., 2023)。

后续版本继续优化同一条 pipeline 在不同硬件上的映射。FlashAttention-2 调整 block 与 warp 的 work partition，提高 occupancy，并减少 shared-memory communication (Dao, 2024)。FlashAttention-3 面向 Hopper 架构，更紧密地重叠 data movement、matrix multiplication 与 softmax (Shah et al., 2024)。FlashAttention-4 又针对 Blackwell 架构重新设计执行流水线 (Zadouri et al., 2026)。贯穿这些版本的核心思路没有变化：让 attention tile 在更靠近计算单元的位置完成从 score 到 output contribution 的全过程。

Training 时 GPU 在做什么

Inference 中的 temporary activation 完成最后一次使用后即可释放。Training 遇到的第一个额外约束正好相反：backward 稍后还需要其中一部分。一个 training step 分成三个连续阶段，forward 计算结果并留下必要记录，backward 计算 gradient，optimizer step 更新 parameter。

Forward 为什么要留下记录

Training forward 在数学上与普通 forward 相同：token embedding 依次经过 attention、MLP、normalization 和 residual path，最后得到 logits 并计算 loss。执行层面仍然使用前面介绍的 kernel、tile、warp 和 memory hierarchy。

变化发生在 intermediate value 的 lifetime。Backward 需要某些 forward input 或 output。Linear layer 的 weight gradient 依赖它在 forward 时看到的 input activation；activation function 的 backward 依赖当时的 input 或 output；normalization 也需要相应统计量。Framework 中负责自动求梯度的 autograd 会保存足够的信息，直到反向传播经过对应 operation。模型更深、sequence 更长、单次处理的 batch 更大时，saved activation 通常随之增长。

Attention 的 backward 也沿用 FlashAttention 的 tiled execution。Forward 保存 online softmax 的归一化统计量，backward 再按 tile 重算需要的 score 与 probability，随后计算 Q、K、V 的 gradient。这样可以用一部分 recomputation 换回更少的 saved intermediate 和 HBM traffic。

每个 forward intermediate 无需全部原样保留。Framework 与 fused kernel 可以保存足以完成 backward 的较小信息，其他 buffer 在最后一次使用后仍能释放。Peak activation memory 最终取决于 model architecture、sequence length、batch、dtype 和具体 implementation。

Backward 与 optimizer step

Loss 产生以后，backward 从最后一层向前推进。每个 operation 接收来自后续 computation 的 output gradient，再产生两类结果：一类继续传向更早的 layer，另一类累积到该 operation 的 trainable parameter 上。Backward 结束后，每个 parameter 都得到当前 batch 对应的 gradient。

Optimizer 随后读取 parameter、gradient 和自己的 persistent state，写出更新后的 parameter。Adam 的 moment 要跨 training step 保留。Optimizer step 因此会读写大量 model-sized array，也带来明显的 HBM bandwidth cost。

相比一次 inference forward，training 增加了两类成本。Capacity 方面，forward record、gradient 和 optimizer state 会同时占用 HBM；compute 与 bandwidth 方面，backward 和 update 又增加一轮 matrix operation 与 model-sized data movement。多块 GPU 训练还要在 update 前合并同一 parameter 的 gradient，这份 gradient 也成为跨卡 communication 的主要 payload 之一。

显存从哪一笔省下来

如果每个 value 占用的 bytes 太多，mixed precision 会让适合低精度的 operation 使用 BF16、FP16 或 FP8，同时让需要更大数值范围或更稳定累积的 path 保持较高 precision。它能缩小部分 weight、activation 和 gradient，也能使用低精度 Tensor Core instruction。BF16 training 往往仍保留更高 precision 的 accumulation 或 optimizer state，HBM 占用不会简单减半 (PyTorch Foundation, 2026)。

如果一个 batch 的 activation 同时存在得太多，microbatching 与 gradient accumulation 会把它拆成几份依次运行。每个 microbatch 完成 forward 和 backward 后释放自己的 activation，parameter gradient 继续累积，等全部 microbatch 完成再执行一次 optimizer step。Peak activation 降低了，每次 matrix operation 的规模也会变小，还会增加 kernel launch。

如果深层模型留下的 forward record 太多，activation checkpointing 只保存选定的 layer boundary。Backward 走到某个区间时，系统重新运行这一段 forward，恢复缺失的 intermediate，再计算 gradient。它把 saved-activation capacity 转成额外 compute，对 weight、gradient 和 optimizer state 没有直接缩减作用 (PyTorch Foundation, 2026)。

Mixed precision 缩小单个 value，microbatching 减少同一时刻存在的 batch state，activation checkpointing 则减少必须留在 HBM 中等待 backward 的 intermediate。显存压力依然超过单块 GPU 时，把 state 或 computation 分到多块 GPU 可以继续扩展 capacity 与 compute，同时也会引入新的数据交换。

一块 GPU 不够时

增加一块 GPU 会带来更多 HBM 和 execution unit，同时增加一道 device boundary。一块 GPU 的 HBM 无法由另一块 GPU 当作 local memory 随意读取。Model state 或 computation 一旦跨卡切分，某些 tensor 就必须通过 interconnect 传输，接收方也可能等待数据到达。

同一 server 内的 GPU 可能通过 NVLink、NVSwitch 或 PCIe 通信，跨 server 时则要经过 InfiniBand、Ethernet 等 network fabric。从这里开始，bandwidth 有了两条边界：HBM bandwidth 描述单块 GPU 内部的数据供应，interconnect bandwidth 描述 GPU 之间的数据传输。选哪种 parallelism，首先取决于究竟是哪一类东西放不下或跑不快。

这些 parallelism 名词的区别，首先在于它们沿哪个 dimension 切：

实际系统经常组合几种切法，因为它们描述的是不同维度。选择哪一种，取决于当前究竟是什么放不下或跑不快。

完整模型放得下时，最直接的扩展方式是复制。Inference serving 把不同 request 交给不同 model replica，同一个 request 在自己的 replica 内完成 decode，通常无需逐层交换 hidden state。Training 中的 data parallelism（DP） 也复制模型，每块 GPU 读取不同的 local batch，独立完成 forward 与 backward，随后同步 parameter gradient (PyTorch Foundation, 2026)。

Training DP 的模型副本会产生不同 gradient。假设两块 GPU 对同一个 parameter 分别得到 $g_0$ 和 $g_1$，一次 sum all-reduce 会让两边都得到 $g_0+g_1$；需要平均 gradient 时，再除以 replica 数量。

如果完整模型的 persistent state 已经放不下，继续复制只会重复这笔开销。ZeRO 与 Fully Sharded Data Parallel（FSDP） 沿 data-parallel group 切分 parameter、gradient 和 optimizer state。某个 layer 即将计算时，GPU 收集所需 parameter shard；backward 产生 gradient 后，再聚合并重新分片。多块卡的 HBM 因而可以共同容纳一份更大的 training state，每层附近也增加了 parameter 与 gradient communication (Rajbhandari et al., 2020; PyTorch Foundation, 2026)。

如果压力来自单个 layer，tensor parallelism（TP） 会切开 layer 内的大 tensor。以 linear layer 为例，几块 GPU 各自保存 weight matrix 的一部分，对同一批 input row 计算不同 output feature，或者分别计算 partial sum。后续 computation 需要完整结果时，这些分散在各卡上的部分还要通过跨卡 communication 重新组合。Transformer 的 attention head 和 MLP dimension 都能沿这种方式切分 (Shoeybi et al., 2019)。这类 communication 往往每个 Transformer block 都会发生，因此参与同一 TP group 的 GPU 通常需要较快的 link。

如果整个 model stack 需要沿 depth 分布，pipeline parallelism（PP） 会把连续 layer 放进不同 stage。Forward 在 stage boundary 发送 activation，backward 再把 activation gradient 传回来。Training batch 通常切成多个 microbatch，让不同 stage 同时处理流水线中的不同位置。流水线开始和结束时仍有部分 stage 空闲，这段空闲称为 pipeline bubble (Narayanan et al., 2021)。

如果压力来自很长的 sequence，context parallelism（CP） 会让不同 GPU 各自保存和处理一段 token position，从而分摊随 sequence length 增长的 activation。Linear 和 normalization 可以直接处理本地的 sequence chunk，attention 却有跨 token 依赖：本地 query 仍要看到整段 sequence 的 K/V。GPU 因此需要在 attention 附近交换 K/V chunk，backward 也要重组对应的 gradient (NVIDIA, 2026)。CP 缩小了每块卡上的 activation footprint，同时增加了与 context 长度相关的 communication。

Mixture-of-Experts（MoE）model 还会遇到另一种结构。Expert parallelism（EP） 把 expert 分布到不同 GPU。Router 为每个 token 选择 expert 后，每张卡都可能需要把不同 token 发往不同目的地，这种交换称为 all-to-all。Expert output 随后再回到原来的 token order，payload 取决于 routing decision 与 load balance (Fedus et al., 2022)。

这些策略背后的 communication 可以收进同一个问题：tensor 切开以后，下一步需要的数据是否仍在本地？如果各卡需要聚合并得到同一结果，会用到 all-reduce；如果需要拼回其他卡持有的 shard，会用到 all-gather；如果聚合后仍希望结果保持分片，会用到 reduce-scatter；如果每张卡要把不同数据发往不同目的地，会用到 all-to-all。它们都属于 collective communication，也就是一组 device 共同参与的数据重组 (NVIDIA, 2026)。DP 的 gradient、FSDP 的 parameter shard、TP 的 activation，以及 EP 的 routed token，只是这些重组动作承载的 payload 不同。

什么时候会 communication-bound

一次 communication 至少支付两类成本：启动与同步需要时间，实际 bytes 还受 link bandwidth 限制。许多小 message 会反复支付前一项；大 payload 更容易受后一项影响。Collective 还要求参与者在合适的时刻全部就绪，某块 GPU 计算较慢或 network path 拥塞时，其他 GPU 也可能等待。

实际实现会尽量重叠 communication 与 compute。Data-parallel backward 从最后一层向前推进，某层 gradient 算好后即可开始同步，同时 GPU 继续计算更早的 layer。Tensor parallelism 与 context parallelism 的 communication 紧贴 layer 内部计算，通常更难隐藏全部时间。Pipeline parallelism 和 expert parallelism 还会受到 microbatch schedule 或 routing balance 的影响。

增加 GPU 后，如果单步时间几乎没有改善，甚至变慢，可以检查三件事：每块 GPU 分到的 compute 是否已经太少，communication payload 与频率是否随 parallel degree 上升，以及 collective 能否与 computation overlap。瓶颈落在跨设备边界时，继续增加理论 FLOPs 无法直接解决。

Optimization 改变了哪一本账

GPU performance 最终可以压进四本账：

一项 optimization 经常会同时改动几本账。Quantization 缩小 weight 或 KV Cache，降低 capacity 与 HBM traffic，同时要求 hardware 和 kernel 真正支持相应的低精度计算。Batching 让一份 weight 服务更多 token row，提高 data reuse 与 throughput，也会增加同时存在的 request state。Kernel fusion 让 intermediate 留在 register 或 shared memory，减少 launch 和 HBM 往返。FlashAttention 改变 attention 的执行顺序，减少 intermediate memory traffic，同时保持 logical attention result。

Training 中，activation checkpointing 减少 saved activation，并在 backward 增加 recomputation。Data、tensor、pipeline、context 和 expert parallelism 把 compute 与 state 分散到更多 GPU，也产生各自的 communication payload。Sharding 可以解决单卡 capacity，却可能把新的瓶颈推到 interconnect 上。Optimization 的名字本身无法说明收益，真正需要追踪的是哪些 tensor 缩小了，哪些 bytes 少搬了，哪些 operation 被省掉或重算，以及哪些 device 需要彼此等待。

下一次看到 .to("cuda")，这只是整条路径的起点。Weight 从 host 进入 HBM；kernel 把 operation 拆成 grid、block、warp 和 thread；memory hierarchy 不断供应当前 working set；inference 维护逐渐增长的 request state，training 保存 backward 与 update 所需的信息；模型跨越多块 GPU 后，部分 tensor 还会通过 interconnect 在 device 之间交换。GPU performance 的大多数讨论，最终都可以放回这张地图，并归入 capacity、bandwidth、compute 与 communication 中的某一本账。