想象你在一个超大的图书馆里找资料。每读一本书,你都要从头到尾重新找一遍相关章节——这就是当前大模型长上下文推理的困境。
现在有人提出了一个简单到让人拍大腿的想法:找一次就够了,为什么每层都要重新找?
长上下文的"三座大山"
现代大模型在长上下文推理时面临三个效率瓶颈:
- 预填充(Prefill)慢:输入越长,第一次处理越慢
- KV 缓存大:每层的键值对都要存,内存吃紧
- 解码慢:每生成一个新 token,都要在全部历史中找相关信息
现有的稀疏注意力方法试图解决第三个问题,但陷入了两难:
- 块稀疏注意力(如 MoBA):结构规整,GPU 友好,加速明显,但粒度太粗,质量损失明显
- Token 稀疏注意力(如 Quest):粒度细,质量好,但每层都要做一次 top-k 路由,路由本身就很慢
核心矛盾在于:路由决策很贵,但每层都在独立做同样的路由决策。
跨层共享路由:一个索引,全模型复用
YOIO(You Only Index Once)的核心思想极其简洁:既然多层读的是同一份 KV 缓存,那路由索引也应该只算一次。
具体来说,YOIO 建立在 YOCO 架构之上。YOCO 把模型分成两部分:
- 自解码器(Self-Decoder):处理输入,生成共享的 KV 缓存
- 交叉解码器(Cross-Decoder):从共享 KV 缓存中检索信息,生成输出
YOIO 在自解码器之上加了一个轻量的"查询感知索引器"(Query-Aware Indexer),一次性计算每个查询 token 应该关注哪些 KV 位置。然后,所有交叉解码器层共享这个索引。
这就像图书馆里配了一个专业索引员,读者不需要每进一个阅览室就重新查一遍目录,直接用索引员给的路标就行。
为什么共享索引不会损失质量?
这是最关键的问题。直觉上,不同层关注的内容应该不同,共享索引会不会"一刀切"?
论文做了详细的注意力模式分析,发现:在共享 KV 缓存的架构下,不同层的路由决策高度一致。 因为它们读的是同一份记忆,需要找的信息源基本相同,只是对信息的处理方式不同(有的层做检索,有的层做变换)。
实验验证了这一点:CLSA 在短上下文和长上下文基准上,与密集注意力基线几乎无损。
效率提升有多猛?
在 128K 上下文长度下:
- 解码吞吐量提升 7.6 倍(相比 Transformer 基线)
- 端到端吞吐量提升 17.1 倍(包括预填充加速)
- KV 缓存占用大幅减少(共享 KV 本身就省空间)
17.1 倍是什么概念?原来需要 17 秒才能完成的推理,现在 1 秒搞定。而且这不是靠牺牲质量换来的——模型在各项基准上的表现几乎不变。
更深层的启示
YOIO 的意义不只是"又一种加速方法"。它揭示了一个重要的架构洞察:在 Transformer 的多层结构中,"去哪里找信息"和"找到信息后怎么处理"是两个独立的决策,前者可以共享,后者保持独立。
这和人类的信息处理方式很像。你在读一篇文章时,眼睛的移动模式(去哪里看)在不同阅读阶段是相似的,但大脑对看到内容的处理方式(怎么理解)则因阶段而异。
YOIO 把"去哪里找"这个昂贵的决策做了一次,让所有层复用;而"怎么处理"这个需要多样性的部分,每层仍然独立。这种分离,恰好抓住了效率和质量的最优平衡点。
一次索引,多层共享。简单,但有效。
📄 论文:You Only Index Once: Cross-Layer Sparse Attention with Shared Routing
💻 相关代码:github.com/THUDM/IndexCache
讨论回复
1 条回复推荐
智谱 GLM-5 已上线
我正在智谱大模型开放平台 BigModel.cn 上打造 AI 应用,智谱新一代旗舰模型 GLM-5 已上线,在推理、代码、智能体综合能力达到开源模型 SOTA 水平。