一句话概括：这本书讲的是「一个模型跑在真实硬件上，时间到底花在哪」。它的价值不在于教你调参，而在于让你在写代码之前就能用纸笔估出：这次训练要多少卡、跑多少天；这个模型自己部署要多少显存、能扛多少并发；某个想法是不是根本不可能快起来。我读下来最大的收获是发现整本书几乎可以由一个式子推出来——所以这份笔记就从那个式子开始讲。
先约定九个词
这几个词后面每一节都会用到，先在这里说清楚，正文里就不重复解释了。已经熟的可以直接跳过。
- token
- 模型处理文字的最小单位，大致相当于一个词或半个词（中文通常一两个字一个）。模型不是一个字一个字读，而是一个 token 一个 token 读。
- 参数 / 权重
- 模型里训练出来的那一大堆数字，本质是很多个矩阵。「70B 模型」就是说它有 700 亿个这样的数。权重是死的、固定的，每次推理都得把它读一遍。
- 激活
- 数据流过模型时产生的中间结果。权重是固定的，激活是随着输入变的——你问的问题不同，激活就不同。
- 梯度
- 训练时算出来的「这个参数该往哪个方向调、调多少」。只有训练有梯度，推理没有。
- FLOPs
- 浮点运算次数，就是「做了多少次加减乘除」。
FLOP/s（带斜杠）是每秒能做多少次，衡量芯片多能算。 - 带宽
- 每秒能搬多少字节数据。注意它和算力是两回事：芯片能算得飞快，但数据喂不进来照样干等——这份笔记基本就是在讲这件事。
- 显存 / HBM
- 显卡上的内存，模型权重就住在这儿。HBM 是它的技术名（高带宽内存）。容量决定你能不能跑，带宽决定跑多快。
- batch / 并发
- 一次同时处理多少条数据。训练时叫 batch（批大小），推理时更常说「并发」——同时在伺候多少个用户。这个数是全篇的主角。
- bf16
- 一种数字格式，每个数占 2 字节。模型现在基本都用它存权重，所以「70B 参数」≈「140 GB」。后面所有字节数都按它算。
0. 先说结论：那个式子
做一次权重矩阵乘（模型里最主要的运算：把输入的一批数据乘上一个权重矩阵，Transformer 里绝大部分时间都花在这上面）：输入是 B 个 token 的激活 [B, D]（B 行 D 列的一张表），权重是 [D, F]。
- 算力开销：每个输出元素要做
D次乘加（乘一下再加到累计值上，矩阵乘就是由一堆乘加组成的），一次乘加算 2 FLOPs，输出有B×F个元素 → 共2·B·D·FFLOPs。 - 搬运开销：bf16 下每个数 2 字节。要搬的是激活
2BD、权重2DF、结果2BF字节。（「搬运」指把数据从显存挪到芯片的计算单元里——数据不搬过去，芯片就没法算。）
推理解码时 B 通常远小于 D 和 F，权重那一项 2DF 完全压倒另外两项。于是算术强度（arithmetic intensity，每搬运一字节能摊到多少次浮点运算）就是：
I = 2·B·D·F / (2·D·F) = B
B。这张图是这份笔记的地基，后面每一节都在用它。这个结果干净得有点吓人：一次权重矩阵乘的算术强度，约等于你这一批喂进去的 token 数。D 和 F 全约掉了——模型多大根本不影响这个比值。
为什么这件事重要？因为每块加速器都有一个临界算术强度 I* = 峰值算力 ÷ 内存带宽。低于它，芯片在等内存；高于它，芯片在算数。把公开规格代进去：
| 芯片 | bf16 峰值 | HBM 带宽 | I*（FLOPs/字节） |
|---|---|---|---|
| NVIDIA A100 80GB | 312 TFLOP/s | 2.0 TB/s | ≈ 156 |
| NVIDIA H100 SXM | 990 TFLOP/s | 3.35 TB/s | ≈ 295 |
| TPU v5e | 197 TFLOP/s | 820 GB/s | ≈ 240 |
| TPU v5p | 459 TFLOP/s | 2.76 TB/s | ≈ 166 |
数量级都在一两百到三百。结合上面那个 I ≈ B，直接得到一条可以贴在显示器上的经验法则：
喂给一次矩阵乘的 token 数如果不到两三百，你买的算力就基本是浪费的——芯片大部分时间在等 HBM。
后面所有内容——为什么要 batching、为什么 GQA 能救命、为什么 decode 比 prefill 难十倍、为什么会有 continuous batching 和投机解码——全都是这一条的推论。这也是我认为这本书最值钱的地方：它把一堆看起来互不相干的工程 trick，还原成同一个约束下的不同解法。
⟶ 原书对应章节：Part 1 · All About Rooflines
1. Roofline：三个上限，不是一个
「性能」这个词太糊，roofline 模型把它拆成三个各自独立的物理上限：
- 算力上限：芯片每秒能做多少次浮点运算（OPs/s）。
- 带宽上限：每秒能在内存和计算单元之间搬多少字节（bytes/s）。
- 容量上限：总共能放下多少字节（bytes）。
前两个决定快不快，第三个决定能不能跑。一个算子（operator，指一次具体的运算，比如「一次矩阵乘」「一次加法」「一次 LayerNorm」——模型就是由成千上万个算子串起来的）的实际耗时是
T ≈ max( FLOPs ÷ 峰值算力 , 字节数 ÷ 带宽 )
取 max 而不是相加，是因为现代芯片能让搬运和计算重叠——理想情况下慢的那个把快的那个完全遮住。画成图就是一条折线：横轴算术强度，左半段是斜率为带宽的上升直线（访存受限），右半段是算力峰值的水平线（算力受限），拐点就是 I*。
B=1 的解码和训练差着两三个数量级的利用率——差别不在芯片，在你喂了多少 token。为什么矩阵乘幸运，逐元素运算不幸
把常见算子的算术强度排一排，差距是数量级级别的：
| 算子 | 算术强度 | 结论 |
|---|---|---|
| 逐元素加 / 乘 | ≈ 1/6（读两个写一个 = 6 字节换 1 FLOP） | 永远访存受限 |
| 激活函数、LayerNorm | O(1) | 永远访存受限 |
权重矩阵乘（batch B） | ≈ B | 取决于 batch |
方阵乘 N×N | ≈ N/3 | N 够大就算力受限 |
逐元素运算（对每个数各做各的，比如整个张量统一加个数、统一过一遍激活函数）的强度是个不随规模变化的小常数——张量（多维数组，模型里的数据和权重都以这个形式存在）放大一万倍，它照样是 1/6。这就是算子融合（fusion）唯一的存在理由：把 norm → matmul → activation → 残差 这一串融成一个 kernel（在芯片上实际跑起来的一段计算程序），中间结果留在片上不落 HBM，几个访存受限的小算子就被合并进一次搬运里。你在 XLA / torch.compile（两个深度学习编译器，分别来自 JAX/TensorFlow 阵营和 PyTorch）里看到的绝大部分优化，本质都是在做这件事。
用它来否定想法
roofline 最实用的用法其实是快速证伪。任何「我加个轻量模块应该不影响速度」的直觉，都可以先算一下这个模块要搬多少字节：如果它是逐元素的、又要完整读写一遍激活，那它的耗时就等于「激活字节数 ÷ 带宽」，跟它算得多不多毫无关系。很多「理论上只增加 2% FLOPs」的改动，实测掉速 30%，原因都在这。
⟶ 原书对应章节：Part 1 · All About Rooflines
2. 硬件的心智模型：为什么长这样
脉动阵列：把 I* 撑高的代价
TPU 的核心是 MXU（矩阵乘单元），一个典型 128×128 的脉动阵列：数据从一边流进去，在阵列里逐级传递、每经过一个单元做一次乘加，从另一边流出结果。它的好处是一个数据一旦进来，会被复用上百次而不需要重新访存——这正是把算术强度撑到几百的硬件基础。
代价是它只擅长矩阵乘。所以芯片旁边还得配一个 VPU（向量单元）处理激活、归一化这些逐元素的活。你的模型里但凡有 MXU 干不了的形状，就得掉到 VPU 上跑，算力差着一两个数量级。这解释了为什么「维度对齐到 128 的倍数」这种听起来很迷信的建议是真有效的——形状不对齐，阵列就有一部分在空转。
内存层级：真正的战场
从快到慢大致是：寄存器 → 片上 SRAM（TPU 叫 VMEM，GPU 叫 SMEM/L2）→ HBM → 主机内存 → 别的机器。每往下一级，容量涨一两个数量级，带宽掉一个数量级。
这几个名字听着唬人，其实就是离计算单元的远近：寄存器紧挨着计算单元、快到几乎不要钱但只放得下几个数；SRAM 是刻在芯片上的高速小内存（TPU 管自己那块叫 VMEM，GPU 叫 SMEM 或 L2 缓存），几十 MB；再往外才是显存 HBM，几十上百 GB 但慢得多。可以类比成手边 → 桌上 → 书架 → 仓库。
S×S 那个中间结果掉到 HBM 这一层。所有高性能 kernel 的写法都是同一套路：把数据切成能塞进片上 SRAM 的 tile（小块），搬一次，在上面把所有能做的计算全做完，再搬下一块。FlashAttention（一种被普遍采用的注意力加速实现）之所以快，不是因为它减少了 FLOPs（它反而增加了一点重算），而是因为它避免了把 S×S 那张注意力矩阵写回 HBM——序列有 S 个 token，每个 token 都要和其它所有 token 算一次相关度，于是中间会产生一张 S×S 的大表；序列一长这张表大得吓人（S=8192 时单头就是 6700 万个数）。FlashAttention 把它拆成小块、一直留在片上算完，从不整张落地。这是 roofline 思维最漂亮的一次应用：用更多的算力换更少的搬运，在访存受限的区域里是稳赚的。
芯片之间：环面（torus）与树
单芯片放不下模型，就得连起来。这里 TPU 和 GPU 走了两条不同的路，理解这个差别对选并行策略很关键。先解释两个词：环面（torus）是把芯片排成网格、再把每一行每一列的首尾接起来形成的环状结构（想象一张纸卷成筒、再把筒的两头接成甜甜圈），好处是任意两点之间的最远距离被砍掉一半、而且没有「边缘节点」；NVLink 是 NVIDIA 用来连同一台机器内几张卡的高速专线，InfiniBand 则是连不同机器的网络，后者比前者慢一个数量级。
| 维度 | TPU Pod（2D/3D 环面） | GPU 集群（节点内全连 + 节点间网络） |
|---|---|---|
| 连接方式 | 每芯片只连上下左右邻居，边缘绕回 | 节点内 NVLink 全互联，节点间 InfiniBand/以太网 |
| 带宽分布 | 相对均匀，没有断崖 | 陡峭：节点内极快，出了节点掉一个数量级 |
| 扩展性 | 加芯片同时加链路，对分带宽随规模增长 | 受限于网络层级与收敛比 |
| 对策略的影响 | 大规模集合通信代价可预测 | 必须把重通信的并行方式关在节点内 |
环面的关键性质是没有「跨节点」这个断崖。GPU 集群上你必须精心设计「什么并行放节点内、什么放节点外」，而 TPU 上更多是在一个连续的 mesh 上切维度。这也是为什么同一套并行策略从 GPU 挪到 TPU（或反过来）常常要重新调——不是软件问题，是拓扑性格不同。
⟶ 原书对应章节：Part 2 · How to Think About TPUs ｜ Part 12 · How to Think About GPUs
3. 通信的代价模型：四个集合操作
一旦模型被切到多个芯片上，就需要集合通信（collective）把碎片拼回去——指的是「一群设备一起参与的、有固定套路的数据交换动作」，跟点对点发消息不同，它是全体同时进场的。只要记住四个，以及一条反直觉的结论。
| 操作 | 做什么 | 环形算法耗时 |
|---|---|---|
| AllGather | 每人有 1/N，最后每人都有全部 | ≈ V/W |
| ReduceScatter | 每人有全量，求和后每人留 1/N | ≈ V/W |
| AllReduce | = ReduceScatter + AllGather | ≈ 2V/W |
| AllToAll | 每人给每人发一份不同的数据 | ≈ V/W 量级 |
反直觉的结论在耗时那一列：N 不在里面。环形算法分 N-1 步，每步只传 1/N 的数据，乘起来 N 就约掉了。也就是说 AllReduce 的代价只取决于数据体积，几乎不取决于有多少台机器参与。这是数据并行能扩到上万卡的根本原因——否则每加一倍机器通信翻一倍，几百卡就到头了。
▸展开：为什么加机器不加通信时间？用 4 台机器走一遍
这个结论第一次看几乎不可信——「大家都要拿到全部数据，人越多不是越慢吗？」关键在于「每人只负责一小段」。设有 4 台机器，每台手里有一份 V 字节的梯度，目标是让每台最后都拿到四份相加的结果。
先把每台机器的 V 切成 4 小块，编号 ①②③④。机器排成一个圈，每台只跟右邻居说话。
第一阶段：ReduceScatter（边传边加）。第 1 拍，机器 A 把自己的①传给 B，B 收到后加到自己的①上；同时 B 把②传给 C、C 把③传给 D、D 把④传给 A——四条线同时在用，谁也没闲着。第 2 拍、第 3 拍重复这个动作。3 拍之后，每台机器手上都有某一块的完整总和（比如 D 手里的①已经是四台的①加在一起了）。
第二阶段：AllGather（把结果转一圈）。现在每台各有一块正确答案，再沿着圈传 3 拍，每台就集齐了四块。
现在算账。每一拍只传一小块，也就是 V/4。两个阶段各 3 拍，共 6 拍，每台机器总共发出 6 × V/4 = 1.5V 字节，约等于 2V。
换成 100 台呢？切成 100 块，每阶段 99 拍——拍数变多了，但每拍传的东西小了 100 倍。总量还是 2 × 99/100 × V ≈ 2V。这就是 N 约掉的原因：机器数同时出现在分子（拍数）和分母（每拍大小）上。
另一个容易忽略的点：所有链路是并行工作的。不是一台机器吭哧吭哧发给所有人，而是每台都只跟邻居收发，全场的线同时在跑。所以决定墙上时间的是「每台设备要吐出多少字节 ÷ 它那条线的带宽」，跟总人数无关。
这也顺带解释了为什么 AllReduce ≈ 2V 而 AllGather 和 ReduceScatter 各是 V——因为 AllReduce 就是上面那两个阶段拼起来的。
（现实里 N 还是会通过两个次要因素影响耗时：延迟项——每一拍都有固定的启动开销，拍数多了就累积；拓扑跳数——数据要多绕几个芯片。消息很小的时候这两项会冒头。但对训练里那些几百 MB 的梯度块，上面的体积项完全主导。）
算并行策略的通信开销时，我习惯只问两个问题：每一步要过线的字节是多少？它正比于参数量还是正比于 token 数？下一节会看到，这个区分直接决定了策略怎么组合。
⟶ 原书对应章节：Part 3 · Sharded Matrices and How to Multiply Them
4. Transformer 的三本账
要估任何东西，先得会算这三笔。先说记号——Transformer 由 L 个结构相同的层堆起来，每层里有注意力和 FFN（前馈网络，就是两个连着的矩阵乘）两个部件：
D= 隐藏维度：每个 token 在模型内部被表示成一个多长的数字向量。常见 4096、8192。这是模型「宽度」。F= FFN 中间维度：FFN 内部先把D撑宽到F再压回来，通常F ≈ 4D。L= 层数，模型「深度」。70B 级别通常 80 层。S= 序列长度：这次处理多少个 token（就是俗称的「上下文长度」）。B= batch：同时处理多少条序列。
第一本：参数量
每层里，注意力的 Q/K/V/O 四个投影各是 D×D，共 4D²（投影就是「乘一个矩阵把向量变成另一个向量」；Q/K/V/O 分别叫 query、key、value、output，下面会讲它们干嘛的）；FFN 上下投影是 2DF，用 SwiGLU 这类门控结构（多用一个矩阵算出一组「闸门」系数去缩放另一路的输出，效果更好但多一个矩阵）则是 3DF。取常见的 F ≈ 4D（门控结构通常把 F 缩到 8D/3 以保持参数量不变），每层大约 12D²，全模型 ≈ 12LD²，另加词表嵌入 V·D（把每个词映射成向量的那张查找表，V 是词表大小，通常十几万）。
第二本：训练算力，也就是 6ND
这个著名的估算可以两行推出来。训练一步分前向（把数据喂进去算出答案）和反向（从答案的误差倒推回去，算出每个参数该怎么调）两趟。前向时每个参数对每个 token 参与一次乘加 = 2 FLOPs；反向要算两个梯度（一个是「误差该继续往前一层传多少」，一个是「这个参数本身该调多少」），工作量约是前向的两倍 = 4 FLOPs。合计每参数每 token 约 6 FLOPs：
训练总算力 ≈ 6 × 参数量 × 训练 token 数
拿它估一次真实训练：70B 模型、15T token → 6 × 70e9 × 15e12 ≈ 6.3e24 FLOPs。假设 8192 张 H100、MFU（实际算力利用率）40%：8192 × 990e12 × 0.4 ≈ 3.2e18 FLOP/s。两者一除，约 23 天。用一个信封背面就得到了和公开报告同量级的答案——这就是这本书想教给你的能力。
注意 6ND 不含注意力里那两个 S×S 的矩阵乘，它们是 O(S²) 的。序列短时可以忽略，序列一长就会反超 FFN 成为主项——长上下文训练贵得不成比例，根子在这。
第三本：显存，以及 KV cache 为什么是主角
▸展开：注意力在干什么，以及 KV cache 是从哪冒出来的
要理解为什么 KV cache 是显存大户，得先知道注意力这一步在算什么。
注意力要解决的问题是：模型读到「它」这个字时，得知道这个「它」指的是前面哪个东西。所以每个 token 都要回头看看前面所有 token，挑出跟自己相关的。
具体做法是给每个 token 算三个向量，各由一个权重矩阵乘出来：
- Q（query，查询）：我想找什么样的信息。
- K（key，键）：我这里有什么样的信息，用来被别人匹配。
- V（value，值）：如果别人匹配上我了，我实际交出去的内容。
可以类比查字典：拿你的 Q 去和前面每个 token 的 K 逐个比对，算出一串「相关度打分」；打分高的说明相关，然后按这个分数加权平均它们的 V，就是这个 token 注意力的输出。前面提到那张 S×S 的大表，装的就是「每个 token 对每个 token 的打分」。
现在关键的一步来了。生成第 500 个 token 时，你需要前面 499 个 token 的 K 和 V。但——第 3 个 token 的 K 和 V，只取决于第 3 个 token 自己，跟你现在生成到第几个毫无关系。也就是说它们算过一次之后永远不会变。
那每生成一个字就把前面几百个 token 的 K、V 重算一遍，就是纯粹的浪费。于是干脆把它们存下来——这就是 KV cache。名字很直白：缓存起来的 K 和 V。
代价是它会一直长：每多生成一个 token 就往里追加一份，每多接一个并发用户就多一整份。所以它和权重完全不同——权重是固定的 140 GB，KV cache 是随着「聊多久 × 多少人在聊」线性膨胀的。长上下文服务真正的显存压力全在这儿。
还有一个细节：注意力不是只做一遍，而是分成多个头（head）并行做，每个头关注不同类型的关系，各自有自己的 Q、K、V。所以 KV cache 要按头数再乘一遍——下面公式里的 H_kv 和 d_head（每个头的向量长度）就是这么来的。
推理时显存分两块：权重（固定，bf16 下 = 2 × 参数量字节）和 KV cache（随并发和上下文线性增长）。后者每 token 的大小是：
2（K 和 V）× L × H_kv × d_head × 2 字节
关键是 H_kv——KV 头的数量。MHA（multi-head attention，多头注意力，最原始的做法）里每个头都有自己的 K 和 V，所以它等于总头数；GQA（grouped-query attention，分组查询注意力）让好几个 Q 头共用同一组 K/V——查询还是各查各的，但被查的那份「资料」共享，于是要缓存的东西成倍减少，效果几乎没损失。这个数就掉下来了。以 Llama-3-70B 的配置（L=80，64 个 Q 头，H_kv=8，d_head=128）实算：
| 方案 | KV 头数 | 每 token | 8K 上下文单条 |
|---|---|---|---|
| 若用 MHA | 64 | 2.5 MiB | 20 GiB |
| 实际用 GQA | 8 | 320 KiB | 2.5 GiB |
8 倍差距。在 8×H100（640 GB）上：权重占 140 GB，剩下约 500 GB 全给 KV，能装下约 190 条 8K 并发；换成 MHA 只剩 23 条。GQA 不是什么精度上的小让步，它是让长上下文服务在经济上成立的前提。
⟶ 原书对应章节：Part 4 · All the Transformer Math You Need to Know
5. 训练怎么切：五种切法与它们的账单
核心矛盾很简单：单卡放不下，多卡要通信。每种并行都是在「省多少显存」和「多花多少通信」之间选一个点。
先用一句话说清五种切法各自在切什么——它们的区别就是「拿刀往哪个方向砍」：
- 数据并行（DP）：模型不切，每张卡都放一整份，切的是数据。四张卡各算四分之一的数据，算完把梯度对一下。
- 完全分片（FSDP / ZeRO-3）：切参数。每张卡只存模型的四分之一，谁要用到哪部分，临时找大家要过来拼齐，用完就扔。省显存最狠。
- 张量并行（TP）：把单个矩阵横着或竖着劈开，四张卡各算一块，拼起来才是完整结果。
- 流水并行（PP）：切层。第 1 张卡算前 20 层，算完把结果传给第 2 张卡接着算——像工厂流水线。
- 专家并行（EP）：只有 MoE（mixture of experts，混合专家：模型里养着一堆「专家」小网络，每个 token 只激活其中一两个，从而参数很多但计算不多）模型才有，切的是专家，每张卡管几个专家，token 要去哪个专家就把它发过去。
下表把它们放在一起对比，关键看最后一列——通信量正比于参数还是正比于 token。表里的「优化器态」指训练时为每个参数额外维护的辅助数据（如动量），用 Adam 时它比参数本身还大一两倍，所以能不能分片它对显存影响很大。
| 策略 | 切什么 | 显存收益 | 每步通信 | 正比于 |
|---|---|---|---|---|
| 数据并行 DP | 切 batch，模型全复制 | 无 | AllReduce 梯度 ≈ 2P | 参数 |
| 完全分片 FSDP / ZeRO-3 | 切参数+梯度+优化器态 | ÷N，收益最大 | ≈ 3P | 参数 |
| 张量并行 TP | 切 D/F 维 | ÷N | 每层 AllReduce 激活 | token |
| 流水并行 PP | 切层 | ÷N | 只传边界激活，极小 | token（很小） |
| 专家并行 EP（MoE） | 切专家 | ÷N | 两次 AllToAll | token |
怎么读这张表
DP 和 FSDP 的通信正比于参数量，与 batch 无关。所以把每卡 batch 调大，通信就被摊薄了——这是数据并行唯一也是最重要的调参方向。反过来，卡多到每卡 batch 只剩几条时，DP 的效率会断崖式下跌。
TP 的通信正比于 token 数，而且是每层都要来一次。这意味着它对链路延迟极其敏感，一旦跨出高速域（GPU 的节点、TPU 的一小片 mesh）立刻崩盘。所以铁律是：TP 只在最快的那一层互联里用，并行度一般不超过单节点卡数（如 8）。
▸展开：一个矩阵劈成两半，为什么还能算对？以及那次 AllReduce 是哪来的
「把矩阵切开分给几张卡」听起来很悬——结果还能对吗？答案是能，而且有两种切法，代价完全不同。这是理解 TP 的关键。
设要算 Y = X · W，X 是输入，W 是权重矩阵。
切法一：竖着切（按列）。把 W 从中间竖着劈成左右两半 W₁、W₂。卡 1 算 X·W₁，卡 2 算 X·W₂。两边算出来的正好是最终结果的左半边和右半边，各自完整、拼起来就完事——全程不需要通信。
切法二：横着切（按行）。把 W 横着劈成上下两半，同时把 X 竖着劈成两半。卡 1 算 X₁·W₁，卡 2 算 X₂·W₂。这次两边算出来的都是整个结果的形状，但都只是「一部分贡献」——真正的答案是这两个加起来。于是必须做一次 AllReduce 把它们相加。
看出门道了吗：竖着切免费，横着切要通信。那为什么不全用竖着切？因为 FFN 是两个矩阵乘连在一起的（先撑宽到 F，再压回 D）。妙处在于：
- 第一个矩阵竖着切 → 免费，输出天然按列分好；
- 这个分好的输出，正好就是第二个矩阵横着切所需要的输入形状 → 直接接上，中间不用重排；
- 只在第二个矩阵之后付一次 AllReduce。
所以一层 FFN 只付一次通信，而不是两次。注意力那边同理，按头切——每张卡管几个完整的头，最后 O 投影时汇总一次。
那为什么每层都躲不掉？因为层与层之间夹着 LayerNorm、激活函数这些非线性运算。非线性的意思是「不能先各算各的再相加」——必须先把完整的结果拼齐，才能往下走。所以每一层结尾都得同步一次，L 层就是 L 次。
这也解释了那条铁律：这些 AllReduce 传的是激活（正比于 token 数），而且一层一次、密集发生，任何一点链路延迟都会被乘上 L 倍放大。放进 NVLink 内部（几百 GB/s）没问题，一旦跨到机器之间（几十 GB/s）就是灾难。
PP 通信最省，但它有独有的问题——流水气泡（bubble，指流水线开头和结尾那段「后面的人还没活干 / 前面的人已经干完」的空转时间）。N 个阶段、M 个微批次（把一批数据再切成更小的小份，一份份地送进流水线——不切的话流水线上永远只有一件工件，后面三个阶段全在等）时，空转比例约 (N-1)/(M+N-1)。要把气泡压到 10% 以下，微批次数得是阶段数的十倍左右，这又反过来要求全局 batch 足够大。PP 省的是带宽，花的是调度复杂度。
N−1 拍），所以唯一的解法是把微批次数堆上去把它摊薄——而这又要求全局 batch 足够大。这就是 PP 那条「省带宽、花调度」的代价具体长什么样。组合起来的次序
实践中这几种是叠着用的（所谓 3D/4D 并行），排布次序几乎总是按互联速度从快到慢来分配：
- 最内层（节点内 / mesh 内小片）→ TP：它最吃带宽，必须占最快的线。
- 中间层 → FSDP：通信量中等且可与计算重叠。
- 最外层（跨机架 / 跨集群）→ PP：通信最少，扛得住慢链路。
- MoE 的 EP 通常和 TP 同层或稍外。
把这个次序记住，比记住任何具体配置都有用——换了硬件，按互联层级重新套一遍就行。
⟶ 原书对应章节：Part 5 · How to Parallelize a Transformer for Training ｜ Part 6 · Training LLaMA 3 on TPUs
6. 推理是另一个世界
这是我认为最容易被低估的一章。推理不是「训练去掉反向」——它内部还分成两个瓶颈完全相反的阶段。
先说这两个阶段是什么。你给模型发一段话，它的处理分成两步：
- Prefill（预填充）：把你输入的整段提示词一次性读进去，建立起前面说的 KV cache。这一步一次处理成百上千个 token，所以矩阵很大、芯片很忙。它决定了你按下回车后多久看到第一个字。
- Decode（解码）：开始一个字一个字往外吐。每吐一个字都要把整个模型重新过一遍，但这一遍只处理 1 个 token。它决定了字往外冒得快不快。
行业里给这两件事各起了个缩写：TTFT（time to first token，首字延迟）和 TPOT（time per output token，每个字的间隔）。前者由 prefill 决定，后者由 decode 决定。
| Prefill（读提示词） | Decode（逐字生成） | |
|---|---|---|
| 一次处理 | 整个 prompt，几百上千 token | 每条序列 1 个 token |
| 算术强度 | 高（≈ prompt 长度） | 低（≈ 并发数） |
| 瓶颈 | 算力受限 | 访存受限 |
| 决定的指标 | 首字延迟 TTFT | 吐字速度 TPOT |
| 加卡有用吗 | 有用 | 只在能提高并发时有用 |
回到开头那个式子 I ≈ B：decode 时的 B 就是当前同时在生成的序列条数。单条请求 B=1，算术强度 1，而芯片需要两三百——算力利用率不到 1%。这时候你每生成一个 token，都要把整个模型的权重完整读一遍，只为了做那么一点点乘法。
由此可以直接推出 decode 的速度上限，而且这个估算准得惊人：
每 token 耗时 ≈ 模型字节数 ÷ HBM 带宽
70B 模型 bf16 = 140 GB，8×H100 聚合带宽约 26.8 TB/s → 约 5.2 ms/token，约 190 token/s，且这个数在并发拉满之前几乎不随并发变化（因为读权重的开销被所有序列共享）。这解释了两件事：为什么量化到 int8/fp8 能几乎线性地提升解码速度（字节数减半，时间减半），以及为什么单用户体验到的速度提不上去时，加卡往往没用——你需要的是把更多请求塞进同一次权重读取里。
于是所有推理优化都指向同一个方向
- Continuous batching（连续批处理）：不等一整批做完，谁生成完了就立刻换新请求进来，让
B时刻保持高位。对比一下：老做法是「凑够 8 个人一起发车，全部到站才发下一趟」，中间有人早早下车了座位也空着；新做法是有人下车就马上让下一个人补上。这是吞吐提升最大的单项优化。 - PagedAttention（分页注意力）：KV cache 按小页分配，而不是一上来就按「这个对话最长可能到 8000 字」预留一大块。后者会造成大量碎片（占着没用满的空间），前者用多少给多少 → 同样显存能装更多并发 →
B更大。思路是从操作系统的虚拟内存分页借来的。 - GQA / MQA：见上一节，KV 小了并发就能大。（MQA 是 multi-query attention，GQA 的极端版——所有 Q 头只共享一组 K/V。）
- 量化：把权重从 bf16 的 2 字节压成 int8 的 1 字节、甚至 fp8/int4 更少（用更粗的精度存同一个数，略损效果换体积）。直接砍掉「模型字节数」这一项，而那正是 decode 耗时公式的分子。
- 投机解码（speculative decoding，也叫推测解码）：用一个小模型先草拟
k个 token，大模型一次性并行验证这几个。它不减少总 FLOPs（反而增加），但把 decode 的B从 1 变成k——在访存受限区里，这些额外的算力是免费的。这是「用算力换搬运」思路的又一次胜利，下面展开讲。 - Prefill / Decode 分离部署：两个阶段瓶颈相反，混在一起会互相拖累（有人发来一篇长文档做 prefill，会把同一台机器上所有人的吐字都卡住）。拆到不同机器上，各自按各自的瓶颈配硬件。
▸展开：投机解码为什么能「白嫖」？——它其实是在花你本来就浪费掉的算力
投机解码第一眼看着像作弊：多算了一堆东西，反而更快了？要理解它，得先接受前面那个结论——decode 时你的芯片算力利用率不到 1%，99% 的时间在等 HBM 把权重搬过来。
⚠️ 这里的「利用率」千万别看 nvidia-smi 那个百分比。那一栏（GPU-Util，等价于 DCGM 的 SM_ACTIVE）量的是「这段时间里有没有 kernel 在跑」，不是「算力用了几成」。decode 时 kernel 一直在跑——只不过它在原地等内存——所以这个数会显示 90~100%，而真实的算力利用率不到 1%。这是整个性能调优里最容易骗人的一个数字，看着满载，其实全在空转。
要判断「还有没有闲置算力」，得看这两个之一：MFU（实测 FLOP/s ÷ 峰值 FLOP/s），或者 profiler 里的 Tensor Core 活跃度（DCGM 的 PIPE_TENSOR_ACTIVE、Nsight 里的 tensor pipe active）。这两个才反映矩阵乘单元真正在干活的比例。
换句话说：你已经付钱买了那些算力，但一直没用上。投机解码就是去把这块闲置的产能利用起来。
它怎么运作：
- 找一个又小又快的「草稿模型」（比如同系列的 1B 小模型），让它连着猜 5 个 token。它便宜，猜 5 个几乎不花时间。
- 把这 5 个候选一次性塞给大模型。注意——大模型这一次前向要处理 5 个 token，而不是 1 个。但它读权重的次数还是一次。
- 大模型顺便就算出了「在每个位置上我自己会选什么」。逐个比对：从头开始，猜对的直接采纳，一直到第一个猜错的地方为止，错的那个用大模型自己的答案替换。
- 假设前 3 个猜对了：这一轮就一口气产出了 4 个 token，而代价只是一次权重读取。
关键的账在这里。回忆 decode 的耗时公式：每步耗时 ≈ 模型字节数 ÷ 带宽——它跟这一步处理几个 token 无关，因为瓶颈是搬权重，不是算数。所以：
- 老办法：一次权重读取 → 1 个 token。
- 投机解码：一次权重读取 → 平均 2～3 个 token。
那多出来的计算量呢？确实多了 5 倍，但它是从那 99% 的闲置算力里出的，不占用瓶颈资源，所以基本不体现在耗时上。这就是「免费」的含义——不是没有成本，而是成本花在了你本来就在浪费的那一项上。
还有一个常被误解的点：输出质量不会下降。因为最终采纳哪个 token 是大模型说了算的，草稿模型只负责提议。猜错了就丢掉重来，所以结果和不用投机解码时在统计上是一致的——它是纯粹的加速，不是精度换速度。
那是不是「只要还有闲置算力就一定赚」？不是。有闲置算力是必要条件，不是充分条件。还有两件事能把收益吃光：
- 接受率 α：草稿模型猜得准不准。设每个 token 被接受的概率是
α、一次草拟k个，那么一轮平均能产出(1−α^(k+1))/(1−α)个 token。α=0.7、k=4时约 2.8 个；α=0.3时只剩 1.4 个——猜不准的话，多算的全白算。 - 草稿模型自己的搬运成本：它也要读自己那份权重，而且是连着跑
k次（自回归地猜）。用 1B 草稿配 70B 大模型，这部分只占4/70 ≈ 6%，基本白送；但用 8B 草稿配 70B，就是4/8.75 ≈ 46%的额外开销——这已经不是零头了。
把两边乘起来：α=0.7 + 1B 草稿 ≈ 2.6× 提速；换成 8B 草稿 ≈ 1.9×；如果 α 掉到 0.3 又配了个 8B 的重草稿，1.43 ÷ 1.46 ≈ 0.98——白忙一场甚至倒亏。所以调它的两个旋钮就是「草稿够不够准」和「草稿够不够轻」。
另一头，什么时候它整体失效？当你的并发已经很高时。回到 I ≈ B：B 逼近脊点那两三百之后，芯片已经接近算力受限，闲置算力没了，投机解码就从「白嫖」变成「真花钱」。所以它主要用在低并发、追求单用户速度的场景——聊天产品有用，离线批量跑基本用不上。
最后是延迟与吞吐的取舍（延迟=单个用户等多久，吞吐=整台机器每秒能服务多少 token）：加大并发能提高吞吐，但排队和单步耗时会让单用户变慢。这两者构成一条帕累托前沿（指一条「已经没法白赚」的边界——落在上面时，想让任何一项变好就必然要牺牲另一项），没有普适最优点——聊天产品往左选低延迟，批量离线任务往右选高吞吐。先明确 SLA（service level agreement，你对用户承诺的服务指标，比如「95% 的请求首字在 1 秒内」），再选工作点。
⟶ 原书对应章节：Part 7 · All About Transformer Inference ｜ Part 8 · Serving LLaMA 3-70B on TPUs
7. 落到代码：profiling 与 JAX
纸面估算的价值，在于它给你一个可以被证伪的预期。真正的工作流是：先估算 → 再测量 → 差距超过 2 倍就说明有你没想到的东西。
profiling（性能剖析）就是用工具录一段程序的运行过程，看每一微秒花在哪个算子上——相当于给程序做一次心电图。看 profile 时我通常按这个顺序问：
- MFU 是多少？（model FLOPs utilization，实测 FLOP/s ÷ 芯片峰值，也就是「这些卡有几成力气真用上了」。）训练能到 40–50% 算健康，低于 25% 一定有事。
- 时间花在哪类算子上？如果大量时间在逐元素算子上，说明融合没做好。
- 通信和计算重叠了吗？看 collective 是否和计算 kernel 并行。理想情况下传数据的同时芯片还在算别的；没重叠说明卡在了某个必须等所有人到齐的同步点上。
- 有没有意外的排布转换？同一份数据在显存里怎么摆（哪一维在前、切在哪张卡上）叫它的排布。当上下游算子对排布的要求对不上时，框架会偷偷插入一次重排把数据搬来搬去——它不出现在你的代码里，却实实在在吃带宽。这类开销最阴险。
▸展开：不装 profiler，30 秒判断自己是不是访存受限
先重复一遍那个陷阱：nvidia-smi 显示 100% 完全不能说明算力用满了。它只表示「有 kernel 在跑」，而 decode 时 kernel 一直在跑、一直在等内存。看到 100% 是正常现象，不是没有余量的证据。
方法一：手算 MFU（不需要任何工具）。推理时每个参数每个 token 约 2 FLOPs（只有前向），所以：
实测 FLOP/s ≈ 2 × 参数量 × 每秒产出 token 数
举例：70B 模型跑在 8×H100 上，实测总吞吐 500 token/s。
实测算力 = 2 × 70e9 × 500 = 7e13 = 70 TFLOP/s。
峰值 = 8 × 990e12 ≈ 7.9e15。
MFU ≈ 0.9% —— 算力有 99% 在闲着，投机解码这类「花算力换搬运」的手段有巨大空间。
方法二：加并发看吞吐（最实用，也不用工具）。把并发翻一倍，看总 token/s：
- 接近翻倍 → 你在访存受限区，权重读取被更多请求摊薄了，还有余量。
- 基本不动 → 撞到算力天花板了，真的满了。
因为在访存受限区，每步耗时由「读一遍权重」决定，跟这步捎带处理几个 token 几乎无关——所以加并发是近乎免费的。这条一测就知道。
方法三：真要看指标，看这两个。DCGM 里同时盯 PIPE_TENSOR_ACTIVE（矩阵乘单元活跃度）和 DRAM_ACTIVE（显存活跃度）：
- 显存高、张量低 → 访存受限。这是 decode 的典型指纹。
- 张量高 → 算力受限，此时才是真的满了。
而 SM_ACTIVE（也就是 nvidia-smi 那个数）在两种情况下都会很高，所以它区分不出这两者，参考价值最低。
方法四：如果你在用 vLLM，它直接把答案报给你。OpenAI 兼容服务默认开着 Prometheus 端点，curl localhost:8000/metrics 就能看到（--disable-log-stats 会关掉）。最值钱的是这一个：
vllm:iteration_tokens_total —— 每个 engine step 实际处理了多少 token 的直方图。这就是本文开头那个 I ≈ B 里的 B，等于直接读出了你的算术强度。把它的均值和脊点那 150~300 一比，是不是访存受限当场就有答案，连算都不用算。
另外几个配着看：vllm:num_requests_running（当前并发）、vllm:generation_tokens（算吞吐用，配合上面手算 MFU）、vllm:kv_cache_usage_perc（KV 还剩多少空间，决定你还能不能往上加并发）。
做投机解码的话还有 vllm:spec_decode_* 一族，可以直接量出接受率而不用猜（不同版本字段名有出入，早期版本有现成的 draft_acceptance_rate，新版是按草稿位置分桶的计数，得自己除一下）。真正的验收还是 A/B 打开前后的 vllm:request_time_per_output_token_seconds——那才是用户体感的吐字速度。
JAX 这边值得学的是它的声明式分片：你先把所有卡组织成一个逻辑网格（mesh，比如把 32 张卡看成 4×8 的一张网，两个轴各起个名字），然后用 jax.sharding 描述「这个张量的这一维切在 mesh 的哪个轴上」，编译器自己推导需要插哪些集合通信。好处是并行策略变成了几行配置而不是一次重写，坏处是你必须能看懂它推导出来的结果——而看懂它，靠的正是前面几节那套代价模型。
⟶ 原书对应章节：Part 9 · How to Profile TPU Programs ｜ Part 10 · Programming TPUs in JAX
8. 怎么配合原书读
原书全部 13 章都在这里免费公开，附带大量图示和带答案的习题——习题是精华，这份笔记完全没有覆盖，强烈建议直接做。章节对照：
| 原书章节 | 本文对应 |
|---|---|
| Part 0 · 引言 | §0 |
| Part 1 · Rooflines | §0、§1 |
| Part 2 · TPUs | §2 |
| Part 3 · 分片矩阵 | §3 |
| Part 4 · Transformer 数学 | §4 |
| Part 5 · 训练并行 | §5 |
| Part 6 · 训练 LLaMA 3 | §4（6ND 实算） |
| Part 7 · 推理 | §6 |
| Part 8 · 服务 LLaMA 3-70B | §6 |
| Part 9 · Profiling | §7 |
| Part 10 · JAX 编程 | §7 |
| Part 11 · 结语与延伸 | — |
| Part 12 · GPUs | §2 |
如果只有一小时，我的建议是：读 Part 1，然后跳到 Part 7。roofline 给你工具，推理那章让你看到这个工具能拆开多少个平时想不明白的现象。中间几章更像参考手册，等真的要动手切模型时再回头查。
以上是我自己的理解与取舍，出错的部分算我的。原书写得非常好，也非常慷慨地免费公开了——真想学明白，请去读原书，顺手给作者一个 star。