← 返回 2026-04-28

随机 KV 路由:实现自适应深度方向缓存共享 Stochastic KV Routing: Enabling Adaptive Depth-Wise Cache Sharing

Anastasiia Filippova, David Grangier, Marco Cuturi, João Monteiro 📅 2026-04-03 👍 8 2026-07-13 08:36
KV缓存优化 LLM推理 Transformer 推理效率 深度方向共享

R-CLA 训练时随机丢弃层缓存,使单一模型自适应多种深度共享策略。

前置知识

KV Cache

自回归 Transformer 推理时为避免重复计算,每一层会缓存先前 token 的 Key 和 Value 张量。

KV 缓存随序列长度线性增长,是长上下文推理的内存与带宽瓶颈,本文方案正是要削减它。

Cross-Layer Attention (CLA)

Brandon 等(2024) 提出让相邻层共享同一组 KV,从而把层 l 的 query 注意力到 l'≠l 的缓存。

R-CLA 是在 CLA 基础上随机化源层选择,读者需先理解跨层注意力的标准写法。

Grouped-Query Attention (GQA)

Ainslie 等(2023) 提出层内多个 query head 共享同一组 KV,通过减少 head 数降低缓存。

GQA 是与本文正交的工程方案,本文的深度共享可与 GQA、量化等叠加获得复合收益。

Bernoulli 采样

以概率 p 取 1、以概率 1-p 取 0 的二值随机变量;R-CLA 训练时对每层每步独立做一次该采样,决定该层是走自注意还是跨层注意,是整个随机化机制的概率论基础。

R-CLA 训练机制的核心是每层每步独立做一次 Bernoulli 采样决定 KV 来源。

研究动机

现有 KV 缓存压缩研究多沿时间轴剔除不重要 token,存在两大痛点。其一是查询相关性:某 token 对当前 query 不重要,对下一个 query 却可能关键,被剔除后必须重算或引发幻觉;其二是开销大:H2O、SnapKV、PyramidKV 等方法需动态重算注意力分数或维护完整索引,反而增加峰值显存与带宽压力。Llama-2-7B 单 token 缓存约 512 KB,相对原始 2-4 字节的 token 整数膨胀了 10 万倍,缓存《爱丽丝梦游仙境》一本书的上下文就需要约 14 GB,几乎与模型权重本身相当。深度方向缓存共享方法(XC-Cache、Layer-Condensed KV Cache、MiniCache)虽证明跨层冗余,但需要额外编码器、增加 prefill 时间或仅事后合并浅层,无法做到部署可配的弹性共享。

