性能建模与 Profiling：算力、带宽与显存账本

所属方向：推理加速 ｜ 难度：进阶 ｜ 前置知识：混合精度与数值稳定性、自注意力机制的计算与显存账本
关键词：性能建模、Profiling、Roofline、算术强度、memory bandwidth、FLOPs、MFU、PyTorch Profiler、Nsight

01. 为什么需要它

先看一组在同一台笔记本（Apple silicon，numpy fp32）上此前实测出来的数字（第 04 节另给本次复核快照），出自文末附录的 machine_probe.py，可以自己复现：

GEMM    : F = 0.998 GFLOP, D = 8.1 MB,   实测 0.498 ms（2003 GFLOP/s）
逐元素  : F = 1.000 GFLOP, D = 12000 MB, 实测 189.649 ms（5 GFLOP/s）
→ FLOPs 基本相同，耗时差 381 倍。

两段计算量的浮点运算次数几乎一样（都约 1 GFLOP），耗时差了 381 倍。如果拿「FLOPs 少的算子更快」这类直觉去做优化决策，在这个数字面前是反的：这里逐元素算子与 GEMM 的 FLOPs 相同，却因为每次都要把数据从内存搬进搬出，被带宽死死卡住。

再比如一个真实场景的优化评审：有人提议「把 LayerNorm 内部换成 fp8 计算单元重写，算力能翻几倍」。查一下账（第 03 节会算）：LayerNorm 的算术强度只有约 1.5 FLOP/Byte（下文不含 beta 偏置的简化 LayerNorm），远低于 A100 的 ridge point 200.6，是典型的带宽受限算子——给它换更快的算力，在只改变算力峰值、访存不变的模型中加速比是 1.00，一分钱收益都没有。反过来，同是 fp8，把它用在 decode 阶段的大权重 GEMM 上却有收益，但收益来自权重字节减半，不是算力翻倍。同一笔投资，用在哪类算子上，结论完全相反。

这篇的目的就是把这套「先算账再动手」的方法补齐：两个账本——速度账（算力 vs 带宽）和容量账（显存四项）——加上一套实测手段（profiler）。量化、算子融合、KV cache 管理、并行切分，所有推理优化节点的收益判断都站在这篇的地基上。

02. 最小可用理解

三句话讲完核心：

1. 任何算子的耗时有一个下界，由两种资源里更慢的那个决定：$T \ge \max(F/P_{\text{peak}},\ D/\beta)$。$F$ 是算子要做的浮点运算数，$P_{\text{peak}}$ 是硬件与该指令类型匹配的峰值算力；$D$ 是算子要搬运的字节数，$\beta$ 是带宽。算得再快也快不过「数据没到」。
2. 算术强度 $I = F/D$ 决定卡在哪：与 ridge point $I^{\ast} = P_{\text{peak}}/\beta$ 比较，$I < I^{\ast}$ 是带宽受限，优化方向是少搬字节（融合、量化、FlashAttention）；$I > I^{\ast}$ 是算力受限，优化方向才是少算（更好算法、更低精度计算单元）。
3. 显存容量是另一本独立的账：权重 + KV cache + 激活 + 额外开销四项加总，才决定一张卡能塞多少并发。「权重放得下就能跑」只覆盖了四项里的第一项。

03. 数学推导

3.1 时间下界为什么取 max

一个算子要做 $F$ 个浮点运算（FLOP，口径：一次乘加记 2 个 FLOP），硬件每秒最多做 $P_{\text{peak}}$ 个——就算计算单元一刻不停，也至少要 $F/P_{\text{peak}}$ 秒。同理，算子要把 $D$ 字节的数据在内存和计算单元之间搬个来回，总线每秒最多搬 $\beta$ 字节——至少要 $D/\beta$ 秒。这两件事用的是不同资源，理想情况下可以完全重叠（算上一批数据的同时搬下一批），所以总时间的下界是两者取 max：

$$T \ge \max\left(\frac{F}{P_{\text{peak}}},\ \frac{D}{\beta}\right)$$

注意这是下界：真实 kernel 还有 kernel launch、同步、缓存未命中、TLB miss 等额外开销，实测只会更慢。后面的实验用经验标定值和估算流量代入，此时算出的只是模型估计；计时波动或缓存会产生超过 100% 的达成率，并不违反物理下界。差距也不能全算成可消除的优化空间。

3.2 算术强度与 ridge point

定义算术强度：

$$I = \frac{F}{D}$$

物理含义：每从内存搬 1 字节数据，能换来多少次浮点运算。它取决于实现与所选存储层级；实际缓存命中和重复加载又与硬件有关，不能视为完全与机器无关。再定义机器的 ridge point：

$$I^{\ast} = \frac{P_{\text{peak}}}{\beta}$$

物理含义：这台机器「算」和「搬」一样快的分界点，单位都是 FLOP/Byte，所以可以比。把 $I$ 与 $I^{\ast}$ 代回 3.1 的下界：

• $I < I^{\ast}$（带宽受限）：$D/\beta$ 那一项更大，$T \approx D/\beta$，模型性能上界 $F/T \le I \cdot \beta$——性能与算力峰值无关，只跟 $I$ 成正比，这就是 roofline 图上那条斜线；
• $I \ge I^{\ast}$（算力受限）：$T \approx F/P_{\text{peak}}$，性能封顶在 $P_{\text{peak}}$，这就是平顶。

$P_{\text{peak}}$ 用哪个口径要非常小心：A100 SXM 的 BF16 dense 是 312 TFLOP/s，A100 不支持原生 FP8 Tensor Core；H100 SXM 的 BF16 dense 约 989、FP8 dense 约 1979 TFLOP/s，差 6.3 倍——口径选错，受限类型的判断直接反掉（第 06 节细说）。

3.3 给几类算子记账

GEMM：$A[M,K] \times B[K,N] \to C[M,N]$。每个输出元素要做 $K$ 次乘加，共 $M N K$ 次，乘加各记一次：

$$F_{\text{gemm}} = 2MNK,\qquad D_{\text{gemm}} = b\,(MK + KN + MN)$$

$b$ 是每元素字节数（bf16 取 2）。读 $A$、读 $B$、写 $C$ 各一遍，统计量这类小东西忽略。$M$ 越大，权重 $B[K,N]$ 被摊得越薄，$I$ 越高——这解释了为什么大 batch 的 GEMM 是算力受限、小 batch 的 GEMM 是带宽受限。

逐元素算子：$n$ 个元素各做 1 次运算，读 2 份写 1 份：

$$F = n,\qquad D = 3 b n,\qquad I = \frac{1}{3b}$$

$I$ 是常数（fp32 下约 0.083，bf16 下约 0.17），与 $n$ 无关。在固定 dtype 与访存模型下，增大 n 本身不会提高 I，因而不会像增加 GEMM 的复用维度那样跨越 roofline 分界。

softmax 的读写口径：普通三阶段实现约 3 读 2 写。若整行能驻留片上存储，融合实现可约 1 读 1 写，得到 2.5 倍的理想流量比；独立的 online normalizer 通常先流式求归一化量，再重读输入写出概率，为 2 读 1 写，不能混为一谈。本文用约 $5nd$ 作为普通 softmax 的操作计数；max、exp 与除法并非都能跑在 Tensor Core 上，online 递推还会增加标量运算。可对照 在线归一化论文与 Triton 整行融合示例的不同前提。

单头注意力（$N$ 是 token 数，$D_{h}$ 是每头维度）：两次 $N \times N \times D_{h}$ 的矩阵乘加 softmax：

$$F_{\text{attn}} = 4N^{2}D_{h} + 5N^{2}$$

朴素实现的分数矩阵 $S$ 和概率矩阵 $P$ 都要落显存，二次项访存约 $4N^{2}$；FlashAttention 让 $N^{2}$ 只留在片上 SRAM，若仅统计每份输入读取一次与输出写回一次，可得到不可避免的数据流量下界；真实分块内核会反复读入 K/V 或 Q 等数据：

$$D_{\text{朴素}} = b\,(4N^{2} + 4ND_{h}),\qquad D_{\text{flash,min}} = b \cdot 4ND_{h}$$

代入 $N=4096$、$D_{h}=128$、bf16（脚本 roofline_model.py 实算）：朴素 $D = 138.41\ \text{MB}$、$I = 62.7$，落在 A100（$I^{\ast} = 200.6$）的带宽受限区；理想最低流量 $D_{\mathrm{flash,min}} = 4.19\ \text{MB}$、$I = 2068$，跳进算力受限区。FLOPs 一动没动（×1.000），roofline 时间下界从 0.089 ms 降到 0.028 ms（×3.2），这是按最低流量计算的理想示例，33 倍不是实际 FlashAttention 访存或速度的测量；有限片上存储下还需分块 IO 模型。FlashAttention 那篇的完整递推在知识树的下一节点展开，这里先用 roofline 把它的收益定位清楚。

图 1：这张图要看三样——散点的横坐标是各算子的算术强度 $I$，点越靠右越「算得过来」；按 I 与 I 的相对位置给点作模型分类；点到上界的距离并不能单独证明真实瓶颈；FlashAttention 会改变访存量和横坐标，无法直接把这段垂直差距视为它的加速收益。*

3.4 优化收益的上限

把算力翻倍（$P_{\text{peak}} \to 2P_{\text{peak}}$），加速比是：

$$S_{\text{算力}} = \frac{\max(F/P_{\text{peak}},\ D/\beta)}{\max(F/(2P_{\text{peak}}),\ D/\beta)}$$

带宽受限时分子分母都是 $D/\beta$，$S_{\text{算力}} = 1$：白花钱。算力受限时加速至多为 2；当 $1<I/I^{\ast}<2$ 时，算力翻倍会遇到带宽上限，加速为 $I/I^{\ast}$。带宽翻倍对称地反一次。roofline_model.py 第四节把第 3.3 节的每个算子都代了一遍（A100 口径）：

算子                              I/I*      算力×2 的加速    带宽×2 的加速
逐元素 add [16M]                  0.001          1.00x          2.00x
LayerNorm [4096, 3072]            0.007          1.00x          2.00x
softmax 朴素三遍 [4096, 4096]     0.002          1.00x          2.00x
GEMM 512x4096x4096                2.041          2.00x          1.00x
attention 朴素 [N=4096,D=128]     0.312          1.00x          2.00x
attention ideal-min [N=4096,D=128]   10.307          2.00x          1.00x

