YOCO_论文与推理数据流详解

YOCO：You Only Cache Once——论文、架构与推理数据流详解

论文：You Only Cache Once: Decoder-Decoder Architectures for Language Models arXiv：2405.05254（v2） 作者：Yutao Sun、Li Dong 等，Microsoft Research / Tsinghua University

本文完整整理 YOCO 的论文方法、架构、Prefill 与 Decode 数据流，以及围绕以下问题的逐步解释：

• YOCO 为什么只需要一份全局 KV Cache？
• Self-Decoder 和 Cross-Decoder 分别负责什么？
• 为什么 Prefill 可以跳过绝大多数 Cross-Decoder 计算？
• Cross-Decoder 的不同层能不能并行？
• YOCO 究竟减少了哪些计算，哪些计算仍然存在？
• 为什么说它在计算拓扑上类似 encoder-decoder，但行为仍然是 decoder-only？
• YOCO 的 Prefill 与 Decode 在 Cross-Decoder 上是否具有相同计算量？
• YOCO 对 Hybrid-aware continuous batching 和 Hybrid state memory manager 有什么启发？

---

1. 一句话理解 YOCO

标准 decoder-only Transformer 是：

每一层都保存自己的历史 Key/Value。

YOCO 是：

前半模型负责把每个 token 编码成一份通用的全局 Key/Value；后半模型的所有层只生成不同的 Query，反复查询同一份全局 Key/Value。

因此它叫 You Only Cache Once。

需要从一开始就明确：YOCO 不是给现有 Transformer 加一个 KV 压缩器，也不是简单删除若干层的 KV Cache。它直接改变模型依赖关系，因此通常需要按照 YOCO 架构重新训练模型。

---

2. 标准 Transformer 的两个主要问题

2.1 KV Cache 容量随层数和上下文共同增长

假设 decoder-only Transformer 有 层，上下文长度为 。每一层都有自己的历史 KV：

KV Cache 大小近似为：

所以复杂度是：

当上下文很长时，KV Cache 会逐渐成为推理显存的主要组成部分。论文举出的例子是：65B 模型即使使用 grouped-query attention 和 8-bit KV，512K token 也会占用大约 86GB KV Cache，超过一张 H100-80GB 的容量。

2.2 Prefill 中每层都要处理整个 prompt

普通 Transformer 的每一个 self-attention 层都需要处理完整的 prompt，并建立该层自己的 KV Cache。全局 attention 的主要计算复杂度为：

FlashAttention 可以显著减少中间张量和 HBM I/O，但不会改变全局 attention 关于 的计算复杂度。

因此长上下文带来的问题不仅是“KV Cache 装不下”，还包括“Prefill 计算太慢”。

---

3. YOCO 总体架构

YOCO 将总共 个 block 分成两部分：

1. 前 层：Self-Decoder；
2. 后 层：Cross-Decoder。

总体架构如下。

flowchart TB X["输入 Token / Embedding X"] subgraph SD["Self-Decoder：前 L/2 层"] direction TB ESA["高效 Causal Self-Attention<br/>gRet / Sliding-Window Attention"] SFFN["RMSNorm + SwiGLU FFN"] ESA --> SFFN end SS[("每层局部状态<br/>gRet：递归状态 S<br/>SWA：窗口 KV")] M["中间表示 M = X^(L/2)"] PROJ["共享投影<br/>K̂ = LN(M)W_K<br/>V̂ = LN(M)W_V"] GKV[("唯一全局 KV Cache<br/>K̂₁…K̂ₙ，V̂₁…V̂ₙ<br/>随 N 线性增长")] CUR["送入 Cross-Decoder 的当前位置<br/>Prefill：最后一个 prompt 位置<br/>Decode：当前新 token"] subgraph CD["Cross-Decoder：后 L/2 层"] direction TB Q1["第 1 层生成自己的 Query：Q¹"] A1["Causal Cross-Attention"] F1["FFN"] MID["⋮"] QL["第 L/2 层生成自己的 Query：Qᴸ"] AL["Causal Cross-Attention"] FL["FFN"] Q1 --> A1 --> F1 --> MID --> QL --> AL --> FL end HEAD["LM Head"] NEXT["生成下一个 Token"] X --> ESA SFFN --> M SS <-->|"每个 token 更新<br/>大小不随 N 增长"| ESA M --> PROJ PROJ -->|"每个 token 追加一对 K/V"| GKV M -->|"只取当前位置的表示"| CUR CUR --> Q1 GKV -->|"所有层共享相同 K̂、V̂"| A1 GKV -->|"再次读取相同 K̂、V̂"| AL FL --> HEAD --> NEXT NEXT -. "下一步自回归" .-> X

给定输入：

首先通过 Self-Decoder：

得到边界表示：

然后只做一组共享 KV 投影：

这就是唯一的 token-level 全局 KV Cache。

第 个 Cross-Decoder 层只生成自己的 Query：

然后查询共享 KV：

不同 Cross-Decoder 层的 Query 和当前隐藏状态不同，但所有层读取相同的 (\hat K,\hat V)。

---

4. Self-Decoder：负责写入历史记忆

Self-Decoder 不能采用普通全局 self-attention，否则它的每一层仍然要保存随 增长的 KV Cache，YOCO 就失去了主要意义。

论文提供两类 Self-Decoder：

1. Gated Retention，默认实现；
2. Sliding-Window Attention。

