为什么 FFN 不需要 KV Cache——兼谈 Prefill 与 Decode 的计算本质

和 Gemini 聊 Transformer 推理的时候，聊到一个问题：为什么 Transformer 有 KV Cache，却没有 FFN Cache？

能答出”FFN 不涉及跨 token 计算”只是及格线。这个问题真正的价值在于——顺着它往下挖，会一路穿过因果掩码的数学等价性、GPU 的计算模式差异，最终落到 Prefill 与 Decode 的本质区别上。

本文来自笔者与 Gemini 的一次讨论，经整理和深度扩展而成。

因果掩码：并行计算如何等价于串行

先把最基础的问题说清楚：[A, B, C] 一次性喂给模型（Prefill），和先喂 A、再喂 B、最后喂 C（逐 token Decode），算出来的 Q、K、V 以及最终特征向量，数字上是否完全一致？

是的，完全一致。 原因在 Causal Mask（因果掩码）。

在自注意力计算中，掩码矩阵的右上角全是 $-\infty$。softmax 之后：

• 第 1 行（token A）：只能 attend 到位置 1
• 第 2 行（token B）：只能 attend 到位置 1、2
• 第 3 行（token C）：只能 attend 到位置 1、2、3

模型通过一次并行计算，同时产出了三个 token 各自的注意力输出，但每个 token 的”视野”严格受限于其位置及之前。这完美模拟了”先有 A、再有 B、最后有 C”的时间自回归顺序。

换句话说，Prefill 没有发明新算法。数学上，它和逐 token 计算做的是同一件事。

FFN 为什么不需要 Cache

Transformer 每一层的结构是：

Input → [Attention → Add&Norm] → [FFN → Add&Norm] → Output

FFN 的计算公式（以 SwiGLU 为例）：

$FFN(x) = W_{down} \cdot (SiLU(W_{gate} \cdot x) \odot (W_{up} \cdot x))$

注意：这里的 $x$ 是单个 token 的隐藏状态向量，$W_{gate}$、$W_{up}$、$W_{down}$ 是层参数。整个计算不涉及任何跨 token 交互。

两个原因

原因一：FFN 是位置无关的 per-token 操作。 FFN 做的事是把一个 $d_{model}$ 维向量升到 $d_{ff}$ 维，激活后再降回来。这个变换对每个 token 独立执行。生成第 $t+1$ 个 token 时，你只需要对它自己的隐藏状态做 FFN——前面 $t$ 个 token 的 FFN 输出根本不会参与计算。

原因二：跨 token 信息已经在 Attention 中汇入，不需要 FFN 再传一次。 生成新 token 时，它的 FFN 输入来自这一层 Attention 的输出。而 Attention 已经通过 $Q_{new}$ 与 $K_{cache}$ 的内积，把历史上下文的语义信息压缩进了新 token 的表示向量中。FFN 拿到的输入本身已经”含有时序上下文”，它只需要对这个向量做非线性变换。

说白了：Attention 负责跨 token 传递信息，FFN 负责对单个 token 做非线性加工。 前者需要历史状态（K、V），后者不需要。

那能不能反过来——把 FFN 输出也缓存起来，加速点什么？不能。Decode 阶段每生成一个新 token，它走的路径是：Attention 融合上下文 → FFN 非线性变换 → 进入下一层 Attention。历史 token 的 FFN 输出永远不会被后续 token 的 FFN 查询或复用。存了也是白存。

Prefill 与 Decode 的计算本质

既然数学上等价，为什么工程上要区分 Prefill 和 Decode？答案在 GPU 的脾气。

Prefill：Compute-bound

长度为 $L_{prompt}$ 的输入一次性进入模型。Attention 中的 $QK^T$ 是 $(L, L)$ 矩阵乘法，FFN 中的 $WX$ 是 $(d_{ff}, d_{model}) \times (d_{model}, L)$ 的大矩阵乘法。全是 GEMM（通用矩阵乘法）。

GPU 只需把模型权重从显存（HBM）里捞一次，几千个 CUDA 核心就能并行处理所有 token。此时瓶颈在计算单元——Compute-bound。

Decode：Memory-bound

Decode 每次只生成 1 个 token。$Q_{new}$ 是 $(1, d)$ 向量，$Q_{new}K_{cache}^T$ 退化为 GEMV（矩阵-向量乘法）。FFN 的 $Wx$ 同样是矩阵-向量乘法。

每次生成一个 token，GPU 仍然需要把全部模型权重（几十 GB）从 HBM 加载一遍，但只用来算 1 个 token。绝大多数 CUDA 核心在等数据送达——Memory-bound。

一句话：Prefill 把 $L$ 次”捞权重、算一个 token”的低效操作，合并成 1 次”捞权重、算 $L$ 个 token”的满载操作。

Decode = $L=1$ 的 Prefill

Prefill 和 Decode 底层跑的是同一套前向计算。唯一的变量是序列长度：

• Prefill：$L = L_{prompt}$，一次性并行处理全部已知 token。把所有层的 K、V 一键写满 KV Cache。产出第一个生成 token。
• Decode：$L = 1$，喂入新生成的 1 个 token，计算它的 Q，从 KV Cache 读取历史 K、V 做 Attention，过 FFN，追加本层 K、V 进 Cache。产出下一个 token。

整个生命周期就是这么简洁的两步。KV Cache 的角色在这两步中也很清楚：Prefill 是”存”（空间换时间），Decode 是”取”（时间换空间）。

收束

讨论 Transformer 推理优化，绕不开两条线：Prefill 怎么更快（吞吐）、Decode 怎么不卡（延迟）。

两者的底层数学从未改变——自始至终在做同一套矩阵运算。Causal Mask 保证了并行与串行的结果一致，KV Cache 把 Decode 的 Attention 从 $O(L^2)$ 压到 $O(L)$，而 Prefill 用 GPU 最擅长的 GEMM 把 prompt 处理从”循环 2000 次”变成”并行一炮”。

至于 FFN——它不需要 cache。不是巧合，是架构设计的必然。跨 token 的信息流动全部收敛在 Attention 里，FFN 只对已经融合了上下文的向量做变换。知道什么该缓存、什么不需要，就是理解 Transformer 推理计算本质的入口。