起因#
这是 CS336 系列的又一篇。
之前那篇讨论了 FLOPs 到 Latency 之间的鸿沟：FLOPs 相同的两个 GEMM，latency 可能差几倍甚至十几倍。那篇文章的核心结论是 FLOPs 是 workload metric，不是 performance metric，并且引入了 Arithmetic Intensity 和 Roofline Model 作为分析工具。
但写完之后我发现，有几个问题其实还没展开讲：
- 为什么 Transformer 训练大量时间花在矩阵乘法上？矩阵乘法到底特殊在哪？
- 为什么一个大 GEMM GPU 可以跑满 Tensor Core，一个小 MLP 推理反而很慢？
- 为什么 LayerNorm、Softmax、Embedding lookup 这些 FLOPs 不高的操作，经常成为性能瓶颈？
- 拿到一个新的 workload，怎么判断它是 compute bound 还是 memory bound？
上一篇更多在讲"FLOPs 为什么不等于 latency"。这一篇想往回退一步，从更基础的地方讲起：性能优化不能只看 FLOPs，需要理解计算和数据搬运之间的关系。
为什么 AI 里到处都是矩阵乘法#
这一节不是线代课，只补三个对后面理解性能分析最关键的直觉。
Vector#
深度学习里到处都是向量，只是叫法不同：
- 一个词的 embedding 是向量
- Transformer 每一层的 hidden state 是向量
- 推荐系统里的 user / item feature 是向量
本质上都是"用一组数字表示一个东西"。维度越高，能表达的信息越丰富。
Dot Product#
两个向量的点积（dot product）：
$$a \cdot b = \sum_{i} a_i b_i$$从代数上看就是逐元素相乘再求和。但它有一个重要的几何含义：
$$a \cdot b = \|a\| \cdot \|b\| \cdot \cos\theta$$其中 \(\theta\) 是两个向量之间的夹角。
也就是说，点积同时编码了两个信息：向量的长度（magnitude）和方向的相似程度（角度）。 如果把长度归一化掉（除以各自的 norm），剩下的就是 cosine similarity——纯粹衡量"方向有多像"。
这就是为什么 Transformer 的 Attention 机制大量使用点积：
$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$\(QK^T\) 本质上就是在算 query 和每个 key 之间的相似度。点积越大，说明这个 key 和 query 越"相关"，attention weight 越高。
Matrix Multiplication = 批量 Dot Product#
矩阵乘法：
$$C = A \times B, \quad A \in \mathbb{R}^{N \times M},\ B \in \mathbb{R}^{M \times K},\ C \in \mathbb{R}^{N \times K}$$教科书上说"行乘列"，但这个描述太机械了。更有用的理解方式是：
矩阵乘法 = 大量 dot product 的结构化批量执行。
\(C\) 的每个元素 \(C_{ij}\) 就是 \(A\) 的第 \(i\) 行和 \(B\) 的第 \(j\) 列做一次 dot product。\(C\) 一共有 \(N \times K\) 个元素，所以一共做了 \(N \times K\) 次长度为 \(M\) 的 dot product。
这就是矩阵乘法的计算量来源：\(N \times K\) 次 dot product，每次 \(M\) 个 multiply-add，所以总共 \(2NMK\) FLOPs。
连接到神经网络#
神经网络中的 Linear layer：
$$Y = XW + b$$其中 \(X \in \mathbb{R}^{B \times D_{\text{in}}}\) 是一个 batch 的输入，\(W \in \mathbb{R}^{D_{\text{in}} \times D_{\text{out}}}\) 是权重。
这就是一个矩阵乘法。每个 sample 的输入向量和 \(W\) 的每一列做 dot product，得到输出的一个维度。Linear layer 本质是一个可学习的空间变换——把输入从 \(D_{\text{in}}\) 维空间映射到 \(D_{\text{out}}\) 维空间。
Transformer 里的 QKV projection、FFN、output projection，全都是 Linear layer，也就是全都是 GEMM。这就是为什么 Transformer 训练的大部分时间花在矩阵乘法上。
从 FLOPs 到 Arithmetic Intensity#
FLOPs 只是计算量#
GEMM 的 FLOPs 是 \(2NMK\)，这个在 上一篇 已经推过了。
但 FLOPs 只回答一个问题：一共需要做多少次运算。
它不回答：这些运算能跑多快。
上一篇用过一个例子：三组 GEMM，\(2MKN\) 几乎一样，但 shape 从 balanced 到 tall-skinny，Achieved FLOPS 差了好几倍。
那到底是什么决定了"同样多的计算，GPU 能以多快的速度做完"？
Arithmetic Intensity#
答案的核心在 Arithmetic Intensity（算术强度）：
$$\text{Arithmetic Intensity (AI)} = \frac{\text{FLOPs}}{\text{Bytes}}$$单位是 FLOPs/Byte。
它回答的问题是：每从内存搬运 1 byte 数据到计算单元，能做多少次计算？
这个 ratio 非常关键，因为它刻画了一个 workload 的"计算密度"——计算和数据搬运之间的比例关系。
高 AI 意味着：
- 数据搬进来以后被反复使用（高复用）
- 计算单元大部分时间在算，不在等数据
- GPU 的算力容易被"喂饱"
低 AI 意味着：
- 搬一次数据只做几次运算就扔掉
- 计算单元大量时间在等内存
- 算力再高也用不上
这就是为什么光看 FLOPs 不够。两个 workload FLOPs 一样，但如果一个 AI 高、一个 AI 低，前者能充分利用 GPU 算力，后者被内存带宽卡住——latency 可能差很多。
为什么 GEMM 通常是 Compute Bound#
这是本文最核心的一节。
先说结论：矩阵乘法最大的优势不是计算量大，而是数据复用率极高。
数据复用：矩阵乘法的核心优势#
回到矩阵乘法 \(C = A \times B\)，\(A \in \mathbb{R}^{N \times M}\)，\(B \in \mathbb{R}^{M \times K}\)。
看 \(A\) 中的一个元素 \(A_{ij}\)：它参与 \(C\) 的第 \(i\) 行所有 \(K\) 个输出元素的计算。
看 \(B\) 中的一个元素 \(B_{jk}\)：它参与 \(C\) 的第 \(k\) 列所有 \(N\) 个输出元素的计算。
也就是说：\(A\) 的每个元素被 reuse \(K\) 次，\(B\) 的每个元素被 reuse \(N\) 次。
具体算一下。以 BF16 为例（每个元素 2 bytes）：
- 读 \(A\)：\(2NM\) bytes
- 读 \(B\)：\(2MK\) bytes
- 写 \(C\)：\(2NK\) bytes
- 总 memory traffic：\(2(NM + MK + NK)\) bytes
- FLOPs：\(2NMK\)
理论 Arithmetic Intensity：
$$\text{AI}_{\text{GEMM}} = \frac{2NMK}{2(NM + MK + NK)} = \frac{NMK}{NM + MK + NK}$$拿一个具体例子感受一下。\(A(1024, 512) \times B(512, 2048)\)：
- FLOPs = \(2 \times 1024 \times 512 \times 2048 \approx 2.1 \times 10^9\)
- Bytes = \(2 \times (1024 \times 512 + 512 \times 2048 + 1024 \times 2048) \approx 8.4 \times 10^6\)
- AI \(\approx 256\) FLOPs/Byte
每搬 1 byte 数据进来，可以做 256 次运算。 这就是为什么 GEMM 能够充分利用 GPU 的计算能力。
再看另一个极端。如果 \(N=1\)（batch=1 推理），\(A(1, 512) \times B(512, 2048)\)：
- FLOPs = \(2 \times 1 \times 512 \times 2048 \approx 2.1 \times 10^6\)
- Bytes = \(2 \times (512 + 512 \times 2048 + 2048) \approx 2.1 \times 10^6\)
- AI \(\approx 1\) FLOPs/Byte
AI 从 256 骤降到 1。因为 \(B\)（weight matrix）有 \(512 \times 2048\) 个元素，但只被用了一次（\(N=1\)，没有 batch 维度来 reuse）。搬了这么多数据进来，每个只用一次就扔掉了。
这就是 batch size 对 GEMM 性能影响巨大的根本原因：batch size 直接决定了 weight 的 reuse 次数，从而决定了 AI。
和 Elementwise 操作对比#
作为对照，看 LayerNorm。
LayerNorm 对每个 token 的 hidden vector 做 normalize：算均值、方差、减均值、除标准差、乘 \(\gamma\) 加 \(\beta\)。
对一个长度为 \(D\) 的 vector：
- 读入 \(D\) 个元素：\(2D\) bytes（BF16）
- 计算量大约是 \(5D\sim8D\) FLOPs（均值 + 方差 + normalize + scale/shift）
- 写回 \(D\) 个元素：\(2D\) bytes
AI 大约是 \(\frac{5D \sim 8D}{4D} \approx 1 \sim 2\) FLOPs/Byte。
每搬 1 byte 只做 1~2 次运算。 这和 GEMM 的几百甚至上千的 AI 完全不是一个量级。
这就是为什么 LayerNorm 虽然 FLOPs 很少，却经常成为性能瓶颈——它不是算不过来，而是搬数据搬不过来。
GPU 如何利用 GEMM 的高 AI#
GPU 做 GEMM 时的 tiling 策略，正是为了最大化数据复用：
- 把 \(C\) 的输出切成 tile
- 每个 tile 对应 \(A\) 的一组行和 \(B\) 的一组列
- 把这些 \(A\)、\(B\) 的子块从 HBM 搬到 shared memory
- 在 shared memory 里反复 reuse，用 Tensor Core 做大量 FMA
- 只有算完了才写回 HBM
flowchart LR
HBM["HBM
（慢，大）"]
SMEM["Shared Memory
（快，小）"]
TC["Tensor Core
（极快）"]
HBM -->|"加载 A/B 的 tile"| SMEM
SMEM -->|"反复 reuse"| TC
TC -->|"算完后写回"| HBM
tiling 的本质就是把 GEMM 天然的数据复用，映射到 GPU 的 memory hierarchy 上。 shared memory 是 SRAM，带宽远高于 HBM，所以只要数据能在 shared memory 里被反复使用，计算单元就不会饿。
这也是为什么 GPU 喜欢大 GEMM：矩阵越大，每个 tile 能做的计算越多，数据复用越充分，Tensor Core 利用率越高。
Roofline Model#
有了 Arithmetic Intensity，就可以用 Roofline Model 来判断一个 workload 的瓶颈。
核心公式：
$$\text{Achievable Performance} = \min(\text{Peak Compute},\ \text{Memory BW} \times \text{AI})$$直觉上很好理解。一个 workload 的实际性能，要么被计算能力限制，要么被数据搬运能力限制——取决于哪个先成为瓶颈。
两个区域#
Memory Bound 区域（AI 低）：
当 AI 低时，\(\text{BW} \times \text{AI}\) 小于 Peak Compute。瓶颈在内存带宽——数据搬不过来，计算单元在等。
优化方向：减少 memory traffic（fusion、quantization）、提高带宽利用率。
Compute Bound 区域（AI 高）：
当 AI 高时，\(\text{BW} \times \text{AI}\) 超过了 Peak Compute。瓶颈在算力——数据来得及，但计算单元忙不过来。
优化方向：用更高效的计算指令（Tensor Core）、更好的 tiling。
Ridge Point#
两个区域的分界点叫 ridge point：
$$\text{AI}_{\text{ridge}} = \frac{\text{Peak Compute}}{\text{Memory BW}}$$当 workload 的 AI 等于 ridge point 时，计算和内存刚好同时饱和。低于这个点是 memory bound，高于是 compute bound。
用一个简化的 Roofline 概念图来理解：
flowchart TB
subgraph Roofline["Roofline Model"]
direction LR
MB["Memory Bound 区域
性能 = BW × AI
（斜线，随 AI 线性增长）"]
RP["Ridge Point
AI = Peak / BW"]
CB["Compute Bound 区域
性能 = Peak Compute
（水平线，不再增长）"]
MB --> RP --> CB
end
Low["低 AI 操作
LayerNorm, Softmax
Embedding Lookup"] --> MB
High["高 AI 操作
大 GEMM, Conv"] --> CB
ridge point 左边，你能做的是减少搬运、提高 AI。ridge point 右边，你该做的是提升计算吞吐。搞反了方向，优化就是在浪费时间。
关于 Roofline 的更多细节（包括它的局限性——它只给 upper bound，不 model parallelism 和 launch overhead），可以看 上一篇。
Accelerator Intensity#
上面说的 ridge point，其实就是 Accelerator Intensity（有些文献也叫 machine balance 或 ops:byte ratio）。
$$\text{Accelerator Intensity} = \frac{\text{Peak Compute (FLOPS)}}{\text{Memory Bandwidth (Bytes/s)}}$$单位也是 FLOPs/Byte，和 Arithmetic Intensity 一样。但含义完全不同：
- Arithmetic Intensity 是 workload 的属性——这个计算任务本身，每 byte 数据能产生多少 FLOPs
- Accelerator Intensity 是硬件的属性——这块芯片，计算能力和搬运能力的比值
判断 bound 类型的规则非常简单：当 workload 的 AI > 硬件的 Accelerator Intensity 时，是 compute bound；反之是 memory bound。
具体硬件的数字#
看几块典型 GPU（BF16 Tensor Core）：
| GPU | Peak BF16 TFLOPS | HBM BW (TB/s) | Accelerator Intensity (FLOPS/Byte) |
|---|---|---|---|
| A100 SXM | 312 | 2.0 | ~156 |
| H100 SXM | 989 | 3.35 | ~295 |
| B200 | 2250 | 8.0 | ~281 |
需要注意的是，这里用的是 peak spec。实际 sustained bandwidth 和 compute throughput 都会低一些，所以真实的 Accelerator Intensity 和上面的数字可能有偏差。但数量级是对的，用来建立直觉足够了。
算力涨得比带宽快#
看上面的表，一个趋势很明显：
- A100 → H100：BF16 算力涨了约 3.2x，HBM 带宽涨了约 1.7x
- Accelerator Intensity 从 ~156 涨到 ~295
算力增长速度远超内存带宽增长速度。
这意味着：随着 GPU 越来越强，workload 需要越来越高的 AI 才能跑进 compute bound 区域。 很多在 A100 上还算 compute bound 的操作，到了 H100 上可能就变成 memory bound 了——因为计算能力更强，但数据供应没跟上。
这是 ML Infra 领域一个结构性趋势：memory wall 问题只会越来越严重。 FlashAttention、kernel fusion、quantization 这些优化技术之所以越来越重要，根本原因就在这里——它们本质上都是在减少 memory traffic 或提高 AI。
常见 ML 操作属于什么 Bound#
有了 AI 和 Accelerator Intensity 这两个工具，就可以把常见的 ML 操作分个类。
| Operation | Typical Bound | 原因 | 常见优化 |
|---|---|---|---|
| GEMM / Linear（大 batch） | Compute Bound | 数据复用率高，AI 几百以上 | Tensor Core、tiling、BF16 |
| Conv（大 feature map） | Compute Bound | im2col 转成 GEMM，同样高复用 | Tensor Core |
| LayerNorm | Memory Bound | 每元素只做几次运算，AI ≈ 1~2 | Kernel Fusion |
| Softmax | Memory Bound | 需要多次遍历 tensor（max、exp、sum、normalize） | FlashAttention / Fusion |
| Embedding Lookup | Memory Bound | random access，几乎没有计算 | 改 layout、prefetch |
| Attention（\(QK^T\) 部分） | Depends | 大 seq → compute bound；decode 阶段 → memory bound | FlashAttention |
| Small MLP（batch=1） | Memory / Latency Bound | weight 搬一次只用一次，AI ≈ 1 | Fusion、CUDA Graph |
| Python-heavy pipeline | Latency Bound | kernel launch / dispatch 开销主导 | Graph capture、C++ runtime |
需要强调的是：一个操作是什么 bound，取决于具体的 shape、batch size 和硬件。 上表只是 typical case。
比如 GEMM，大 batch 时是 compute bound，batch=1 时可能变成 memory bound（前面算过，AI 从 256 降到 1）。Attention 在 prefill 阶段（长 sequence）和 decode 阶段（逐 token）的 bound 类型完全不同。
还有一类不太容易归类的：Latency Bound。 这种情况下，瓶颈不在计算也不在内存，而在各种固定开销——Python 解释器、CUDA kernel launch、CPU-GPU 同步、operator dispatch。对于极小的 workload（比如 batch=1 的一个小 MLP），每个 kernel 的计算可能只要几微秒，但 launch + sync 的开销可能更大。
联系实际 ML Infra 场景#
把上面的分析框架应用到三个典型场景。
大模型训练#
特点：
- batch size 大（几千甚至几万 tokens per step）
- GEMM 维度大且规整
- 计算占主导
这是最"幸福"的场景——大 batch 意味着高 AI，GPU 的 Tensor Core 容易被喂饱。
主要优化方向：
- 充分利用 Tensor Core：用 BF16 / TF32，选合适的 tile size
- Distributed training：模型大到一张卡放不下，需要 Tensor Parallel / Pipeline Parallel，但通信开销引入新的瓶颈
- 减少非 GEMM 开销：LayerNorm、Softmax 这些 memory bound 操作通过 fusion 减少 HBM 读写
大模型推理（LLM Serving）#
这里需要区分两个阶段。
Prefill 阶段：输入 prompt 一次性算完。行为类似训练的 forward pass——sequence 长度大，GEMM 规整，通常 compute bound。
Decode 阶段：逐 token 生成。每步只产生 1 个 token，等价于 batch=1 的一系列 GEMM。weight matrix 每个 token 都要从 HBM 重新读一遍，但只做一次乘加。AI 极低，典型的 memory bound。
再加上 KV Cache——每一步 decode 需要读取之前所有 token 的 K、V，随着 sequence 增长，memory traffic 持续增加。
主要优化方向：
- FlashAttention：减少 Attention 的 HBM 读写，把多次 pass 合并成一次
- KV Cache 优化：PagedAttention（vLLM）、GQA/MQA 减少 KV 数量
- Speculative Decoding：用小模型 draft 多个 token，大模型一次验证，把 decode 的多步 memory bound 变成一步 compute bound 的 verify
- Quantization：weight 从 BF16 → INT8/INT4，直接减少搬运量
低延迟小模型推理#
这是我自己工作中最常接触的场景。量化金融里的 MLP / GRU 模型，特点：
- 模型很小（参数量可能只有几百 K 到几 M）
- batch=1（逐笔数据来了就要预测）
- FLOPs 极少，但对 latency 要求极高（微秒级）
这种场景下：
- GEMM 太小，Tensor Core 根本跑不满（tile 都填不满几个 SM）
- weight 搬一次只用一次，AI ≈ 1，memory bound
- 每个 kernel 的实际计算可能就几微秒，但 kernel launch、CUDA driver overhead、Python dispatch 可能和计算本身一个量级——latency bound
优化方向完全不同：
- Kernel Fusion：把多个小 kernel 合成一个，减少 launch 次数和中间结果的 HBM 读写
- CUDA Graph：把一系列 kernel launch 录制下来一次性回放，消除逐个 launch 的开销
- CPU-side 优化：SIMD、cache locality，甚至绕过 GPU 直接在 CPU 上跑
- 避免 Python overhead：用 C++ inference runtime，或者 TorchScript / torch.compile
flowchart TB
subgraph Training["大模型训练"]
T1["Batch 大，GEMM 大"] --> T2["Compute Bound"]
T2 --> T3["优化：Tensor Core
BF16，Distributed"]
end
subgraph Serving["大模型推理"]
S1["Prefill: Compute Bound
Decode: Memory Bound"] --> S2["瓶颈在 KV Cache
和 Weight Loading"]
S2 --> S3["优化：FlashAttention
KV Cache，Quantization"]
end
subgraph LowLatency["低延迟小模型推理"]
L1["Batch=1，模型小"] --> L2["Memory / Latency Bound"]
L2 --> L3["优化：Fusion
CUDA Graph，SIMD"]
end
建立性能分析的心智模型#
写到最后，收一个实用的判断框架。拿到一个 workload，不要先急着优化，先问三个问题：
1. 主要操作的 AI 大概是多少？
大 GEMM 的 AI 通常在几百以上。elementwise / reduction 操作通常在个位数。batch size 对 GEMM 的 AI 影响巨大。
2. 目标硬件的 Accelerator Intensity 是多少？
对现代 GPU（BF16 Tensor Core），大概在 150~300 之间。AI 低于这个值就是 memory bound。
3. 还有没有 latency bound 的因素？
kernel launch overhead、Python dispatch、CPU-GPU 同步。对极小 workload 这些可能是主导因素。
然后根据判断选择优化方向：
- Compute Bound → 提升计算效率：Tensor Core、更好的 tiling、BF16/INT8 提高吞吐
- Memory Bound → 减少搬运：kernel fusion、FlashAttention、quantization、KV Cache 优化
- Latency Bound → 减少固定开销：CUDA Graph、operator fusion、C++ runtime
搞反了方向，就是在 memory bound 的 workload 上拼命优化计算（没用），或者在 compute bound 的 workload 上做 fusion（没必要）。
参考资料#
- CS336: Language Modeling from Scratch
- Williams, Waterman, Patterson. Roofline: An Insightful Visual Performance Model for Multicore Architectures. Communications of the ACM, 2009.
- 从 FLOPs 到 Latency：GPU 推理性能到底由什么决定？（本系列上一篇）
- 从 IEEE 754 到 BF16：理解 ML Infra 中的浮点精度选择（姊妹篇）
- Dao et al. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS, 2022.
- NVIDIA. CUDA C++ Programming Guide — Warp-level Matrix Multiply-Accumulate (WMMA) and Tensor Cores.
- NVIDIA. H100 Tensor Core GPU Architecture Whitepaper.