这张表就是「投资之前先看图 2」的数字版：你的 kernel 在分界线哪一侧，决定哪类投资是零收益。

图 2：这张图要看什么——横轴 $I/I^{\ast}=1$ 那条竖线就是分界线：线左边算力翻倍的理想收益为 1，带宽翻倍收益在 1 到 2 之间；只有 I/I≤0.5 时完整得到 2 倍。线右边对称，I/I≥2 才完整得到算力翻倍的 2 倍。投入硬件或投入算子融合之前，先看自己在哪一侧。

3.5 显存的容量账

速度账之外另有一本容量账。下面以未分块 prefill、同时处理 BS 个 token 为例，显存分四项：

$$M_{\text{total}} = N_{\text{params}} \cdot b_{w} + 2BSLd_{\text{kv}}b + BS \cdot a + \rho \cdot M_{\text{sum}}$$

逐项说物理含义：$N_{\text{params}}$ 是参数量、$b_{w}$ 是每参数字节数（fp16 为 2）——这一项与并发无关，是常数；$B$ 是并发序列数、$S$ 是序列长度、$L$ 是层数、$d_{\text{kv}}$ 是每层 KV 总维度（GQA 模型用实际的 KV head 数乘头维度）——KV cache 每个 token 每层都要存一份 K 和一份 V，所以是 $2BSLd_{\text{kv}}b$ 字节，随并发线性增长；$a$ 是每 token 的激活峰值（推理不保留整层中间结果，但当前层十几份临时张量要同时活着，$d_{\text{model}}=4096$、fp16 时取约 128 KiB/token 是经验值，随实现差距很大）；$\rho$ 是额外开销率，$M_{\text{sum}}$ 是前三项之和。逐 token decode 时通常只有 B 个活跃 token，应将激活项 BS·a 改成 B·a；KV cache 仍随 BS 增长。分块 prefill 则按实际活跃 chunk 计。第 04 节把未分块 prefill 的教学账本代进一个 7B 模型。

04. 代码实现

三个脚本全部只用 numpy，因为 roofline 的方法论不依赖 GPU：同一台机器、同一套口径，把「峰值」和「落点」都实测出来，预测和实测的差距才看得见。本次复核实测峰值：$\beta = 71.8\ \text{GB/s}$、$P = 1635\ \text{GFLOP/s}$、$I^{\ast} = 22.8\ \text{FLOP/Byte}$（注意这是「numpy 能摸到的上限」，不是芯片标称值——方法论可比的前提是口径一致）。

4.1 标定两个峰值

带宽用两输入向量加法模式标定（不是含乘法的 STREAM triad）：用足够大的工作集降低缓存影响；128 MB 是否超过目标机器缓存仍需核对，测得的是该访问模式的有效带宽。

def measure_bandwidth(n: int = 32_000_000, repeat: int = 8) -> float:
    """z = x + y：读 2n、写 n，共 3n 个 fp32 元素。"""
    x = np.ones(n, dtype=np.float32)
    y = np.ones(n, dtype=np.float32)
    z = np.empty(n, dtype=np.float32)
    dt = _best(lambda: np.add(x, y, out=z), repeat)
    return 3 * n * FP32 / dt          # Byte/s，_best 取 repeat 次最快

算力用足够大的方阵乘标定（4096 的方阵乘访存被摊薄，$F = 2n^{3}$）。实跑输出：

本机标定（arm64 / numpy 2.1.3 / fp32，2026-10-02 02:21:20）
  实测可达带宽 beta  =    71.78 GB/s
  实测可达算力 P     =  1635.04 GFLOP/s
  ridge point  I*    =     22.8 FLOP/Byte

4.2 把真实算子打上 roofline

关键测试对象是 matmul-softmax-matmul 微基准；此处省略 1/√D 缩放，不是完整的生产 attention。下面这七行就是「账本」本身——每一行右边标了它读写了几个 $N^{2}$ 量级的遍数，$D$ 就是这么数出来的，不是拍脑袋：

def attention_naive():
    np.matmul(q, kt, out=S)                        # 写 S            1 遍
    np.max(S, axis=-1, keepdims=True, out=rowmax)  # 读 S            2 遍
    np.subtract(S, rowmax, out=S)                  # 读写 S          4 遍
    np.exp(S, out=S)                               # 读写 S          6 遍
    np.sum(S, axis=-1, keepdims=True, out=rowsum)  # 读 S            7 遍
    np.divide(S, rowsum, out=S)                    # 读写 S          9 遍
    np.matmul(S, v, out=O)                         # 读 S 写 O      10 遍

实跑结果（本机口径，$I^{\ast} = 22.8$）：

算子                              I   模型分类  T估计    T实测   达成率    实测算力
逐元素 add（1 遍）                   0.08  带宽   5.35ms   5.07ms  106%      6.3 GF/s
add+relu 两遍（预分配）               0.10  带宽   8.92ms  16.76ms   53%      3.8 GF/s
add+relu 写成一行（有中间数组）           0.10  带宽   8.92ms  26.80ms   33%      2.4 GF/s
LayerNorm [8192,4096] 预分配      0.15  带宽  18.70ms  32.56ms   57%      6.2 GF/s
attention 朴素 [N=2048,D=128]   12.61  带宽   2.40ms  11.41ms   21%    190.1 GF/s
GEMM 512x4096x4096           204.80  算力  10.51ms  16.18ms   65%   1061.5 GF/s
GEMM 4096x4096x4096          682.67  算力  84.06ms  84.82ms   99%   1620.3 GF/s

逐行解读，三种达成率各说明一件事：

• 逐元素 add 106%：与标定时相同访问模式，已接近本次有效带宽。超过 100% 来自经验标定与实际计时的差异，不能理解成超过硬件物理峰值。
• GEMM 大矩阵 99%：已接近本机同类 GEMM 标定值；这不能证明其他实现没有改进空间。
• 朴素 attention 21%：实测约为模型估计时间的 4.8 倍，可能涉及缓存、指令吞吐、线程和内核调度；仅凭总时间不能分离原因，也不能直接解释成多搬了几倍字节。

图 1 使用完整脚本中的 8 个算子（正文节选了 7 行）。点到 roofline 的距离表示相对模型上界的性能差距；它本身不能区分额外访存、指令开销或同步等原因。

4.3 显存账本与 batch 拐点

未分块 prefill 的容量教学账代一个 7B 模型（32 层、$d_{\text{kv}} = 4096$、fp16）：

模型：7B，fp16 权重 = 13.0 GiB，32 层，d_kv = 4096
KV cache 单价：512 KiB / token （一条 4096 长的序列 = 2.00 GiB）

 batch    权重   KV cache    激活     额外开销     合计    假设 78 GiB 可用预算
     1   13.0G     2.00G    0.50G    4.66G   20.20G     装得下
     8   13.0G    16.00G    4.00G    9.91G   42.95G     装得下
    16   13.0G    32.00G    8.00G   15.91G   68.95G     装得下
    24   13.0G    48.00G   12.00G   21.91G   94.95G     OOM
    32   13.0G    64.00G   16.00G   27.91G  120.95G     OOM

两个读数：权重是常数，batch 再大都是 13.0 GiB；KV cache 按 512 KiB/token 线性涨，batch 16 时（32 GiB）已经是权重的 2.5 倍。此处 30% 是人为设定的额外开销/前三项总量比率，不是实测额外开销率；它也不同于 PagedAttention 论文中 KV cache 已分配容量的浪费比例。同一张卡只把本教学模型的额外开销率改成 4%，batch 24 从 94.95 GiB 降到 75.96 GiB，从 OOM 变装得下：不换卡、不改模型，这个假设账本跨过一个 batch 档位；真实分页收益需只对可优化的 KV 分配项建模，并检查实际可用显存。

容量账决定了「能开多大」，速度账决定「该开多大」。固定一份 64 MiB 的权重矩阵，扫 batch（每个 batch 位置喂一个 token），本机实测：

   B    F (GFLOP)   D (MB)     I   模型分类    延迟 ms   吞吐 K/s
   1       0.034    67.14   0.50    带宽       1.718       0.6
   2       0.067    67.17   1.00    带宽       8.232       0.2
   4       0.134    67.24   2.00    带宽       8.161       0.5
   8       0.268    67.37   3.98    带宽       4.477       1.8
  16       0.537    67.63   7.94    带宽       4.436       3.6
  32       1.074    68.16  15.75    带宽       4.461       7.2
  64       2.147    69.21  31.03    算力       4.865      13.2
 128       4.295    71.30  60.24    算力       6.362      20.1
 256       8.590    75.50 113.78    算力       8.612      29.7

延迟：B=1 时 1.718 ms，B=256 时 8.612 ms，涨 5.0 倍；吞吐涨 51.1 倍。理想带宽模型中，加大 batch 可摊薄权重读取；本机表格并不呈严格平坦段。B=32 到 64 之间越过模型分界 $I^{\ast}=22.8$，这不保证实测曲线在同一处出现锐利拐点。B=256 时 64 MiB 的权重读取被摊到每样本 0.25 MiB。

B=2 的延迟是 B=8 的 1.84 倍，可能涉及 BLAS 内核选择、线程调度或计时波动；仅凭这张时间表不能确证原因。roofline 看不见这类事情，定 batch 必须结合实测。

图 3：这张图要看什么——左边是实测延迟，右边是实测吞吐，曲线不保证出现理想平坦段；竖虚线是本次标定值代入模型后的分类交界；B=2 处那个尖是 roofline 模型解释不了的库行为，标出来是为了让你对「模型给方向、实测给结论」有体感。

图 4：这张图要看什么——左图是显存四项的堆叠：蓝色权重是常数，红色 KV cache 随并发线性涨，越过黑色可用线就是 OOM；右图是只改一个参数（额外开销率 30%→4%）的效果，batch 24 从 OOM 变装得下。

05. 工业级实现对照

最小实现是方法论，生产里测量走的是另一套工具，但问的是同一组问题。

PyTorch Profiler（pytorch/pytorch · torch/profiler/profiler.py · profile，以 2026-09 的实现为准）是最常用的第一站：

from torch.profiler import profile, ProfilerActivity

