✨ 要点🔬 技术摘要
想象你是一位厨师,正在为一系列客人准备一顿复杂的多道菜大餐。在大语言模型(LLM)的世界里,“大餐”是生成回复的过程,而“食材”则是模型已经处理过的单词(token)。
旧方法:“全有或全无”的厨房
传统上,当一位新客人(一个新请求)到来时,厨师会检查他们点的菜是否与上一位客人相似。
如果点了完全相同的前菜 :厨师会重用整盘菜。
如果点了略有不同的菜 :厨师会扔掉整盘前菜,从头开始烹饪,即使前 90% 的食材完全相同。
用技术术语来说,这被称为密集缓存(dense caching) 。系统保存每一步(每个 token)的副本以供日后重用。这对标准模型效果很好,但对于一种名为**混合模型或循环模型(Hybrid or Recurrent Model)**的新型模型,这种方法就像为了读一句话而携带整个图书馆的书籍。它太沉重,占用内存过多。
新想法:“检查点”策略
这篇论文提出了一种更聪明的方法来处理这些特定模型。将模型的内存想象成一种思维状态 ,而不是存储每个单词的图书馆。
想象你正在阅读一部长篇小说。
旧方法 :你在每一页 上都贴一张便利贴,以便能瞬间跳回。(便利贴太多了!)
新方法(稀疏前缀缓存) :你只在第 1 页、第 100 页、第 200 页 等位置贴便利贴。
如果一位新读者想从第 150 页 继续读故事:
你不会扔掉整本书。
你找到最后一张便利贴(第 100 页)。
你快速重读第 101 页到 149 页的故事,以回到当前状态。
然后,从第 150 页继续。
因为模型是“循环”的(它逐步演变其状态),所以它不需要整个历史,只需要特定时间点的状态。这篇论文将这些便利贴称为检查点(checkpoints) 。
问题:把便利贴贴在哪里?
现在到了棘手的部分。你的便利贴预算(内存)有限。你应该把它们放在哪里才能最节省时间?
“平衡”策略 :均匀地放置便利贴(每 100 页一张)。这很稳妥,但可能不是最快的。
“智能”策略(本文所做的) :观察读者的习惯。
如果大多数人在第 50 页左右停止阅读,你就在那里贴一张。
如果人们通常一直读到结尾,你就在结尾附近贴。
如果人们经常在第 200 页停止,你就在那里贴一张。
作者创建了一个数学公式(一种“动态规划”),它就像一位超级智能的图书管理员 。它分析过去的请求,预测未来的读者最可能在哪里停止。然后,它将便利贴精确地放置在最有用之处,而不是均匀分布。
结果:节省时间和内存
这篇论文在现实场景中测试了这种方法,例如:
QuALITY :一篇长文档,人们针对同一文本提出不同的问题。
系统提示词(System Prompts) :一套长指令,后跟许多不同的用户问题。
他们的发现:
更少内存,同等速度 :通过根据人们实际停止的位置“智能”地放置检查点,他们可以使用比标准“均匀间隔”方法更少的便利贴(检查点) ,同时仍能节省相同的烹饪时间。
短预算下的巨大优势 :最大的改进发生在他们只有极少量便利贴可用的时候。在这些紧张的情况下,“智能”放置比随意猜测或均匀间隔要好得多。
精确结果 :与某些猜测答案的捷径不同,这种方法保证输出与从头开始工作100% 完全相同 。它只是通过跳过已知部分来加快速度。
总结
这篇论文提出了一种方法,使使用“循环”内存的 AI 模型更加高效。它不是保存每一步,也不是什么都不保存,而是保存模型大脑的几个战略性“快照”。通过利用数学,根据人们实际使用 AI 的方式,精确计算出应该在哪里保存这些快照,系统可以运行得更快并使用更少的内存,特别是在许多用户针对同一份长文档提出相似问题时。
这就像拥有一个 GPS,它不仅显示整张地图,还知道你最可能走哪些转弯,因此它只保存这些特定转弯的路线。
技术摘要:面向混合与循环大语言模型服务的稀疏前缀缓存
1. 问题陈述
现有的大语言模型(LLM)服务系统严重依赖密集前缀缓存 ,即存储共享前缀中每个 token 的 Key/Value(KV)状态,以避免冗余的预填充计算。这种方法在 Transformer 架构中是标准做法。然而,状态空间模型(SSMs)和 混合架构 (结合注意力机制与循环层)呈现出结构性不对称:循环层通过 h t = F ( h t − 1 , x t ) h_t = F(h_{t-1}, x_t) h t = F ( h t − 1 , x t ) 演化其状态,这意味着它可以仅从单个存储状态 h c h_c h c 恢复计算,而无需整个中间状态历史。
对循环状态进行每个 token 的密集缓存通常是不必要的,且代价高昂。对于 Qwen-3.5 等模型,每个头的循环状态规模按 d h e a d 2 d_{head}^2 d h e a d 2 增长,而注意力 KV 存储则呈线性增长。因此,存储每个 token 的循环状态所消耗的内存远多于注意力 KV 缓存,但相比存储稀疏的检查点集合,并未带来额外收益。
本文解决的核心问题是:给定长度为 N N N 的前缀和 M M M 个检查点槽位的内存预算,在特定的未来请求重叠深度分布下,应将检查点放置在何处以最小化预期的重计算成本? 该优化问题不同于缓存准入/淘汰策略;它专注于在单个保留条目内哪些 token 应被设为检查点,以最大化部分前缀重叠的效用。
2. 方法论
理论表述
作者将该问题形式化为直线上的单边加权 k-中值问题 。
设 C = { c 1 , … , c M } C = \{c_1, \dots, c_M\} C = { c 1 , … , c M } 为检查点位置集合。
对于重叠前缀至深度 t t t 的请求,可复用的最深检查点为 ℓ ( t ; C ) = max ( { 0 } ∪ { c ∈ C : c ≤ t } ) \ell(t; C) = \max(\{0\} \cup \{c \in C : c \le t\}) ℓ ( t ; C ) = max ({ 0 } ∪ { c ∈ C : c ≤ t }) 。
成本为需重计算的 token 数量:r ( t ; C ) = t − ℓ ( t ; C ) r(t; C) = t - \ell(t; C) r ( t ; C ) = t − ℓ ( t ; C ) 。
目标是在给定重叠深度分布 T T T (概率为 p t p_t p t )的情况下,最小化预期重计算成本 E [ r ( T ; C ) ] E[r(T; C)] E [ r ( T ; C )] 。
算法方法
最优平衡放置(均匀/最坏情况): 作者证明,在均匀重叠分布或最坏情况下,平衡间距 (均匀分布的检查点)是最优的。
分布感知动态规划(DP): 针对非均匀的现实世界重叠分布,作者推导出了精确的 $O(NM)$ 动态规划 算法。
递推关系为:d p [ m , j ] = min 1 ≤ s ≤ j ( d p [ m − 1 , s − 1 ] + w ( s , j ) ) dp[m, j] = \min_{1 \le s \le j} (dp[m-1, s-1] + w(s, j)) d p [ m , j ] = min 1 ≤ s ≤ j ( d p [ m − 1 , s − 1 ] + w ( s , j )) ,其中 w ( s , j ) w(s, j) w ( s , j ) 是将 s s s 用作深度 s … j s \dots j s … j 的最深检查点的成本。
该问题利用前缀和的单调性,通过单调凸包技巧 高效求解。
处理分布漂移: 由于真实的重叠分布未知且必须从历史数据中估计,本文提供了理论保证(通过引理 1 和定理 3),即价值函数关于重叠分布是利普希茨连续 的。这确保了使用经验直方图(或针对漂移分布的指数加权直方图)可产生具有有界次优性的近优调度。
实施策略
检查点设置: 系统在稀疏位置存储精确的循环状态。发生缓存命中时,系统从存储的最深检查点恢复,并精确重计算后缀。
精确性: 由于循环更新是确定性的,从检查点恢复并重放后缀所产生的输出与完整预填充的结果在比特级完全一致。
混合集成: 对于混合模型,该方法存储完整的注意力 KV 缓存(密集),但仅对循环层应用稀疏检查点。
3. 主要贡献
理论最优性: 证明了在均匀重叠和最坏情况下,平衡间距是最优的(定理 1)。
精确算法: 将稀疏前缀缓存形式化为加权 k-中值问题,并提供了一个精确的 $O(NM)$ 动态规划解法(定理 2)。
鲁棒性保证: 证明了最优调度在分布误设下是稳定的。次优性由经验分布与真实重叠分布之间的总变差距离界定(定理 3),并通过指数重加权扩展至分布漂移场景(定理 4)。
实证验证: 在真实数据集(QuALITY, NarrativeQA, System Prompts)上验证了该方法,显示出相对于固定预算基线的一致改进。
4. 实验结果
数据集与设置
数据集:
QuALITY: 单文档包含多个问题(在接近全长处具有高重叠)。
NarrativeQA: 结构类似,具有广泛的模态重叠。
System Prompts: 真实系统提示与用户查询的组合(重叠集中在较短长度)。
基线: 与“无缓存”、"KV 仅”、“平衡”(均匀间距)、“块”(固定块大小)、“对数”和"L \sqrt{L} L 间距”策略进行比较。
硬件: 原型在 NVIDIA RTX 2080 Super 上运行,使用 Qwen-3.5-0.8B 模型。
发现
Token 节省: DP 最优调度在循环工作减少因子方面,始终优于确定性基线(平衡、块、对数)。
低预算优势: 在低检查点预算 下,性能差距最大,此时重叠分布最不均匀。例如,在 System Prompts(重叠较短)上,DP 方法在分歧峰值附近密集放置检查点,使用更少的检查点即可达到与块缓存相同的运行时。
挂钟时间: 在 QuALITY 和 System Prompts 上,DP 导出的调度在代表性层组上产生了可测量的挂钟时间加速。运行时间的减少与节省的 token 数量呈线性相关。
开销: 从 CPU 加载检查点到 GPU 的开销相对于因减少重计算而获得的节省而言很小,从而实现了净运行时改进。
5. 意义与主张
本文主张,稀疏前缀缓存 解锁了混合与循环 LLM 服务的一个此前不可用的新设计点。
内存/计算权衡: 它提供了内存使用与重计算之间的权衡,这是密集 KV 缓存无法提供的,特别是针对那些状态可以被精确提取和恢复的架构。
精确性: 与近似的 KV 压缩技术不同,该方法在不改变循环计算内核的情况下保留精确输出 。
适用性: 它最有利于请求共享实质性但不完全相同前缀 的工作负载(例如,带有不同问题的 RAG,或不同的系统提示)。对于严格的仅追加聊天工作负载(最后一个状态已足够),它不提供额外收益,但在此类情况下也不会产生开销。
未来工作: 作者承认工程挑战,特别是生产运行时中需要高效的状态提取/恢复 API,以及将此放置策略与准入/淘汰策略(如 Marconi 等人提出的方法)集成。他们还指出,当前的单前缀理论可以扩展到通用前缀树。
每周获取最佳 machine learning 论文。
受到斯坦福、剑桥和法国科学院研究人员的信赖。
请查收邮箱确认订阅。
出了点问题,再试一次?
无垃圾邮件,随时退订。