第六章 从模型计算图到系统成本

前五章回答了模型为什么这样设计、如何学习和生成。本章把同一张计算图映射到参数、运行时状态、运算和数据传输。后续章节的公式由这张图展开，而不是另起一套需要背诵的知识。

1 三类对象有不同生命周期

权重是训练学到、推理时复用的矩阵和向量。一次请求结束后，服务中的权重通常继续存在，供后续请求使用。

临时激活是输入经过计算后产生的中间结果，例如 Q、FFN 中间向量和 logits。在不需要反向传播的推理中，某个结果被最后一次使用后，其缓冲区就可能复用。

KV Cache也由前向计算产生，但因为后续生成会读取历史 K/V，需要跨 Decode step 保留。它属于运行时状态，而不是训练参数。

因此“新增 token”产生的是一份新的计算与状态，并不会复制整个模型。对固定模型的普通推理，参数个数保持不变；请求相关缓存随有效上下文增长。

2 一个线性层揭示三种不同的计数

设输入 $X\in\mathbb{R}^{n\times a}$，权重 $W\in\mathbb{R}^{a\times b}$，输出 $Y=XW$。

权重有 $ab$ 个元素。每个输出元素由 a 项乘积累加得到，共 nb 个输出，按一次乘加 2 FLOPs 的常见口径，主项计算量是 $2nab$。若数据各用 s bytes 存储，理想单次读取输入、权重并写出结果的字节量为：

$$ D_{\rm ideal}=s(na+ab+nb) $$

参数量 $ab$ 与 token 数 n 无关；FLOPs 随 n 增长；权重读取能否被 n 个 token 有效复用，则决定数据移动效率。真实 kernel 受片上容量和分块影响，可能多次读取数据，因此上式只是理想 IO 模型。

忽略输入输出 IO，算术强度约为：

$$ I=\frac{F}{D}\approx\frac{2nab}{sab}=\frac{2n}{s} $$

当 n 较小，花很大成本读入权重，却只执行少量复用；n 增大后，同一份权重服务更多输入行。这里已出现 Prefill 与 Decode 性能差异的基本原因。

3 将一层拆成两类工作

Q/K/V 投影、输出投影和 FFN 都属于对 token 行进行矩阵变换，计算主项随新处理 token 数增长。

QK 和 AV 则把 query 位置与 key 位置关联起来。完整 Prefill 中，有 S 个新 query，每个 query 读取合法的前缀；总配对数为 $S(S+1)/2$。普通 Decode 中，每条序列只有一个新 query，读取长度约为 T 的历史，配对数约为 T。

因此线性层与注意力需要分别计数：前者主要跟“本次新增多少输入行”有关，后者还跟“每一行能看到多少历史位置”有关。这解释了为什么不能永远用模型参数量乘一个常数代替全部 FLOPs。

4 为什么 FLOPs 小也可能慢

假设某个算子几乎不做算术，却必须搬运大量张量，它可能被内存带宽限制。反之，一个数据复用充分的大矩阵乘法，可能主要受计算速率限制。

设工作量为 F FLOPs、HBM 传输为 D bytes，设备峰值算力为 C FLOP/s、带宽为 $\beta$ bytes/s，至少需要满足：

$$ t\ge F/C,\qquad t\ge D/\beta $$

合起来得到理想资源下界：

$$ t\ge\max(F/C,D/\beta) $$

这个模型没有包括不同算子串行、低利用率、通信和调度成本，却已经能解释：两个参数量相近的模型，或同一个模型的两个执行阶段，为什么可能有完全不同的吞吐。

它也说明优化应先对应瓶颈。减少不必要的 HBM 中间结果，可能比减少少量算术更有效；增加 batch 可以提升权重复用，但同时增加 KV 容量与读取压力。

5 将数学注意力映射到存储层次

直接实现注意力，会先形成分数矩阵，再做 Softmax，再乘 Value。对长度 S 的序列，中间分数或概率矩阵包含 $S^2$ 个位置对，可能造成很大 HBM IO。

FlashAttention 的关键思路是分块计算，并维护在线 Softmax 所需的归一化信息，避免把完整中间注意力矩阵物化到 HBM。这改变执行计划和数据移动，而不是把完整注意力换成另一种近似关系。FlashAttention

在线归一化可以用下面的合并关系理解。对某个 query，已处理一组分数，其最大值为 m，归一化分母为 $l=\sum_j e^{z_j-m}$，未归一化输出为 $u=\sum_je^{z_j-m}v_j$。新块的对应量为 $m_b,l_b,u_b$。令 $m'=\max(m,m_b)$，则：

$$ l'=e^{m-m'}l+e^{m_b-m'}l_b $$

$$ u'=e^{m-m'}u+e^{m_b-m'}u_b $$

全部块处理结束后，输出为 $u'/l'$。同一份 Softmax 加权和可以分块得到，因此不必保留全部分数。浮点运算顺序变化可能造成数值舍入差异，不要求逐 bit 相同。

6 容量优化与带宽优化有不同作用

KV 分页管理减少连续预分配浪费，并支持按块共享，主要改善物理内存管理；GQA 减少每个 token 的 KV 宽度；KV 量化减少每个元素的字节数；前缀复用避免部分历史重算。这些机制影响不同的乘数或执行路径。

例如，两条请求共享同一个前缀时，可以减少重复存储的 KV 块，但不同 Query 仍需计算各自对前缀的注意力。节省容量不等价于按相同比例节省计算时间。

同样，MoE 减少单 token 的专家执行数量，却保留完整专家池，并可能引入路由通信。总参数描述容量，激活路径描述部分计算，真实传输还取决于整个 batch 触达的专家并集。

7 接下来的资源推导采用什么结构

后续统一采用一个明确的现代模型作为分析对象：decoder-only、Pre-Norm、每层两个 RMSNorm、final RMSNorm、普通 RoPE、无 bias 的线性层、SwiGLU，以及标准 MHA/GQA/MQA。MoE 只在 FFN 位置替换模块；不同层不共享参数。

这是为了使每个公式有确定含义，不是在规定所有 Transformer 都必须这样实现。原始 Encoder-decoder、多种混合结构、MLA 和滑动窗口，需要按其真实计算图调整。

资源推导分成四步：先统计可学习矩阵，再追踪运行时状态，然后统计各算子的运算量，最后用硬件与工作负载将工作量转换为时间。参数量、显存、FLOPs 和吞吐之间有联系，但没有可以跳过所有中间条件的一步换算。

下一章：模型容量与参数量推导。需要集中核对 shape 时，可查看张量形状参考。