Roofline：算力、带宽与算术强度

同样是 \(10^{12}\) 次运算，一个矩阵乘法在 H100 上可能 1 毫秒就做完，一串逐元素运算却可能慢几十倍。只知道 FLOPs 不够，差别在于数据要搬多少、从哪里搬。roofline 分析回答的就是「一个运算的时间到底花在计算上，还是花在搬数据上」，所有并行策略的分析都建立在它之上。

做 roofline 估算的原材料是硬件的规格数字：TPU 和 GPU 的计算单元有多快，存储层次和网络的每条通道有多宽，决定了一个运算会卡在哪里。

1. 时间花在哪：计算、搬数据、通信

先问一个朴素的问题：为什么一个算法要 50 ms，而不是 50 s 或 5 ms？时间只可能花在两类事情上。

计算。深度学习模型基本就是一堆矩阵乘法，每个由浮点乘和加组成。硬件每秒能做的浮点运算数（FLOP/s）是固定的，于是：

例如 H100 的 BF16 稠密峰值约 \(9.89\times 10^{14}\) FLOP/s，TPU v6e 约 \(9.2\times10^{14}\)。做 \(10^{12}\) 次运算在 H100 上大约 1.01 ms。这是按标称峰值算的理想值：H100、B200 上的矩阵乘法通常只能达到标称峰值的 80–85%，TPU 正常使用时可以接近 95%。

通信，分两种：

• 芯片内：张量要在显存（HBM，high bandwidth memory）和计算核心之间搬运，这条路的带宽叫 HBM 带宽（NVIDIA 也叫 memory bandwidth）。H100 约 3.35 TB/s，TPU v6e 约 1.64 TB/s。
• 芯片间：模型分布在多块芯片上时，张量要在芯片之间传，走 NVLink、ICI、InfiniBand、DCN、PCIe 等不同的链路，带宽各不相同。

两种都按同一个式子估算：

计算和通信通常可以重叠：在算这一块的同时，下一块数据已经在路上了。所以总时间有一个下界和一个上界：

因为 \(T_\text{math} + T_\text{comms} \le 2\max(T_\text{math}, T_\text{comms})\)，上下界最多差 2 倍。实际估算时按 \(\max\) 来，因为代数简单，而且好的实现确实能接近它；之后再根据 profiling 补上重叠不完美的部分。

假设完全重叠：\(T_\text{math} > T_\text{comms}\) 时硬件算力被用满，叫计算受限（compute-bound）；反之叫通信受限（communication-bound；瓶颈在 HBM 时也常叫 memory-bound 或 bandwidth-bound），一部分算力在等数据。

2. 算术强度与临界强度

判断一个运算会受哪个限制，看它的算术强度（arithmetic intensity，也叫 operational intensity）：

单位是 FLOPs/byte，含义是「每搬一个字节，做多少次运算」。推导计算受限的条件只需要一行：

右边这个只取决于硬件的量叫临界算术强度（critical / peak arithmetic intensity）。它的意思是：硬件要跑满算力，每从 HBM 读一个字节，就必须做这么多次运算。几个芯片（BF16 矩阵单元峰值 / HBM 带宽，稠密值）：

表里都是稠密值。读 NVIDIA 规格表时要小心：Tensor Core 的数字通常是带结构化稀疏的，正好是稠密值的 2 倍。例如 H100 SXM 规格表上 BF16 写的是 1,979 TFLOPS（带星号），稠密值是 989；HGX B200 写 8 卡 BF16 36 PFLOPS（稀疏），折合每卡稠密 2.25 PFLOPS，HBM 带宽每卡 8 TB/s。

可以看到，主流加速器的临界强度大多在 150–600 之间，而且有意思的是，几代硬件之间它变化不大：算力和带宽大体同步增长。

3. Roofline 图

把上面的关系画出来就是 roofline 图：横轴是算术强度，纵轴是可达到的吞吐（FLOP/s），两个轴都取对数。对强度为 \(I\) 的算法，能达到的吞吐是

图形像一个屋顶：左边是一条斜率为 1 的斜线（对数坐标下，斜线的高度由带宽决定），右边是一条水平线（峰值算力）。两线交点就是临界强度，也叫 ridge point。点落在斜线下方的区域是带宽受限，落在水平线下方是计算受限。提升性能只有两种方式：让算法的强度变大（点往右移），或者换更高的带宽（斜线往上移）。

选芯片和运算，拖动参数看点怎么移动。除了 HBM 的屋顶，还画了芯片间网络的屋顶（虚线），和片上高速存储的屋顶（点线，有数据时才画）。

几个值得自己拖出来的现象：