with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
             record_shapes=True, profile_memory=True, with_flops=True) as prof:
    model(x)
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=15))
prof.export_chrome_trace("trace.json")

和本文方法的对应关系：key_averages().table 按算子事件键聚合 CPU/CUDA 耗时，而非直接按 GPU kernel 名称聚合，回答「时间花在哪」；with_flops=True 只为支持的算子（如矩阵乘与二维卷积）估算 FLOPs，不能视为全模型所有操作的完整计数，除以耗时和相应算力峰值得到该算子的利用率估计，并不是通常定义的模型级 MFU（Model FLOPs Utilization，模型理论 FLOPs ÷ 墙钟时间 ÷ 硬件峰值算力）；profile_memory=True 抓张量级显存分配，对应容量账的激活项。训练报告里常见的整体 MFU 是同一口径的粗化：拿模型理论 FLOPs 除以墙钟时间和卡数峰值，这个数在 decode 型负载里天然上不去，原因见第 08 节第 4 条。

Nsight 全家桶是更细的一层：Nsight Systems（nsys）看时间线——kernel 之间的空隙、同步等待、通信重叠，对应「下界假设完美重叠」不成立的部分；Nsight Compute（ncu）看单个 kernel，它的 Speed Of Light 面板直接给出 Compute Throughput 和 Memory Throughput 两个百分比——那就是 roofline 的粗版：应结合指令流水线、occupancy、缓存层级和 stall 原因判断，不能仅比较两个百分比便确定瓶颈。实操顺序通常是：nsys 找到热点和空隙，ncu 对热点 kernel 看 SOL 定受限类型，再决定投算力还是投带宽。

和框架选择的关系：SDPA 会根据输入与硬件分派后端，不能断言所有 diffusers/transformers 都默认使用 FlashAttention。本文 2068 FLOP/byte 来自每份 Q/K/V 只读一次的理想下界；真实 FlashAttention 还需考虑分块重读。小 batch GEMM 的权重读取成本通常很高，但具体量化收益与核实现、反量化和 batch 都有关。

06. 代价与边界

roofline 是模型，模型有假设。四条主要假设和不成立时的样子：

1. 假设计算与访存完全重叠。真实 kernel 在算和搬之间来回切换，还有 launch 和同步的开销。小 kernel（本机实测 B=1 约 1.7 ms 的场景）里这些固定开销占比不小，模型系统性偏乐观。
2. 假设峰值算力是一个数。实际有 fp16/bf16/fp8/稀疏好几档，差 6 倍以上；按 BF16 dense 口径，H100 SXM 相比 A100 40GB SXM 的 ridge point 从约 201 升到 295——换新卡后更多 kernel 会落进带宽受限区，拿旧卡的直觉做判断会错。
3. 假设 $D$ 与缓存无关。账本里的 $D$ 按落盘遍数数，但缓存命中会让有效带宽远大于 DRAM 标称值。锚点论文 Hierarchical Roofline（arXiv:2009.05257）就是把单一 roofline 扩展成每级缓存一条，用于定位数据移动发生在哪一层。
4. 假设算子孤立。decode 阶段 GEMM 的「权重」每层都要重新读一遍，全局账（整个模型、整个请求）和单算子账结论可能不同；通信算子（allreduce）的账本里延迟和消息数占大头，照搬本文公式会算错。

什么时候不用它：动态 shape、算子间强耦合（融合边界在变）、通信密集的分布式场景——这些先上 profiler 看时间线，roofline 只对「单 kernel、口径清晰」的问题给下界。坦诚标注：本文所有「实测」都来自一台笔记本的 numpy，数字本身不可迁移，理想分段趋势也可能被缓存、内核选择和调度打破。

07. 经典论文脉络

• Roofline: An Insightful Visual Performance Model（Williams et al., CACM 2009，未挂 arXiv）：提出算术强度与 ridge point，一根折线把「算力受限/带宽受限」变成可判定的题。一切性能建模的原点。
• Hierarchical Roofline Performance Analysis for Deep Learning Applications（Yang et al., 2020）：把 roofline 按缓存层级展开，回答「多搬的字节发生在哪一级存储」，是本文锚点论文，也补了单一 roofline 在深度学习负载上最大的盲区。
• FlashAttention（Dao et al., 2022）：IO-aware 的代表作——FLOPs 一动不动，靠 tiling + online softmax 减少 N×N 中间量落地及相应 IO，把注意力从带宽受限拉进算力受限。roofline 视角下「优化带宽」的教科书案例。
• Mixed Precision Training（Micikevicius et al., 2017）：换 dtype 同时改两本账——$b$ 变小省字节，算力单元换挡提峰值。哪半边有收益取决于受限类型（已发长文专门讲数值稳定那半）。
• Efficient Memory Management with PagedAttention（Kwon et al., 2023）：在论文比较的服务负载中显著减少 KV cache 分配浪费；这些比例不能直接乘到权重、激活和所有显存上，虚拟内存的分页思想搬进显存管理。知识树里 KV cache 一篇的主角。

五篇连起来是一条线：先有判定工具（roofline），再按受限类型各给一把钥匙——算力侧（混合精度）、带宽侧（FlashAttention）、容量侧（PagedAttention）。

08. 常见误解

1. 「FLOPs 少的算子更快」。本文开头的实测：同样约 1 GFLOP，GEMM 0.498 ms，逐元素 189.6 ms，差 381 倍。FLOPs 只在算力受限区才和耗时挂钩；带宽受限区里，决定耗时的是字节数。
2. 「显存够放权重就能跑」。7B fp16 权重只要 13 GiB，但 batch 24 时 KV cache 48 GiB + 额外开销 21.9 GiB，假设有 78 GiB 可用预算的设备照样 OOM。容量账要四项加总，KV cache 那一项随并发线性涨，batch 16 就反超权重了。
3. 「新卡算力翻倍，我的推理一定提速」。带宽受限的 kernel 加速比是 1.00（第 3.4 节表格里的 1.00x）。H100 SXM 的 FP8 dense 峰值约为 A100 BF16 的 6.3 倍，但实际换卡同时还改变带宽与内核；带宽受限时应估算字节数/实际带宽。权重量化会减少读取量，单独提高计算峰值则未必有用。
4. 「MFU 低就是实现烂」。decode 阶段每个 token 过一遍全部权重，$I$ 天然低于 ridge point，MFU 高不了——这是负载形状决定的，不是代码烂。看 MFU 前先分清 典型 prefill 与小 batch decode 的负载形状；足够大的 decode batch、长上下文注意力或通信可能改变瓶颈。
5. 「账本算出来的就是实际」。朴素 attention 模型估计 2.40 ms，实测 11.41 ms（差 4.8 倍）；batch 扫描里 B=2 延迟是 B=8 的 1.84 倍（原因需进一步 profile）。模型给方向和上限，实测给结论——两个都要，缺一个都会做出错误决策。

09. 动手验证

三个都能在笔记本上跑（附录有完整代码）：

1. python machine_probe.py——预期：逐元素 add 达成率 ≈100%，GEMM 大矩阵 ≈100%，朴素 attention 明显低于 50%。如果你的机器上 attention 达成率反而很高，多半是缓存把 $N^{2}$ 矩阵装下了，把 N 调大一倍再看。
2. python memory_ledger.py——观察实际延迟/吞吐，不预设理想三段式；按 I 与 I* 标记模型分类交界（本机在 32→64 之间）。改 batch_sweep 里的 B 列表，看吞吐什么时候不再涨。
3. 打开 roofline_model.py，仅把 A100 示例的 peak_flops 改成 1979e12、其余不变，重跑第五节（这是控制变量实验，不代表实际 H100，真实换卡还需更新带宽与 dtype）——预期：GEMM 拐点从 M∈(128, 256] 右移，更多算子被判为带宽受限。这一步会让你体感 ridge point 抬高对优化决策的影响。

10. 延伸阅读

按知识树的依赖关系，从这篇出发有三个方向：

• 往上游：自注意力机制的计算与显存账本——本文 3.3 节那笔注意力账的完整推导；混合精度与数值稳定性——换 dtype 的另一本账（数值范围与累加精度）。
• 往下游（受限类型各一把钥匙）：FlashAttention（带宽侧）、KV Cache 与自回归视频生成（容量侧）、算子融合与 CUDA Graph（launch 开销侧，已排期）、量化（字节侧，已排期）。
• 目录：完整知识树见博客目录页 AIGC 基本功知识树，按依赖顺序排好了先修课。

继续阅读 FlashAttention 为什么不需要存下注意力矩阵：本文 3.3 节的最低流量示例省略了真实分块重读，那篇把 online softmax 的递推式一步步推出来。

附录：完整代码

09 节用到的脚本全文如下（machine_probe.py、roofline_model.py、memory_ledger.py、make_figures.py）。复制到本地存成同名文件，按各脚本开头的依赖说明准备环境后即可运行。

machine_probe.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""在同一台机器上把 roofline 的「峰值」和「落点」都实测出来。

只依赖 numpy。直接 `python machine_probe.py` 即可运行，约 20 秒。

为什么要自己测一遍
------------------
roofline_model.py 里的 A100/H100 峰值是厂商标称值，读者手边不一定有 GPU。
这个脚本用同一套口径（F 记 2·MAC、D 记读写字节）在同一台机器上先标定
峰值算力与峰值带宽，再把真实算子打到这张 roofline 上——所以你能看到
「预测的下界」和「实测」差多少，而不只是相信规格表。

每个算子的 D 都按实现里实际的读写遍数来数：
    LayerNorm 预分配版    10 遍（mean / 减均值 / 平方 / 再 mean / 除 / 乘 gamma）
    attention 朴素预分配   10 遍（写 S / max / 减 / exp / sum / 除 / 读 S 写 O）
遍数不是拍脑袋来的，是照着下面每一行 numpy 数出来的。

