第九章 Prefill 与 Decode 计算量

本章计算矩阵乘法的主项，统一按一次乘加 2 FLOPs。Norm、Softmax、RoPE、激活函数、路由选择、采样和数据搬运并未被这些矩阵公式完整覆盖；它们参数少，却可能影响实际耗时。

1 线性层为什么近似是 2P

对一个 token，$[1,a]\times[a,b]$ 的每个输出元素需要约 $a$ 次乘加；一共 $b$ 个输出，因此是 $2ab$ FLOPs。矩阵有 $ab$ 个参数，所以一次使用它约为 2 倍参数量。

将各 block 实际执行的线性层参数加起来，记为 $P_{\rm lin}$。对 Dense SwiGLU、$d_q=d$：

$$ P_{\rm lin}=L(2d^2+2dd_{kv}+3dm) $$

对同构 MoE，无共享专家：

$$ P_{\rm lin}=L(2d^2+2dd_{kv}+3kdm_e+dE) $$

$n$ 个 token 的 block 线性层计算约为 $2nP_{\rm lin}$。这排除了输入 Embedding 查表，并把 LM Head 单独处理；也还没有算 Attention 的 QK 和 AV。

“每 token 约 $2P$”只适合线性权重主导的粗略估计。长上下文、超大词表、MoE、权重共享或反复复用模块时，都要回到实际算子。

2 完整 Prefill 的注意力

对每个 Q head，$QK^\top$ 为 $2S^2d_h$ FLOPs，$AV$ 同样如此；乘上 batch、head 数和层数：

$$ F_{\rm attn,prefill,full}=4LBh_qd_hS^2=4LBd_qS^2 $$

这是按完整方阵计算的口径。因果注意力真正合法的 query-key 配对只有：

$$ 1+2+\cdots+S=\frac{S(S+1)}{2} $$

每对在 QK 与 AV 中合计约 $4d_h$ FLOPs，因此理想跳过无效三角区域时：

$$ F_{\rm attn,prefill,causal}=2LBd_qS(S+1) $$

实际 kernel 按 tile 工作，还可能计算部分被 mask 的位置，因此不能只看到 causal mask 就宣称硬件实际工作量严格减半。回答时可以同时给出完整方阵口径和因果有效工作量口径，并说明使用哪个。

3 Prefill 总 FLOPs 与输出层

若只对每条 prompt 最后一个位置计算 logits：

$$ \boxed{F_{\rm pre}\approx2BSP_{\rm lin}+2LBd_qS(S+1)+2BdV} $$

如果对所有 prompt 位置计算 logits，则末项变成 $2BSdV$。这可能出现在训练、需要所有位置 logprobs 或某些实现中。LM Head 的权重是否绑定不影响此处 FLOPs。

变长 batch 不应盲目用最大 $S$ 乘 $B$。在能去掉 padding 计算的实现中，线性项按 $\sum_iS_i$，注意力项按 $\sum_iS_i(S_i+1)$；如果 kernel 真的计算 padding，则要按实际执行形状计费。

4 带缓存前缀的 chunked Prefill

已有 $T$ 个前缀 token 的 KV，本次追加 $C_t$ 个已知输入 token。第 $j$ 个新 query 可以读取 $T+j$ 个 key，因此合法配对数为：

$$ C_tT+\frac{C_t(C_t+1)}{2} $$

相应的理想有效注意力 FLOPs 为：

$$ 4LBd_q\left(C_tT+\frac{C_t(C_t+1)}{2}\right) $$

线性层只处理 $BC_t$ 个新 token。前缀缓存复用免去了旧 token 的重新投影和旧位置的 Attention，但新 query 对旧前缀的读取与注意力仍然存在。

把无缓存 prompt 拆成多个 chunk，并不会自动改变所有合法配对的总数；收益主要来自调度、峰值显存和延迟安排，代价可能包括重复读取与更多调度开销。

5 一个 Decode step

step 开始前有 $T$ 个旧 token，本次处理一个新输入 token。新 Query 读取 $T+1$ 个 K/V，包括当前位置：

$$ F_{\rm attn,dec}=4LBd_q(T+1) $$

加上线性层与每序列一次 LM Head：

$$ \boxed{F_{\rm dec,step}\approx2BP_{\rm lin}+4LBd_q(T+1)+2BdV} $$

一个 step 总计产出 $B$ 个预测 token，因此平均每输出 token 的主项是上述数值除以 $B$。这里不包含推测解码一次验证多个候选 token 的情况。

固定模型下，线性项与 $T$ 无关，Attention 项随 $T$ 线性增长。所以“Decode 每 token 固定 2P”会低估长上下文成本。

6 一整段生成

若 Prefill 输入长度为 $S$，要生成 $N$ 个输出 token，后续有 $N-1$ 次 Decode。第 $j$ 次 Decode 看到 $S+j$ 个 key，$j=1,\ldots,N-1$。

线性层和 LM Head 的主项是：

$$ 2B(N-1)(P_{\rm lin}+dV) $$

Decode 注意力总量是：

$$ 4LBd_q\left((N-1)S+\frac{N(N-1)}{2}\right) $$

再加 Prefill 即得请求计算主项。输出长度增大时，不仅步数增多，每步看到的上下文也逐渐变长。

7 什么时候 Attention 计算显著

在单序列 Decode 中，忽略输出层，用 $2P_{\rm lin}=4LdT$ 求两项相等的上下文长度：

$$ T_{\rm equal}\approx\frac{P_{\rm lin}}{2Ld} $$

这是 FLOPs 的交点，不是耗时的交点。Attention 可能在 FLOPs 仍较小时，就因为 KV 带宽成为延迟瓶颈。

Prefill 的线性项随 $S$ 增长，注意力项随 $S^2$ 增长；Decode 每步只有一个新 Query，因此其注意力是 $O(T)$。若没有 KV Cache，每一步都重新处理整个前缀，才会重复支付大量前缀计算。

8 训练 FLOPs 的边界

对大矩阵乘法，反向传播通常还要计算输入梯度和权重梯度，所以前向加反向粗略约为前向的三倍，形成常见的每训练 token 约 $6P$ 估算。它仍是权重主导近似：注意力、重计算、输出层、稀疏路由和优化器步骤应按需要另计。不要把 $6P$ 用于普通推理解码。

下一章：吞吐与延迟估算。