• BF16 矩阵乘法里，强度几乎只由 \(B\) 决定，\(D,F\) 大时 \(I \approx B\)。把 \(B\) 拖过临界值（H100 约 295）点就爬上屋顶。
• \(D=F=1024\) 时，要达到同样的强度需要的 \(B\) 明显更大（\(BF\)、\(BD\) 两项不再可以忽略）。
• 点积和逐元素加法的强度是常数（1/2 和 1/12），无论 \(N\) 多大都在屋顶的斜坡上。
• 注意力核心在生成阶段（\(T=1\)）强度约等于分组大小 \(G\)，永远到不了屋顶；在训练/prefill（\(T=S\)）强度随 \(T\) 线性增长，很快就计算受限。

4. 低强度的运算：点积与逐元素运算

点积

计算 BF16 向量的点积 \(x\cdot y\)：\(x, y\) 各 \(N\) 个元素，读 \(2N + 2N\) 字节，写回 2 字节；做 \(N\) 次乘、\(N-1\) 次加：

每读一个字节只做半次运算。拿它和矩阵单元的临界强度（几百）比没有意义，因为点积本来就不在矩阵单元上跑，而在向量单元上（TPU 的 VPU、GPU 的 CUDA core）。但向量单元的算力低得多，它们的临界强度也低：TPU v5p 的 VPU 每核约 \(7\times10^{12}\) FLOP/s，整芯片两个核约 \(1.4\times10^{13}\)，临界强度约 5；H100 的 CUDA core 不算 FMA 时约 33.5 TFLOP/s，除以 3.35 TB/s，临界强度约 10。即便如此，1/2 依然远低于它们，所以点积是带宽受限的。

逐元素运算

两个 FP32 向量相加：读 \(2\times 4N\)、写 \(4N\) 字节，做 \(N\) 次加法，强度 \(N/12N = 1/12\)，糟糕至极。在 H100 上：

在 H100 上用 JAX 实测一下：\(N = 1024^3\) 时 roofline 预测约 3.84 ms，实测 4.1 ms，很接近；但 \(N=65{,}536\) 时预测 0.23 μs，实测约 1.5 μs——数据太小时，内核启动和延迟才是主要成本。

这个结论的实际意义很大：LayerNorm、激活函数、残差加、dropout、softmax……这些逐元素或归约运算的 FLOPs 在总账里可以忽略，但每一个都要把整个张量从 HBM 读一遍、写一遍。所以要么把它们融合（fuse）进相邻的矩阵乘法内核里（读进来顺手算掉，不单独写回），要么至少把几个逐元素运算融合成一个内核。融合内核（fused kernel）的价值、FlashAttention 不把注意力矩阵写回 HBM 的原因，背后都是这一张图。

5. 矩阵乘法：为什么 batch 要过 ~240 个 token

推导

\(X\) 是 \(\text{bf16}[B, D]\)，\(Y\) 是 \(\text{bf16}[D, F]\)，求 \(Z = XY\)，\(Z\) 是 \(\text{bf16}[B,F]\)。要读 \(2BD + 2DF\) 字节，写 \(2BF\) 字节（输出在累加时是 FP32，但写回 HBM 前一般转成 BF16），做 \(2BDF\) 次运算：

在 Transformer 里，\(X\) 是激活（\(B\) 是 token 数），\(Y\) 是权重（\(D, F\) 是模型维度）。如果 \(B\) 远小于 \(D\) 和 \(F\)，分母里 \(DF\) 占主导：

于是计算受限的条件变成一个极其简单的规则：

直觉上：权重 \(DF\) 从 HBM 读一次的成本是固定的，它要被 \(B\) 个 token 复用；每个 token 对每个权重做 2 次运算（乘和加）而权重占 2 字节，所以强度就是 \(B\)。token 越多，读一次权重摊到的计算越多。

这就是训练和推理里一系列选择的来源：

• 训练：每个副本的本地 token batch 通常 \(\gt 1024\)，而 \(D, F\gt 8000\)，矩阵乘法舒服地计算受限。但如果数据并行开得太大，本地 batch 掉到几百以下，就开始浪费算力。
• 推理的生成阶段：每步每条序列只出一个 token，\(B\) 就是并发的序列数。batch 小时，每一步都要把全部权重从 HBM 读一遍，而计算只用到一点点。这时 roofline 给出一个延迟下界：权重字节数 / HBM 带宽。举个例子：200B 参数的 BF16 模型切到 32 块 TPU v4 上，每块约 12.5 GB，以约 1.2 TB/s 读完需要约 10 ms，所以单步采样不可能快于 10 ms。

精确的临界值与形状的影响