本文的目标是本文目标是训练一个能容忍任意深度方向缓存共享策略的单一模型。形式上,对每个 Transformer 层 l 与每个前向步,作者独立采一个以 p 为参数的 Bernoulli 变量 d:若 d=1 则 Output_l = Attn(Q_l, K_l, V_l) 走标准自注意;若 d=0 则从 {1,...,l-1} 均匀采 l',改为 Output_l = Attn(Q_l, K_{l'}, V_{l'}) 走跨层注意。推理时切换到确定性共享策略 S⊆{1,...,L},通过 μ(l)=max{j∈S|j<l} 把任意 l 路由到最近前层 leader 的缓存,无需重训即可在 100% 到 25% 缓存保留率之间灵活切换,同时不损失甚至提升完整缓存下的质量,从而让同一模型可部署在从高配集群到边缘设备的全光谱硬件环境中。

与已有工作不同的是,作者指出深度维度是与时间轴正交却长期被低估的优化方向。已有工作要么把跨层共享做成固定模式(Wu et al. 2025 需要为每种模式重新预训练),要么仅在浅层事后融合(MiniCache),要么需要昂贵外部编码器(XC-Cache)。本文独特之处在于把跨层选择做成训练期间的随机事件而非推理期的固定规则,模拟结构级 cache fault,使查询投影 Ql 学习对多种 KV 分布鲁棒,并在大模型数据受限微调场景下观察到类似正则化的额外收益。

核心方法

方法核心是 Random Cross-Layer Attention (R-CLA),把每层每步的 KV 来源当成一个随机变量。训练时对每个 Transformer 层 l 独立掷一枚以 p 为参数的硬币:正面(概率 p)走标准自注意 $\mathrm{Output}_l = \mathrm{Attn}(Q_l, K_l, V_l)$,反面(概率 1-p)从 $\{1, \ldots, l\!-\!1\}$ 中均匀采样 l',改为 $\mathrm{Output}_l = \mathrm{Attn}(Q_l, K_{l'}, V_{l'})$。这一过程模拟硬件上的结构性 cache miss,迫使 Q 投影学习与多种历史 KV 分布兼容。推理时切换到确定性共享:选定策略 $S \subseteq \{1,\ldots,L\}$ 表示保留缓存的层集合,未在 S 中的层通过 $\mu(l) = \max\{j \in S \mid j < l\}$ 复用一个最近前层的缓存,每组只 load/update 一次 K, V,可彻底省掉非 leader 层的 K/V 投影与更新开销。

R-CLA 与已有跨层方案的本质区别在于训练期间的不确定性而非推理时的不确定性。Brandon 等(2024) 的 CLA 把每两层固化成一对,Wu et al.(2025) 系统枚举固定共享模式,MiniCache、CommonKV、KVSharer 都是后处理无训练。本文的关键洞察是:只有当训练时模型就被暴露给大量不同的共享模式(既包括完全自注意也包括各种跨层组合),推理时无论选哪种固定策略 S 都不会破坏 Q 与 KV 的特征对齐。此外,作者用控制变量法分离了"共享"与"随机性"两个因素——固定 CLA@2/4 在 100% 缓存下表现甚至略优于 R-CLA,但只要进入 50% 或 25% 缺失缓存区,随机暴露的训练分布就显示出压倒性优势,这印证了正则化效应来自跨层 KV 共享造成的信息瓶颈。

方法步骤详情

第一步定义跨层注意力算子:对层 l 计算 Q_l,但到 (K_{l'}, V_{l'}) 上做注意力,其中 l'≤l,l'=l 时退化为自注意。第二步定义缓存共享策略 S:S⊆{1,...,L} 是允许持有 KV 的层集合,目标是最小化 |S| 同时保持性能;任意 l 通过映射 μ(l)=l(若 l∈S)或 max{j∈S|j<l}(若 l∉S)确定访问的缓存。第三步实现 R-CLA 训练:每步每层独立采 d~Bernoulli(p),d=1 走 (K_l,V_l),d=0 从 {1,...,l-1} 均匀采 l' 后走 (K_{l'},V_{l'});前向过程中维持一个 CurrentCache 指针,按 S 在合适位置插入 load 与 update。第四步推理时按部署硬件约束选 S,例如 50% 保留即 S 取每隔一层、25% 保留则每四层取一个 leader;非 leader 层直接跳过 K/V 投影,复用最近 leader 的 K, V 矩阵。第五步(可选)将 R-CLA 与 GQA、KV 量化或时间轴剪枝叠加以获得复合收益。

技术新颖性

技术新颖性体现在三方面:其一,把 cache sharing 视作训练时随机扰动而非推理期工程技巧,借鉴了结构化 dropout 的随机性思想但保持每层都参与计算,仅随机化 KV 来源;其二,单模型可同时服务 100%、50%、25% 多种保留率,无需为每种硬件重训;其三,论文系统地用 RD-CLA 消融把"共享的容量压缩效应"与"随机训练的正则化效应"分离,发现正则化主要来自 KV 共享本身而非随机性本身,纠正了对正则化来源的直觉。此外,论文提供了完整的算法 1/2 伪代码与 Group g∈{1,2,4,8} 的端到端推理基准,使 R-CLA 具备落地到 vLLM/TGI 类系统的工程基础。

Behavior of a model implementing R-CLA.
Figure 4: Behavior of a model implementing R-CLA.

实验结果

预训练 Qwen3-1.7B、34B tokens、Chinchilla-optimal 设定下,评估损失从 p=0 的 2.424 仅升到 p=0.75 的 2.461,退化 <2%,训练稳定;等效缓存预算下 28 层 R-CLA + 跨层共享始终优于 7/14/21 层浅层 Transformer(图 5)。微调 p=0.6 训练 Qwen3-8B、Mistral-7B、Llama-3.1-8B,QA 五数据集上:25% 保留率下基线几乎崩溃,如 Qwen3-8B 在 HotpotQA 从 0.233 跌到 0.011,R-CLA 维持 0.098;100% 保留率下 Llama-3.1-8B 在 SQuAD v2 从 0.655 升到 0.758(+15.7%),印证正则化效应。消融 Table 3:100% 保留下固定 CLA@2/4 略胜 R-CLA,50%/25% 缺缓存区 R-CLA 系统性领先。推理:8K g=4 让 KV 1170→293 MB(4×)、吞吐 34.0→41.6 tok/s(+22%)、batch=16 基线 OOM、g=4 仍可服务。

Pre-training evaluation loss for Qwen3-1.7B trained with R-CLA at varying cross-layer attention probability p.
Table 1: Pre-training evaluation loss for Qwen3-1.7B trained with R-CLA at varying cross-layer attention probability p.
Effect of R-CLA on F1 Score for models fine-tuned for Question-Answer.
Table 2: Effect of R-CLA on F1 Score for models fine-tuned for Question-Answer.
Ablations on top of Llama-3.1-8B.
Table 3: Ablations on top of Llama-3.1-8B.
Inference efficiency across input lengths.
Table 4: Inference efficiency across input lengths.
Batch size scaling at 8,192-token context.
Table 5: Batch size scaling at 8,192-token context.
Performance comparison under varying cache retention rates for base models and their counterparts trained with R-CLA.
Figure 2: Performance comparison under varying cache retention rates for base models and their counterparts trained with R-CLA.
Cache size and prediction performance trade-offs for different models.
Figure 3: Cache size and prediction performance trade-offs for different models.
Cache size vs. evaluation loss for Qwen3-1.7B.
Figure 5: Cache size vs. evaluation loss for Qwen3-1.7B.
查看结构化数据
任务指标本文基线提升
Qwen3-1.7B 预训练 (OpenWeb, 34B tokens) Eval Loss p=0.75 时 2.461 p=0 标准训练 2.424 退化 <2%,训练稳定
HotpotQA @ Llama-3.1-8B, 25% 缓存 F1 0.237 0.080 +196.2%
HotpotQA @ Qwen3-8B, 25% 缓存 F1 0.098 0.011 +826.6%
SQuAD v2 @ Llama-3.1-8B, 100% 缓存 F1 0.758 0.655 +15.7%(正则化效应)
RepLiQA @ Mistral-7B, 50% 缓存 F1 0.547 0.291 +88.2%
TriviaQA @ Qwen3-8B, 25% 缓存 F1 0.131 0.005 +2287.3%
Qwen3-8B 推理, 8K context, g=4 Decode Throughput 41.6 tok/s 34.0 tok/s +22%
Qwen3-8B 推理, 8K context, g=4 KV Cache 内存 293 MB 1170 MB 4× 缩减
Qwen3-8B 推理, 8K context, batch=16 可服务性 60.3 GB 峰值,8.0 tok/s Out of Memory 基线 OOM,g=4 可服务

局限与改进

作者明确指出四点局限。其一,R-CLA 需训练阶段介入,post-hoc 方法或轻量适配器能否诱导类似跨层鲁棒性尚未验证;其二,实验在 Chinchilla-optimal 计算最优设定下进行,未考察过训练(overtraining)情形;其三,未在 MoE 架构上评估,理论上因只改注意力 KV 来源应可直接迁移但缺少证据;其四,QA 任务虽是 context retention 的良好代理,但缺少更广泛任务(代码、对话、推理)的覆盖。读者也应注意到:Table 2 中 100% 保留率下 Qwen3-8B 在 HotpotQA 的 0.233→0.357 提升达 +53.1%,这远超"正则化"通常量级,提示可能与 Qwen 预训练分布与 QA 数据集存在某种分布匹配,而不能完全归因于 R-CLA 本身;以及 Table 4 的实现是"跳过 K/V 投影"但未真正把 K, V 驻留 SRAM 做层间复用,作者明确说这只是保守下界。

独立分析的弱点

独立审视可发现若干可改进点。其一,R-CLA 训练时每步每层都需采样 l' 并可能跨多组做注意力,相比标准 self-attention 的并行友好性会引入额外同步开销,工程实现复杂;可探索固定 stride 的近似随机以减少运行时分支。其二,所有消融与基准都基于 group size g∈{1,2,4,8} 这种规则结构,对不规则或非均匀 S(如前密后疏)下的鲁棒性未充分验证,理论上 R-CLA 应能应对但缺乏实验。其三,p=0.6 是单一超参选择,论文未给出在 1.7B vs 8B 不同规模下 p 的最佳值扫描,限制实际部署可参照性;可补充 p-scaling law。其四,Table 3 中 RD-CLA 在 100% 保留率下超过 R-CLA 提示确定性的容量压缩可能更优,未来工作可研究"训练期高随机 + 推理期可调确定性"的组合策略。其五,25% 保留下 F1 仍与 100% 有显著差距(HotpotQA 0.237 vs 0.306),说明深度共享并非完全无损,存在信息损失下界。

未来方向

作者提出的方向包括:自适应共享策略(根据 query 复杂度或当前生成步动态调整缓存深度)、与 GQA / SSM / MoE 架构创新的复合、与时间轴剪枝与 KV 量化的组合研究、自适应 p-schedule(前期低 p 稳定训练后期高 p 强化共享)。基于结果还可延伸:把 R-CLA 思想用于 diffusion 模型的 cross-attention KV 共享;研究 MoE 路由与 KV 路由的协同设计;将 R-CLA 与 speculative decoding 结合(被共享层 KV 也能被 draft 模型复用);探索 post-hoc 蒸馏路径——用一个全缓存大模型蒸馏出 R-CLA 训练目标下的小模型,避免完整 R-CLA 训练成本;以及把 R-CLA 与 KV eviction oracle 拼接,构造查询感知 + 深度感知两阶段缓存系统。

复现评估

复现难度中等偏低:方法本身仅改前向计算路径,伪代码在 Algorithm 1/2 中完整给出;超参 p=0.6、AdamW β=(0.9, 0.95)、lr 5e-6 线性 warmup 1.5% 都明确说明;预训练用 Qwen-1.7B 风格解码器、OpenWeb 子集、2K 上下文、34B tokens 在 H100 上训练;微调使用 HotpotQA/SQuAD v2/MSMarco/TriviaQA/RepLiQA 五任务 50K 步 batch 128、8K 上下文。论文未声明代码仓库是否开源(截至摘要撰写时未见官方 repo),需读者自行实现 R-CLA forward 路径。最大复现成本是 8B 模型的 50K 步微调与 1.7B 预训练所需的 H100 算力;Table 4/5 的端到端推理基准需要 80GB 单卡 GPU(如 A100/H100 80G)并复现 g=4 的非 leader 层 K/V 投影跳过逻辑;QA 数据集均公开可下载,整体可复现性良好但计算资源门槛较高。