☰ DSv4.1 Flash 架构导读

3. 层维度(上):CED 把 prefill 减半

3.1 为什么先动 prefill

barm 第 2 章把一次生成拆成 prefill 和 decode 两段,其中 prefill 是"把整段提示词过一次模型"。在两本书的场景里,提示词通常几百个 token,prefill 只占一次调用的零头。

Agent 工作流改变了这个比例。一次工具调用返回几千 token,模型要重新读一遍(或部分重读)上下文;一轮任务可能调用几十次工具。于是prefill 从零头变成主要开销,而且它是纯计算:提示词没变、缓存命中不了的部分必须从头算。

这里有一个简单但容易忽略的事实:prefill 的成本正比于"token 数 × 层数"。1.2 节的 KV 账是"层数 × token 数 × 条目字节",prefill 的算力账是"层数 × token 数 × 每层 FLOP"。两者共享同一个乘数 层数。CED 就是在这个乘数上动手。

3.2 YoCo:让 KV 只被生产一次

CED 的直接来源是 YoCo(You Only Cache Once)。它的观察很朴素:标准 Transformer 里,第 \(l\) 层的 KV 由第 \(l\) 层的隐状态算出,所以 40 层就有 40 份不同的历史表示。可是这 40 份里有多少是真正独立的?

YoCo 的答案:把模型竖着切成两半。下半部分负责把历史压成一份表示,上半部分只负责读这份表示。

标准 Transformer
L40  Q40 K40 V40
L39  Q39 K39 V39
...
L1   Q1  K1  V1

YoCo
L40  Q40   读共享 KV
L39  Q39   读共享 KV
...
L21  Q21   读共享 KV
L20  K20 V20  产出共享 KV
L19  K19 V19
...
L1   Q1  K1  V1

下半部分仍要算 KV(因为要读自己的历史),上半部分只算 Q。于是"层维度"上的 KV 份数从 \(L\) 降到 \(L/2+1\),同时上半部分的层在 prefill 时可以只看最后一个位置——历史已经压进共享 KV 里了。

这个思路在工业界还有别的名字,比如 KV-Mirror;报告也引用了腾讯 WeLM 的工作。它们的共同点是跨层共享一份全局 KV,但投影怎么共享、Decoder 的局部支路怎么处理、prefill 什么时候提前退出,各家并不相同。所以这里只把它们当作同一条路线,不当作同一个实现。

3.3 CED 的具体做法

CED(因果编码器-解码器,causal encoder-decoder)把 \(L=40\) 层切成 \(E=D=20\):\(l\le 20\) 是 Encoder,\(l>20\) 是 Decoder。对 Decoder 的全局注意力,KV 不再由本层隐状态 \(H_l\) 导出,而是从第 20 层的隐状态投影。报告式 (1) 写的是:

\[ C_l=H_{20}\,W_l^{KV}, \qquad Z_l=H_{20}\,W_l^{Z}, \] \[ l>\tfrac{L}{2}, \]

其中 \(H_{20}\in\mathbb{R}^{N\times d}\) 是第 20 层(Encoder 最后一层)输出的全部位置隐状态,\(W_l^{KV}\)(\(d\to d_c\))是逐层的投影权重,\(C_l\) 是该层要用的主 KV;\(Z_l\) 是配套的索引侧表示(报告称其为 corresponding compression weights)。源层只有一个,但投影可以逐层不同——这就是"同一份记忆、不同读法"。

CSA2 再往前走一步。按 2.3.1 的配对规则,Decoder 里只有被指定为 Full 的那一层真正用 \(W_l^{KV}\) 产出全局缓存(这里是 layer_idx=20),其余层走 Reindex / Reuse,直接复用这一份。于是缓存份数从 40 降到等效 2.5(4.4 节);代价是 Decoder 各层只能靠自己的 Q、局部窗口、输出投影与 FFN 区分彼此。注意这两种说法并不冲突:式 (1) 定义的是"每个 Decoder 层可以怎么从 \(H_{20}\) 取得自己的 KV",CSA2 的调度决定的是"实际只有一层去算"。

和第 1 章的 MLA 对照一下,会发现这是同一个手法的另一面。MLA 把"每层的 K、V 由本层隐状态算出"改成"由一个低维潜在向量算出";CED 把"每层的潜在向量由本层隐状态算出"改成"由第 20 层的隐状态算出"。前者压的是每个条目的宽度,后者压的是条目的份数。

共享一份 KV,会不会让 20 层 Decoder 退化成"同一层重复 20 次"?不会。共享的只是全局历史这一项输入;每层仍有自己的 Q、自己的局部 SWA KV、输出投影、MoE 与残差,层输入也随深度变化。报告 2.2 节对 SWA 的说法正好相反——SWA KV 仍然逐层从本层隐状态生成,这让局部表示的"计算深度"反而变深了(3.6 节)。共享真正限制的是"随层更新全局记忆内容"的自由度——这正是第 5 章要算清楚的事:顶点固定之后,改 Q 还能在凸包里走多远(5.5 节)。

3.4 计算量:为什么正好是一半

设 Encoder 深度 \(E=L/2\),Decoder 深度 \(D=L/2\),窗口 \(W\),序列长 \(N\)。