\(B \ll D, F\) 是近似。不做近似，直接解 \(2BDF/P = (2BD + 2DF + 2BF)/W\)：H100 上，\(D=F=8192\) 时需要 \(B \ge 319\)，\(D=F=4096\) 时 \(B \ge 345\)，\(D=F=1024\) 时要 \(B\ge 698\)，几乎翻倍。再如 int8 的 \([B, 4096]\times[4096, 16384]\) 在 v5e 上（int8 峰值 3.94e14），精确解出 \(B \gt 259\)，略大于 240，因为把 \(D, F\) 的有限大小也算进去了。矩阵越小，越难吃饱。

分块（tiling）

大矩阵乘法在硬件里会被拆成小块，放进片上的高速存储（TPU 的 VMEM、GPU 的 SMEM/TMEM），块会被重复读取，所以并不严格是只读 \(O(N^2)\) 字节。设 \((m,k)\cdot(k,n)\) 的矩阵乘法用块大小 \(b_m, b_k, b_n\)，块数 \(t_m = m/b_m\) 等，总 FLOPs 是 \(2\,t_m t_n t_k\, b_m b_n b_k\)，总字节约 \(2\,t_m t_n\big(t_k(b_m b_k + b_k b_n) + b_m b_n\big)\)。忽略最后一项，强度约为 \(\frac{b_m b_n}{b_m + b_n}\)，形式和整体的式子一样。这说明块要足够大，片上存储的容量因此很重要。

量化

• INT8（或 FP8）全程：每个数 1 字节，强度变成 \(2B\)，但硬件的 INT8 峰值也是 BF16 的 2 倍（v5e 3.94e14，临界强度 480），条件仍是 \(2B > 480\)，即 \(B > 240\)。基本不变。
• INT8 权重 + BF16 激活和计算：只有权重变成 1 字节，强度约 \(2BDF/DF = 2B\)，而算力仍按 BF16 的 240，于是 \(B > 120\)。权重量化在小 batch 推理里很划算，因为它直接减少了要搬的字节。
• H100、B200 的 FP8 峰值正好是 BF16 的 2 倍，临界强度翻倍到约 590 和 562；如果权重也用 FP8 读入，实际门槛不变。

强度永远很低的两种情况

• 每个样本一套权重：\(\text{int8}[B, D]\cdot\text{int8}[B, D, F]\)。FLOPs 不变（\(2BDF\)），但要读 \(BDF\) 字节的权重，强度约为 2，与 \(B\) 无关。没有权重复用，就永远带宽受限。
• 生成阶段的注意力：每条序列有自己的 KV cache，本质上同上。GQA 下（\(N\) 个 Q 头、\(K\) 个 KV 头、每组 \(G = N/K\)，\(T\) 个 query、\(S\) 个 key），\(QK^\top\) 和乘 V 各 \(2BTSKGH\) FLOPs，共 \(4BTSKGH\)；要读写 Q 和输出各 \(2BTKGH\) 字节、K 和 V 各 \(2BSKH\) 字节，共 \(4BHK(TG + S)\)。强度为 \(\frac{TSG}{TG + S}\)。训练/prefill 时 \(T = S\)，强度 \(\frac{SG}{G+1}=O(T)\)，序列稍长就计算受限；生成时 \(T = 1\)，强度 \(\frac{SG}{G+S}\to G\)，永远带宽受限。增大 \(G\)（更少的 KV 头）会让它更接近屋顶，这是 GQA/MQA 在推理上的意义之一。

MoE

MoE 有 \(E\) 个专家，每个 token 选 \(k\) 个。每个权重矩阵要读 \(E\cdot DF\) 个数，但 FLOPs 只有 \(2k\cdot BDF\)。以 INT8 权重、BF16 计算、v5e（临界 240）为例，计算受限需要 \(2kBDF/(EDF) > 240\)，即 \(B > 120\,E/k\)。DeepSeek-V3（\(E=256, k=8\)）得到 \(B > 3840\) 个 token——在生成阶段这是很大的 batch。换句话说：要让每个专家都吃饱，平均每个专家分到的 token 数要过临界值。

6. 交互：矩阵乘法耗时计算器

计算器把 \(T_\text{math}\)、\(T_\text{HBM}\) 和芯片间通信时间 \(T_\text{comm}\) 放在一起算，可以选择把矩阵乘法切到多块芯片上：

• 沿 \(D\) 切（缩并维）：每块芯片拿 \(X\) 的 \(D/n\) 列和 \(Y\) 的 \(D/n\) 行，算出完整形状 \([B,F]\) 的部分和，然后需要把部分和加起来（AllReduce）。
• 沿 \(F\) 切（输出维）：每块芯片拿完整的 \(X\) 和 \(Y\) 的 \(F/n\) 列，各算一部分输出，不需要通信（但 \(X\) 要在每块芯片上都有一份）。