4.1 Gated Retention

Gated Retention 可以表示为带数据依赖遗忘门的线性递归：

其中：

• 是压缩历史信息的 recurrent state；
• (\gamma_t) 是由输入决定的遗忘门；
• (K_t^{\mathsf T}V_t) 将当前 token 信息写入状态；
• 从状态中读取与当前 Query 有关的信息。

关键性质是：

不会随着序列长度 增长。

但是不能把这理解成“YOCO 最终把所有历史 token 都压成一个 recurrent state”。Gated Retention 只负责 Self-Decoder 内部的信息传播。Self-Decoder 最顶层仍会为每一个 token 生成边界表示 ，这些表示会被投影为 token-level 的全局：

因此 YOCO 同时具有：

• 前半段线性 attention/recurrent state 的固定长度状态；
• 后半段 softmax attention 的 token-level 全局检索能力。

这正是它的 hybrid 性质。

4.2 Gated Retention 的三种等价计算形式

论文说明 gated retention 可以使用三种等价形式：

1. Parallel representation；
2. Recurrent representation；
3. Chunkwise recurrent representation。

训练或 Prefill 通常使用 parallel/chunkwise 形式，以充分利用 GPU 矩阵计算；逐 token Decode 使用 recurrent 形式，以获得相对于上下文长度为常数的状态。

Recurrent 形式

每次只输入一个 token：

Chunkwise 形式

序列按长度 切成 chunk。计算被分成：

• Inner-chunk：chunk 内 token 使用矩阵方式并行计算；
• Cross-chunk：历史 chunk 被压缩成一个边界状态 。

可以写成：

chunk 边界更新一次 recurrent state：

论文 profiling 中将 chunk size 设为 256。

4.3 Sliding-Window Attention

另一种选择是使用固定窗口 ：

每个 Self-Decoder 层只保留最近 个 token 的局部 KV，因此状态复杂度为：

只要 固定，它相对于上下文长度 就是常数。

---

5. Cross-Decoder：只读共享历史记忆

可以将 YOCO 的两个部分理解为：

• Self-Decoder：Memory Writer；
• Cross-Decoder：Memory Reader。

Self-Decoder 将每个 token 写入共享记忆：

Cross-Decoder 对当前位置产生不同层的 Query：

但反复读取同一份：

Cross-Decoder 不会为每一层重新生成一套 token-level KV Cache，也不会把历史位置在 Cross-Decoder 中的隐藏状态继续写回共享记忆。

它的历史记忆接口是只读的；只有当前位置的隐藏状态沿层方向继续向上传递。

---

6. “Only Cache Once”并不等于整个模型只有一个状态

YOCO 推理时实际存在三类状态。

状态	是否随上下文 增长	主要访问模式
Self-Decoder 的 gRet state (S^{})	否	每个 token、每层读写一次
Self-Decoder 的滑窗 KV（如果使用 SWA）	不随 ，随窗口	读取最近窗口并更新
全局 (\hat K,\hat V)	是，	每 token 追加一次，被所有 Cross 层读取

因此论文脚注也说明，“once”指只保存一份全局 token-level KV。Self-Decoder 仍然需要保存每层的常数大小状态。

论文将总体缓存复杂度写为：

更具体地说：

• gRet 版本：一份 全局 KV，加 个固定大小 recurrent state；
• SWA 版本：一份 全局 KV，加 局部 KV。

相比普通 Transformer：

当 足够长时，KV 容量理论上接近减少 倍。

---

7. YOCO Prefill：完整数据流与计算过程

假设 prompt 为：

7.1 Prefill 架构图

flowchart TB P["Prompt：x₁ ... xₙ"] EMB["Embedding<br/>X₁ ... Xₙ"] subgraph SELF["Self-Decoder：前 L/2 层"] direction TB S1["第 1 层 Efficient Self-Attention<br/>gRet chunkwise / Sliding Window"] F1["FFN"] SDOT["⋮"] SM["第 L/2 层 Efficient Self-Attention + FFN"] S1 --> F1 --> SDOT --> SM end STATE[("Prefill 结束后保留<br/>gRet：每层最终状态 Sₙ<br/>SWA：每层最后 C 个 KV")] M["所有位置的边界表示<br/>M₁ ... Mₙ"] KP["共享投影 W_K"] VP["共享投影 W_V"] GKV[("唯一全局 KV Cache<br/>K̂₁...K̂ₙ<br/>V̂₁...V̂ₙ")] SKIP["M₁ ... Mₙ₋₁<br/>不进入 Cross-Decoder"] LAST["只取最后位置 Mₙ"] subgraph CROSS["第一次 next-token 计算"] direction TB Q1["第 1 个 Cross 层生成 Qₙ¹"] A1["Qₙ¹ 查询全部 K̂₁:ₙ、V̂₁:ₙ"] F2["FFN"] DOT["⋮"] AL["最后一个 Cross 层<br/>再次查询同一份全局 KV"] Q1 --> A1 --> F2 --> DOT --> AL end LOGIT["LM Head<br/>得到 p(xₙ₊₁ | x₁:ₙ)"] P --> EMB --> S1 SM --> M SM --> STATE M --> KP --> GKV M --> VP --> GKV M -.-> SKIP M -->|"选择 Mₙ"| LAST LAST --> Q1 GKV -->|"所有 Cross 层共享"| A1 GKV -->|"再次扫描"| AL AL --> LOGIT

