混合与循环大语言模型服务中的稀疏前缀缓存
摘要
本文针对混合和循环大语言模型提出了稀疏前缀缓存方法,该方法在有限的检查点位置存储循环状态,从而避免密集缓存,同时最小化重计算量。在真实数据上,该方法优于标准启发式方法,尤其是在请求共享大量但非完全相同的前缀时。
arXiv:2605.05219v1 公告类型:新
摘要:前缀缓存是自回归大语言模型服务中的关键延迟优化技术,但现有系统假设每个token的键/值可以被密集复用。状态空间模型改变了问题的结构:循环层可以从单个存储的状态恢复,而不需要整个token历史。这种不对称性在“不复用”和“密集缓存”之间开辟了一个新的设计点:在稀疏的检查点位置存储精确的循环状态,并在缓存命中时从最深的存储检查点恢复,然后精确地重新计算剩余的后续部分。
我们将稀疏前缀缓存形式化为在重叠深度分布下的检查点放置问题,并得到了一个精确的O(NM)动态规划。对于请求共享非平凡前缀的用例(例如,针对单个长文档提出不同的问题),我们表明我们的方法持续改进了真实数据上标准启发式方法所绘制的帕累托前沿。在QuALITY和系统提示上,基于分布感知的放置方法在测量的层组帕累托前沿上优于所有固定预算的基线,并且与最强的启发式方法(块缓存)相匹配或更优,同时通常使用更少的检查点。在检查点预算较低且重叠分布最不均匀的情况下,增益最大。该方法最适用于许多请求在保留的缓存条目内共享大量但非完全相同的前缀的情况。它保持精确输出,不改变循环计算本身,也不需要新的循环更新内核,适用于隐藏状态可以被精确提取和恢复的循环/SSM层,对于混合模型,可以与现有的KV缓存压缩技术结合使用。
查看缓存全文
缓存时间: 2026/05/08 06:29
# 混合与循环LLM服务中的稀疏前缀缓存技术 来源:https://arxiv.org/html/2605.05219 ###### 摘要 前缀缓存技术是实现自回归LLM低延迟服务的关键优化手段,然而现有系统均假设逐token键/值(key/value)可被密集复用。状态空间模型从根本上改变了问题结构:循环层只需恢复单个存储状态即可继续推理,无需保留完整的历史token序列。这种非对称性催生了一个介于"完全不复用"与"密集缓存"之间的新设计空间——在*稀疏*的检查点位置上存储精确的循环状态,当缓存命中时,从最深匹配的检查点恢复状态,再精确重算后续缺失部分。本文将稀疏前缀缓存形式化为在交叠深度分布约束下的检查点放置问题,并给出精确的O(NM)动态规划解法。对于请求共享非平凡前缀(例如围绕同一篇长文档提出不同问题)的应用场景,实验表明,在真实数据集上,我们的方法持续优于标准启发式方法所描绘的帕累托前沿。在QuALITY和System Prompts数据集上,基于分布感知的放置策略在测得的层组帕累托前沿上全面超越了所有固定预算基线,与最强启发式方法(块缓存)性能持平或更优,且通常使用更少的检查点。在交叠分布最不均匀的低预算区域,增益尤为显著。本方法最适用于在缓存条目内多个请求共享显著但不完全一致前缀的场景。它保证输出精确性,不改变循环计算本身或引入新的循环更新内核,适用于可精确提取和恢复隐藏状态的循环/SSM层,对于混合模型还可与现有KV缓存压缩技术结合使用。
## 1 引言
前缀缓存通过复用前后请求共同前缀的计算结果,避免冗余的预填充计算。在以Transformer为核心的服务系统中,可复用状态由逐token的密集键/值张量构成。现代运行时系统通过分页技术[14 (https://arxiv.org/html/2605.05219#bib.bib4)]、基数树共享[38 (https://arxiv.org/html/2605.05219#bib.bib5)]以及分布式提示调度[29 (https://arxiv.org/html/2605.05219#bib.bib6)],使这一机制得以大规模落地。近来,状态空间模型与线性循环模型[11 (https://arxiv.org/html/2605.05219#bib.bib1), 7 (https://arxiv.org/html/2605.05219#bib.bib2), 26 (https://arxiv.org/html/2605.05219#bib.bib12), 33 (https://arxiv.org/html/2605.05219#bib.bib10)]的发展催生了新的架构:其循环层维护固定大小的隐藏状态,而非逐token的键/值对。
将注意力机制与循环或状态空间层结合的混合架构[27 (https://arxiv.org/html/2605.05219#bib.bib11)]使情况更为复杂。由于循环层的状态演化遵循 ht = F(ht-1, xt),从 token t 处继续推理只需要 ht,而非完整的中间状态历史。因此,对循环状态进行密集逐token缓存既无必要,又代价高昂:例如对于Qwen-3.5模型,其循环层(GatedDeltaNet)的每头状态规模与 d_head^2 成正比,而注意力KV缓存每头的存储量仅随 d_head 线性增长。
混合模型与SSM服务不必在"全量缓存"与"完全不缓存"之间二选一;它们可以选择在*稀疏*的检查点位置上存储循环状态,并在从最深存储的检查点出发,计算到实际分叉点之间的缺失段时进行按需重算。由此衍生出如下优化问题:给定一个长度为N的前缀和M个检查点槽位的内存预算,为了在未知未来请求交叠结构的情况下最小化期望重算量,应该将这些检查点放置在何处?这个问题与Marconi[23 (https://arxiv.org/html/2605.05219#bib.bib8)]所研究的缓存准入与淘汰策略正交——后者决定保留哪些条目。我们这里关注的是:对于每个缓存条目,如何选择哪些token处设置检查点,从而让*部分前缀*的交叠对循环层变得有用。当请求共享很长但不完全匹配的前缀时(例如连续请求共享很长的系统提示、以RAG检索文本开头、或围绕同一文档提出不同问题),这种方法尤其有益。对于聊天类工作负载,其中对话内连续请求总是包含前一个请求作为子串,显然最优策略是只存储最后一个状态,因此我们的方法不带来额外收益(也不会引入任何开销)。
该问题让人联想到训练中的激活检查点技术[5 (https://arxiv.org/html/2605.05219#bib.bib42)]:在前向传播过程中存储一部分中间激活,在反向传播时重算其余部分。关键区别在于:激活检查点针对的是已知的单一计算图进行优化,而我们这里需要考虑跨请求的随机交叠结构、共享缓存容量以及交叠分布的在线估计。
#### 我们的贡献。
1. 我们证明,在均匀交叠情形及最差情况下,均衡间隔是最优的(定理1 (https://arxiv.org/html/2605.05219#Thmtheorem1))。
2. 我们将稀疏前缀缓存形式化为一个单侧加权k-中位数问题,并给出精确的O(NM)动态规划解法(定理2 (https://arxiv.org/html/2605.05219#Thmtheorem2))。
3. 进一步,为支持使用历史请求的经验交叠深度分布代替未来真实分布,我们证明了价值函数在交叠律上是Lipschitz的(引理1 (https://arxiv.org/html/2605.05219#Thmlemma1)),因此其次优性可由经验分布与真实分布之间的总变差界定(定理3 (https://arxiv.org/html/2605.05219#Thmtheorem3))。我们还给出了交叠深度分布漂移情形下的类似界(定理4 (https://arxiv.org/html/2605.05219#Thmtheorem4))。
4. 我们在真实数据集(QuALITY、NarrativeQA、System Prompts)上验证了这些结果,并表明:与我们所考虑的每种固定预算基线相比,我们的方法在每个检查点上都节省了更多重算量。在QuALITY和System Prompts上,DP导出的调度方案在代表性层组上获得了原型级端到端加速,且在交叠分布最不均匀的低预算区域增益最大。
论文其余部分组织如下:第2节 (https://arxiv.org/html/2605.05219#S2) 回顾前缀缓存、KV缓存压缩、混合模型服务以及相关检查点与设施选址公式的先前工作。第3节 (https://arxiv.org/html/2605.05219#S3) 引入稀疏前缀缓存问题及其分布形式化定义。第4节 (https://arxiv.org/html/2605.05219#S4) 呈现主要理论结果,包括均匀交叠下的最优均衡放置、精确动态规划算法,以及在平稳和漂移交叠律下的插件式保证。第5节 (https://arxiv.org/html/2605.05219#S5) 描述实验设置、数据集、基线方法与实现约束。第6节 (https://arxiv.org/html/2605.05219#S6) 报告token节省量及原型墙钟时间结果。第7节 (https://arxiv.org/html/2605.05219#S7) 总结全文,附录包含额外证明与补充图表。
## 2 相关工作
#### 服务系统与密集前缀复用。
vLLM中的PagedAttention机制表明,在实践中,LLM服务高度依赖于灵活的KV缓存共享与低碎片化内存管理[14 (https://arxiv.org/html/2605.05219#bib.bib4)]。SGLang引入了RadixAttention用于结构化多调用应用程序,允许跨程序自动共享前缀[38 (https://arxiv.org/html/2605.05219#bib.bib5)];Preble则研究了联合优化共享与负载均衡的分布式提示调度[29 (https://arxiv.org/html/2605.05219#bib.bib6)]。CacheBlend解决了检索增强生成中非连续缓存块的精确复用问题[34 (https://arxiv.org/html/2605.05219#bib.bib7)]。ChunkAttention提出了一种两阶段分区方案用于前缀感知解码,进一步减少了冗余注意力计算[35 (https://arxiv.org/html/2605.05219#bib.bib13)]。所有这些系统都假设了以token索引的密集KV状态;然而我们的场景不同,因为循环层可以从稀疏的精确检查点恢复。
#### KV缓存压缩与选择。
大量工作致力于通过淘汰、量化或稀疏化来减少密集注意力缓存的内存占用。Scissorhands识别出对未来注意力"关键"的token并淘汰其余部分[17 (https://arxiv.org/html/2605.05219#bib.bib14)]。H2O仅保留累积了最多注意力权重的heavy-hitter KV条目[37 (https://arxiv.org/html/2605.05219#bib.bib15)]。PyramidKV和PyramidInfer利用后层需要的缓存token更少的观察,以层自适应金字塔形状分配预算[4 (https://arxiv.org/html/2605.05219#bib.bib16), 31 (https://arxiv.org/html/2605.05219#bib.bib17)]。SnapKV对注意力特征进行聚类从而为每层选择代表性KV条目[16 (https://arxiv.org/html/2605.05219#bib.bib18)],DMC则提出了一个可学习的策略来决定哪些token应被合并、淘汰或保留[22 (https://arxiv.org/html/2605.05219#bib.bib19)]。与这些方法不同,我们的方法针对的是循环状态而非KV缓存,并且保持输出精确性。
#### 混合模型的前缀缓存。
最接近的先前工作是Marconi等人的研究[23 (https://arxiv.org/html/2605.05219#bib.bib8)]。该工作指出,混合模型使用原地循环更新使得部分回滚困难,从而迫使服务走向精确匹配条目,产生许多体积大但价值低的缓存候选。作者提出了改进的准入和淘汰策略。我们的工作在此处是互补的:我们并不改进精确循环缓存条目的管理,而是通过存储稀疏精确检查点来改变复用单元本身,从而使得部分前缀交叠变得有用。
#### 训练中的激活检查点。
稀疏前缀缓存在结构上类似于激活检查点[5 (https://arxiv.org/html/2605.05219#bib.bib42)]:前向传播时存储一部分层激活,以减少反向传播时的峰值内存。两个问题都提出:沿序列计算过程中,为了在内存预算下最小化重算量,检查点应放置在哪里?关键区别在于:(i) 激活检查点针对的是已知拓扑的单一确定性计算图进行优化,而我们则要在随机交叠的跨请求场景下优化;(ii) 在服务中,检查点放置与缓存准入和淘汰相互影响,这在训练中没有对应项。
#### 设施选址与k-中位数问题。
我们研究的问题在离线形式上是一维线上的单侧加权k-中位数问题。经典的k-中位数在一般度量空间中是NP难的[20 (https://arxiv.org/html/2605.05219#bib.bib20)],但一维情形存在高效动态规划解法[12 (https://arxiv.org/html/2605.05219#bib.bib21), 10 (https://arxiv.org/html/2605.05219#bib.bib22)]。我们的凸壳技巧求解器利用前缀和的单调结构,达到了O(NM)时间复杂度,这与相关区间调度问题中基于SMAWK的优化方法[1 (https://arxiv.org/html/2605.05219#bib.bib23)]类似。
## 3 问题形式化
考虑一个长度为N的缓存前缀和一个检查点集合 C = {c1, c2, ..., cM},满足 1 ≤ c1 < ... < cM ≤ N。检查点 ci 处的隐藏状态(对于循环层而言)被存储。现在一个具有交叠深度 d 的新请求到达(即在之前缓存条目的前 d 个 token 处匹配)。为了恢复推理,我们必须从最深不超过 d 的检查点开始。如果 d = 0(完全不匹配),我们什么也节省不了。如果 d ≥ 1,令 j = argmax{ci : ci ≤ d}(即 ≤ d 的最大检查点)。那么从该检查点恢复之后,需要重算从 cj + 1 到 d 的 token。注意:如果检查点 ci 精确位于 d 处,则无需重算,跳步为零。
在我们的模型下,假设所有请求都共享一个长度为 N 的公共前缀,并且每个请求在某个深度 d ~ p 上与缓存条目中的一个(随机选择的)条目重叠,其中 p 是 [N] ∪ {0} 上的概率质量函数。在下一节中我们将放宽该模型以适应现实场景,但此刻该设置足以清楚说明问题。
我们的目标是选择由检查点之间的间隔向量 g1, ..., gK(其中 K = M+1 表示间隔数,满足 sum(gi) = N+1 且每个 gi ≥ 1)所刻画的检查点 C,以最小化由下面函数给出的期望重算 token 数 R(p, g):
R(p, g) = E_{d~p}[f_g(d)] 其中 f_g(d) = d - max{i: sum_{j≤i} g_j ≤ d} (若 d ≥ 1);若 d=0 则为 0。
其中 f_g(d) 是从最近检查点位置到交叠深度的距离。
[辅助图插入]
图1:稀疏前缀缓存示意图。检查点(阴影方块)存储精确的循环状态。当请求的新前缀与缓存前缀重叠到深度 d 时,系统从最深不超过 d 的检查点恢复,并重算从该检查点之后到 d 的剩余 token(虚线路径)。不需要一直回退到最开头。
注意 p({0}) 实际上是匹配失败的概率;在那种情况下,所需的重算量是整个前缀(最多 N 个 token),因为找不到合适的缓存条目。目前我们将该情况视为损失常数 N,但不严重,因为我们是从静态分布的角度考量。在更实际的设置中,d 表示与之前某个缓存条目的重叠深度,具体机制取决于缓存准入和淘汰策略;这里我们假设一个抽象分布 p。
## 4 理论结果
现在我们转向分析该优化问题并表述主要定理。
**定理 1(均匀交叠下的均衡放置是最优的):** 如果 p 是在 {1,...,N} 上的均匀分布(且 p(0) = 0),那么任何能使期望函数 R(p,g) 达到最小的间隔向量 g 必须满足所有间隔 gi 至多相差 1。
*证明:*
(⇒) 假设存在两个间隔 gi 和 gj(i < j)相差至少 2:设 gi ≥ gj + 2。将 gi 减 1 而 gj 加 1,得到新的向量 g'。我们证明这一局部操作能严格减小目标函数。由于除 i 与 j 外其他间隔不变,只需比较这两个间隔的贡献。记相邻检查点之间的长度为 li = gi - 1(可重算 token 段长度)。原方案中,对于落在间隔 i 内的交叠深度,其重算距离从 1 到 li 分布;类似地,对于间隔 j 内的交叠其重算距离从 0 到 lj。在均匀分布下,这些交叠深度在 [1, N] 上等概率。将 1 个 token 从间隔 i 移到间隔 j 之后,间隔 i+ 和 j+ 的期望重算量变化为:φ(gi) + φ(gj) - φ(gi - 1) - φ(gj + 1) = gi - gj - 1 > 0。因此,任何两个间隔相差至少 2 的可行向量都可以被严格改善。因此,每个最优解的所有间隔至多相差 1。
(⇐) 令 K = M + 1,记 N+1 = qK + ρ,0 ≤ ρ < K。那么若所有间隔均为 q 或 q+1,该向量即为均衡向量,其上界可以由上述局部优性推导得出。具体构造略。
**定理 2(对任意分布 p 的 O(NM) 动态规划):** 对于定义在 [N] 上的任意离散分布 p,最小化期望重算量 R(p, g) 的问题可以转化为一个具有凸代价函数的加权 k-中位数问题,并可通过标准凸壳技巧在 O(NM) 时间内求解。
*证明提纲:* 定义 DP[t][k] 表示只考虑前缀 t 且设置 k 个检查点时的最小期望重算量(每个检查点对应一个间隔结束)。转移方程为 DP[t][k] = min_{s < t} { DP[s][k-1] + C(s+1, t) },其中 C(i, j) 表示在位置 i 到 j 的区间内放置一个检查点(其前一个检查点位于 i-1)所带来的期望重算代价。这里的代价函数是凸的,因为重算距离与累积质量成特定关系。通过预处理前缀和以及凸壳技巧,可达到 O(NM)。
**引理 1(价值函数的 Lipschitz 性):** 定义 V(p, M) = min_g R(p, g)。对于任意两个分布 p 和 q,有 |V(p, M) - V(q, M)| ≤ N * TV(p, q),其中 TV 表示总变差距离。
*证明:* 期望重算量 R(p,g) = Σ_d p(d) f_g(d)。由于 f_g(d) ≤ N,该函数在 p 上是线性且合同系数有界。通过最优性条件可得该界。
**定理 3(经验分布的插件式保证):** 令 p̂ 为基于 T 个历史请求的 d 的经验分布。若使用 p̂ 计算出的最优检查点集 ĝ,则在真实分布 p 下的期望重算量满足:R(p, ĝ) ≤ V(p, M) + 2N * TV(p̂, p)。若经验样本独立同分布于 p,则 TV(p̂, p) 依概率收敛到 0。
*证明:* 由引理 1 及最优性:R(p,ĝ) ≤ V(p,M) + |R(p,ĝ) - V(p,M)| ≤ V(p,M) + N * TV(p̂,p) + N * TV(p,p̂) = V(p,M) + 2N * TV(p̂,p)。
**定理 4(分布漂移下的鲁棒性):** 假设交叠深度分布是随时间变化的,在长度为 T 的窗口内经验平均为 p̂_T,而下一时刻的真实分布为 p_T。如果分布漂移满足 TV(p_T, p̂_T) ≤ ε,则用 p̂_T 计算的检查点 ĝ_T 在 p_T 下的期望重算量至多为 V(p_T, M) + 2Nε。
*证明:* 与定理 3 类似,只需替换分布。
## 5 实验设置
**数据集与分布估计。** 我们使用三个数据集来模拟交叠深度分布:
- **QuALITY**:关于长文档的问答数据集,每个文档对应多个问题。我们将同一文档的问题组成请求后缀,并测量共享前缀长度。
- **NarrativeQA**:关于叙事文的长文档问答,同样构造共享前缀。
- **System Prompts**:来自实际 LLM 服务日志,包含不同系统提示重复使用场景的规模测量。
对于每个数据集,我们抽取 10000 个请求,并与一个缓存条目(例如文档或系统提示)比较,记录交叠深度。然后归一化为分布 p̂。
**基线方法。** 我们比较以下检查点放置策略:
- 无缓存:不存储任何检查点,所有请求均需完全重算。
- 均匀间隔:检查点均匀放置(间隔至多相差 1)。
- 块缓存:将前缀划分为固定长度的块,仅在块边界存储检查点。
- 随机:随机放置检查点。
- 贪心:逐步选择能最大程度减少当前估计的期望重算量的位置。
- 我们的 DP 方法:使用经验分布 p̂ 通过动态规划找到最优检查点集。
**评估指标。** 主要指标是每个检查点节省的期望重算 token 数。我们也测量在特定层组(由循环层和注意力层组成)上的墙钟时间,用于原型实现。
**实现细节。** 使用 Python 实现 DP 算法;利用凸壳技巧达到 O(NM)。对于墙钟测试,我们在 HuggingFace 的混合模型(如 Qwen3.5-7B)上实现原型,通过对每组循环层(每层对应一个 GatedDeltaNet 状态)进行稀疏存储。
## 6 实验结果
**Token 节省量随检查点预算的变化。** 图 2 显示了在 QuALITY 数据集上,随着检查点数量 M 的增加,每种策略节省的期望重算 token 数。DP 方法始终优于或等于所有其他方法。低预算时增益最大,因为均匀分布假设偏差最大。
[图2:QuALITY 上不同 M 下的节省 token 数]
**帕累托前沿。** 我们将控制变量从检查点数量改为总存储预算(考虑循环状态与 KV 缓存的不同规模)。在混合模型中,每个循环检查点占用 O(d_head^2) 内存,而注意力层 KV 缓存是每 token 线性增长。因此我们绘制节省 token 数与总内存的关系。我们的方法在低内存区域主导前沿。
**墙钟加速。** 在系统提示数据集上,我们模拟一个包含 32 层混合模型(12 层注意力,20 层循环)的服务流程。表 1 报告了平均每个请求的预填充时间。DP 方法在 M=16 时比均匀间隔快 18%,比块缓存快 12%。
[表1:墙钟时间对比]
**关于超参数的影响。** 检查点预算 M 的选择对性能影响较大;我们建议通过类似曲线拐点的方式选择 M。交互式的在线更新检查点可能通过滑动窗口分布实现。
**分布漂移实验。** 在连续请求中分布缓慢变化(例如不同日期的对话)时,我们使用最近 100 个请求的滑动窗口更新分布。在一周时间跨度内,DP 方法仍然保持优于任何固定检查点集的性能,TV 距离从未超过 0.1,保证了低次优性。
## 7 结论与未来工作
本文提出了稀疏前缀缓存,即在不改变混合模型输出准确性的前提下,通过在循环层稀疏存储状态来利用部分前缀交叠。我们将该问题形式化为一个加权 k-中位数问题,给出了均衡放置的理论保证,并由对任意分布最优的 O(NM) 动态规划。在真实工作负载上的实验表明,该方法每个检查点节省了更多重算量,且在存储受限场景下主导帕累托前沿。未来工作包括将检查点放置自适应地集成到缓存准入策略中,利用预测器去预估分布漂移,以及为在连续服务中联合优化检查点放置与淘汰提供更强理论保证。
## 致谢
我们将本文的 arXiv 源码公开,以方便社区复现与扩展。相关工作资源链接见论文网站。相似文章
并非所有令牌都值得缓存:学习语义感知的LLM前缀缓存驱逐策略
一种针对LLM前缀缓存的新型语义自适应驱逐策略,学习不同令牌类型间的令牌重用模式,相比现有策略实现了1.4倍至2.7倍的TTFT提升。
为扩散语言模型启用共享前缀的KV缓存
本文提出BiCache,一种面向扩散语言模型共享前缀的新型KV缓存技术,通过动态重用浅层中缓存的键和值来避免精度崩溃,并实现36.3%–98.3%的吞吐量提升。
KV Packet: 免重计算的上下文无关KV缓存用于大语言模型
KV Packet 提出了一种免重计算的缓存复用框架,用于大语言模型。该框架使用可训练的软标记适配器来弥合上下文不连续性,消除了开销,同时在 Llama-3.1 和 Qwen2.5 上的性能与完全重计算基线相当。
知识卸载:将大语言模型分解为稀疏骨干网络与记忆模块
提出KOFF框架,利用结构化剪枝和LoRA适配器将预训练大语言模型分解为稀疏共享骨干网络与领域特定外部记忆,实现12%稀疏度且无明显性能损失。
持续LLM升级循环:一种用于从稠密到稀疏LLM的预测器门控按组稀疏训练方案
本文提出了一种用于大语言模型的从稠密到稀疏的持续训练方法,采用预测器门控的按组稀疏性实现4倍FFN稀疏度,并在Qwen2.5-8B上通过长上下文训练进行了验证。