7. 网络的 roofline

前面的屋顶都是芯片内的：分母是 HBM 带宽。但同样的推理对芯片间的通信完全适用，而且分析并行策略时，关心的几乎都是这一种。

看一个最简单的例子：\(X\sim\text{bf16}[B,D]\)、\(Y\sim\text{bf16}[D,F]\)，沿 \(D\) 均分到 2 块 TPU 上。每块芯片算一半：\(Z_0 = X[:, :D/2]\,Y[:D/2, :]\)，\(Z_1\) 同理，然后把部分和发给对方相加。设每个方向能传 \(4.5\times10^{10}\) 字节/秒，每块芯片算力 \(1.97\times10^{14}\)：

计算受限的条件是

和 HBM 的 roofline 不同，这里的临界条件只和 \(D\) 有关，和 \(B\) 无关。原因：FLOPs 和要传的字节（部分和 \([B, F]\)）都和 \(B\)、\(F\) 成正比，二者相除只剩下 \(D\)——每块芯片在缩并维上做了多少工作，决定了它能「掩护」多少通信。推广到 \(n\) 块芯片做环形 AllReduce，条件大约是 \(D > 2(n-1)\cdot P/W\)，切得越多越难。

每一层带宽都有自己的临界强度。拿 BF16 算力分别除以各层带宽，就能看到越往外走，需要的强度越高：

同一个道理也适用于 PCIe。假设把 \(\text{bf16}[D, F]\) 的权重和 \(\text{bf16}[B, D]\) 的激活都放在主机内存里，在一块 TPU v6e 上做矩阵乘法（\(F = 4D\)，\(B\ll D\)），要不被 PCIe（1.6e10 B/s）卡住，需要 \(B > 9.18\times10^{14}/1.6\times10^{10}\approx 57{,}000\)。这就是为什么权重卸载（offload）到 CPU 内存只在 batch 极大时才划算。

8. TPU 是什么样的

TPU 的结构比 GPU 简单。

概括地说，TPU 基本就是一个专门做矩阵乘法的计算核心（TensorCore），挂在一堆高速显存（HBM）上。TensorCore 里有三个关键单元：

• MXU（Matrix Multiply Unit）：TensorCore 的核心。大多数代际里，它每 8 个周期完成一次 \(\text{bf16}[8,128]\,@\,\text{bf16}[128,128]\to\text{f32}[8,128]\) 的矩阵乘法，用的是脉动阵列（systolic array）。TPU v6e（Trillium）的 MXU 是 256×256，之前各代都是 128×128。一个 128×128 的 MXU 每周期做 16,384 次乘加，即 32,768 FLOPs；v5e 的 TensorCore 有 4 个 MXU，频率 1.5 GHz，合计 \(4\times 32768\times 1.5\times10^9\approx 2\times10^{14}\) BF16 FLOP/s。低精度更快：v5e 的 INT8 是 \(4\times10^{14}\) OP/s。
• VPU（Vector Processing Unit）：做 ReLU、逐元素加乘、归约等一般数学运算。它是一个形状为 (8, 128) 的二维 SIMD 机器：128 那一维叫 lane，8 那一维叫 sublane，每个 (lane, sublane) 位置在 v5 上有 4 个独立的 ALU。于是 v5p 每核 VPU 约 \(8\times128\times4\times1.75\text{e}9 \approx 7\times10^{12}\) FLOP/s，大约是 MXU 的 1/30。沿 sublane 的归约便宜（有一个周期左右的 shuffle），跨 lane 的归约要经过专门的 XLU（cross-lane unit），又慢又贵。
• VMEM（Vector Memory）：TensorCore 里的片上暂存，比 HBM 小得多（v5e 上 128 MiB），但到 MXU 的带宽高得多。它像 CPU 的 L1/L2 缓存，但大得多，而且由程序员/编译器显式管理。HBM 里的数据必须先拷到 VMEM 才能参与计算。

此外还有标量单元（scalar core）：它取指、发射所有指令、执行 HBM 到 VMEM 的传输。它是单线程的，所以每个 TensorCore 每周期只能发起一个 DMA 请求；一个标量核控制着 4,096 个 VPU ALU、4 个 MXU、2 个 XLU 和多个 DMA 引擎。「控制很少、计算很多」是 TPU 硬件效率的来源，也限制了它做依赖数据的灵活计算。

流水线与 VMEM 的强度