输出的每个数字都来自 `time.perf_counter()` 的真实计时。
"""

import json
import platform
import time
from pathlib import Path

import numpy as np

FP32 = 4  # 本机跑 fp32，numpy 在 Apple 上走 Accelerate

SNAPSHOT = Path(__file__).resolve().parent / "probe_results.json"


def _best(fn, repeat: int):
    """跑 repeat 次取最快的一次：避开冷启动和调度抖动。"""
    fn()  # 预热：第一次要分配 / 触发缺页
    best = float("inf")
    for _ in range(repeat):
        t = time.perf_counter()
        fn()
        best = min(best, time.perf_counter() - t)
    return best


# ── 一、标定可达带宽：二元 add 模式（读两份写一份）───────────────────
def measure_bandwidth(n: int = 32_000_000, repeat: int = 8) -> float:
    """返回实测可达带宽（Byte/s）。

    每份 fp32 数组 128 MB；应核对目标机器缓存大小，这里测量此访问模式的有效带宽。
    z = x + y：读 2n、写 n，共 3n 个元素。
    """
    x = np.ones(n, dtype=np.float32)
    y = np.ones(n, dtype=np.float32)
    z = np.empty(n, dtype=np.float32)
    dt = _best(lambda: np.add(x, y, out=z), repeat)
    return 3 * n * FP32 / dt


# ── 二、标定峰值算力：大矩阵乘 ────────────────────────────────────
def measure_matmul_peak(n: int = 4096, repeat: int = 5) -> float:
    """返回实测可达算力（FLOP/s）。

    4096 的方阵乘足够大，访存被摊薄，测出来的是算力上限。
    FLOPs = 2·n³（口径同 roofline_model.py：一次乘加记 2 个浮点运算）。
    """
    rng = np.random.default_rng(0)
    a = rng.standard_normal((n, n), dtype=np.float32)
    b = rng.standard_normal((n, n), dtype=np.float32)
    c = np.empty((n, n), dtype=np.float32)
    dt = _best(lambda: np.matmul(a, b, out=c), repeat)
    return 2 * n ** 3 / dt


def calibrate(verbose: bool = True):
    bw = measure_bandwidth()
    fl = measure_matmul_peak()
    peak = dict(peak_flops=fl, peak_bw=bw)
    if verbose:
        print("=" * 80)
        print(f"本机标定（{platform.machine()}，numpy {np.__version__}，fp32）")
        print("=" * 80)
        print(f"  实测可达带宽 beta  = {bw / 1e9:8.2f} GB/s")
        print(f"  实测可达算力 P     = {fl / 1e9:8.2f} GFLOP/s")
        print(f"  ridge point  I*    = {fl / bw:8.1f} FLOP/Byte")
        print("  注：这是「用 numpy 能摸到」的上限，不是硬件标称值；")
        print("      换 BLAS 后端、换 dtype、换线程数都会变。口径一致才有可比性。")
    return peak


def build_ops():
    """构造待实测的算子。每个元素 = (名字, F, D, 函数, repeat)。"""
    rng = np.random.default_rng(0)
    ops = []

    # ── 逐元素：1 遍 vs 2 遍 vs「看起来融合了」 ──
    n = 32_000_000
    a = rng.standard_normal(n).astype(np.float32)
    b = rng.standard_normal(n).astype(np.float32)
    o1 = np.empty(n, dtype=np.float32)
    tmp = np.empty(n, dtype=np.float32)

    def add_relu_two_pass():          # 5n 字节：tmp 读写各一次
        np.add(a, b, out=tmp)
        np.maximum(tmp, 0, out=o1)

    def add_relu_one_line():          # 写法像融合了，numpy 照样分配中间数组
        np.maximum(np.add(a, b), 0, out=o1)

    ops.append(("逐元素 add（1 遍）", 1.0 * n, FP32 * 3 * n,
                lambda: np.add(a, b, out=o1), 8))
    ops.append(("add+relu 两遍（预分配）", 2.0 * n, FP32 * 5 * n,
                add_relu_two_pass, 8))
    ops.append(("add+relu 写成一行（有中间数组）", 2.0 * n, FP32 * 5 * n,
                add_relu_one_line, 8))

    # ── LayerNorm：预分配的 10 遍版 ──
    ln_n, ln_d = 8192, 4096
    x = rng.standard_normal((ln_n, ln_d)).astype(np.float32)
    g = rng.standard_normal(ln_d).astype(np.float32)
    mu = np.empty((ln_n, 1), dtype=np.float32)
    var = np.empty((ln_n, 1), dtype=np.float32)
    xc = np.empty_like(x)
    xc2 = np.empty_like(x)

    def layernorm_buffered():
        np.mean(x, axis=-1, keepdims=True, out=mu)      # 读 x            1 遍
        np.subtract(x, mu, out=xc)                      # 读 x 写 xc      3 遍
        np.multiply(xc, xc, out=xc2)                    # 读 xc 写 xc2    5 遍
        np.mean(xc2, axis=-1, keepdims=True, out=var)   # 读 xc2          6 遍
        np.sqrt(var + 1e-5, out=var)                    # 小量
        np.divide(xc, var, out=xc)                      # 读 xc 写 xc     8 遍
        np.multiply(xc, g, out=xc)                      # 读 xc 写 xc    10 遍

    ops.append((f"LayerNorm [{ln_n},{ln_d}] 预分配", 6.0 * ln_n * ln_d,
                FP32 * 10 * ln_n * ln_d, layernorm_buffered, 5))

    # ── 朴素注意力：S 落盘，10 遍 N² 访存 ──
    N, D = 2048, 128
    q = rng.standard_normal((N, D)).astype(np.float32)
    k = rng.standard_normal((N, D)).astype(np.float32)
    v = rng.standard_normal((N, D)).astype(np.float32)
    kt = np.ascontiguousarray(k.T)
    S = np.empty((N, N), dtype=np.float32)
    O = np.empty((N, D), dtype=np.float32)
    rowmax = np.empty((N, 1), dtype=np.float32)
    rowsum = np.empty((N, 1), dtype=np.float32)

    def attention_naive():
        # 本微基准省略 1/sqrt(D) 缩放，只测 matmul-softmax-matmul 的执行开销。
        np.matmul(q, kt, out=S)                              # 写 S          1 遍
        np.max(S, axis=-1, keepdims=True, out=rowmax)        # 读 S          2 遍
        np.subtract(S, rowmax, out=S)                        # 读写 S        4 遍
        np.exp(S, out=S)                                     # 读写 S        6 遍
        np.sum(S, axis=-1, keepdims=True, out=rowsum)        # 读 S          7 遍
        np.divide(S, rowsum, out=S)                          # 读写 S        9 遍
        np.matmul(S, v, out=O)                               # 读 S 写 O    10 遍

    ops.append((f"attention 朴素 [N={N},D={D}]", 4.0 * N * N * D + 5.0 * N * N,
                FP32 * (10 * N * N + 4 * N * D), attention_naive, 5))

    # ── GEMM 三个规模 ──
    for m in (512, 2048, 4096):
        ma = rng.standard_normal((m, 4096)).astype(np.float32)
        mb = rng.standard_normal((4096, 4096)).astype(np.float32)
        mc = np.empty((m, 4096), dtype=np.float32)
        ops.append((f"GEMM {m}x4096x4096", 2.0 * m * 4096 * 4096,
                    FP32 * (m * 4096 + 4096 * 4096 + m * 4096),
                    lambda ma=ma, mb=mb, mc=mc: np.matmul(ma, mb, out=mc), 5))
    return ops


def probe(peak: dict, verbose: bool = True):
    ops = build_ops()
    I_star = peak["peak_flops"] / peak["peak_bw"]
    if verbose:
        print("\n" + "=" * 80)
        print("三、真实算子打到这张 roofline 上：下界 vs 实测")
        print("=" * 80)
        hdr = (f"{'算子':<32s}{'I':>9s}{'受限':>8s}{'T下界':>10s}{'T实测':>10s}"
               f"{'达成率':>9s}{'实测算力':>14s}")
        print(hdr)
        print("-" * len(hdr))

    rows = []
    for name, F, D, fn, repeat in ops:
        dt = _best(fn, repeat)
        I = F / D
        bound = "算力" if I >= I_star else "带宽"
        t_pred = max(F / peak["peak_flops"], D / peak["peak_bw"])
        eff = t_pred / dt
        rows.append(dict(name=name, I=I, bound=bound, t_pred=t_pred, t_real=dt,
                         eff=eff, flops=F, bytes=D, gflops=F / dt / 1e9))
        if verbose:
            print(f"{name:<32s}{I:>9.2f}{bound:>8s}{t_pred * 1e3:>8.2f}ms"
                  f"{dt * 1e3:>8.2f}ms{eff * 100:>8.0f}%"
                  f"{F / dt / 1e9:>11.1f} GF/s")
    if verbose:
        print("\n  达成率 = 下界 / 实测，回答的是「这个实现离 roofline 还有多远」。")
        print("  低达成率可能来自指令、缓存、调度或同步；单凭此比值不能诊断额外访存。")
    return rows


def demo_same_flops():
    """同样 1 GFLOP 的算力，GEMM 和逐元素算子差多少时间。"""
    print("\n" + "=" * 80)
    print("四、同样的 FLOPs，两种算子差多少时间")
    print("=" * 80)
    rng = np.random.default_rng(3)
    target = 1.0e9

    # GEMM：2·M·K·N = target，取 K=N=1024
    M = int(target / (2 * 1024 * 1024))
    ma = rng.standard_normal((M, 1024)).astype(np.float32)
    mb = rng.standard_normal((1024, 1024)).astype(np.float32)
    mc = np.empty((M, 1024), dtype=np.float32)
    t_gemm = _best(lambda: np.matmul(ma, mb, out=mc), 10)
    f_gemm = 2.0 * M * 1024 * 1024

    # 逐元素：n 个元素各算 1 次，F = n
    n_elem = int(target)
    a = rng.standard_normal(n_elem).astype(np.float32)
    b = rng.standard_normal(n_elem).astype(np.float32)
    o = np.empty(n_elem, dtype=np.float32)
    t_elem = _best(lambda: np.add(a, b, out=o), 5)

    print(f"  GEMM    : F = {f_gemm / 1e9:.3f} GFLOP, "
          f"D = {FP32 * (M * 1024 + 1024 * 1024 + M * 1024) / 1e6:.1f} MB, "
          f"实测 {t_gemm * 1e3:.3f} ms（{f_gemm / t_gemm / 1e9:.0f} GFLOP/s）")
    print(f"  逐元素  : F = {n_elem / 1e9:.3f} GFLOP, "
          f"D = {FP32 * 3 * n_elem / 1e6:.1f} MB, "
          f"实测 {t_elem * 1e3:.3f} ms（{n_elem / t_elem / 1e9:.0f} GFLOP/s）")
    print(f"  → FLOPs 基本相同，耗时差 {t_elem / t_gemm:.0f} 倍。")
    print("    同样运算量不代表同样时间：这里逐元素算子的运算数并未")
    print(f"    减少，却因为 I 只有 {1 / (FP32 * 3):.2f} FLOP/Byte 而完全被带宽卡住。")


def main():
    peak = calibrate()
    rows = probe(peak)
    demo_same_flops()

    # 把这次标定存成快照：make_figures / memory_ledger 直接读它，
    # 保证「图上的数字」和「正文引用的数字」来自同一次运行。
    SNAPSHOT.write_text(json.dumps(
        dict(peak=peak, rows=rows,
             note=f"{platform.machine()} / numpy {np.__version__} / fp32",
             timestamp=time.strftime("%Y-%m-%d %H:%M:%S")),
        ensure_ascii=False, indent=1))
    print(f"\n标定快照已存到 {SNAPSHOT.name}，make_figures / memory_ledger 会复用它。")
    print(f"这台机器的 roofline：P = {peak['peak_flops'] / 1e9:.0f} GFLOP/s, "
          f"beta = {peak['peak_bw'] / 1e9:.1f} GB/s, "
          f"I* = {peak['peak_flops'] / peak['peak_bw']:.1f} FLOP/Byte")


if __name__ == "__main__":
    main()

roofline_model.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Roofline 模型的最小实现：把「算力账」和「带宽账」合成一条上界曲线。

只依赖 numpy（其实只是用来排版，算法本身是纯标量运算）。
直接 `python roofline_model.py` 即可运行，约 1 秒。

符号与文章第 03 节一一对应
--------------------------
    F        浮点运算数（FLOP）。口径：一次乘加（MAC）记 2 FLOP
    D        访存字节数（Byte）。口径：读 + 写，缓存命中不重复计
    I        算术强度 I = F / D，单位 FLOP/Byte
    P_peak   峰值算力（FLOP/s）
    beta     峰值带宽（Byte/s）
    I_star   ridge point = P_peak / beta，I 低于它就是带宽受限

关键结论（代码算完会打印）：
    T >= max(F / P_peak, D / beta)
    把 F 降为一半或把 D 降为一半时的模型比值，见 optimize_headroom()
"""