7.2 第一步：全部 prompt token 经过 Self-Decoder

所有 个 token 进入前 层：

如果采用 gated retention，Prefill 使用 chunkwise 形式：

• chunk 内并行；
• chunk 间传递 recurrent state；
• 得到与逐 token recurrence 相同的最终结果。

Prefill 结束时，还需要保留每个 Self-Decoder 层的最终状态：

后续 Decode 从这些状态继续递推。

如果采用 Sliding-Window Attention，则每层保留最后 个 token 的局部 KV。

7.3 第二步：建立唯一的全局 KV Cache

Self-Decoder 顶部输出：

通过共享投影生成：

这里只有一份全局 KV，而不是每个 Cross-Decoder 层一份。

“Cache Once”的准确含义是：

每个 token 只生成并存储一次 global K/V，而不是整个序列只做一次投影，也不是整个模型只有一个状态对象。

7.4 第三步：只让最后位置经过 Cross-Decoder

为了预测第 个 token，最终只需要位置 的 logits。

第 个 Cross-Decoder 层对最后位置计算：

历史位置 (1,\ldots,N-1) 的 Cross-Decoder 输出不需要为未来生成新的 KV，因此不必计算。

7.5 Prefill 的复杂度

标准 Transformer 的 attention Prefill：

YOCO 中：

• Self-Decoder 的 gRet/SWA 关于 近似线性；
• Cross-Decoder 只有一个 Query，但每层需要查询 个全局 KV。

因此论文将 attention 部分写成：

这只是 attention 模块复杂度。FFN、RMSNorm、线性投影等计算并没有消失；YOCO 只是让大量 prompt 位置不再经过上半部分 block。

---

8. 为什么可以跳过绝大多数 Cross-Decoder Prefill？

这是 YOCO 最关键、也最容易被误解的地方。

准确说法不是“Cross-Decoder 没有用”，而是：

Cross-Decoder 是只读历史记忆的模块。它从当前位置产生 Query，读取 Self-Decoder 已经生成的共享 KV，但不会为历史 token 再生成供未来使用的分层 KV。

8.1 普通 Transformer 为什么不能跳过历史位置？

假设 prompt 有 4 个 token，我们需要使用最后位置预测第 5 个 token。

普通 Transformer 第 1 层得到：

第 2 层的 KV 来自所有位置的第 1 层表示：

所以为了计算 ，必须先得到：

因为这些表示分别负责产生：

到了第 3 层，又必须计算所有位置的第 2 层输出，以生成第 3 层的历史 KV。

依赖关系是：

第 1 层：H₁¹  H₂¹  H₃¹  H₄¹
          │    │    │    │
          └────┴────┴────┘
                    ↓ 生成第 2 层的 K/V
第 2 层：H₁²  H₂²  H₃²  H₄²
          │    │    │    │
          └────┴────┴────┘
                    ↓ 生成第 3 层的 K/V
第 3 层：              H₄³ → 预测 token 5

虽然最终只需要 ，它却依赖所有历史位置在前面各层产生的 KV。

所以普通 Transformer Prefill 必须完成一个 (N\times L) 的计算矩形：

             Token 1  Token 2  Token 3  Token 4
Layer 1         ✓        ✓        ✓        ✓
Layer 2         ✓        ✓        ✓        ✓
Layer 3         ✓        ✓        ✓        ✓

严格来说，推理系统一般缓存的不是 本身，而是由它投影出来的：

但为了得到这些 KV，必须先计算 。所以“缓存隐藏状态”只是便于理解的说法，真正被缓存的是逐层 KV。

8.2 YOCO 如何切断这个依赖？

YOCO 先让 4 个 token 经过 Self-Decoder：

只生成一份共享 KV：

形成：

Cross-Decoder 第 1 层只需要用 产生 Query：

然后查询共享 KV：

这里不需要计算：

因为第 2 个 Cross-Decoder 层不会从这些历史位置的 生成新的 KV。

第 2 层只使用当前位置的上一层输出：

但 K/V 仍然是同一份：

依赖路径变成：

M₁ ─┐
M₂ ─┼──→ 唯一全局 K̂₁:₄、V̂₁:₄
M₃ ─┤                │
M₄ ─┘                ├──→ Cross Layer 1
 │                   │        ↑
 │                   ├────────┘
 └──→ 当前 Query ────┘
              │
              ▼
             H₄¹
              │ 产生新的 Query
              ▼
        Cross Layer 2
              ↑
              └──── 再读同一份 K̂₁:₄、V̂₁:₄
              │
              ▼
             H₄² → 预测 token 5

Cross-Decoder 只需要计算最后位置的一条“竖线”：

             Token 1  Token 2  Token 3  Token 4
Cross Layer 1                              ✓
Cross Layer 2                              ✓
Cross Layer 3                              ✓

8.3 用依赖公式概括

普通 Transformer：

为了计算 ，必须拥有全部：

YOCO Cross-Decoder：

共享 KV 在进入 Cross-Decoder 前已经完整准备好。因此到了上半模型，只依赖：

• 当前位置的 ；
• 固定不变的共享 (\hat K,\hat V)。

不依赖其他位置的 Cross-Decoder 输出。

8.4 Cross-Decoder 的输出是否完全不影响未来？