TPU 上所有操作都是流水线化、相互重叠的。做 \(X\cdot A\to Y\) 时，先把 \(A\) 和 \(X\) 的块从 HBM 拷到 VMEM，再送进 MXU 按 8×128（\(X\)）和 128×128（\(A\)）的块相乘，结果逐块写回 HBM；拷贝和 MXU 的工作是重叠的，所以 MXU 不必等数据——只要整体是计算受限的。如果 HBM→VMEM 比 MXU 慢，就是带宽受限：MXU 在「挨饿」。

VMEM 的带宽约为 HBM 的 22 倍，这意味着如果输入输出都能放在 VMEM 里，强度只要 10–20 就能吃满 MXU。还是上面那个 \(\text{int8}[B, 4096]\times\text{int8}[4096,16384]\) 的例子：在 v5e 上权重从 HBM 读时需要 \(B > 259\)；如果从 VMEM 读（带宽按 HBM 的 22 倍算），只需 \(B > 11\)（实际上 VMEM 带宽不能全用来读权重，更接近 20）。这就是 VMEM 预取的价值：例如在算注意力的时候，把下一个 FFN 的权重提前搬进 VMEM，就能把权重加载的成本藏起来——前提是一层的权重切得足够小，放得进 VMEM。可惜 VMEM 很小，这往往是个挑战。

脉动阵列是怎么工作的

MXU 的核心是一个 128×128（=16,384 个）乘加单元组成的二维网格。权重（RHS，128×128）从上方预先装入，每个单元固定持有一个权重；输入（LHS，8×128）从左边斜着送进来，每个周期往右移一格；部分和从上往下流，每经过一个单元就加上「输入 × 该单元的权重」，从底部流出的就是结果。开始有一个装载权重和输入的流水线气泡，之后新输入可以连续送入，没有额外气泡。所以张量的维度至少要和 MXU 的边长一样大（128，v6e 上 256），否则要补零浪费算力；有多个 MXU 时，切块的维度要大于 128 × MXU 个数。

一个 3×3 的权重驻留型（weight-stationary）阵列，逐周期运行：

芯片、托盘与主机

一块 TPU 芯片通常有两个 TensorCore，共享显存，可以当成一个算力加倍的大加速器（叫 megacore），v4、v5p 都是这样（TPU v7 去掉了 megacore，改成两个核之间的高带宽连接）。v3 及更早的芯片两个核显存独立，被当作两个加速器；面向推理的 v5e 和 v6e 每芯片只有一个核（v6e 的核里是 2 个 256×256 的 MXU）。芯片 4 个一组装在一个托盘（tray）上，经 PCIe 连到 CPU 主机；v5e 每台主机挂 2 个托盘、8 块芯片。PCIe 带宽有限，例如 v4 是每个方向 16 GB/s，比 HBM 慢约 100 倍：可以把数据放到主机内存，但搬得不快。

9. TPU 的网络：ICI、torus、pod 与 DCN

同一个 pod 里的芯片通过 ICI（inter-chip interconnect）直连，不经过主机。v2/v3、推理芯片 v5e 和 v6e 的每块芯片连 4 个最近邻，构成二维 torus；v4 和 v5p 连 6 个最近邻，构成三维 torus。

• Torus 与回绕：环形回绕把任意两点的最大距离从 \(N\) 降到 \(N/2\)；还有「扭曲 torus」（twisted torus），像莫比乌斯带一样回绕，进一步缩短平均距离。
• Pod 和 superpod：v4 最大 16×16×16，v5p 最大 16×20×28（8,960 块芯片；单个切片最多 6,144 块，更大的规模用 multislice 跨切片拼起来）。大 pod 由 4×4×4 的立方体组成，立方体之间用可重构的光交换（optical switch）连起来，它带宽和 ICI 一样，只是让连接可以重配并保留回绕。
• 小切片没有回绕：申请 2×2×1、2×2×4 之类小于一个整立方体的拓扑，没有回绕链路，大多数通信的时间会翻倍；整立方体的倍数（4×4×4、4×4×8）才有。v5e 和 v6e 的 pod 是单个 16×16 的二维 torus，长度为 16 的轴才有回绕（所以 8×16 在长轴上有回绕）。
• 切片（slice）与 DCN：一组由 ICI 连起来的 TPU 叫一个切片。切片之间通过数据中心网络（DCN）相连，例如跨 pod 训练（multislice）。DCN 是主机到主机的：数据要先过 PCIe 到主机，出网卡，到对端主机，再过 PCIe 进 HBM。

带宽的量级（v5p）：HBM 每芯片 2.77e12 B/s；ICI 每轴 9e10 B/s（单向，每芯片 3 个轴）；DCN 每芯片约 6.25e9 B/s（v5e 是 3.125e9，v6e 和 TPU7x 是 12.5e9）；PCIe 每芯片约 1.6e10。ICI 每条链路的单向/双向带宽：v5e 4.5e10/9e10，v5p 和 v6e 9e10/1.8e11。双向带宽是指沿一个轴、两个方向加起来能发出的字节数，只有在有回绕的轴上（能构成环）才能用满，环形集合通信的公式里经常用它。

