YOCO 完全拆解：Decoder-Decoder 架构如何让 KV Cache 只存一次

路线一：滑动窗口注意力#

滑动窗口注意力只允许位置 $i$ 看见最近 $C$ 个位置：

其中因果窗口掩码 $B$ 为：

解码到第 $n$ 个 token 时，每个 Self-Decoder 层只保留区间 $[n-C+1,n]$ 的 K/V，旧位置可以环形覆盖，所以状态从 $\mathcal{O}(N)$ 变成 $\mathcal{O}(C)$。论文的缩放实验把窗口设为 1024。它的长处是保留 softmax 注意力熟悉的归一化与精确局部匹配，kernel 也容易由局部版 FlashAttention 实现；短处是单层看不到窗口以外的位置。

“局部窗口会不会使模型失去全局上下文”是理解 YOCO 的关键问题。答案分两层：Self-Decoder 堆叠会扩大有效感受野，但这还不是完整答案；更重要的是，它在边界处为所有历史 token 都生成一个向量 $M_i$，并把所有 $M_i$ 投影成全局 $\hat K_i,\hat V_i$。Cross-Decoder 的每个生成位置可以对这份全长记忆做因果全局注意力。局部混合负责把邻域加工好，全局读取则被集中到后半段。

路线二：Multi-Head Gated Retention#

论文默认采用 Gated Retention（gRet）。它继承 RetNet 将注意力写成递推状态的思路，并让遗忘率由当前数据决定。对一个 head，在第 $n$ 个位置先计算 $Q_n,K_n,V_n$ 以及门值 $\gamma_n$：

$W_\gamma$ 是可学习投影，$\tau$ 是温度；论文解释，温度用于让 $\gamma$ 更倾向于接近 1，从而避免状态过快遗忘。状态与输出递推为：

$S_n$ 是当前 head 的矩阵状态，$K_n^\top V_n$ 把新 token 的键值外积写入状态；$\gamma_n$ 则把旧状态整体衰减。把递推式展开，可以看到任意历史位置 $m$ 对位置 $n$ 的贡献：

括号中的连乘是从 $m$ 到 $n$ 的存活系数。若途中门值都接近 1，这段信息长期保留；若某些位置给出较小门值，之前的信息会被快速淡化。它不是“只记最后几个 token”，而是把任意长度历史压入固定形状的充分统计量。状态大小依赖 head 维度而不依赖 $N$，因此增量解码是 $\mathcal{O}(1)$ 状态。

YOCO 采用 head-wise 而不是 element-wise 的衰减门。同一 head 内共享一个标量衰减看起来表达力较弱，却能让 $Q_nS_n$ 和外积更新保持规则矩阵形状，更容易利用 Tensor Core；若每个状态元素都有不同门控，算子会变得碎片化，实际硬件效率可能抵消理论收益。

多头输出还要经过逐 head 归一化与输出门：

GroupNorm 防止不同 head 的状态尺度漂移，$\operatorname{swish}(XW_G)$ 再按元素控制哪些通道写回残差流。这里的门与 $\gamma$ 分工不同：$\gamma$ 管历史状态保留多久，输出门管当前 block 向残差流暴露哪些特征。

Prefill 与 Decode 为什么使用两种计算形态#

递推形式适合一次一个 token 的 Decode，却不适合长 Prompt：若真的在 GPU 上循环 100 万次，状态虽省，串行依赖仍会让 Prefill 很慢。gRet 因而还提供等价的分块形式。块内一次性计算并行 retention，块间只传递末端状态 $R_i$：

$O_{[i]}^{\mathrm{inner}}$ 表示第 $i$ 块内部所有位置之间的贡献，$O_{[i]}^{\mathrm{cross}}$ 表示此前块压缩状态对当前块的贡献。官方评测的 Prefill chunk size 为 256；生成阶段则切回递推形式，只更新 $S_n$。这不是近似：在相同数值精度下，并行、递推和分块递推是同一公式的不同结合顺序。工程上选择不同形态，是为了让 Prefill 吃到矩阵并行、Decode 避免保存长序列。

可以用一个块的边界状态证明这种等价性。设第 $i$ 块从位置 $a$ 到 $b$，块内第 $j$ 个位置的累计衰减为：

从上一块传来的状态记为 $S_{a-1}$。块内位置 $a+j$ 的状态可拆成两项：