Encoder 要处理全部 \(N\) 个位置,成本正比 \(NE\)。Decoder 只需要处理"最后 \(B=\min(N,W)\) 个位置"(原因见 3.5 节的 bounded replay),成本正比 \(BD\)。合计

\[ C_{\text{CED}}\approx NE+BD . \]

没有 CED 时全部 \(N\) 个位置要走完 \(L\) 层,成本 \(NL\)。两者相比,代入 \(E=D=L/2\):

\[ \frac{C_{\text{CED}}}{C_{\text{full}}} \approx\frac{NL/2+BL/2}{NL} =\frac12+\frac{B}{2N}. \]
结论

\(1/2\) 来自"层数减半",\(B/2N\) 是 Decoder 为恢复局部状态付的税。序列越长,税越轻;\(N\gg W\) 时比例趋近 50%。

代入报告的配置 \(L=40,W=128\),取一段 \(N=32768\) 的冷启动提示词:

如果把 \(N\) 换成 \(10^6\):\(0.5+128/(2\times10^6)=50.006\%\)。税几乎消失。这也对应报告里那句"prefill 每 token 激活 8B、decode 激活 16B":prefill 只走 20 层,参数激活量正好是 decode 的一半。

3.5 Bounded replay:那个 \(B\) 从哪来

Decoder 里除了全局 KV,每层还有一份 SWA 局部 KV(最近 128 个位置)。问题来了:Encoder 只处理了"全部位置",Decoder 只处理"最后 128 个",那中间那些位置的局部 KV 就没生成过。要不要补?

如果严格补,代价比想象的大。用玩具窗口(\(W=4\),只为看清结构)举例:要让 Decoder 第 20 层在位置 100 的局部表示正确,需要它下面一层在位置 97–100 的表示;再往下,位置 97 又需要 94–97,依此类推。每下一层,依赖区间向前扩 \(W-1\),所以向前追溯的总跨度是

\[ \begin{aligned} R&=1+D(W-1)=1+20\times127\\ &=2541\approx DW . \end{aligned} \]

注意这是依赖跨度,不是"必须让 2541 个位置逐个过完 20 层":逐层的依赖区间层层收缩,整体形状是三角形而不是矩形。但即便按最坏的矩形估,也要把约 2541 个位置送进 Decoder;每轮工具调用只新增几十个 token 时,这笔恢复成本足以吃掉提前退出的收益。

报告的做法叫 SWA Bounded Replay(有界重放):只重放最后 \(n_{\text{win}}=128\) 个位置,不再向前追溯。报告的估计是"精确恢复要重放 \(L\times n_{\text{win}}\) 个 token"——对 Decoder 那一半就是 \(20\times128=2560\)(精确到 \(1+D(W-1)=2541\));bounded replay 只付 \(128\),把召回范围截成 \([\max(s,i-W+1),i]\)。代价是重建出来的局部 KV 是近似的——重放片段的开头缺少更早的局部依赖。报告的判断是"对回答质量几乎无影响",依据是既有工作对 SWA 有效感受野的观察(实际远小于理论上的 \(n_{\text{win}}\times L/2\))。为了让训练和推理的分布一致,后训练阶段还会模拟一遍同样的 replay。

报告里其实有两个 bounded replay,容易混:Decoder SWA Bounded Replay(本节)把 decoder 前向限制在 \(n_{\text{win}}\) 个 token,几乎把 prefill 计算砍半;Encoder SWA Bounded Replay(3.2.2)用在"持久化缓存里没有 SWA KV"的情形,重放缓存前缀的最后 \(n_{\text{win}}\) 个 token 来近似恢复——这是第 8 章 1/8 那个倍数的前提。

常见错误

被截断的是局部 SWA 状态的重建范围,不是全局记忆。重放最后 128 个位置时,Decoder 仍可以按稀疏规则读更长的全局 KV。所以这个近似不会直接删掉全局缓存里的历史;但局部状态的误差会顺着后续隐状态影响主查询与 Indexer 查询,从而间接改变全局信息怎么被读取。"不影响这段提示词讲了什么"因此只是待验证的假设,不是结论。

3.6 一个反直觉的细节:SWA 反而增加了"深度"

CED 里全局 KV 只来自第 20 层,看起来"深度变浅了"。但 SWA 不是:每一层都用自己的隐状态算局部 KV,包括 20 层 Decoder。于是最近 128 个位置的表示,经过的是完整 40 层;而更早的位置经过的是 20 层。

这不是缺陷而是设计:解码时新 token 就在窗口内,它需要最深的表示;远处的历史只需要"被压缩过的记忆"。于是同一份上下文里,近处深、远处浅——这正是"注意力该有的样子"。

3.7 容易算错的地方

常见错误
  • 把 CED 当成"40 层变 20 层"。参数和层数都没变,变的是"哪些层参与 prefill 的逐位置计算"。
  • 用 \(B=0\) 算 50%。\(B=\min(N,W)\),小提示词下 \(B/2N\) 不可忽略:\(N=256\) 时比例是 \(0.5+128/512=75\%\),CED 在这个量级上几乎不省。
  • 把 2541 当成"要重放的量"。2541 是精确恢复的依赖跨度(上界估算用 \(DW=2560\));bounded replay 只重放 128,剩下的是被接受的近似。
  • 认为 Decoder 层不产生 KV。它们仍产生局部 SWA KV,只是不再各自产生全局 KV。