import json
from pathlib import Path


def _local_spec():
    """本机实测峰值：优先读 machine_probe.py 存下的标定快照。

    快照不存在时退回兜底常数（先跑一次 machine_probe.py 更准）。
    这样 SPECS 里的「本机实测」永远是同一次标定，不会和正文引用的数字分叉。
    """
    snap = Path(__file__).resolve().parent / "probe_results.json"
    if snap.exists():
        p = json.loads(snap.read_text())["peak"]
        return dict(peak_flops=p["peak_flops"], peak_bw=p["peak_bw"])
    return dict(peak_flops=1680e9, peak_bw=77e9)


# ── 厂商标称峰值。来源：NVIDIA A100 / H100 SXM 白皮书，bf16 dense（不含稀疏）
#    注意 H100 BF16 dense 算力增幅大于带宽增幅，ridge point 比 A100 高：
#    带宽涨 2.15 倍、bf16 算力涨约 3.17 倍 → 更多 kernel 落在带宽受限区。
SPECS = {
    "A100-40GB SXM (bf16)": dict(peak_flops=312e12, peak_bw=1555e9),
    "H100-80GB SXM (bf16)": dict(peak_flops=989e12, peak_bw=3350e9),
    "H100-80GB SXM (fp8)": dict(peak_flops=1979e12, peak_bw=3350e9),
    # 下面这一条不是标称值，是 machine_probe.py 在同一台机器上实测出来的
    "本机实测（见 machine_probe.py）": _local_spec(),
}

BF16 = 2  # bytes per element


def roofline_time(F: float, D: float, peak_flops: float, peak_bw: float):
    """roofline 时间下界：算力时间与带宽时间取 max（假设两者完美重叠）。"""
    t_compute = F / peak_flops
    t_memory = D / peak_bw
    return max(t_compute, t_memory), t_compute, t_memory


def bound_of(F: float, D: float, peak_flops: float, peak_bw: float) -> str:
    I = F / D
    I_star = peak_flops / peak_bw
    return "算力受限" if I >= I_star else "带宽受限"


def optimize_headroom(F: float, D: float, peak_flops: float, peak_bw: float):
    """把运算量 F 或搬运量 D 减半，理想模型可快多少？

    答案与算术强度有关：带宽受限时砍算力收益为 0，算力受限时最多 2 倍。
    返回 (F 减半的加速比, D 减半的加速比)，分别等价于只将峰值算力或带宽翻倍。
    """
    def speedup(F2, D2):
        t0, _, _ = roofline_time(F, D, peak_flops, peak_bw)
        t1, _, _ = roofline_time(F2, D2, peak_flops, peak_bw)
        return t0 / t1

    return speedup(F / 2, D), speedup(F, D / 2)


# ── 几类算子的 F / D 账本 ──────────────────────────────────────────
def gemm(M: int, K: int, N: int, b: int = BF16):
    """矩阵乘 [M,K] @ [K,N] -> [M,N]。读 A、读 B、写 C。"""
    F = 2.0 * M * K * N
    D = b * (M * K + K * N + M * N)
    return F, D


def elementwise(n: int, b: int = BF16, n_pass: int = 1, n_read: int = 1):
    """逐元素算子：读 n_read 份、写 1 份；n_pass 表示这样读写几轮。

    二元运算（如 add）是 n_read=2：读两份输入、写一份输出。
    """
    F = 1.0 * n * n_pass          # 每个元素算 1 次
    D = b * (n_read + 1) * n * n_pass
    return F, D


def layernorm(n: int, d: int, b: int = BF16):
    """简化 LayerNorm（无 beta）：两次归约、减均值、平方、除标准差、乘 gamma，约 6 次/元素。
    访存 = 读 x 一遍 + 写 y 一遍（统计量是 O(n) 的小量，忽略）。"""
    F = 6.0 * n * d
    D = b * 2 * n * d
    return F, D


def softmax(n: int, d: int, b: int = BF16, fused: bool = True):
    """softmax：exp / 减最大值 / 归一化，约 5 次运算/元素。

    fused=True  假设整行可驻留片上存储，读 1 遍写 1 遍；不是独立 online normalizer 的通用 IO
    fused=False 朴素三遍（求 max、求 exp 和、归一化），读 3 遍写 2 遍
    """
    F = 6.0 * n * d
    D = b * 2 * n * d if fused else b * (3 * n * d + 2 * n * d)
    return F, D


def attention(N: int, D: int, b: int = BF16, flash: bool = False):
    """单头注意力。N 是 token 数，D 是每头维度。

    flash=False（朴素）：分数矩阵 S 和权重 P 都要落 HBM
        S = Q K^T    写 N^2
        P = softmax(S) 读 N^2 写 N^2
        O = P V      读 N^2
        → 二次项访存 ≈ 4 N^2 b
    flash=True（强缓存假设下的最低流量，不是真实 FlashAttention IO）：Q/K/V 读进来、O 写出去，N^2 项只留在 SRAM
        → 访存 ≈ 4 N D b
    """
    F = 4.0 * N * N * D + 5.0 * N * N          # 两个 N×N×D 的矩阵乘 + softmax
    if flash:
        D_bytes = b * (3 * N * D + N * D)      # 读 Q K V，写 O
    else:
        D_bytes = b * (4 * N * N + 4 * N * D)  # 上面再加 S/P 的读写
    return F, D_bytes


def fmt(x: float, unit: str = "") -> str:
    for u, s in (("T", 1e12), ("G", 1e9), ("M", 1e6), ("K", 1e3)):
        if x >= s:
            return f"{x / s:.2f} {u}{unit}"
    return f"{x:.2f} {unit}"