第一项只依赖块入口状态，可以让整块所有 Query 一次性与 $S_{a-1}$ 相乘；第二项只涉及块内的三角因果矩阵，可以用常规矩阵乘并行计算。到块末尾，再把两项合并成 $S_b$ 传给下一块。这样，跨块方向仍有顺序依赖，但顺序步数从 $N$ 降到约 $N/B$，每个块内部的 $B$ 个位置则并行展开。

分块大小 $B$ 因而是硬件旋钮，不是模型语义参数。$B$ 太小，跨块循环次数多，GPU 启动与同步开销占比上升；$B$ 太大，块内三角矩阵和临时张量膨胀，片上空间与访存压力增加。论文实现选 256 是对应其 Triton kernel 与 H100 评测的工程点，迁移到其他 head 维度、精度或加速器时应重新测量，而不应把 256 当作公式要求。

两条 Self-Decoder 路线的差异可以压缩成一张表：

两者都满足 YOCO 的架构契约：Self-Decoder 只维护有界状态，并在中间边界产出一份可供全局查询的完整 KV。

增量生成时究竟追加什么#

假设 Prompt 长度为 $N$，现在生成第 $N+1$ 个 token。它先经过每一个 Self-Decoder block：滑窗版本向每层环形缓冲追加一组局部 K/V，gRet 版本更新每层的 $S_{N+1}$。到达中间边界后，只为这个新 token 计算一次：

然后把它们追加到唯一的全局缓存。后半段每层生成各自的 $\hat Q_{N+1}^l$，对 $[\hat K_1,\ldots,\hat K_{N+1}]$ 和对应 Value 做注意力。新 token 仍完整穿过全部 $L$ 层，因此 YOCO 没有用“少跑一半网络”换取 Decode 质量。

如果 Cross-Decoder 使用 GQA，缓存形状还可以从 Query 头数 $h_q$ 解耦为更小的 KV 头数 $h_{kv}$。以论文 3B 配置为例，Query 有 24 个头，KV 只有 8 个头；三个 Query 头共享一组 K/V。于是 YOCO 与 GQA 的收益相乘：GQA 在头维上少存，YOCO 在层维上少存。

共享为何可能成立，又会损失什么#

跨层 K/V 并非天然等价。标准 Transformer 的第 2 层与第 30 层面对的是不同抽象层级的隐藏状态，各自投影的 K/V 理论上可以编码不同关系。YOCO 主动施加一个信息瓶颈：所有后半层必须从同一份中间记忆中取信息。它之所以不等同于把层“复制粘贴”，是因为后半层保留了独立 Query 与 FFN；它之所以仍可能有损，则是因为共享 K/V 无法为每个深度重新编码整段历史。

论文用从 160M 到 13B 的等参数缩放曲线验证这种瓶颈没有破坏总体可扩展性：YOCO 与采用 Llama 式 RMSNorm、SwiGLU 的 Transformer 基线相当，gRet 版本优于论文中的 SWA 版本。这个结果只说明论文设置下质量—效率折中可行，不能推出任意已训练 Transformer 都能无损改成 YOCO。因为 K/V 共享位置、Self-Decoder 的 mixer 与后半段读取方式都改变了，YOCO 必须作为原生架构获得相匹配的参数；它不是推理时对普通模型做一次配置开关。

用 3B 配置手算一遍#

论文 3B 模型有 26 层，隐藏维度 3072，24 个 Query 头、8 个 KV 头，因此每个 Cross-Attention head 的维度是 $3072/24=128$。若标准 26 层 GQA Transformer 用 BF16 保存缓存，单 token、单请求的理论 KV 字节数为：

YOCO 的全局缓存只做一份同样的 8-head 投影：

因此，仅比较随上下文增长的 global KV 主项，恰好是 26 倍差距。32K 上下文对应约 3.25GiB 对 128MiB，1M 上下文则约为 101.6GiB 对 3.9GiB。这里还没有加入 gRet 的固定状态：13 个 Self 层、24 个 head、每个 $128\times128$ 状态，若按 BF16 粗算约为 9.75MiB/请求。它在 32K 已明显小于 global cache，在 1M 更接近可忽略常数，但高并发时仍要计入调度预算。

这次手算也解释了“约省 $L$ 倍”为什么不是永远等于总层数。若比较对象的层数、KV 头数、精度或 Self/Cross 切分改变，比例也会改变；如果采用 SWA，固定项还会变成每个 Self 层的 $C$ 个 K/V。论文 Figure 7 中 65B 配置达到 80 倍，来自该配置更深的层数，而不是 YOCO 对大模型使用了另一种压缩公式。