它会影响未来，但影响方式不是保存连续隐藏状态。

当前位置 的最终 Cross-Decoder 输出：

决定下一个 token：

下一轮把生成的 token 重新变为 embedding，再送入 Self-Decoder：

Hₙᴸ
 ↓
LM Head
 ↓
生成 token xₙ₊₁
 ↓
Embedding(xₙ₊₁)
 ↓
Self-Decoder
 ↓
Mₙ₊₁
 ↓
追加 K̂ₙ₊₁、V̂ₙ₊₁

因此 Cross-Decoder 通过“生成出的 token ID”影响未来，但不会把：

作为历史连续状态缓存到下一步。

8.5 Early Exit 的适用边界

Early exit 只适用于自回归推理中“只需要预测下一个 token”的场景。

以下情况不能跳过全部历史位置的 Cross-Decoder：

• 训练时需要每个位置的 next-token loss；
• 用户需要所有 prompt token 的最终 logits；
• 下游任务需要每个 prompt 位置的顶层 embedding；
• 某种算法显式消费所有历史位置的 Cross-Decoder 表示。

所以 early exit 的准确含义是：

---

9. Cross-Decoder 的不同层能不能并行？

不能。Cross-Decoder 的层间仍然存在严格的串行依赖。

对于当前位置 ：

下一层的 Query 是：

必须先得到 (H_t^{})，才能执行第 层。

当前位置 Mₜ
    │
    ▼
Cross Layer 1
生成 Qₜ¹ → 读取共享 KV → Attention → FFN
    │
    ▼ Hₜ¹
Cross Layer 2
生成 Qₜ² → 读取共享 KV → Attention → FFN
    │
    ▼ Hₜ²
Cross Layer 3
生成 Qₜ³ → 读取共享 KV → Attention → FFN
    │
    ▼
最终输出

YOCO 消除的是历史 token 在上半模型中的位置依赖，不是当前 token 的层间依赖。

依赖关系	普通 Transformer	YOCO Cross-Decoder
当前 token 的层间依赖	有	有
历史 token 是否经过所有上层	需要	不需要
每层是否生成独立 KV	是	否
是否共享同一全局 KV	否	是
Cross 层能否直接并行	否	否

“共享 KV”主要减少 KV 构建、存储和状态管理，并不会自动让 Cross-Decoder 层并行。

---

10. YOCO Decode：完整数据流与计算过程

假设已经处理到位置 ，当前输入 token 是 ，这一轮需要预测 。

10.1 Decode 架构图

flowchart TB XT["当前 token xₜ"] EMB["Embedding Xₜ"] OLDSTATE[("上一步 Self 状态<br/>Sₜ₋₁¹ ... Sₜ₋₁ᴸᐟ²")] OLDKV[("已有全局 KV<br/>K̂₁:ₜ₋₁、V̂₁:ₜ₋₁")] subgraph SELF["Self-Decoder：只处理一个 token"] direction TB R1["第 1 层<br/>更新 Sₜ¹，产生 hₜ¹"] RF1["FFN"] RDOT["⋮"] RL["第 L/2 层<br/>更新 Sₜᴸᐟ²，产生 Mₜ"] R1 --> RF1 --> RDOT --> RL end NEWSTATE[("保存新状态<br/>Sₜ¹ ... Sₜᴸᐟ²")] KNEW["K̂ₜ = LN(Mₜ)W_K"] VNEW["V̂ₜ = LN(Mₜ)W_V"] NEWKV[("更新后的全局 KV<br/>K̂₁:ₜ、V̂₁:ₜ")] subgraph CROSS["Cross-Decoder：当前位置逐层执行"] direction TB CQ1["Qₜ¹ = LN(Mₜ)W_Q¹"] CA1["Attention 1<br/>扫描全部 K̂₁:ₜ、V̂₁:ₜ"] CF1["FFN → 当前隐藏状态 hₜ²"] CDOT["⋮"] CQL["生成 Qₜᴸ"] CAL["Attention L/2<br/>再次扫描相同全局 KV"] CFL["FFN"] CQ1 --> CA1 --> CF1 --> CDOT --> CQL --> CAL --> CFL end HEAD["LM Head"] SAMPLE["采样 xₜ₊₁"] XT --> EMB --> R1 OLDSTATE -->|"读取并更新"| R1 RL --> NEWSTATE RL --> KNEW RL --> VNEW OLDKV --> NEWKV KNEW -->|"append"| NEWKV VNEW -->|"append"| NEWKV RL --> CQ1 NEWKV -->|"第 1 次读取"| CA1 NEWKV -->|"第 L/2 次读取"| CAL CFL --> HEAD --> SAMPLE SAMPLE -. "下一轮 Decode" .-> XT

10.2 第一步：Self-Decoder 递归更新

Decode 时不再重新处理全部历史 token，只输入当前 token 。

以 gated retention 为例，第 层执行：

每一层：

• 读取旧状态 (S_{t-1}^{})；
• 写入新状态 (S_t^{})；
• 不读取所有历史 token 的 KV；
• 与上下文长度有关的状态容量保持不变。

经过前半模型后得到：

10.3 第二步：向全局 KV 追加一个位置

只执行一次共享 KV 投影：

然后 append：

这一轮只新增一个 token 的 KV，不会为每个 Cross-Decoder 层分别追加一份。

10.4 第三步：Cross-Decoder 逐层扫描共享 KV