def main():
    print("=" * 78)
    print("一、四台「机器」的峰值与 ridge point")
    print("=" * 78)
    print(f"{'设备':<28s}{'P_peak':>14s}{'beta':>14s}{'I* = P/beta':>16s}")
    ridges = {}
    for name, s in SPECS.items():
        I_star = s["peak_flops"] / s["peak_bw"]
        ridges[name] = (s, I_star)
        print(f"{name:<28s}{fmt(s['peak_flops'], 'FLOP/s'):>16s}"
              f"{fmt(s['peak_bw'], 'B/s'):>14s}{I_star:>12.1f} FLOP/B")

    a100 = SPECS["A100-40GB SXM (bf16)"]
    h100 = SPECS["H100-80GB SXM (bf16)"]
    print(f"\n  A100 -> H100：带宽 ×{h100['peak_bw'] / a100['peak_bw']:.2f}，"
          f"算力 ×{h100['peak_flops'] / a100['peak_flops']:.2f}，"
          f"ridge point {a100['peak_flops'] / a100['peak_bw']:.0f} -> "
          f"{h100['peak_flops'] / h100['peak_bw']:.0f}（更高 = 更多算子落入带宽受限区）")

    # ── 二、算子账本 ────────────────────────────────────────────────
    N, D = 4096, 128          # 一个典型的 DiT / LLM 单头规模
    cases = [
        ("逐元素 add [16M]", elementwise(16_000_000, n_read=2)),
        ("LayerNorm [4096, 3072]", layernorm(4096, 3072)),
        ("softmax fused [4096, 4096]", softmax(4096, 4096, fused=True)),
        ("softmax 朴素三遍 [4096, 4096]", softmax(4096, 4096, fused=False)),
        ("GEMM 512x4096x4096", gemm(512, 4096, 4096)),
        ("GEMM 4096x4096x4096", gemm(4096, 4096, 4096)),
        ("attention 朴素 [N=4096,D=128]", attention(N, D, flash=False)),
        ("attention ideal-min [N=4096,D=128]", attention(N, D, flash=True)),
    ]

    print("\n" + "=" * 78)
    print("二、算子账本（bf16，A100 口径）")
    print("=" * 78)
    hdr = (f"{'算子':<32s}{'F (FLOP)':>14s}{'D (Byte)':>14s}"
           f"{'I':>10s}{'受限':>10s}{'T_roof':>12s}")
    print(hdr)
    print("-" * len(hdr))
    for name, (F, Db) in cases:
        I = F / Db
        T, tc, tm = roofline_time(F, Db, a100["peak_flops"], a100["peak_bw"])
        print(f"{name:<32s}{fmt(F):>14s}{fmt(Db, 'B'):>14s}"
              f"{I:>10.2f}{bound_of(F, Db, a100['peak_flops'], a100['peak_bw']):>10s}"
              f"{T * 1e3:>10.3f} ms")

    # ── 三、FlashAttention 把算术强度抬了多少 ────────────────────────
    print("\n" + "=" * 78)
    print("三、物化 S/P 与理想最低 IO 对比（单头 N=4096, D=128, bf16）")
    print("=" * 78)
    F_naive, D_naive = attention(N, D, flash=False)
    F_flash, D_flash = attention(N, D, flash=True)
    T_naive, _, _ = roofline_time(F_naive, D_naive, a100["peak_flops"], a100["peak_bw"])
    T_flash, _, _ = roofline_time(F_flash, D_flash, a100["peak_flops"], a100["peak_bw"])
    print(f"  FLOPs    : {fmt(F_naive)} -> {fmt(F_flash)}  "
          f"（×{F_flash / F_naive:.3f}，几乎没变）")
    print(f"  访存     : {fmt(D_naive, 'B')} -> {fmt(D_flash, 'B')}  "
          f"（×{D_flash / D_naive:.4f}，省 {D_naive / D_flash:.1f} 倍）")
    print(f"  算术强度 : {F_naive / D_naive:.1f} -> {F_flash / D_flash:.1f} FLOP/B "
          f"（×{(F_flash / D_flash) / (F_naive / D_naive):.1f}）")
    print(f"  受限类型 : {bound_of(F_naive, D_naive, **a100)} -> "
          f"{bound_of(F_flash, D_flash, **a100)}")
    print(f"  roofline : {T_naive * 1e3:.3f} ms -> {T_flash * 1e3:.3f} ms "
          f"（×{T_naive / T_flash:.2f}）")
    print(f"  注意：加速来自「少搬 {D_naive / D_flash:.0f} 倍字节」，"
          "不是「少算」——此处仅比较相同主导 FLOPs 与理想最低访存，不是实测内核。")

    # ── 四、砍算力 vs 砍带宽，谁的收益大 ────────────────────────────
    print("\n" + "=" * 78)
    print("四、同一个 kernel，只翻倍峰值算力或带宽的理想加速")
    print("=" * 78)
    print(f"{'算子':<32s}{'I/I*':>10s}{'算力×2 的加速':>16s}{'带宽×2 的加速':>16s}")
    I_star = a100["peak_flops"] / a100["peak_bw"]
    for name, (F, Db) in cases:
        s_flops, s_bw = optimize_headroom(F, Db, a100["peak_flops"], a100["peak_bw"])
        print(f"{name:<32s}{(F / Db) / I_star:>10.3f}"
              f"{s_flops:>14.2f}x{s_bw:>14.2f}x")
    print("\n  读法：带宽受限的行（I/I* < 1）里「算力×2」那一列几乎都是 1.00——"
          "\n  在只改变计算峰值且访存不变的理想模型里，该列为 1。")

    # ── 五、GEMM 什么时候从带宽受限翻到算力受限 ─────────────────────
    print("\n" + "=" * 78)
    print("五、GEMM [M,4096] x [4096,4096]：M 多大才开始算力受限")
    print("=" * 78)
    K = Nout = 4096
    print(f"{'M':>8s}{'I (FLOP/B)':>14s}{'受限':>12s}")
    prev = None
    for M in (1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 4096):
        F, Db = gemm(M, K, Nout)
        b = bound_of(F, Db, a100["peak_flops"], a100["peak_bw"])
        print(f"{M:>8d}{F / Db:>14.2f}{b:>12s}")
        if prev == "带宽受限" and b == "算力受限":
            print(f"         ↑ 拐点：M 在 ({M // 2}, {M}] 之间，"
                  f"I 越过 ridge point {I_star:.0f}")
        prev = b
    print("\n  推论：小 batch 推理（M=1~8）里的 GEMM 是带宽受限的，"
          "这时候做 fp8 量化省的是字节、不是算力。")


if __name__ == "__main__":
    main()

memory_ledger.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""显存账本 + batch 扫描：容量这本账怎么算，以及它怎么反过来决定速度。

只依赖 numpy（第二部分要真跑计时）。直接 `python memory_ledger.py` 即可运行，约 10 秒。

第一部分是「容量账」：权重 / KV cache / 激活 / 额外开销四项加总，看 假设 78 GiB 可用预算
能塞多少条并发请求。第二部分是「速度账」：batch 变大之后，同一份权重被更多
token 摊薄，算术强度抬上去，kernel 从带宽受限翻到算力受限——这个拐点用真实
计时测出来，顺便告诉你 batch 该开多大。
"""

import json
import time
from pathlib import Path

import numpy as np

GIB = 1024 ** 3


def _load_peak():
    """优先读 machine_probe.py 存下的标定快照（保证口径一致）；
    没跑过 machine_probe.py 时退回兜底常数，并照常工作。"""
    snap = Path(__file__).resolve().parent / "probe_results.json"
    if snap.exists():
        peak = json.loads(snap.read_text())["peak"]
        return peak["peak_flops"], peak["peak_bw"]
    return 1680.64e9, 77.21e9      # 兜底：先跑一次 machine_probe.py 更准


P_PEAK, BETA = _load_peak()
I_STAR = P_PEAK / BETA


def kv_cache_bytes(n_tokens: int, n_layers: int, d_kv: int, b: int = 2) -> int:
    """KV cache 字节数：每个 token 每层都要存一份 K 和一份 V。

    n_tokens  序列长度（或并发请求的总 token 数）
    n_layers  层数 L
    d_kv      每层的 KV 总维度 = n_kv_heads × d_head（GQA 时用实际的 KV head 数）
    b         每个元素的字节数（fp16/bf16 = 2，fp8 = 1）
    """
    return 2 * n_tokens * n_layers * d_kv * b


def ledger(n_params: float, b_w: int, n_layers: int, d_kv: int,
           batch: int, seq: int, act_per_token: int, frag_rate: float):
    """推理显存账本的四项，单位 GiB。

    act_per_token  每个 token 的激活峰值（字节）。推理不保留整层的中间结果，
                   但当前层的十几份临时张量要同时活着：d_model=4096、fp16 时
                   一份 [d] 张量 8 KiB，取 16 份 ≈ 128 KiB / token。
                   这是经验值，随实现（融合程度、是否分块）差距很大。
    frag_rate      教学额外开销 / (权重 + KV + 激活)，不是实测碎片率。
                   本函数按未分块 prefill 计 batch*seq 个活跃 token；
                   逐 token decode 的激活应改为 batch*act_per_token，KV 仍取 batch*seq。
    """
    w = n_params * b_w
    kv = kv_cache_bytes(batch * seq, n_layers, d_kv)
    act = batch * seq * act_per_token
    frag = (w + kv + act) * frag_rate
    return dict(weights=w / GIB, kv=kv / GIB, act=act / GIB, frag=frag / GIB,
                total=(w + kv + act + frag) / GIB)


def part1():
    print("=" * 80)
    print("一、显存账本：一张 假设 78 GiB 可用预算，能塞多少条并发请求")
    print("=" * 80)
    # Llama-2-7B 的配置：32 层、32 个 KV head、d_head=128 → d_kv = 4096
    N_PARAMS, N_LAYERS, D_KV = 7e9, 32, 4096
    SEQ = 4096
    per_token = kv_cache_bytes(1, N_LAYERS, D_KV)
    print(f"  模型：7B，fp16 权重 = {N_PARAMS * 2 / GIB:.1f} GiB，"
          f"{N_LAYERS} 层，d_kv = {D_KV}")
    print(f"  KV cache 单价：{per_token / 1024:.0f} KiB / token "
          f"（一条 {SEQ} 长的序列 = {per_token * SEQ / GIB:.2f} GiB）")
    print()
    hdr = (f"{'batch':>6s}{'权重':>9s}{'KV cache':>10s}{'激活':>9s}"
           f"{'额外开销':>9s}{'合计':>9s}{'78GiB预算':>10s}")
    print(hdr)
    print("-" * len(hdr))
    for batch in (1, 4, 8, 16, 24, 32):
        # 朴素预分配：按最大长度全预留，教学假设：以权重+KV+激活总和的 30% 估计额外开销，不是实测碎片比例
        naive = ledger(N_PARAMS, 2, N_LAYERS, D_KV, batch, SEQ,
                       act_per_token=131072, frag_rate=0.30)
        ok = "装得下" if naive["total"] < 78 else "OOM"
        print(f"{batch:>6d}{naive['weights']:>8.1f}G{naive['kv']:>9.2f}G"
              f"{naive['act']:>8.2f}G{naive['frag']:>8.2f}G"
              f"{naive['total']:>8.2f}G{ok:>10s}")

    # 换成低开销假设（教学假设，非 PagedAttention 实测）：额外开销率降到 4%
    print("\n  同一张卡，把额外开销率从 30% 降到 4%（低开销假设）：")
    print(f"{'batch':>6s}{'朴素合计':>11s}{'低开销假设':>11s}{'多出来的并发':>16s}")
    for batch in (8, 16, 24, 32):
        naive = ledger(N_PARAMS, 2, N_LAYERS, D_KV, batch, SEQ, 131072, 0.30)
        paged = ledger(N_PARAMS, 2, N_LAYERS, D_KV, batch, SEQ, 131072, 0.04)
        extra = ""
        if naive["total"] > 78 >= paged["total"]:
            extra = "从 OOM 变装得下"
        print(f"{batch:>6d}{naive['total']:>10.2f}G{paged['total']:>10.2f}G"
              f"{extra:>16s}")
    print("\n  读法：权重那一项是常数，batch 再大也不变；KV cache 才是随并发")
    print("        线性增长的那一项，batch 16 时它已经比权重还大。")
    print("        「显存够放权重就能跑」这句话漏掉了后面三项。")


# 用 roofline_model 里同款口径：下界 = max(F/P, D/beta)。
# P_PEAK / BETA / I_STAR 已在文件顶部从标定快照读出，这里不再覆盖，
# 否则快照刷新后正文引用的数字会和实际跑出来的对不上。


def batch_sweep(repeats: int = 15):
    """固定一份「权重」，扫 batch，返回每一档的实测延迟/吞吐/算术强度。"""
    rng = np.random.default_rng(0)
    d = 4096
    W = rng.standard_normal((d, d)).astype(np.float32)   # 64 MiB 的「权重」
    W_bytes = W.nbytes
    FP32 = 4

    def run(B):
        x = rng.standard_normal((B, d)).astype(np.float32)
        y = np.empty((B, d), dtype=np.float32)
        np.matmul(x, W, out=y)                            # 预热
        best = float("inf")
        for _ in range(repeats):
            t = time.perf_counter()
            np.matmul(x, W, out=y)
            best = min(best, time.perf_counter() - t)
        return best

    rows = []
    for B in (1, 2, 4, 8, 16, 32, 64, 128, 256):
        F = 2.0 * B * d * d
        D = W_bytes + 2 * B * d * FP32
        I = F / D
        dt = run(B)
        rows.append(dict(B=B, F=F, D=D, I=I, lat=dt, thr=B / dt,
                         bound="算力" if I >= I_STAR else "带宽",
                         t_pred=max(F / P_PEAK, D / BETA),
                         eff=max(F / P_PEAK, D / BETA) / dt))
    return rows, W_bytes


def part2():
    """batch 扫描实测：延迟、吞吐、算术强度，以及拐点在哪。"""
    print("\n" + "=" * 80)
    print("二、batch 扫描实测：吞吐什么时候不再涨")
    print("=" * 80)
    rows, W_bytes = batch_sweep()
    (Path(__file__).resolve().parent / "batch_results.json").write_text(
        json.dumps(dict(rows=rows, weight_bytes=W_bytes,
                        peak=dict(peak_flops=P_PEAK, peak_bw=BETA)), indent=2))
    print(f"  固定「权重」W 形状 [4096, 4096]，fp32 = {W_bytes / 2 ** 20:.0f} MiB；"
          f"batch B 就是一次喂进去的 token 数")
    print(f"  峰值口径（来自 machine_probe 的标定快照）："
          f"P = {P_PEAK / 1e9:.0f} GFLOP/s，beta = {BETA / 1e9:.1f} GB/s，"
          f"I* = {I_STAR:.1f}")
    print()
    hdr = (f"{'B':>6s}{'F (GFLOP)':>12s}{'D (MB)':>10s}{'I':>9s}{'受限':>8s}"
           f"{'延迟 ms':>10s}{'吞吐 K/s':>11s}{'下界 ms':>10s}{'达成率':>9s}")
    print(hdr)
    print("-" * len(hdr))

    prev_bound = None
    for r in rows:
        print(f"{r['B']:>6d}{r['F'] / 1e9:>12.3f}{r['D'] / 1e6:>10.2f}"
              f"{r['I']:>9.2f}{r['bound']:>8s}{r['lat'] * 1e3:>10.3f}"
              f"{r['thr'] / 1e3:>9.1f}K{r['t_pred'] * 1e3:>10.3f}"
              f"{r['eff'] * 100:>8.0f}%")
        if prev_bound == "带宽" and r["bound"] == "算力":
            print(f"         ↑ 拐点：B 从 {r['B'] // 2} 到 {r['B']} 之间"
                  f"越过 ridge point {I_STAR:.1f}")
        prev_bound = r["bound"]

    # 反常检测：batch 变大反而变慢，说明库换了代码路径（roofline 看不见这件事）
    for r0, r1 in zip(rows, rows[1:]):
        if r1["B"] == 2 * r0["B"] and r1["lat"] > r0["lat"] * 1.5:
            print(f"         ! 反常：B={r0['B']} 只要 {r0['lat'] * 1e3:.2f} ms，"
                  f"B={r1['B']} 却要 {r1['lat'] * 1e3:.2f} ms")

    base = rows[0]
    top = rows[-1]
    print(f"\n  延迟：B=1 时 {base['lat'] * 1e3:.3f} ms，B={top['B']} 时 "
          f"{top['lat'] * 1e3:.3f} ms（涨 {top['lat'] / base['lat']:.1f} 倍）")
    print(f"  吞吐：B=1 时 {base['thr'] / 1e3:.1f} K/s，B={top['B']} 时 "
          f"{top['thr'] / 1e3:.1f} K/s（涨 {top['thr'] / base['thr']:.1f} 倍）")
    print(f"  单样本成本：B={top['B']} 时把 64 MiB 权重的读取摊薄到 {top['B']} 个样本，"
          f"每个样本只摊 {W_bytes / top['B'] / 2 ** 20:.2f} MiB")
    print("  → 增大 batch 可摊薄权重访问；延迟是否平坦、吞吐能增加多少仍需实测。")
    print("  → 上面标 ! 的行是 roofline 看不见的东西：同样的公式、更大的 batch，"
          "可能涉及内核选择、调度或计时波动，单凭本表不能确诊。")


if __name__ == "__main__":
    part1()
    part2()

make_figures.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""生成「性能建模与 Profiling」一文的四张解释图。

图里的计时优先读取已保存的实测快照，避免重绘时图文使用不同测量：
  * roofline.png     峰值与落点来自 machine_probe.calibrate() / probe()
  * batch_scaling.png 来自 memory_ledger.batch_sweep()
  * memory_ledger.png 来自 memory_ledger.ledger()
  * opt_gain.png      来自 roofline_model.optimize_headroom()

只依赖 numpy + matplotlib，直接 `python make_figures.py` 即可运行，约 40 秒。
改了另外三个脚本，这里要重跑，避免图与正文数字对不上。
"""