容量减少也不等于每步只从 HBM 读 4KiB/token。13 个 Cross 层都会读取同一份 4KiB/token 全局缓存，层与层之间还隔着输出投影和 FFN；若缓存无法长期驻留片上，它可能被反复从 HBM 取回。YOCO 消除了 26 份物理存储，并把全局扫描集中到后半层，但不会自动把 13 次逻辑读取合并成一次。这正是容量复杂度和带宽复杂度必须分开分析的原因。

总显存：权重不变，缓存斜率改变#

3B 配置的总推理显存。32K 时 YOCO 约省 1.95 倍，1M 时差距扩大到 9.38 倍；YOCO 在 1M 时为 12.4GB。图源：YOCO 论文 Figure 6。

32K 时，模型权重和中间激活仍占显著比例，因此理论上的缓存缩减不会等比例变成总显存缩减，图中只有约 1.95 倍。上下文继续增长后，Transformer 的逐层 KV 成为主导，YOCO 曲线则只缓慢上升；到 1M，论文报告 YOCO 总占用 12.4GB，Transformer 是其 9.38 倍。这个趋势比单个倍数更重要：YOCO 改变的是长上下文下显存增长的斜率。

每 token 缓存：模型越深，层共享越值钱#

按论文模型配置估算的每 token KV 占用。模型从 1.2B 扩到 65B，YOCO 相对 Transformer 的缓存缩减从 24 倍扩大到 80 倍。图源：YOCO 论文 Figure 7。

图中 Transformer 从约 40KB/token 增长到约 640KB/token，YOCO 的增长小得多。原因并非“大模型每层 KV 更容易压缩”，而是更大模型通常层数更多，标准缓存付出的层数乘数更大；共享一次的价值随深度增加。论文据此给出一个容量化例子：在 65B 配置上，1GB KV 预算约能容纳 128K 个 YOCO token，而 GQA Transformer 约为 1.6K 个 token，对应 80 倍差距。这个数字只比较缓存预算，不包含权重、激活、碎片和服务框架元数据。

Prefill：提前退出与线性 mixer 叠加#

输入长度从 32K 增至 1M 时的 Prefill 延迟。Transformer 的全局注意力接近平方增长，YOCO 的 gRet Self-Decoder 近似线性增长。图源：YOCO 论文 Figure 8。

32K 时 YOCO 已有 2.87 倍加速，因为 bulk Prompt 不进入后半段；上下文增长后，全局注意力的 $N^2$ 项开始主导，差距扩大。512K 时基线约 180 秒，YOCO 不到 6 秒，对应约 30.3 倍；1M 时图中从约 300 秒降到约 4 秒，论文标注 71.82 倍。图中 YOCO 在超长区间没有随长度明显上升，除了线性复杂度之外，也受到固定开销、并行效率和测量尺度影响，不能据此推断 Prefill 在现实中与长度完全无关。

吞吐：不只是单请求 kernel 更快#

吞吐同时计入 Prefill 和 Generation。上下文越长，YOCO 相对优势越大；512K 时从 4.5 token/s 提升到 43.1 token/s。图源：YOCO 论文 Figure 9。

512K 上下文时，论文报告 Transformer 为 4.5 token/s，YOCO 为 43.1 token/s，即 9.56 倍。这个收益有两条来源：第一，单请求的 Prefill 工作量下降；第二，缓存变小后，同样显存可以放入更大的 batch，GPU 更容易维持足够并行度。后者说明 KV 优化常有“二阶收益”：省下的容量不只避免 OOM，还会反过来提高吞吐。

把四组数据放在一起，可以避免两个常见误判。其一，理论上的约 $L$ 倍 KV 缩减不能直接写成 $L$ 倍总显存或吞吐，因为权重、激活和计算仍在；其二，论文报告的 71.82 倍 Prefill 与 9.56 倍吞吐属于特定 3B、H100、超长输入和实现条件，不应外推为所有请求长度与硬件上的固定加速比。

推理状态应分成两类#

一个服务引擎至少要维护以下状态：

1
RequestState
2
├── self_state[0 ... L_self-1]
3
│   ├── gRet: last_hidden_state
4
│   └── SWA:  prev_key[C], prev_value[C]
5
└── global_cache
6
    ├── key[0 ... current_length-1]
7
    └── value[0 ... current_length-1]

self_state 按 Self-Decoder 层各有一份，但长度有界；global_cache 随序列增长，却只有一份并被所有 Cross-Decoder 层引用。若把后者错误地在每层各注册一份张量，即使内容相同，也会把 YOCO 的核心收益重新复制掉。工程实现应传递同一 storage 的视图或句柄，而不是逐层拷贝。