第一层用 产生 Query：

查询：

得到新的当前位置隐藏状态。

第二层根据第一层输出生成新的 Query：

但仍然查询相同的：

一直重复到最后一层。

所以 Cross-Decoder 的性质是：

Query：每层不同
K/V：所有层相同
当前隐藏状态：逐层演化

10.5 Decode 完成后究竟保存什么？

长期状态为：

不需要保存：

• 每个 Cross-Decoder 层自己的 KV；
• 历史 token 的 Cross-Decoder 隐藏状态；
• 历史 token 在 Cross-Decoder 各层产生的 Query。

状态	标准 Transformer	YOCO
每层历史 KV	份	无
共享全局 KV	无	1 份
Self recurrent state	无	个固定状态
随 增长的主要状态

---

11. Prefill 与 Decode 的本质差异

阶段	Self-Decoder	全局 KV	Cross-Decoder
Prefill	处理全部 个 prompt token	一次性建立 对 KV	仅最后位置需要完整执行
Decode	每轮只处理一个 token	每轮 append 一对 KV	每层用一个 Query 扫描全部历史 KV
主要优势	线性/近似线性 Prefill	只存一份	不需要上层 per-layer KV
主要瓶颈	recurrent state 更新或窗口计算	长度随 增长	同一份 KV 被多个层重复读取

最关键的系统结论是：

一次 Decode 的全局 KV 读取量近似为：

也就是说：

• KV 只写一份；
• HBM 中只存一份；
• 但后 个 Cross 层会依次读取多次。

---

12. YOCO 究竟减少了哪些计算？

不能简单概括为“减少 KV Cache 和 Cross-Decoder 的串行计算”。准确来说，YOCO 有三类收益，但没有消除 Cross-Decoder 的层间串行性。

12.1 第一类收益：减少 KV Cache 容量和写入

标准 Transformer：

YOCO：

因此不仅减少容量，还减少：

• 每层 KV projection；
• 每层新 KV 的写入；
• KV allocation 和 page-table 管理；
• 分布式环境中重复的 per-layer KV 通信或搬运。

12.2 第二类收益：Prefill 跳过上半模型的大量位置

标准 Transformer 的 Prefill 计算网格：

                Prompt token
             1   2   3   ...   N
Layer 1      ✓   ✓   ✓         ✓
Layer 2      ✓   ✓   ✓         ✓
Layer 3      ✓   ✓   ✓         ✓
  ...
Layer L      ✓   ✓   ✓         ✓

YOCO：

                Prompt token
             1   2   3   ...   N
Self 1       ✓   ✓   ✓         ✓
Self 2       ✓   ✓   ✓         ✓
  ...
Self L/2     ✓   ✓   ✓         ✓
----------------------------------
Cross 1                        ✓
Cross 2                        ✓
  ...
Cross L/2                      ✓

跳过的不只是 Cross-Attention，还包括这些历史位置在上半 block 中的：

• Query projection；
• Attention；
• Output projection；
• RMSNorm；
• FFN；
• residual 等计算。

Prefill 的主要加速来源是：

1. 下半模型使用 gRet/SWA，避免所有层的全局 attention；
2. 上半模型只计算最后位置，而不是全部 个位置。

12.3 第三类收益：Decode 中只有后半模型扫描长 KV

标准 Transformer 每生成一个 token：

Layer 1：扫描自己的 N 个 KV
Layer 2：扫描自己的 N 个 KV
...
Layer L：扫描自己的 N 个 KV

YOCO：

前 L/2 层 Self-Decoder：
    更新固定大小 recurrent state
    或访问固定窗口
    不扫描全部 N 个全局历史 KV

后 L/2 层 Cross-Decoder：
    扫描共享的 N 个全局 KV
    仍然逐层串行

因此：

• 扫描长上下文的全局 attention 层数从 减为约 ；
• K/V projection 从每层一次减少为边界处一次；
• 前半模型从 历史扫描变成固定状态或固定窗口；
• 所有 个 block 的 FFN 在逐 token Decode 时仍需要执行；
• 后 个 Cross 层仍然串行；
• 共享全局 KV 仍会被扫描 次。

标准 Transformer 的单步 Decode 可以近似写为：

YOCO：

所以 YOCO 的 Decode 仍然会随上下文长度线性增长：

只是需要扫描全局长上下文的层数减少了。

12.4 YOCO 没有减少的部分

YOCO 没有消除：

• 模型总层数；
• 当前 token 必须依次经过所有层的串行依赖；
• 每层的 FFN；
• 每个 Cross 层自己的 Query projection；
• Cross-Decoder 对长全局 KV 的重复扫描；
• 单步 Decode 关于上下文长度 的线性依赖。

最准确的总结是：

---

13. Prefill 和 Decode 对 Cross-Decoder 的计算量是否相同？

基本正确，但必须区分“纯 Prefill”和“生成第一个 token”。

YOCO 的实际过程可以分为：

Prompt：x₁ ... xₙ
        │
        ▼
① Bulk Self-Prefill
   N 个 token 经过 Self-Decoder
        │
        ▼
② 建立全局 KV₁:ₙ
        │
        ▼
③ First-token Cross Pass
   只有 Mₙ 经过 Cross-Decoder
   扫描 KV₁:ₙ
        │
        ▼
   生成 xₙ₊₁
        │
        ▼
