GQA，MLA之外的另一种KV Cache压缩方式：动态内存压缩（DMC）

最新推荐文章于 2025-07-29 00:12:12 发布

原创

最新推荐文章于 2025-07-29 00:12:12 发布 · 1.7k 阅读

17 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #论文阅读

0x0. 前言

在openreview上看到最近NV的一个KV Cache压缩工作：https://openreview.net/pdf?id=tDRYrAkOB7 ，感觉思路还是有一些意思的，所以这里就分享一下。

简单来说就是paper提出通过一种特殊的方式continue train一下原始的大模型，可以把模型在generate过程中的KV Cache分成多个段，并且每个token都会学出来一个0或者1的标记，对于标记为1的token的KV Cache可以累加到上一个标记为1的token的KV Cache上（需要这两个token连续才可以累加），而不是concat，这样就可以高效压缩KV Cache，效果也是不错的，且可以配合GQA联合使用。此外，它在continue train或者推理prefill的时候仍然可以用上Flash Attention，推理的decode阶段可以用上Paged Attention。

但是读完方法部分发现这个方法也有几个缺陷，估计会失去工程应用的机会，只能当一篇paper读一下。第一点就是要对整个模型做全量参数的continue train，虽然训练的数据可以很少，但是能把大模型加载起来的成本已经非常高了，普通人肯定这一步就卡死了。第二，这种方法不能from scratch训练，这样就更阻碍了这种方法的广泛应用，毕竟MLA的成功一个重要原因就是因为Deepseek2直接用这个架构from scratch训出来的model并开源。另外，这里的开源链接目前是空的，还不能看到代码细节。

0x1. 摘要

Transformer 已经成为大型语言模型 (LLM) 的核心。然而，由于需要在内存中存储过去token的key value的缓存（KV Cache），而缓存的大小与输入序列长度和batch大小线性相关，因此生成仍然效率低下。为了解决这个问题，paper提出了动态内存压缩 (DMC)，一种在推理时对KV Cache进行在线压缩的方法。最重要的是，模型学习在不同的注意力头和layer中应用不同的压缩率。paper将预训练的 LLM（如 Llama 2（7B、13B 和 70B））改造为 DMC Transformer，在 NVIDIA H100 GPU 上进行自回归推理时，实现了高达 ~3.7 倍的吞吐量提升。DMC 通过在原始数据的极小百分比上进行continue pretrained来应用，而无需添加任何额外的参数。paper发现，DMC 在高达 4 倍的缓存压缩的情况下，保留了原始的下游性能，优于经过微调的分组查询注意力 (GQA) 和key value驱逐策略 (H2O、TOVA)。GQA 和 DMC 可以结合起来获得复合收益。因此，DMC 在任何给定的内存预算内都适合更长的上下文和更大的batch。我们在 https://github.com/NVIDIA/Megatron-LM/tree/DMC 上发布了 DMC 代码和模型。

0x2. 介绍

首先还是提到了大模型推理的时候KV Cache会伴随着序列长度和Batch大小不断增长，这个问题在长文本生成（例如，对话和故事）或为大量用户查询提供服务时更加明显。然后为了缓解这个问题，GQA，Key-Value驱逐策略比如H20，TOVA等被提出，然后Paper说这些方法往往会牺牲预训练模型的精度。另外Flash-Attention等IO-aware或者子平方注意力算法等都无法改善KV Cache。

因此paper提出了动态内存压缩（DMC）方法对KV Cache进行压缩，如图 1 所示，在每个时间步长，DMC 会决定是将当前的key-value表示添加到缓存中，还是对它们与缓存中顶部的项进行加权平均。DMC 中的内存以亚线性方式增长，虽然比不上Linear Attention Transformer推理时的内存恒定，但明显好于Transformer。

在这里插入图片描述

作者团队使用了部分预训练数据（2%对应2倍压缩，4%对应4倍压缩）对应用了DMC的LLama 2（7B，13B）模型进行continue training。 paper在一些下游任务上评估了我们的 DMC 模型，例如 MMLU 用于事实性，QA 数据集用于常识推理，以及 HumanEval 用于代码。作者发现，DMC LLM 保持了与原始 LLM 相似的下游性能，而基线（如 GQA、H2O 和 TOVA）在高压缩率下会造成显著的性能下降。最后，作者表明 DMC 可以与 GQA 混合，使得它们的压缩率相乘。对于使用 GQA 8 倍预训练的 Llama 2 70B，DMC 2 倍可以实现总共 16 倍的压缩。

作者验证了 KV 缓存压缩在实践中可以转化为更有效的生成。最后测量到，DMC 4 倍在不损失性能的情况下，将 Llama 2 7B 和 13B 在 NVIDIA H100 或 A100 GPU 上的推理吞吐量提高了 340% 到 370%。事实上，它使大模型能够在给定的内存预算中容纳更大的batch和更长的序列。

0x3. 动态内存压缩方法（DMC）

LLM 的推理通常受内存限制，而不是计算能力限制。减少 KV Cache的大小可以降低延迟并提高 GPU 利用率。DMC 是一种简单且廉价的在线压缩 KV Cache的方法。通过continue pretrain，可以教会预训练的 LLM 使用 DMC。

0x3.1 推理

考虑自回归推理过程中注意力层的正向传播。在普通 Transformer 中，在每个时间步 $t$ ， $k_t$ 和 $v_t$ 都会被追加到 KV Cache中。另一方面，在 DMC 中，KV Cache更新过程有所不同，如算法 1 所示。首先，预测一个决策变量 $αt∈\alpha_t \in$ {0, 1} (只能取0和1) 和一个重要性变量 $ω_t ∈ [0, 1]$ 。为了避免添加新的参数，我们分别重用 $k_t$ 和 $q_t$ 中的第一个神经元来提取这两个分数。根据 $αt\alpha_t$ ，决定是将 KV 表示 $k_t$ 和 $v_t$ 追加到缓存中，还是将其与缓存的最后一个元素累加。

在这里插入图片描述

具体来说，对于累加，paper根据对当前token预测的重要性分数 $ω$ 和自上次预测 $α=0\alpha = 0$ 以来所有 token 的重要性分数之和 $z_t$ 进行加权平均。事实上， $α\alpha$ 变量有效地将输入序列分割：每个决策决定当前段是否应该继续（ $α=1\alpha = 1$