import json
import sys
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))

import machine_probe as MP          # noqa: E402
import memory_ledger as ML          # noqa: E402
import roofline_model as RM         # noqa: E402

OUT = HERE.parent / "figures"
try:
    # exist_ok=True 是必须的：目录已存在时 pathlib 会抛 FileExistsError，
    # 少数沙箱环境连 exist_ok=True 的 mkdir 也一并拦，这里再兜一层。
    OUT.mkdir(exist_ok=True)
except PermissionError:
    if not OUT.is_dir():
        raise

plt.rcParams.update({
    "font.sans-serif": ["PingFang SC", "Arial Unicode MS", "DejaVu Sans"],
    "axes.unicode_minus": False,
    "figure.dpi": 160,
    "savefig.dpi": 160,
})

C_MEM = "#e05263"      # 带宽受限：红
C_CMP = "#3f7fbf"      # 算力受限：蓝
C_ACC = "#2f9e6f"      # 好结果：绿
C_GREY = "#94a3b8"


# ── 图 1：roofline 曲线 + 实测落点 ────────────────────────────────
def fig_roofline(peak, rows):
    P = peak["peak_flops"] / 1e9     # GFLOP/s
    beta = peak["peak_bw"] / 1e9     # GB/s
    I_star = P / beta                # 因为两个都除过 1e9，比值不变

    I = np.logspace(-2, 4, 400)
    perf = np.minimum(P, beta * I)

    fig, ax = plt.subplots(figsize=(8.4, 5.4))
    ax.loglog(I, perf, color="k", lw=2.2)
    ax.axvline(I_star, color=C_GREY, ls="--", lw=1.4)
    ax.text(I_star * 1.15, 3, f"ridge point I* = {I_star:.0f}\nFLOP/Byte",
            color="#475569", fontsize=9, va="bottom")
    ax.text(0.012, 25, f"峰值算力 P = {P:.0f} GFLOP/s", fontsize=9,
            color="#475569")
    ax.text(0.012, 8.5, "斜线 = 带宽天花板，beta = "
            f"{beta:.0f} GB/s", fontsize=9, color="#475569")

    # 落点：用实测时间反推实际性能
    short = {
        "逐元素 add（1 遍）": "逐元素 add",
        "add+relu 两遍（预分配）": "add+relu 两遍",
        "add+relu 写成一行（有中间数组）": "add+relu 一行写法",
        "LayerNorm [8192,4096] 预分配": "LayerNorm",
        "attention 朴素 [N=2048,D=128]": "朴素 attention",
        "GEMM 512x4096x4096": "GEMM M=512",
        "GEMM 2048x4096x4096": "GEMM M=2048",
        "GEMM 4096x4096x4096": "GEMM M=4096",
    }
    offsets = {                      # 手工调过：右边缘和重叠的标签让位
        "GEMM M=4096": (-10, 10),
        "GEMM M=2048": (10, -4),
        "GEMM M=512": (12, -14),
        "朴素 attention": (7, -3),
    }
    # 左下角四个点挤在一起，用箭头把标签拉到右边的空地上
    arrows = {
        "逐元素 add": (0.42, 9.5),
        "LayerNorm": (0.42, 5.6),
        "add+relu 两遍": (0.42, 3.1),
        "add+relu 一行写法": (0.42, 1.75),
    }
    for r in rows:
        name = short.get(r["name"], r["name"])
        y = r["gflops"]
        ax.scatter([r["I"]], [y], s=52,
                   color=C_MEM if r["bound"] == "带宽" else C_CMP,
                   zorder=5, edgecolor="white", linewidth=0.8)
        if name in arrows:
            tx, ty = arrows[name]
            ax.annotate(name, (r["I"], y), xytext=(tx, ty),
                        fontsize=8.5, color="#334155",
                        arrowprops=dict(arrowstyle="-", color=C_GREY, lw=0.9),
                        va="center")
            continue
        dx, dy = offsets.get(name, (7, -3))
        ha = "right" if dx < 0 else "left"
        ax.annotate(name, (r["I"], y), textcoords="offset points",
                    xytext=(dx, dy), fontsize=8.5, color="#334155", ha=ha)

    ax.set_xlabel("算术强度 I = F / D  (FLOP/Byte)")
    ax.set_ylabel("实测性能  (GFLOP/s)")
    ax.set_title("图 1：本机的 roofline —— 斜线是带宽，平顶是算力，点是实测落点",
                 fontsize=11)
    ax.set_xlim(0.01, 5000)
    ax.set_ylim(0.3, 5000)
    ax.grid(alpha=0.25, which="both", ls=":")
    handles = [plt.Line2D([], [], marker="o", ls="", color=C_MEM,
                          label="模型分类：I < I*"),
               plt.Line2D([], [], marker="o", ls="", color=C_CMP,
                          label="模型分类：I >= I*")]
    ax.legend(handles=handles, loc="lower right", fontsize=8.5, framealpha=0.9)
    fig.tight_layout()
    fig.savefig(OUT / "roofline.png")
    plt.close(fig)