④ 普通 Decode
   每轮一个 token：
   Self recurrence → append KV → Cross Pass

13.1 严格定义的纯 Prefill

如果将 Prefill 定义为：

处理 prompt，并建立后续生成所需的状态和 KV Cache。

那么 Cross-Decoder 计算量是：

因为所有 prompt token 只需要经过 Self-Decoder，并生成共享 KV。

13.2 如果将第一个 token 的 logits 计算计入 TTFT

实际服务中的 TTFT 需要得到第一个输出 token，因此最后位置 要经过一次完整 Cross-Decoder：

其中：

这里的 ((1,t)) 表示：

• Query 长度为 1；
• KV 长度为 。

13.3 与一次 Decode 对比

生成到位置 时，一次 Decode 的 Cross-Decoder 同样是：

因此：

它们的计算形态相同：

一个当前位置 Query
        │
Cross Layer 1 ──扫描全部 global KV
        │
Cross Layer 2 ──再次扫描全部 global KV
        │
       ...
        │
Cross Layer L/2 ──再次扫描全部 global KV

后续 Decode 的上下文会逐渐增长：

因此后续每轮会稍微更贵。

13.4 整个 Prefill 不等于一次 Decode

整个 Prefill 还有大规模 Self-Decoder 阶段：

一次 Decode：

因此准确结论是：

但：

从系统角度，YOCO 的推理最好拆成三阶段：

Self-Prefill
    ↓
Global-KV Finalize
    ↓
Cross-Decode（包括首 token 与后续 token）

• Self-Prefill：处理大量 token，使用 chunkwise gRet，较适合矩阵并行；
• Global-KV Finalize：批量生成并写入一份全局 KV；
• Cross-Decode：单 Query 扫描长 KV，形态与普通 Decode 类似，容易受到 HBM 带宽限制。

---

14. 为什么说 YOCO 的计算拓扑像 Encoder-Decoder？

说 YOCO 与 encoder-decoder“本质完全一样”不够严谨。准确说法是：

YOCO 与 encoder-decoder 的 Memory Writer/Memory Reader 计算拓扑相似，但模型语义、因果性和缓存方式不同。

二者共有的基本结构是：

14.1 普通 Encoder-Decoder

以机器翻译为例：

源文本："I love you"
      │
      ▼
   Encoder
      │
      ▼
Encoder Memory：M_source
      │
      ├─────────────┐
      ▼             ▼
Decoder Layer 1   Decoder Layer 2 ...
      │             │
      └── Cross-Attention
              │
              ▼
        生成目标 token

Encoder 先得到：

Decoder 生成目标 token 时，用当前隐藏状态产生 Query，再查询 Encoder Memory：

Encoder 是记忆生成器，Decoder 是记忆读取器。源 token 不需要逐个经过 Decoder。

14.2 YOCO 的对应关系

YOCO 的 Self-Decoder 相当于 Memory Writer：

建立共享记忆：

Cross-Decoder 相当于 Memory Reader：

结构对应关系如下：

Encoder-Decoder：

Source tokens
      │
   Encoder
      │
Encoder Memory
      │
   Decoder ← Target token
      │
   Next token


YOCO：

所有历史 tokens
      │
Causal Self-Decoder
      │
共享 Global KV Memory
      │
Cross-Decoder ← 当前 token
      │
   Next token

因此二者共同点是：下半模块把一组 token 编码成 memory，上半模块不重新编码这些历史 token，而是让当前 Query 读取 memory。

这也是 YOCO 可以跳过历史位置 Cross-Decoder 计算的直观原因：类似于 encoder-decoder 中，源 token 只经过 Encoder，不需要逐个经过 Decoder。

14.3 区别一：是否存在两个独立序列

传统 encoder-decoder 有两个序列：

• Source：输入序列；
• Target：生成序列。

建模目标为：

YOCO 只有一条由 prompt 和 generated token 组成的连续序列：

所以从外部任务接口看，它仍是 decoder-only language model。

14.4 区别二：传统 Encoder 通常是双向的，YOCO Self-Decoder 是因果的

普通 Encoder 的位置 可以看到未来位置：

如果在末尾添加新 token，旧位置的表示可能发生变化：

因此双向 Encoder 通常不能通过简单 append 支持无限自回归生成。

YOCO Self-Decoder 使用 causal mask：

追加 后，旧的：

都保持有效，只需计算：

并追加：

14.5 区别三：Encoder Memory 通常固定，YOCO Memory 会增长

在翻译等传统 encoder-decoder 任务中，Encoder Memory 在生成阶段通常保持不变。

YOCO 的 memory 会随着生成增长：

第 N 步：     KV₁:ₙ
第 N+1 步：   KV₁:ₙ₊₁
第 N+2 步：   KV₁:ₙ₊₂

每个生成 token 都会先经过 Self-Decoder，然后写入全局记忆：

因此 YOCO 的 memory 是 growing causal memory。

14.6 区别四：普通 Encoder-Decoder 通常仍有 per-layer KV 投影

标准 encoder-decoder 的不同 Decoder 层虽然读取同一份 Encoder 输出 (M_{\text{src}})，但通常使用各自的投影：

因此系统可能仍然缓存：

YOCO 更进一步：

所有 Cross-Decoder 层共享完全相同的 (\hat K,\hat V)，只有 Query 是每层独立的：

