MosaicKV：使用动态二维KV缓存压缩服务长上下文LLM

# MosaicKV：面向超长上下文LLM服务的动态二维KV缓存压缩

**来源：** https://arxiv.org/html/2607.00760

**强盛**
上海交通大学 并行与分布式系统研究所，上海，中国
qiangsheng@sjtu\.edu\.cn (https://arxiv.org/html/2607.00760v1/mailto:[email protected])

**陈瑞伟**
上海交通大学 并行与分布式系统研究所，上海，中国
chenruiwei@sjtu\.edu\.cn (https://arxiv.org/html/2607.00760v1/mailto:[email protected])

**吴银鹏**
上海交通大学 并行与分布式系统研究所，上海，中国
wyp1536481268@foxmail\.com (https://arxiv.org/html/2607.00760v1/mailto:[email protected])

**顾金宇**
上海交通大学 并行与分布式系统研究所，上海，中国
gujinyu@sjtu\.edu\.cn (https://arxiv.org/html/2607.00760v1/mailto:[email protected])

**华志超**
上海交通大学 并行与分布式系统研究所，上海，中国
zchua@sjtu\.edu\.cn (https://arxiv.org/html/2607.00760v1/mailto:[email protected])

**夏虞斌**
上海交通大学 并行与分布式系统研究所，上海，中国
xiayubin@sjtu\.edu\.cn (https://arxiv.org/html/2607.00760v1/mailto:[email protected])

**臧斌宇**
上海交通大学 并行与分布式系统研究所，上海，中国
byzang@sjtu\.edu\.cn (https://arxiv.org/html/2607.00760v1/mailto:[email protected])

**陈海波**
上海交通大学 并行与分布式系统研究所，上海，中国
haibochen@sjtu\.edu\.cn (https://arxiv.org/html/2607.00760v1/mailto:[email protected])

###### 摘要. 长上下文LLM服务现已支持数十万至数百万token的提示词，这使得键值（KV）缓存成为首要的服务成本。由于缓存随上下文长度线性增长，会耗尽GPU内存、迫使批量变小、降低服务吞吐量。现有的KV缓存压缩技术通常仅针对序列维度或通道维度，随着上下文窗口扩展，其压缩空间有限。同时压缩两个维度有望实现更高的内存缩减，但直接应用这两种压缩会导致显著的精度损失。本文介绍MosaicKV，一种面向极长上下文服务的动态二维KV缓存压缩系统。MosaicKV采用*动态二维压缩*来应对精度挑战，利用KV缓存中元素重要性分布的非均匀性。MosaicKV并非全局应用单一压缩模式，而是为每个KV向量识别重要元素，并以KV缓存段为粒度选择压缩策略。为了解决性能挑战（即细粒度稀疏性和压缩管理开销可能抵消压缩带来的增益），MosaicKV引入了*压缩KV缓存管理*。该机制利用未充分利用的GPU和CPU资源来维护压缩后的KV缓存并加速注意力计算。在H800 GPU上对多种LLM进行的评估表明，与未压缩的基线相比，MosaicKV实现了高达16×\times注意力加速、4.8×\times更低的解码延迟和7.3×\times更高的吞吐量。同时，它将内存使用量减少了3×\times，并且在LongBench和RULER上仅导致平均1.76%的精度损失。

## 1. 引言

大语言模型（LLM）越来越多地部署在需要极长上下文的场景中，例如仓库级代码生成(Zhang et al.,2023a (https://arxiv.org/html/2607.00760#bib.bib57),2024 (https://arxiv.org/html/2607.00760#bib.bib58))、多文档问答(Wang et al.,2024 (https://arxiv.org/html/2607.00760#bib.bib46); Cheng et al.,2025 (https://arxiv.org/html/2607.00760#bib.bib18))以及多轮智能体工作流(Li et al.,2026 (https://arxiv.org/html/2607.00760#bib.bib31); Wei et al.,2025 (https://arxiv.org/html/2607.00760#bib.bib47); cla,2026 (https://arxiv.org/html/2607.00760#bib.bib5); ope,2026 (https://arxiv.org/html/2607.00760#bib.bib9))。为此，前沿模型(noa,2026b (https://arxiv.org/html/2607.00760#bib.bib8),a (https://arxiv.org/html/2607.00760#bib.bib7),2025b (https://arxiv.org/html/2607.00760#bib.bib4); ant,2026 (https://arxiv.org/html/2607.00760#bib.bib10); noa,2025a (https://arxiv.org/html/2607.00760#bib.bib3))已迅速将其支持的上下文窗口扩展到1M token，例如GPT-5.4(noa,2026a (https://arxiv.org/html/2607.00760#bib.bib7))和Gemini 3.1 Pro(noa,2025b (https://arxiv.org/html/2607.00760#bib.bib4))。然而，长上下文通过键值（KV）缓存造成了严重的内存瓶颈，其大小在推理过程中随序列长度线性增长。例如，在128K token下同时服务8个LLaMA-3.1-8B请求，仅KV缓存就需要128 GB——已经达到模型权重的8×\times；在1M token时，这一需求增长到1 TB，是模型权重的64×\times。这种内存压力会耗尽GPU容量，迫使批量大小急剧减小，并显著增加每步解码延迟，最终降低服务吞吐量。

先前的工作通过KV缓存压缩(Zhang et al.,2023b (https://arxiv.org/html/2607.00760#bib.bib60); Tang et al.,2024 (https://arxiv.org/html/2607.00760#bib.bib45); Xu et al.,2025 (https://arxiv.org/html/2607.00760#bib.bib51); Lee et al.,2024 (https://arxiv.org/html/2607.00760#bib.bib29); Sun et al.,2025 (https://arxiv.org/html/2607.00760#bib.bib44))来缓解这一瓶颈，该过程沿两个维度进行：*序列压缩*和*通道压缩*。序列压缩仅保留历史token的一个子集用于注意力计算。例如，Quest(Tang et al.,2024 (https://arxiv.org/html/2607.00760#bib.bib45))在运行时执行查询感知的token选择，以实现高精度和低延迟。相比之下，通道压缩保留每个token KV向量中的一部分通道。例如，ThinK(Xu et al.,2025 (https://arxiv.org/html/2607.00760#bib.bib51))根据预填充键缓存计算统一的通道掩码，并将其应用于选择所有键向量中的重要元素。然而，每个单独维度无法从另一维度受益，随着上下文窗口继续扩展，很难进一步提高压缩率。因此，很自然的下一步是将两个压缩维度结合起来，以实现对极长上下文服务的更高压缩率。

**挑战1：精度损失。** 最近的工作尝试结合两个压缩维度。例如，ThinK(Xu et al.,2025 (https://arxiv.org/html/2607.00760#bib.bib51))首先对键向量应用通道压缩，然后选择重要的token用于注意力计算。然而，它观察到当同时对KV缓存应用通道和序列压缩时，精度会显著下降。因此，ThinK仅对键向量应用二维压缩，而值向量仅进行序列压缩。其他现有工作，包括InfiniGen(Lee et al.,2024 (https://arxiv.org/html/2607.00760#bib.bib29))和ShadowKV(Sun et al.,2025 (https://arxiv.org/html/2607.00760#bib.bib44))，也无法在token选择和注意力计算的整个过程中对键和值向量同时应用二维压缩。根本原因是，简单地将这种全生命周期的二维压缩应用于KV缓存会显著降低模型精度。在我们的评估中，在序列压缩基础上再应用70%的通道压缩率会导致82.8%的精度损失（图3 (https://arxiv.org/html/2607.00760#S3.F3)）。

参见标题图1.不同的KV缓存压缩方法。

*观察与方法1：动态二维压缩。* 我们观察到KV缓存中元素的重要性分布是非均匀的。现有工作未考虑这种变化，而是对整个KV缓存应用“一刀切”的压缩策略，导致在二维压缩下出现严重的精度下降。为了解决这些局限性，本文介绍了MosaicKV，一种新颖的框架，利用KV缓存差异化的重要性分布，以低内存使用、高解码性能和低精度损失来服务长上下文LLM推理。为了克服精度挑战，MosaicKV引入了动态二维压缩，该压缩采用动态压缩策略执行二维压缩。MosaicKV利用了两个层次的差异化。首先，即使在SVD旋转之后，各个KV向量的重要性分布也不同，如图4 (https://arxiv.org/html/2607.00760#S3.F4)所示。即使使用仅256 token的小上下文，只有62.28%的前25%重要元素位于全局前25%重要通道中。与对所有KV向量应用统一通道掩码的现有工作不同，MosaicKV引入了*逐向量元素选择*，即为每个KV向量单独识别重要元素。其次，KV缓存区域之间的元素分布特征也有所不同。例如，在某些区域，元素幅度方差较大，存在明显的离群点；而在其他区域，元素分布更均匀，幅度相似。基于这一观察，MosaicKV引入了*动态压缩策略*，将KV缓存划分为多个段，并为每个段应用合适的策略。动态二维压缩方法使得MosaicKV能够在整个解码阶段对KV缓存应用二维压缩，从而提高压缩率同时保持模型精度。

**挑战2：性能下降。** 尽管MosaicKV的动态二维压缩实现了高KV缓存压缩率，但它引入了细粒度的稀疏内存访问模式和大量的内存管理开销。这两者都会损害推理性能，甚至可能比未压缩的基线更差。对于前者，逐向量元素选择会产生细粒度的稀疏KV缓存矩阵，具有较差的缓存局部性。访问这样的矩阵会显著增加GPU缓存未命中率，抵消压缩可能带来的加速。对于后者，动态压缩策略将大量的KV缓存管理开销（包括内存划分、SVD计算和重新压缩）引入关键的解码阶段，进一步降低性能。

*观察与方法2：压缩KV缓存管理。* 我们观察到，在长上下文解码过程中，注意力计算在内存使用和延迟方面均占主导地位。而且，解码阶段的注意力表现出高度不平衡的工作负载：GPU内存带宽完全饱和，而GPU计算单元（尤其是CUDA核心）和CPU即使在FlashAttention(Dao et al.,2022 (https://arxiv.org/html/2607.00760#bib.bib20))下也处于显著未利用状态。在A800 GPU上对LLaMA-3.1-8B的256K token请求进行FlashInfer(Ye et al.,2025 (https://arxiv.org/html/2607.00760#bib.bib54))性能分析时，我们发现内存带宽利用率达到90.5%，而CUDA核心利用率仅为10.35%，几乎所有CPU核心都处于空闲状态。为了克服性能挑战，MosaicKV引入了一种新颖的**压缩KV缓存管理**方法，利用这些未充分利用的资源来管理压缩后的KV缓存并加速注意力计算。MosaicKV首先引入了*紧致稀疏注意力*，以加速对二维压缩KV缓存的注意力计算。它将稀疏的KV缓存矩阵打包成GPU内存中的稠密格式，并利用未充分利用的CUDA核心（而非张量核心）在不解包的情况下执行注意力计算。这种方法显著降低了GPU缓存未命中率和GPU内存使用量，从而加速了注意力计算。此外，MosaicKV引入了*异构双缓冲区压缩*，以减少压缩KV缓存管理的开销。GPU端缓冲区存储临时压缩的、新生成的KV向量，这些向量用于中间的解码计算。CPU端缓冲区用于执行异步的KV缓存管理操作，包括划分KV缓存段、计算SVD旋转矩阵、生成压缩策略以及重新压缩新段。一种切换机制在不解码阶段阻塞的情况下，用最终压缩版本替换临时压缩的KV段。

我们实现了MosaicKV的原型，并在H800 GPU上对多个LLM模型进行了评估。与未压缩的基线相比，MosaicKV实现了高达16×\times注意力加速、4.8×\times更低的解码延迟和7.3×\times更高的吞吐量，同时将GPU内存使用量减少了3×\times。尽管有这些性能提升，MosaicKV仍保持了模型精度，在LongBench(Bai et al.,2024 (https://arxiv.org/html/2607.00760#bib.bib13))和RULER(Hsieh et al.,2024 (https://arxiv.org/html/2607.00760#bib.bib25))基准测试上仅导致平均1.76%的精度损失。本文做出了以下贡献：

- • 一种动态二维压缩方法，在整个解码阶段对KV缓存同时应用通道和序列压缩。它采用逐向量元素选择和动态压缩策略，以在高压缩率下保持精度。
- • 一种压缩KV缓存管理系统，利用未充分利用的GPU和CPU资源来管理压缩后的KV缓存并加速注意力计算。
- • 原型实现与评估。MosaicKV实现了高达16×\times注意力加速、4.8×\times更低的解码延迟和7.3×\times更高的吞吐量，同时与未压缩基线相比仅导致平均1.76%的精度损失。

## 2. 背景与动机

### 2.1. 长上下文LLM推理中的大规模KV缓存

大语言模型（LLM）越来越多地部署在需要极长上下文的场景中，包括仓库级代码生成(Zhang et al.,2023a (https://arxiv.org/html/2607.00760#bib.bib57),2024 (https://arxiv.org/html/2607.00760#bib.bib58))、多文档问答(Wang et al.,2024 (https://arxiv.org/html/2607.00760#bib.bib46); Cheng et al.,2025 (https://arxiv.org/html/2607.00760#bib.bib18))以及多轮智能体工作流(Li et al.,2026 (https://arxiv.org/html/2607.00760#bib.bib31); Wei et al.,2025 (https://arxiv.org/html/2607.00760#bib.bib47))。前沿模型(noa,2026b (https://arxiv.org/html/2607.00760#bib.bib8),a (https://arxiv.org/html/2607.00760#bib.bib7),2025b (https://arxiv.org/html/2607.00760#bib.bib4); ant,2026 (https://arxiv.org/html/2607.00760#bib.bib10); noa,2025a (https://arxiv.org/html/2607.00760#bib.bib3))现在支持1M token的上下文窗口，例如Claude 4.6(ant,2026 (https://arxiv.org/html/2607.00760#bib.bib10))和GPT-5.4(noa,2026a (https://arxiv.org/html/2607.00760#bib.bib7))。LLM推理包括一个处理输入token的预填充阶段和一个自回归生成输出token的解码阶段。为了避免冗余重新计算，现代推理框架维护一个KV缓存，用于在每个注意力层存储所有历史token的键（K）和值（V）向量。KV缓存驻留在GPU内存中以最小化计算延迟。然而，它的大小随上下文长度线性增长，成为主要的内存瓶颈。对于在128K token下同时服务8个LLaMA-3.1-8B请求，KV缓存需要128 GB（模型权重的8×\times）；在1M token时，这增加到1 TB（模型权重的64×\times）。这种内存压力也增加了注意力延迟：注意力计算