总结一下：TPU 很简单。大多数情况下可以把它看成一个矩阵乘法单元，连着显存（超快）、通过 ICI 连着邻居芯片（相当快）、通过 DCN 连着数据中心的其余部分（还算快）。让通信量和各条通道的速度成比例，就不会卡住 MXU。

10. GPU 是什么样的

现代 ML GPU（H100、B200）从高处看和 TPU 惊人地相似：一堆专做矩阵乘法的计算单元（叫 SM，streaming multiprocessor），连着一块高速显存（HBM）。区别是 TPU 有 1–2 个大 TensorCore，而 GPU 有一百多个小 SM（H100 132 个，B200 148 个）。每个 SM 比 TPU 的 TensorCore 弱得多，但彼此基本独立，所以 GPU 可以同时做几百件不同的事。

SM 内部

每个 SM 分成 4 个相同的子分区（SM subpartition），每个子分区有：

• Tensor Core：专用的矩阵乘法单元，相当于 TPU 的 MXU（注意命名冲突：GPU 的 Tensor Core 是 SM 里的矩阵子单元，TPU 的 TensorCore 是包含 MXU、VPU 等的整个核心）。它提供了绝大部分算力：H100 上 Tensor Core 是 989 TFLOP/s BF16，CUDA core 只有 67 TFLOP/s FP32。按 132 个 SM、约 1.83 GHz 的频率反推，每个子分区的 Tensor Core 每周期做 1,024 次 BF16 FLOPs，相当于一次 8×8×8 的矩阵乘法；同样的口径下 A100 是 512，B200 约 2,048，每代翻一倍。
• CUDA core 与 Warp Scheduler：32 个 FP32 的 ALU（就是 NVIDIA 说的 CUDA core）加上调度它们的单元，做 SIMD/SIMT 的向量运算，相当于 TPU 的 VPU：ReLU、逐元素运算、归约都在这里。每个 ALU 每周期做一次运算；FMA（乘加融合）算两次，NVIDIA 的规格数字就是按 FMA 算的。
• 寄存器堆：16,384 个 32 位字。

算下来，H100 共有 \(132\times4\times32 = 16{,}896\) 个 FP32 CUDA core，B200（148 个 SM）有 18,944 个；而 TPU v5p 一块芯片有 \(2\times4\times8\times128 = 8{,}192\) 个独立 ALU。

SIMT、warp 与调度

GPU 的线程以 32 个为一组，叫 warp，同一 warp 的线程同时执行同一条指令，处理不同的数据。warp 再组成更大的线程块（block，例如 512 或 1,024 个线程），每个块分配给一个 SM。一个 SM 可以同时驻留多个块和最多 64 个 warp，但每个 warp scheduler 每个周期只执行一个 warp 的指令；它在活跃的 warp 之间切换，以此掩盖访存延迟。相比之下 TPU 基本是单线程的。

CUDA core 比 TPU 的 VPU 灵活：它用 SIMT（single instruction, multiple threads）模型，每个线程有自己的指令指针，可以独立编程；而 VPU 是纯 SIMD。代价是：如果同一个 warp 里的线程走了不同的分支（warp divergence），硬件会把两条分支都执行一遍，把不该执行的线程屏蔽掉，性能悄悄下降。

在 GPU 上运行的程序叫内核（kernel），可以用 CUDA 或 Triton 写，编译成 PTX。写内核大致有四条路：纯 PyTorch（简单但慢）、torch.compile（简单、快、不够灵活）、Triton（更难、更快、更灵活）、CUDA（最难、最快、最灵活）。常见的优化有：访存合并（memory coalescing，让一个 warp 的 32 个线程访问连续地址，从而合并成一次 DRAM 突发读取）、分块（tiling，把数据搬进共享内存反复使用）、线程粗化和减少分支分歧。这些都是在给定内核里提高实际达到的带宽或算力，对应 roofline 上「离屋顶有多远」。内核细节见 GPU 内核与 FlashAttention。

存储层次