官方 CrossDecoder 先对边界隐藏状态做一次 RMSNorm 和 K/V 投影，再把结果写入 incremental_state["prev_key"] 与 incremental_state["prev_value"]；随后循环所有 Cross layer，把同一 key, value 参数传入。它还分别对 K 与 Query 应用 RoPE，因此共享并没有取消位置信息。CrossAttention 内部使用 FlashAttention 接口，并支持 Query 头多于 KV 头的 GQA 布局。

Prefill 的实现骨架#

下面的伪代码只保留推理依赖，不对应仓库中的逐行复制：

1
def prefill(prompt, state):
2
    # 1. 整段 Prompt 通过 Self-Decoder。
3
    hidden = embed(prompt)
4
    for layer_id, layer in enumerate(self_layers):
5
        hidden, state.self_state[layer_id] = layer.prefill(
6
            hidden,
7
            state.self_state[layer_id],
8
        )
9

10
    # 2. 只投影一次全局 K/V，并写入唯一缓存。
11
    k, v = boundary_kv_project(rms_norm(hidden))
12
    state.global_cache.append(k, v)
13

14
    # 3. Prompt 的非前沿位置不跑 Cross-Decoder。
15
    frontier = hidden[:, -1:, :]
16
    for layer in cross_layers:
17
        frontier = layer(
18
            frontier,
19
            state.global_cache.key,
20
            state.global_cache.value,
21
            causal=True,
22
        )
23

24
    return lm_head(final_norm(frontier)), state

gRet 路径会先把 Prompt padding 到 256 的倍数以适配分块 kernel，再用 mask 排除填充位置；同时计算每个 Self 层 Prefill 结束后的 last_hidden_state，供下一次递推 Decode 继续使用。SWA 路径则把各层最后 $C$ 个 K/V 留在固定缓冲区。两者都不能只保留边界层状态，因为新 token 仍要逐层穿过 Self-Decoder。

论文图把 Prefill 阶段的 Cross-Decoder 标为 skipped，代码中也有 skip_cross_decoder 开关。官方 needle 评测的生成循环先用该开关处理区间 $[0,N-1)$，特意留下最后一个 Prompt token；随后在 start_pos=N-1 处让这个 token 同时通过 Self-Decoder 与 Cross-Decoder，得到首个 next-token logits。另一种等价组织是先让 Self-Decoder 处理完整 Prompt，再单独把边界的最后位置送入 Cross-Decoder。无论采用哪种接口，都不能漏掉首个输出所需的前沿计算。可以把它视作 Prefill 收尾，也可以视作 Generation 的第零步；性能计时必须统一这个边界，否则不同实现的 TTFT 会少算或多算一段。

单 token Decode 的实现骨架#

1
def decode_one(token, position, state):
2
    hidden = embed(token)
3

4
    # 每层只更新有界状态。
5
    for layer_id, layer in enumerate(self_layers):
6
        hidden, state.self_state[layer_id] = layer.decode_one(
7
            hidden,
8
            position,
9
            state.self_state[layer_id],
10
        )
11

12
    # 新 token 的全局 K/V 只生成、追加一次。
13
    k_new, v_new = boundary_kv_project(rms_norm(hidden))
14
    state.global_cache.append(k_new, v_new)
15

16
    # 所有 Cross 层共享读取完整缓存，各自计算 Query。
17
    for layer in cross_layers:
18
        hidden = layer(
19
            hidden,
20
            state.global_cache.key,
21
            state.global_cache.value,
22
            causal=True,
23
        )
24

25
    logits = lm_head(final_norm(hidden))
26
    return sample(logits), state

对 gRet，每个 Self 层的核心更新就是 current_state = exp(gate) * old_state + outer(k, v)；官方实现把门控 logit 除以 16，并对 K 乘 $d_h^{-1/2}$ 的缩放。对 SWA，代码预分配窗口大小的 prev_key、prev_value，超过窗口后只保留尾部。对全局缓存，参考实现按 max_seq_len 预分配连续张量并按 start_pos 写入，逻辑简单但可能浪费保留空间。

接入推理引擎时的四个关键点#

第一，全局缓存仍然适合分页。YOCO 消掉的是层维复制，不是请求长度不确定性；服务端依旧可能遇到不同请求长度、抢占、beam 分叉与前缀共享。把唯一的 global cache 接到块式分配器，可以叠加 PagedAttention 的低碎片优势。