普通 Encoder-Decoder：
同一个 Encoder 表示
    → Decoder Layer 1 投影 K¹/V¹
    → Decoder Layer 2 投影 K²/V²
    → Decoder Layer 3 投影 K³/V³

YOCO：
Self-Decoder 边界表示
    → 只投影一次 K̂/V̂
    → 所有 Cross 层共享

14.7 为什么叫 Decoder-Decoder？

Self-Decoder 虽然承担类似 Encoder 的“建立 memory”作用，但它必须：

• 保持 causal；
• 支持增量计算；
• 保证旧表示不会因新 token 到来而改变；
• 支持自回归生成；
• 维护 recurrent state 或滑窗缓存。

这些都是 Decoder 的性质。

因此 YOCO 可以概括为：

也就是：结构上像 encoder-decoder，行为上像 decoder-only，并额外要求所有 Cross 层共享同一份 projected KV。

---

15. YOCO 论文里的两种 Chunk

论文中出现两个不同层面的 chunk，需要区分。

15.1 Chunkwise Gated Retention

这是 linear attention/retention 意义上的计算分块：

• chunk 内用矩阵运算并行计算；
• chunk 间只传递 recurrent state；
• 主要目的是降低循环次数并利用 Tensor Core。

它属于算子级计算边界。

15.2 Chunk Parallelism

附录 A 将超长序列切成多个 chunk，分配到不同 GPU：

GPU 1：token 1 ... B
GPU 2：token B+1 ... 2B
GPU 3：token 2B+1 ... 3B

Self-Decoder 中：

• gated retention 只需向相邻设备传递边界 recurrent state；
• sliding-window attention 只需处理窗口边界通信。

Self-Decoder 输出全局 KV 后，Cross-Decoder 只进行一次 KV all-gather，因为所有 Cross 层共享同一份 KV，而不需要每层重新构造和通信一次。

因此两者的区别是：

• Chunkwise retention：算子计算边界；
• Chunk parallelism：分布式序列切分边界。

它们可以采用相同的序列切分，但优化目标不同。

---

16. 论文实验结果

16.1 语言模型质量

YOCO-3B 的主要配置：

• 26 层；
• hidden size 3072；
• 24 个 Query head；
• 8 个 KV head；
• 使用 grouped-query attention；
• Self-Decoder 使用 gated retention；
• 训练到 1.6T token。

论文报告的 Eval Harness 八任务平均分：

模型	平均分
OpenLLaMA-3B-v2，1T token	0.619
StableLM-base-alpha-3B-v2，1T token	0.612
YOCO-3B，1T token	0.634
YOCO-3B，1.6T token	0.636
YOCO-3B-1M	0.645

这些结果说明 YOCO 没有因为共享 KV 而出现明显的短任务性能下降。

不过这些模型的训练数据、tokenizer、参数量和训练细节并非完全一致，因此不能把这张表当作严格、完全控制变量的架构消融。

16.2 模型规模扩展

论文训练了 160M、400M、830M、1.4B、2.7B、6.8B 和 13B 等不同规模模型，并拟合 validation loss scaling curve。

YOCO-gRet、YOCO-SWA 与 Llama-style Transformer 的 scaling 趋势接近，YOCO-gRet 略优。

但每个规模只训练了 10B token。对于 13B 模型而言，这远低于充分训练所需 token，因此该实验主要证明架构可稳定扩展，而不是证明 13B YOCO 的最终能力已经优于成熟 Transformer。

16.3 1M 上下文扩展

YOCO-3B 逐步扩展上下文：

训练长度	训练 token	RoPE theta
64K	6B	640K
256K	4B	5M
1M	1.5B	80M

在 1M Needle-in-a-Haystack 测试中，YOCO 获得接近满分的单针检索结果。论文对相同深度和长度组合运行 10 次，并报告平均准确率。

128K 多针测试：

Needle 数量	YOCO-3B-1M 准确率
1	0.98
2	0.98
4	0.84
8	0.56

结果表明 YOCO 保留了全局检索能力，但多目标检索仍明显下降。单针 Needle 测试本身也较简单，不能等价为复杂的 1M-token 组合推理能力。

论文还报告了 book 和 repository-level code 数据上的累计平均 NLL，随着可用上下文增长而下降，说明模型能够从更长历史中获得一定语言建模收益。

16.4 推理效率

推理实验使用 H100-80GB。Transformer baseline 已加入：

• grouped-query attention；
• Flash-Decoding；
• kernel fusion。

3B 模型在 1M context 下：

• Transformer：约 116GB；
• YOCO：12.4GB；
• 总体推理显存降低约 9.38 倍。

512K Prefill：

• Transformer：约 180 秒；
• YOCO：小于 6 秒；
• 加速约 30.3 倍。

1M Prefill 加速约 71.82 倍。

512K 综合吞吐：

• Transformer：4.5 token/s；
• YOCO：43.1 token/s；
• 约 9.56 倍。

需要注意，这里的 throughput 综合了长 prompt Prefill 和随后 1024-token generation，还包含 YOCO 因显存更小而能够使用更大 batch 的收益，不是相同 batch 下的纯 Decode kernel 对比。

---

17. 论文的主要局限

17.1 不是现有 Transformer 的无损推理优化

YOCO 改变了模型结构和依赖图。不能直接把现有 Llama 每层 KV 删除后变成 YOCO，通常需要重新训练或进行专门的模型转换与适配。

