IndexCache：解决稀疏注意力的隐藏瓶颈

处理 20 万 Token 的大型语言模型既昂贵又缓慢：上下文越长，成本增长越快。清华大学和 Z.ai 的研究团队开发了一种名为 IndexCache 的技术，可以削减稀疏注意力模型中高达 75% 的冗余计算。

DSA 的隐藏问题

大型语言模型依赖自注意力机制，模型计算上下文中每个 Token 与所有先前 Token 的关系来预测下一个 Token。

核心限制：自注意力的计算复杂度随序列长度二次增长。对于需要长上下文窗口的应用（大文档处理、多步代理工作流、长链式推理），这种二次缩放导致推理速度缓慢和显著的内存成本。

稀疏注意力解决方案 不再计算每个 Token 与所有先前 Token 的关系，稀疏注意力优化了这一过程——让每个查询只选择和关注最相关的 Token 子集。

DeepSeek 稀疏注意力（DSA） 在每一层引入轻量级”闪电索引器模块”，对所有先前 Token 评分并选择一小批给主核心注意力机制处理。这将核心注意力计算从二次方削减到线性。

遗留问题 研究团队发现：DSA 索引器本身在每一层仍以二次方复杂度运行。虽然索引器计算量小于主注意力，但随着上下文长度增长，模型运行这些索引器的时间急剧增加。

IndexCache 解决方案

研究团队发现了 DSA 模型处理数据的一个关键特征：索引器选择的重要 Token 子集在连续层之间保持惊人稳定。对 DSA 模型的实证测试显示，相邻层共享 70% 到 100% 的选定 Token。

层分类 IndexCache 将模型层分为两类：

  • 完整层（F）：保留索引器，主动评分 Token 并选择最重要的进行缓存
  • 共享层（S）：不执行索引，复用最近 F 层缓存的索引

推理流程

  • 到达 F 层：计算并缓存新索引
  • 到达 S 层：跳过计算，直接复制缓存数据

与 KV 缓存压缩的区别 IndexCache 不是传统的 KV 缓存压缩或共享技术。它通过跨层复用索引消除冗余，减少计算而非仅内存占用。这与现有方法互补，可以结合使用。

两种部署方式

无训练方法 适用于开箱即用的 DSA 模型：

  • 使用贪婪层选择算法
  • 运行小型校准数据集
  • 自动确定 F 和 S 层的最优配置
  • 无需权重更新

实证表明，贪婪算法可安全移除 75% 的索引器，同时保持原始模型的下游性能。

训练感知方法 适用于预训练或大量微调自有基础模型的团队：

  • 引入多层蒸馏损失
  • 训练时强制每个保留索引器学习选择对所有后续层高度相关的共识 Token 子集

实测性能

GLM-4.7 Flash（300 亿参数）

  • 200K 上下文：prefill 延迟从 19.5 秒降至 10.7 秒（1.82 倍加速）
  • 生成吞吐量：从 58 TPS 提升至 86 TPS（1.48 倍加速）
  • 服务器内存饱和时：总解码吞吐量提升高达 51%

质量保持

  • 移除 75% 索引器后，长上下文基准平均分：49.9 vs 原始 50.2
  • AIME 2025 数学推理：优化模型 92.6 分 vs 原始 91.0 分

GLM-5（7440 亿参数）

  • 100K+ Token 上下文：至少 1.3 倍加速
  • 长上下文任务质量基本保持

企业价值

“在 ROI 方面，IndexCache 在各种场景下提供一致收益，但在长上下文工作负载（如 RAG、文档分析、代理管道）中收益最明显，”论文合著者 Yushi Bai 表示，”在这些情况下，我们观察到部署成本至少降低约 20%，用户感知延迟也有类似改善。”

生产部署

开源补丁已在 GitHub 提供，支持主流推理引擎：

  • 可应用于 vLLM 或 SGLang
  • 最小配置更改即可启用
  • 推荐使用领域特定数据作为校准集

适用模型 IndexCache 仅适用于使用 DSA 架构的模型：

  • DeepSeek 最新模型
  • GLM 最新系列

未来展望

研究团队的哲学指向 AI 行业模型设计方法的更广泛转变：”未来的基础模型可能会从一开始就考虑到下游推理约束来设计。这意味着不仅要考虑模型规模的可扩展性，还要针对真实世界的吞吐量和延迟进行优化，而不是将这些作为事后考虑。”