文章摘要
本文介绍面向大模型Transformer解码阶段的KV缓存优化技术YOCO，它针对大模型Prefill与Decode任务不对称、需从层维度压缩KV缓存两大痛点，通过多层共享KV缓存降低存储与计算开销，衍生出多个优化变种，目前DeepSeek V4.1、Gemma4、WeLM等模型均已借鉴该思路落地相关优化。

随着近期大模型推理优化技术的迭代，YOCO架构再次受到广泛关注。作为一款针对Transformer解码阶段KV缓存优化的经典工作，YOCO不仅本身设计简洁高效，还衍生出了多种实用的改进版本，同时也为后续众多主流模型的优化方案提供了核心思路。

YOCO的设计初衷主要源于两个关键行业痛点：

第一，Prefill与Decode阶段的任务不对称。在大模型的Prefill阶段，核心目标是构建上下文的KV缓存，只有最后一个Prompt Token需要生成初始的logits，前面的Token其实无需单独计算输出。传统的Transformer架构中每个层都需要独立维护自己的KV缓存，很难充分利用这种不对称性；而YOCO通过让后续的Cross-Decoder共享全局的共享KV缓存，使得长Prompt的前置Token在完成共享KV的构建后就可以提前退出计算，极大放大了这种不对称性带来的性能收益。

第二，从Layer维度压缩KV缓存。KV缓存的空间开销通常由三个维度决定：单Token的KV维度、上下文序列长度、模型层数。此前的优化已经从另外两个维度进行了探索：比如通过GQA/MQA减少KV头数、MLA压缩单Token的KV表示维度来优化Entry维度；通过CSA、HCA、Linear等注意力机制减少对超长上下文缓存的依赖来优化Sequence维度；而YOCO这类方法则聚焦于Layer维度，通过让多层共享同一份KV缓存来实现压缩。

YOCO基础架构

YOCO的核心架构非常直观，将标准的仅解码器Transformer拆分为上下两个部分。前半部分为Self-Decoder，部分优化版本甚至会采用门控注意力或者滑动窗口注意力，将整个上下文压缩为一份全局共享的KV缓存。后半部分为Cross-Decoder，每一层不再单独保存历史的K/V值，而是通过自身的Query去交叉注意力读取这份共享的KV缓存。

这种设计将标准Transformer的O(LND)级别的KV缓存开销降低到了O((N+L)D)级别，更关键的是，在Prefill阶段完全不需要运行后半部分的Cross-Decoder，这在超长上下文的Agent场景中尤为实用。

YOCO-U变种

YOCO-U在Self-Decoder部分引入了循环迭代机制，通过对前置的Self-Decoder进行三次循环迭代来生成最终的共享KV缓存；而在Cross-Decoder部分则采用了NOPE（无位置编码）的设计来进一步优化推理效率。

YOCO-Sparse变种

YOCO-Sparse的优化思路更加直接：既然YOCO后半部分的多个Cross-Decoder层本身就在读取同一份KV缓存，那么可以进一步共享Sparse Attention的Routing Index，实现了更进一步的计算和存储节省，堪称极致的优化方案。

DeepSeek-V4.1-Flash的落地实现

作为在KV缓存优化上最为激进的团队之一，DeepSeek的V4.1版本将多种优化方案结合到了一起，其注意力机制整体基于YOCO/CED架构，结合了CSA2和滑动窗口注意力（SWA）技术。

YOCO/CED：减少Prefill计算量

全局KV缓存采用YOCO作为核心框架，将40层的模型拆分为20层因果编码器和20层解码器，也就是CED架构。在Prefill阶段只运行前20层，解码器的全局KV缓存直接由编码器最后一层的输出通过各层独立投影生成，避免了让整个Prompt完整运行后20层的冗余计算。

CSA2：跨层复用全局KV与Top-K检索

在此基础上，全局注意力通过CSA2实现跨层复用，分为三种不同的层类型：

  • Full层：负责生成新的全局KV缓存和Top-K检索集合
  • Reindex层：复用已有的KV缓存，但重新选择Top-K的检索位置
  • Reuse层：直接复用已有的KV缓存和Top-K检索结果

解码器中只有第一个Full层会生成全局KV缓存，后续层共享这份缓存，并每隔几层通过Reindex更新检索位置，最终全局稀疏注意力只会读取Top-512个远程KV缓存，进一步压缩了计算开销。

SWA：短期缓存与局部重建优化

为了进一步节省存储开销，DeepSeek不再将滑动窗口的KV长期存入持久缓存，仅将全局主KV缓存使用FP4量化，而滑动窗口KV保留FP8格式，仅放在短TTL的主机内存池中。如果滑动窗口KV丢失，通常需要重新回放128乘以层数的Token来重建，但DeepSeek V4.1直接只回放最后128个Token，极大压缩了重建所需的计算量。在解码阶段，新生成的Token会正常经过全部40层，同时使用Top-512的全局稀疏KV和128的本地滑动窗口KV。

Gemma 4 E2B / E4B的优化方案

Gemma4系列模型在架构上也采用了类似YOCO的优化思路，后半段的大量层不再计算自己的K/V值，仅保留独立的Query，复用前面最后一个同类型注意力层的KV缓存。具体参数如下：

模型 总层数 后段共享范围 Sliding Attention复用来源 Full / Global Attention复用来源
Gemma 4 E2B 35层 L15–L34 L13 的KV L14 的KV
Gemma 4 E4B 42层 L24–L41 L22 的KV L23 的KV

WeLM的KV-Mirror方案

WeLM也采用了类似YOCO的KV共享操作，但其核心目标不是节省KV缓存，而是减少Prefill阶段的计算量。他们采用了U形的KV共享策略：将前1/3的层镜像到后1/3的层进行共享，不过复用的不是最终的KV缓存，而是KV投影前的隐藏状态，再通过目标层的投影重新计算KV，这样就节省了Prefill阶段后1/3层的完整注意力和MoE等计算。这种设计的灵感来自于KVSharer的研究发现：共享差异更大的层的KV缓存，反而能获得更好的模型性能。

以上内容不代表本平台立场，仅供读者参考