这与 KV quantization、PagedAttention、prefix caching 或 cache eviction 不同。

17.2 减少容量不等于同比例减少读取流量

尽管只保存一份全局 KV，Cross-Decoder 的每一层仍要查询它：

因此：

• 存储容量从 下降到 ；
• 但 Decode 仍可能逐层扫描长全局 KV；
• 单 token Decode 的全局 attention 仍随 线性增长。

YOCO 解决了大部分“放不下”问题，但没有彻底解决“每轮读多少数据”的问题。

17.3 边界表示可能成为信息瓶颈

所有 Cross-Decoder 层都依赖 Self-Decoder 顶部产生的同一份 KV。

Cross-Decoder 对历史 token 做出的上层语义修正不会重新写入历史 memory。模型必须学会让中间表示 成为对所有 Cross 层都足够通用的记忆表示。

17.4 长上下文证据仍以检索为主

论文的 1M 上下文证据主要包括：

• 单 Needle；
• 多 Needle；
• book/code 长序列 NLL。

缺少更复杂的跨段组合推理、真实多轮对话、代理轨迹和长上下文 serving workload。

17.5 Prefill 复杂度表主要针对 Attention

论文将 Prefill 从 降到 ，主要描述 attention 部分。

FFN、线性层、归一化和其他算子仍然存在；模型总时间并不会严格按照 attention 渐进复杂度同比例下降。

---

18. 对 Hybrid-aware Continuous Batching 的启发

YOCO 同时具有两类性质完全不同的请求状态。

18.1 固定大小、频繁读写的 Self State

对于 gRet：

具有以下特点：

• 大小不随上下文增长；
• 每生成一个 token 都要读取和写回；
• 访问频率高；
• 计算通常包含矩阵状态更新。

18.2 随上下文增长、追加一次、反复读取的 Global KV

全局 KV：

具有以下特点：

• 容量随 线性增长；
• 每个 token 只追加一次；
• 每个 Cross 层都会读取；
• 长上下文下容易成为 HBM bandwidth bottleneck。

因此只判断“请求状态是否装得下”是不够的。一个 YOCO 请求可能只占较少显存，但因为 Cross-Decoder 要多次读取很长的 global KV，单步时间仍然很长。

可以建立近似模型：

HBM 流量近似为：

这正好支持 State-Traffic-Aware Continuous Batching 的核心动机：

Memory footprint 小不代表执行时间短；YOCO 只缓存一次，但同一份缓存可能被读很多次。

18.3 推理调度可以进一步细分阶段

传统 serving engine 通常只区分 Prefill 和 Decode。YOCO 更适合拆成：

1. Self-Prefill；
2. Global-KV Finalize；
3. Cross-Decode，包括 first-token Cross Pass；
4. 后续循环的 Self recurrent update + KV append + Cross Pass。

不同阶段具有不同资源特征：

阶段	Query 数量	KV 长度	主要特征
Self-Prefill		recurrent/chunkwise	计算并行度高
Global-KV Finalize		写入 对 KV	projection 与 HBM write
First-token Cross Pass	1		Decode-like、带宽敏感
后续 Cross Decode	1		逐步变长、带宽敏感

潜在系统方向包括：

• 按 global KV 长度和 recurrent-state 大小联合组 batch；
• 分别管理 append-once/read-many 的全局 KV 与 read-write-every-step 的 recurrent state；
• 在 Prefill/Decode 解耦架构中，只传输一份 global KV 和固定大小 Self state；
• 对共享 KV 只建立一次压缩、量化或检索索引，供所有 Cross 层复用；
• 将 TTFT 预测拆为 Self-Prefill、KV finalize 和 first-token Cross scan；
• 针对长短请求混合，预测各请求的 Cross-Decoder HBM 流量，而不只预测 cache allocation；
• 为 recurrent state 和 global KV 设计不同的驻留、换入换出与优先级策略。

---

19. 最终总结

YOCO 最核心的创新不是某一个 kernel，而是改变模型的历史依赖图：

1. 前半 Self-Decoder 负责以 recurrent state 或局部窗口传播历史；
2. Self-Decoder 顶部为每个 token 写入一份共享全局 KV；
3. 后半 Cross-Decoder 的不同层只产生不同 Query；
4. 所有 Cross 层读取同一份 KV；
5. 历史 prompt 位置无需经过 Cross-Decoder；
6. 只有当前位置需要依次经过所有 Cross 层；
7. Cross 层之间仍然串行，因为下一层 Query 依赖上一层输出；
8. Prefill 的 first-token Cross Pass 与一次普通 Decode Cross Pass 具有相同计算形态；
9. 整个 Prefill 仍包含处理全部 prompt 的 Self-Decoder 和 global KV 建立过程；
10. KV 容量显著降低，但长 global KV 被 Cross 层重复扫描的读取流量仍然存在。

可以用一句话概括：

因此它同时具备：

• 类似 encoder-decoder 的 Memory Writer/Reader 拓扑；
• decoder-only 的单序列 causal 行为；
• linear/recurrent Self-Decoder 的固定状态；
• global softmax Cross-Attention 的精确 token-level 检索。

对推理系统而言，它把主要矛盾从“每层 KV Cache 都装不下”，进一步转化为：

一份更容易装下的 global KV，如何在多个串行 Cross 层之间被高效、可预测地重复读取。