# ── 图 3：batch 扫描的延迟与吞吐 ──────────────────────────────────
def fig_batch(rows):
    B = [r["B"] for r in rows]
    lat = [r["lat"] * 1e3 for r in rows]
    thr = [r["thr"] / 1e3 for r in rows]

    fig, ax1 = plt.subplots(figsize=(8.4, 5.0))
    ax1.plot(B, lat, "o-", color=C_CMP, lw=2, label="延迟（左轴）")
    ax1.set_xscale("log", base=2)
    ax1.set_xlabel("batch B（一次喂进去的 token 数）")
    ax1.set_ylabel("延迟 (ms)", color=C_CMP)
    ax1.tick_params(axis="y", labelcolor=C_CMP)

    ax2 = ax1.twinx()
    ax2.plot(B, thr, "s--", color=C_ACC, lw=2, label="吞吐（右轴）")
    ax2.set_ylabel("吞吐 (K token/s)", color=C_ACC)
    ax2.tick_params(axis="y", labelcolor=C_ACC)

    # 拐点：I 越过 ridge point 的地方
    cross = next((r for r in rows if r["bound"] == "算力"), None)
    if cross:
        ax1.axvline(cross["B"], color=C_GREY, ls=":", lw=1.5)
        ax1.text(cross["B"] * 1.05, max(lat) * 0.92,
                 f"越过 ridge point\nB ≈ {cross['B']}", fontsize=9,
                 color="#475569")

    # 库的反常：延迟不随 batch 单调
    b8 = next((r["lat"] for r in rows if r["B"] == 8), None)
    spike = next((r for r in rows if r["B"] in (2, 4) and
                  b8 is not None and r["lat"] > 1.5 * b8), None)
    if spike:
        ax1.annotate(f"B={spike['B']} 延迟为 B=8 的 {spike['lat']/b8:.1f} 倍\n（原因需 profile）",
                     (spike["B"], spike["lat"] * 1e3),
                     textcoords="offset points", xytext=(14, -6),
                     fontsize=8.5, color=C_MEM)

    ax1.set_title("图 3：batch 扫描 —— 实测延迟、吞吐与 roofline 模型交界",
                  fontsize=11)
    ax1.grid(alpha=0.25, ls=":")
    lines1, lab1 = ax1.get_legend_handles_labels()
    lines2, lab2 = ax2.get_legend_handles_labels()
    ax1.legend(lines1 + lines2, lab1 + lab2, loc="upper left", fontsize=9)
    fig.tight_layout()
    fig.savefig(OUT / "batch_scaling.png")
    plt.close(fig)


# ── 图 4：显存账本 ────────────────────────────────────────────────
def fig_ledger():
    batches = [1, 4, 8, 16, 24, 32]
    naive = [ML.ledger(7e9, 2, 32, 4096, b, 4096, 131072, 0.30) for b in batches]
    paged = [ML.ledger(7e9, 2, 32, 4096, b, 4096, 131072, 0.04) for b in batches]

    fig, (ax, ax2) = plt.subplots(1, 2, figsize=(11.2, 4.8),
                                  gridspec_kw={"width_ratios": [1.35, 1]})
    x = np.arange(len(batches))
    keys = [("weights", "权重", "#3f7fbf"), ("kv", "KV cache", "#e05263"),
            ("act", "激活", "#e8a33d"), ("frag", "额外开销", "#94a3b8")]
    bottom = np.zeros(len(batches))
    for k, label, color in keys:
        vals = np.array([d[k] for d in naive])
        ax.bar(x, vals, bottom=bottom, label=label, color=color, width=0.62,
               edgecolor="white", linewidth=0.6)
        bottom += vals
    ax.axhline(78, color="k", ls="--", lw=1.3)
    ax.text(len(batches) - 0.4, 79, "教学假设：78 GiB 可用预算", fontsize=8.5,
            ha="right", color="#334155")
    for i, d in enumerate(naive):
        if d["total"] > 78:
            ax.text(i, d["total"] + 2, "OOM", ha="center", fontsize=9,
                    color=C_MEM, fontweight="bold")
    ax.set_xticks(x)
    ax.set_xticklabels(batches)
    ax.set_xlabel("并发序列数 batch")
    ax.set_ylabel("显存 (GiB)")
    ax.set_title("图 4a：未分块 prefill 教学账本（额外开销率 30%）", fontsize=10.5)
    ax.legend(fontsize=8.5, ncol=2)
    ax.grid(axis="y", alpha=0.25, ls=":")

    w_naive = [d["total"] for d in naive]
    w_paged = [d["total"] for d in paged]
    ax2.bar(x - 0.2, w_naive, width=0.4, label="朴素", color="#e05263")
    ax2.bar(x + 0.2, w_paged, width=0.4, label="低开销假设", color="#2f9e6f")
    ax2.axhline(78, color="k", ls="--", lw=1.3)
    ax2.set_xticks(x)
    ax2.set_xticklabels(batches)
    ax2.set_xlabel("并发序列数 batch")
    ax2.set_ylabel("合计显存 (GiB)")
    ax2.set_title("图 4b：只把额外开销率从 30% 降到 4%", fontsize=10.5)
    ax2.legend(fontsize=8.5)
    ax2.grid(axis="y", alpha=0.25, ls=":")

    fig.suptitle("图 4：显存账本的四项 —— 权重是常数，KV cache 随并发线性增长",
                 fontsize=11)
    fig.tight_layout()
    fig.savefig(OUT / "memory_ledger.png")
    plt.close(fig)


# ── 图 2：砍算力 vs 砍带宽，各有多少钱 ────────────────────────────
def fig_opt_gain():
    a100 = RM.SPECS["A100-40GB SXM (bf16)"]
    ratio = np.logspace(-2, 1.2, 300)
    sp_flops, sp_bw = [], []
    for r in ratio:
        # 造一个算术强度恰好是 r × I* 的算子：固定 D，F 由 r 决定
        D = 1e8
        F = r * (a100["peak_flops"] / a100["peak_bw"]) * D
        s_f, s_b = RM.optimize_headroom(F, D, a100["peak_flops"], a100["peak_bw"])
        sp_flops.append(s_f)
        sp_bw.append(s_b)

    fig, ax = plt.subplots(figsize=(8.4, 4.8))
    ax.semilogx(ratio, sp_flops, color=C_CMP, lw=2.2,
                label="算力翻 2 倍能拿到的加速")
    ax.semilogx(ratio, sp_bw, color=C_MEM, lw=2.2,
                label="带宽翻 2 倍能拿到的加速")
    ax.axvline(1.0, color=C_GREY, ls="--", lw=1.4)
    ax.text(1.05, 1.02, "I = I*：分界线", fontsize=9, color="#475569")
    ax.fill_between(ratio, 0.98, 2.02, where=np.array(ratio) < 1,
                    color=C_MEM, alpha=0.07)
    ax.fill_between(ratio, 0.98, 2.02, where=np.array(ratio) >= 1,
                    color=C_CMP, alpha=0.07)
    ax.text(0.05, 1.9, "带宽受限区：\n换更快的算力单元 = 0 收益",
            fontsize=9, color=C_MEM)
    ax.text(6, 1.9, "算力受限区：\n加带宽 = 0 收益", fontsize=9, color=C_CMP)
    ax.set_xlabel("算术强度 / ridge point   (I / I*)")
    ax.set_ylabel("能拿到的加速比")
    ax.set_ylim(0.95, 2.1)
    ax.set_title("图 2：投资之前先看这张图 —— 你的 kernel 在分界线哪一侧",
                 fontsize=11)
    ax.grid(alpha=0.25, ls=":")
    ax.legend(fontsize=9, loc="center right")
    fig.tight_layout()
    fig.savefig(OUT / "opt_gain.png")
    plt.close(fig)


def main():
    # --only <图名>：只重画指定一张（roofline / batch / ledger / gain），
    # 避免为了改一张图的标签把所有实测算子重跑一遍、数字跟正文引用分叉。
    only = None
    if "--only" in sys.argv:
        only = sys.argv[sys.argv.index("--only") + 1]

    # 优先读 machine_probe.py 存的标定快照，保证图上数字与正文引用同源
    if MP.SNAPSHOT.exists():
        snap = json.loads(MP.SNAPSHOT.read_text())
        peak = snap["peak"]
        print(f"读取标定快照（{snap['timestamp']}，{snap['note']}）：")
    else:
        print("没找到标定快照，现场标定一遍…")
        peak = MP.calibrate(verbose=False)
        MP.SNAPSHOT.write_text(json.dumps(
            dict(peak=peak, rows=MP.probe(peak, verbose=False),
                 note="fresh", timestamp="")))
    print(f"  P = {peak['peak_flops'] / 1e9:.0f} GFLOP/s, "
          f"beta = {peak['peak_bw'] / 1e9:.1f} GB/s, "
          f"I* = {peak['peak_flops'] / peak['peak_bw']:.1f}")
    if only in (None, "roofline"):
        rows = snap.get("rows") if MP.SNAPSHOT.exists() and "snap" in locals() else None
        if not rows:
            rows = MP.probe(peak, verbose=False)
        fig_roofline(peak, rows)
    if only in (None, "batch"):
        batch_snap = HERE / "batch_results.json"
        if batch_snap.exists():
            sweeps = json.loads(batch_snap.read_text())["rows"]
        else:
            sweeps, weight_bytes = ML.batch_sweep()
            batch_snap.write_text(json.dumps(dict(rows=sweeps, weight_bytes=weight_bytes, peak=peak), indent=2))
        fig_batch(sweeps)
    if only in (None, "ledger"):
        fig_ledger()
    if only in (None, "gain"):
        fig_opt_gain()
    if only and only not in ("roofline", "batch", "ledger", "gain"):
        raise ValueError("unknown figure: " + only)
    print(f"\n图已写入 {OUT}")


if __name__ == "__main__":
    main()