• 寄存器：每 SM 256 KB。每个线程最多用 256 个寄存器；如果每个线程都用满，一个 SM 同时只能驻留 8 个 warp。
• SMEM / L1：每 SM 256 KB 的片上存储，可以当作程序员控制的共享内存，也可以当硬件缓存，用来存 Tensor Core 矩阵乘法的输入。全卡合计约 33 MB，加上寄存器约 66 MB，大约是 TPU VMEM（约 120 MB）的一半。
• TMEM（Blackwell 新增）：B200 的 Tensor Core 大到输入（尤其是累加器）放不进寄存器和 SMEM，所以加了每 SM 256 KB 的 Tensor Memory。相关地，Ampere 的 Tensor Core 可以由一个 warp 驱动，Hopper 需要整个 SM（warpgroup），Blackwell 由 2 个 SM 一起驱动。
• L2：所有 SM 共享约 50 MB（H100；实际分成两半，各 25 MB），带宽 NVIDIA 没公布，实测约 5.5 TB/s。它不由程序员控制，所以要靠调整访存模式间接用好它；所有 SM 共享 L2 也迫使各 SM 以相当协调的方式运行。
• HBM：主显存，存权重、梯度、激活。容量从 Volta 的 32 GB 涨到 Blackwell 的 192 GB；带宽 H100 3.35 TB/s，B200 约 8 TB/s。

比较一下两种设计思路。GPU 起源于图形渲染，设计目标是通用，新任务上更容易「直接能跑」，对编译器的依赖小；但正因为硬件管得多，也更难推理和逼近 roofline——一个内核可能因为打爆 L2、访存没合并而性能很差，而你很难看出来。TPU 把更多负担交给编译器（单线程控制、只能做整个 VPU 宽度的指令，所有搬运要编译器排好流水线），反过来更容易接近峰值。另一个差异是 TPU 的快速片上存储（VMEM）多得多，如果能把权重放进或预取进 VMEM，LLM 推理可以更快。

11. GPU 的网络：NVLink、NVSwitch 与 InfiniBand

网络是 GPU 和 TPU 差别最大的地方。TPU 是只连邻居的 torus；GPU 用的是更传统的分层交换网络。

节点内（scale-up）：NVLink 与 NVSwitch

一个节点通常是 8 块 GPU（GB200 NVL72 可以到 72 块），彼此通过 NVLink 在一跳之内全互联，带宽高、延迟低。节点里有几个 NVSwitch 负责在本地 GPU 之间交换数据。H100 节点有 4 个 NVSwitch，每块 GPU 有 18 条 NVLink 4.0 链路，每条单向 25 GB/s（全双工，两个方向互不影响），于是每块 GPU 的出口带宽是 \(18\times25 = 450\) GB/s（NVIDIA 规格页写的 900 GB/s 是双向合计）。B200 的 NVLink 5 每条链路翻倍到 50 GB/s，每卡 900 GB/s 单向。

一个 8×H100 节点：NVSwitch 一侧理论上能处理 \(4\times64\times25\) GB/s = 6.4 TB/s，但每块 GPU 最多 450 GB/s，所以节点的峰值是 \(8\times450\) GB/s = 3.6 TB/s，这也是它的对分带宽（把 8 块卡分成两半，两半之间能流过的总带宽）。

理论值和实测有差距：在 8×H100 节点内实测 AllReduce，总线带宽最多接近 370 GB/s，而且要 10 GB 级别的大消息才能达到；像 Llama 3 70B 一个 MLP 权重那么大的数组（8 路切分后约 58 MB），只能跑到约 150 GB/s。Hopper 开始的 NVSwitch 支持 SHARP（交换机内归约），理论上能让 AllReduce 的成本减半，实测提升约 30%。

节点间（scale-out）：InfiniBand 胖树

节点之间用带宽低一些的 InfiniBand（IB）或以太网相连，每块 GPU 配一张网卡。NVIDIA 的参考设计 DGX SuperPod（1,024 块 H100）是这样搭的：

• 每个节点 8 张 400 Gbps 的 ConnectX-7 网卡，即每节点 \(8\times 400/8 = 400\) GB/s 的出口带宽（每卡 50 GB/s）；
• 32 个节点（256 块 GPU）挂在 8 台叶（leaf）交换机下，叫一个 SU（scalable unit）；
• 4 个 SU 再由 16 台脊（spine）交换机连起来，共 1,024 块 GPU。

整个 IB 网络是一棵胖树（fat tree），线缆和交换机的带宽经过设计，保证节点之上有完整的对分带宽：任意把节点分成两半，每个节点都能以 400 GB/s 向另一半发送。这意味着跨节点的 AllGather/ReduceScatter 有大致恒定的带宽，无论跨多少节点。这只是参考设计，不同客户和云厂商会改；例如 Meta 训练 Llama 3 用的是以太网、三层交换和顶层超额订阅（oversubscribed）的网络。

