Skip to content
团子云技术 Lite 1.048596
Go back

为什么 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 之后:

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

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

FFN 为什么不需要 Cache

Transformer 每一层的结构是:

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

FFN 的计算公式(以 SwiGLU 为例):

FFN(x)=Wdown(SiLU(Wgatex)(Wupx))FFN(x) = W_{down} \cdot (SiLU(W_{gate} \cdot x) \odot (W_{up} \cdot x))

注意:这里的 xx单个 token 的隐藏状态向量WgateW_{gate}WupW_{up}WdownW_{down} 是层参数。整个计算不涉及任何跨 token 交互。

两个原因

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

原因二:跨 token 信息已经在 Attention 中汇入,不需要 FFN 再传一次。 生成新 token 时,它的 FFN 输入来自这一层 Attention 的输出。而 Attention 已经通过 QnewQ_{new}KcacheK_{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

长度为 LpromptL_{prompt} 的输入一次性进入模型。Attention 中的 QKTQK^T(L,L)(L, L) 矩阵乘法,FFN 中的 WXWX(dff,dmodel)×(dmodel,L)(d_{ff}, d_{model}) \times (d_{model}, L) 的大矩阵乘法。全是 GEMM(通用矩阵乘法)

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

Decode:Memory-bound

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

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

一句话:Prefill 把 LL 次”捞权重、算一个 token”的低效操作,合并成 1 次”捞权重、算 LL 个 token”的满载操作。

PrefillDecode
输入长度LpromptL_{prompt}(数百~数万)1
核心算子GEMMGEMV
瓶颈计算(Compute-bound)显存带宽(Memory-bound)
Attention 复杂度O(L2d)O(L^2 \cdot d)O(Lcached)O(L_{cache} \cdot d)
KV Cache写入(一次性填满)读取 + 追加 1 个位置
GPU 利用率高(>80% 可行)低(通常 <10%)

Decode = L=1L=1 的 Prefill

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

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

收束

讨论 Transformer 推理优化,绕不开两条线:Prefill 怎么更快(吞吐)、Decode 怎么不卡(延迟)。

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

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


Share this post on:

Previous Post
【转载】00年互联网泡沫,半导体都发生了什么?悲剧重演?历史已给出答案!
Next Post
从 Softmax 梯度消失到 KV Cache 的深度解密:拆解 Transformer 的时空内幕