第二，调度器要分别核算两类状态。每个请求的预算可以近似写成“所有 Self 层的常量状态 + 一份随 token 增长的 global KV”。若仍沿用 num_layers × bytes_per_token 估算，准入控制会严重低估可服务并发。

第三，Cross-Decoder 的长 KV 读取仍是带宽热点。缓存只存一份不等于每步只读一遍：每个 Cross 层都会扫描或分块读取同一份 K/V。容量复杂度从 $L$ 份降到 1 份，但 Decode 的全局注意力计算仍大致乘 Cross 层数。FlashAttention、GQA、量化或稀疏检索仍有叠加空间。

第四，Prefix Cache 的共享单位会变得更简单。同一模型、同一 Prompt 前缀只需要缓存一套边界 K/V，加上能够恢复各 Self 层递归状态或窗口状态的快照。相比逐层前缀块，这降低了长前缀存储量；但如果只保存全局 K/V 而没有 Self 状态，续写时无法从前缀末端正确推进 Self-Decoder。

它不是已有权重的无损插件#

普通 Transformer 的每层都有自己的 K/V 投影和历史表征，YOCO 则要求在模型中间形成一个可被所有后半层读取的共享接口。直接删掉已有模型的逐层缓存、把某一层 K/V 广播给所有层，会同时遇到维度、语义和质量问题。部署 YOCO 需要原生 YOCO checkpoint，或者经历针对新架构的适配过程；不能像切换 FlashAttention kernel 一样只改推理后端。

它省容量，但 dense Decode 仍反复扫描长缓存#

每个 Cross-Decoder 层都读取同一份全局 K/V。物理存储只有一份，但若上下文为 $N$，每个新 Query 的 dense attention 仍需在每个 Cross 层处理 $N$ 个位置。粗略地说，Decode 的注意力读取与计算仍为：

所以 YOCO 首先解决“缓存复制 $L$ 次”和“Prompt 的 $N^2$ Prefill”，并没有让 dense 长上下文 Decode 变成与 $N$ 无关。长输出、长历史的任务最终仍可能受 HBM 带宽限制。这也解释了为什么后续工作会在共享缓存之上继续做稀疏索引，而不是认为 YOCO 已经结束了长上下文优化。

有界状态不是免费状态#

SWA 每个 Self 层仍保存 $C$ 个位置，$C$ 过大会抬高常量；gRet 每个 head 的递归矩阵状态与 head 维度相关，同样要占显存并执行外积更新。$\mathcal{O}(1)$ 只表示不随序列长度 $N$ 增长，不表示字节数为零。对于短 Prompt，权重与这些固定状态可能占主导，YOCO 的总显存优势自然小于超长上下文。

一百万上下文不等于一百万有效推理能力#

论文的单针检索在 1M 长度上接近满分，多针实验中 YOCO-3B-1M 在 128K 长度、针数 $1/2/4/8$ 时准确率分别为 $0.98/0.98/0.84/0.56$。它证明模型能在极长输入中访问远处信息，也同时显示任务变复杂后准确率会下降。Needle-in-a-Haystack 主要测试精确召回，不能代替跨章节综合、代码库推理或多跳归因等长上下文能力评估。

与其他 KV 优化不是替代关系#

这张表对应前文公式 $2BNLh_{kv}d_hs$ 的不同因子。GQA 减少头数，量化减少字节数，驱逐减少 token 数，YOCO 减少长缓存的层数副本，PagedAttention 则降低物理分配浪费。它们瞄准的轴不同，因而不是只能择一。关于低秩缓存接口可参考 MLA 完全拆解，关于 2-bit KV 可参考 KIVI 完全拆解。

最适合与不太适合的负载#

YOCO 最适合长输入、共享前缀多、KV 容量限制并发、输出相对较短的服务，例如长文档问答、代码库检索增强和多轮会话。此时 Prefill early exit 与单份全局缓存都会直接命中瓶颈。

如果上下文很短、batch 已经能把 GPU 算满，权重与 FFN 计算占主导，节省 KV 未必带来显著端到端收益。如果任务要求在百万历史上持续生成极长结果，dense Cross-Decoder 的反复全局扫描又会成为新瓶颈。最后，现有部署生态大量围绕标准 Transformer 的逐层 paged KV、连续批处理和成熟 fused kernel 建设；原生 YOCO 还需要调度器、缓存核算和 kernel 路径的专门支持，论文原型性能不能自动等价为所有框架里的生产性能。