做估算时可以记：节点内每卡 450 GB/s，跨节点每节点 400 GB/s，后者在 AllToAll 这种「每卡对每卡」的通信里会退化得很厉害（跨 2 个节点时相比节点内慢 4 倍以上）。GB200 NVL72 把 72 块 GPU 放进一个 NVLink 域，每卡 900 GB/s，节点出口带宽也同比提高到 3.6 TB/s，这显著改变了 GPU 上的通信 roofline。

Scale-up 与 scale-out

这两个词描述的正是上面两层：scale-up 网络（NVLink 域）是少量芯片之间的高带宽、低延迟全互联，不为大规模和容错设计（有人形容它是「加强版的 PCIe」）；scale-out 网络（IB/以太网）可以连成千上万个节点，带宽每卡低一个数量级，更像传统的数据中心网络。TPU 的 ICI 介于两者之间：它在整个 pod（数千芯片）范围内提供每芯片恒定的带宽（v5p 每链路 90 GB/s、三个轴共 540 GB/s 出口），但只能用于规整的通信模式。并行策略的核心约束就在这里：通信量最大的并行方式（例如张量并行）要放在快的那一层里。

12. 对比与规格表

部件对照

芯片规格

GPU 的互联是每卡单向；TPU 是每链路单向。B200 芯片上有 192 GB HBM3e，HGX/DGX B200 里每卡可用 180 GB（8 卡共 1,440 GB）。价格方面，Google Cloud 上 H200 的标价约每小时 10 美元、TPU v5p 约 4 美元，单块 H200 的算力接近 v5p 的 2 倍、显存 1.5 倍——GPU 单卡更强，TPU 更依赖把很多芯片连起来。

13. 小结

• 任何运算的时间都可以用 \(T_\text{math} = \text{FLOPs}/\text{峰值}\) 和 \(T_\text{comms} = \text{Bytes}/\text{带宽}\) 来估算，完全重叠时取 \(\max\)，不重叠时取和，二者最多差 2 倍。
• 算术强度 = FLOPs / Bytes。大于硬件的临界强度（峰值 / 带宽）就计算受限，否则带宽受限。主流芯片的 HBM 临界强度在 150–600 之间。
• BF16 矩阵乘法的强度约等于每副本的 token batch \(B\)；在 H100 上要 \(B \gtrsim 300\)、TPU v5e 上 \(B\gtrsim 240\) 才能计算受限。只量化权重时门槛减半。
• 逐元素运算、点积、生成阶段的注意力、每个样本独有的权重：强度是常数，永远带宽受限，只能靠融合和减少字节来加速。
• 网络也有 roofline：沿缩并维切分的矩阵乘法，临界条件取决于每块芯片上的 \(D\) 而不是 \(B\)。越外层的网络临界强度越高（H100：HBM 约 300，NVLink 约 2,200，IB 约 2 万）。
• TPU = 少数大核心 + 大片上存储 + 只连邻居的 torus；GPU = 上百个小 SM + 小缓存 + 节点内全互联、节点间胖树。

有了每条链路的带宽，就能算具体的通信操作（AllGather、ReduceScatter、AllReduce、AllToAll）各要多少时间，以及矩阵乘法被切分后需要哪些通信：见集合通信与分片矩阵乘法。

参考与延伸阅读

这一篇的 roofline 和硬件部分主要整理自 How to Scale Your Model 的第 1、2、12 章，GPU 内核优化的部分来自 Ultra-Scale Playbook。想动手练习，推荐把第 1 章和第 12 章末尾的习题做一遍；芯片规格以 NVIDIA 和 Google Cloud 的官方页面为准。

• Jacob Austin 等，How to Scale Your Model：Part 1: A Brief Intro to Roofline Analysis（Where Does the Time Go、Visualizing rooflines、Matrix multiplication、Network communication rooflines 及习题）；Part 2: How to Think About TPUs（What Is a TPU、TPU Networking、TPU specs、附录 A/B）；Part 12: How to Think About GPUs（What Is a GPU、Memory、GPUs vs. TPUs、Networking 及习题）；Part 4 的习题 4、8（注意力与 MoE 的强度）。
• Nouamane Tazi 等，The Ultra-Scale Playbook：A primer on GPUs、Improving performance with kernels（memory coalescing、tiling、thread coarsening、control divergence）、Fused kernels。
• NVIDIA 规格页：H100、HGX、DGX B200。
• Google Cloud TPU 文档：v5e、v5p、v6e。
• Llama Team, Meta，The Llama 3 Herd of Models（2024）：第 3.3 节的训练集群网络（RoCE、三层 Clos、顶层超额订阅）。
• Samuel Williams、Andrew Waterman、David Patterson，Roofline: An Insightful Visual Performance Model for Multicore Architectures（CACM 2009），roofline 模型的原始论文。