核心发现
方法论
论文定义映射函数kv(i),表示第i层查询所使用的KV来源层;kv(i)=i时为KV层,否则复用目标层KV。作者组合三种分组方式:pizza、sandwich、lasagna,以及bottom、top、middle三种目标位置,共九类配置,并覆盖LCKV、YOCO和CLA。top/middle配置采用m=7、b=2的迭代训练与预填充近似。
关键结果
- 在1.1B模型、RTX 3090上,短提示5+2043时,九种配置均明显超过标准Llama,且KV层越少吞吐越高。长提示512+1024时,top/middle因迭代预填充吞吐显著下降,部分低于基线;bottom配置仍保持明显优势。
- 在MiniPile的1.7B token训练中,110M模型训练1个epoch、1.1B模型训练2个epoch;KV层保留一半时,多数配置困惑度接近标准Transformer。KV层继续减少时,bottom退化最明显,top/middle相对稳健。
- 在SlimPajama 100B token训练及Hellaswag、OpenBookQA、WinoGrande、ARC、BoolQ、PIQA、SciQ评测中,sandwich-top通常优于两类bottom配置;KV层很少时,sandwich-middle综合表现最佳。
研究意义
该研究将分散的跨层KV共享方法放入同一坐标系,揭示缓存容量、目标层位置、预填充成本与任务性能之间的可解释权衡。结果为服务商选择配置提供实证依据:若只需约2倍缓存压缩,应优先考虑pizza-bottom或lasagna-bottom;若追求更激进压缩且生成远长于提示,则可接受sandwich-middle的训练和预填充开销。
技术贡献
核心贡献是形式化kv(i)映射及pizza/sandwich/lasagna×bottom/top/middle设计空间,并系统测试此前未研究的lasagna-top及middle系列。针对top/middle的循环依赖,作者采用逐token迭代近似:前n-b次用m次近似,最后b次反向传播;推理预填充同样迭代。研究还明确指出,KV层数量直接决定WK、WV参数量和缓存容量。
新颖性
创新不在于提出单一新缓存算法,而在于首次以统一框架横向比较主要跨层KV共享方案及其新变体。论文把LCKV、YOCO、CLA分别对应到sandwich-top、pizza-bottom、lasagna-bottom,并揭示目标层位于底部、顶部或中部时,速度与质量的系统差异。
局限性
- 实验规模有限,主要使用110M和1.1B模型,较大模型、更多数据和更长上下文上的结论仍需验证。
- top/middle因迭代训练和预填充产生额外成本,并通过屏蔽注意力对角线处理循环依赖,可能造成与标准Transformer的性能差距。
未来方向
后续可扩展至更大模型和真实服务负载,优化迭代预填充;作者特别提出为每个token单独计算自注意力以弥补对角线屏蔽损失,并探索quarter、three-quarter位置、预训练初始化及更高效的训练方案。
AI 总览摘要
大型语言模型的生成速度常受KV缓存限制。每层都保存Key和Value会占用大量显存,已有GQA、滑动窗口和重型token保留等方法缓解这一问题,但跨层共享的具体策略尚未被系统比较。Wu、Wu与Tu将LCKV、YOCO和CLA统一为一个映射框架,考察不同KV层位置与分组方式。
框架用kv(i)指定第i层查询读取哪一层KV,并组合pizza、sandwich、lasagna三种分组与bottom、top、middle三种目标位置。bottom可直接训练;top和middle存在循环依赖,因此采用迭代训练与预填充,实验设置m=7、b=2。结果显示,在1.1B模型和RTX 3090上,短提示5+2043时九种配置都比标准Llama快,且KV层越少吞吐越高;长提示512+1024时,top/middle因迭代编码显著变慢,而bottom仍具优势。
质量实验使用MiniPile 1.7B token及SlimPajama 100B token,并测试Hellaswag、OpenBookQA、WinoGrande、ARC、BoolQ、PIQA和SciQ。保留约一半KV层时,多数配置仍接近基线;更激进压缩则bottom退化明显,sandwich-middle更有潜力。论文因此给出务实结论:2倍压缩优先选pizza-bottom或lasagna-bottom;若生成长度远大于提示且能承担额外训练和预填充成本,可选择sandwich-middle。
深度分析
研究背景
Transformer解码会为每层、每个历史token保存K和V,长上下文与大批量服务因此受到显存带宽和容量限制。GQA/MQA减少头间复制,H2O和StreamingLLM减少token缓存;LCKV、YOCO、CLA则进一步跨层复用KV。但这些方法的层选择和共享方向不同,缺少统一比较。
核心问题
研究问题是:在固定KV层数量下,哪些分组方式和目标层位置能同时降低缓存、提高吞吐并保持语言建模与下游能力?难点在于top/middle共享会形成跨层、跨token的顺序依赖,增加训练和预填充成本;不同上下文长度也会改变收益。
核心创新
作者提出kv(i)统一表示,并构造九种配置。pizza保留前l-1层为KV层,sandwich在底部和顶部保留KV层,lasagna把全部层均匀分组。每种结构再选择bottom、top或middle目标层;首组固定使用底层目标,以避免首层非KV造成明显性能损失。
方法详解
- �� 输入:L层Transformer和l个KV层预算;输出是kv(i)映射。
- �� KV层计算自身WK、WV,非KV层删除对应参数并复用目标层KV,因此KV层数同时控制模型参数与缓存。
- �� bottom配置按标准方式训练;top/middle使用迭代训练,前n-b次以m次近似,最后b次反向传播。
- �� 预填充阶段在最后KV层后可early exit;top/middle用m+b次迭代并行近似。对角线注意力被屏蔽以打破循环依赖。
- �� 实现采用HuggingFace Transformers、FlashAttention 2、融合RMSNorm、交叉熵和SwiGLU。
实验设计
模型为110M和1.1B参数,分别含12和22层;1.1B配置为hidden size 2048、32个注意力头、4个KV头。吞吐在RTX 3090上测试,比较5+2043与512+1024。小规模训练使用MiniPile 1.7B token;大规模使用SlimPajama 100B token。指标包括最大吞吐、困惑度和LM Eval Harness零样本准确率。
结果分析
短提示下九类配置均快于标准Llama,减少KV层可继续提高吞吐;长提示下只有bottom稳定,因为top/middle需要迭代编码。保留一半层时多数配置质量接近基线;继续压缩时bottom损失最大。大规模实验中sandwich-top通常优于bottom配置,sandwich-middle在KV层很少时最好。
应用场景
适用于GPU显存受限的聊天机器人、长文本生成、批量推理和边缘部署。短提示、高并发场景可采用bottom配置;若输出远长于输入,可用sandwich-middle换取更小缓存。部署前需依据提示长度、生成长度、显存预算和是否能重新训练模型进行选择。
局限与展望
论文只验证到1.1B参数和100B训练token,不能保证在更大模型上保持相同排序。top/middle的迭代过程增加训练与prefill延迟,且对角线mask可能损伤质量。吞吐依赖RTX 3090、实现融合和批量设置;未来需测试更长上下文、更多硬件、预训练模型初始化及专门的自注意力补偿。
通俗解读 非专业人士也能看懂
把模型想成一家工厂。标准Transformer让每个车间都保存自己的“订单记录”和“库存记录”,所以记录很多、查找很慢。跨层KV共享则让若干车间共用一个记录柜:每个车间仍可提出问题,但不必重复制作同样的记录。
pizza方案把记录柜集中放在前面,sandwich方案把柜子放在前后两端,lasagna方案则把工厂分成几段,每段设置一个柜子。柜子越少,省下的空间越多,工厂通常也能更快出货;但如果后面的车间必须等待尚未完成的柜子,反而会拖慢开工,这就是top和middle方案在长订单上的问题。
实验说明,订单较短时几乎所有方案都比标准工厂快;订单很长时,前置柜或底部柜更可靠。若只想把记录空间减半,pizza-bottom和lasagna-bottom是稳妥选择;若订单输出特别长,可以接受额外准备时间,sandwich-middle可能更划算。
简单解释 像给14岁少年讲一样
想象你在玩一个超大型游戏。游戏角色每走一步,系统都要记住以前看到过的地图、敌人和道具。普通做法是每个关卡都保存一整套记忆,所以游戏越长,电脑越吃力。
这篇论文想了个办法:不同关卡可以共享同一个“记忆背包”。作者研究了背包放在前面、后面还是中间,也研究了把关卡分成几组。这样能少占显存,还可能让游戏跑得更快!
结果很有意思:如果开场提示很短、后面玩很久,几乎所有共享方式都比普通方法快;如果开场提示本身很长,把背包放在后面或底部更好,因为放在中间、顶部需要反复等待。保留一半记忆层时,回答质量通常接近普通模型;删得太多,底部方案容易变差。
所以没有“永远最强”的方案。想稳妥省一半空间,用pizza-bottom或lasagna-bottom;想进一步压缩,而且愿意训练更久、等待更长的开场处理,可以试sandwich-middle。未来还要在更大的模型和更长的上下文中继续测试!
术语表
KV cache(键值缓存)
模型保存历史token的Key和Value,生成下一个token时可直接复用,避免重复计算。它会随层数、序列长度和KV头数增长。
论文研究跨层复用KV,以减少缓存容量和访存成本。
Cross-layer KV sharing(跨层KV共享)
多个Transformer层的查询使用同一目标层计算的Key和Value。非KV层因此可以移除WK、WV参数。
LCKV、YOCO、CLA及九种统一框架配置均属于此类。
LCKV
Layer-Condensed KV Cache,只在顶部层计算KV,并让多层查询共享。它能大幅压缩缓存,但可能增加训练约束。
对应本文的sandwich-top配置。
YOCO
You Only Cache Once,在中间层计算KV,并将其提供给上半部分查询;原方法还对下半部分使用高效注意力。
论文中的pizza-bottom与YOCO共享思想相近,但使用标准注意力。
CLA
Cross-Layer Attention把相邻层均匀分组,每组查询共享底层KV。它在缓存节省与实现简洁之间取得平衡。
对应本文的lasagna-bottom配置。
Prefilling(预填充)
模型一次性处理用户已有提示并建立KV缓存的阶段。提示越长,该阶段的时间越重要。
top/middle配置因迭代编码而在长提示下吞吐下降。
开放问题 这项研究留下的未解疑问
- 1 更大模型和超长上下文是否保持相同配置排序尚不清楚;需要跨参数规模、硬件和真实服务负载验证。
- 2 如何消除top/middle的循环依赖,同时保留其强压缩能力,仍缺少低成本、无质量损失的训练与预填充算法。
应用场景
近期应用
显存受限的批量生成
服务商可在1.1B级模型上采用pizza-bottom或lasagna-bottom,将KV缓存约压缩2倍,并在短提示、高并发场景获得高于标准Llama的吞吐。部署前应按RTX 3090级硬件和实际批量重新基准测试。
长输出对话服务
当生成长度显著超过提示长度时,可评估sandwich-middle。它需要额外训练和预填充迭代,但在KV层很少时质量更稳,适合显存比首token延迟更重要的应用。
远期愿景
自适应KV布局引擎
未来推理系统可依据提示长度、生成长度、批量和显存动态选择bottom、top或middle布局,实现速度、容量与质量的自动折中。
原文摘要
Recently, sharing key-value (KV) cache across layers has been found effective in efficient inference of large language models (LLMs). To systematically investigate different techniques of cross-layer KV sharing, we propose a unified framework that covers several recent methods and their novel variants. We conduct comprehensive experiments on all the configurations of the framework, evaluating their generation throughput and performance in language modeling and downstream tasks. We find that when reducing the size of the KV cache by 2$\times$, most configurations can achieve higher throughput than standard transformers while maintaining competitive performance. When further reducing the size of the KV cache, however, pairing queries of all layers with KVs of upper layers performs better, at the expense of additional training cost and prefilling latency. We hope that this work will help users make more informed choices of cross-layer KV sharing approaches and facilitate future research on efficient LLM inference.