DiffKV: Differentiated Memory Management for Large Language Models with Parallel KV Compaction

发表时间: 2025-10 · arXiv:2412.03131 (SOSP 2025)

Yanqi Zhang, Yuwei Hu, Runyuan Zhao (Huawei), John C.S. Lui (The Chinese University of Hong Kong), Haibo Chen (Shanghai Jiao Tong University)

速读

一句话结论
DiffKV 提出了一个针对大语言模型 KV Cache 的差异化压缩与并行内存管理框架，在保持近乎无损精度的前提下将 KV Cache 压缩了 2.7 到 5.7 倍，并使推理吞吐量提升了 1.9 到 5.4 倍。

要解决什么问题
大语言模型在推理时需要缓存历史的键值向量（KV Cache）以避免重复计算，这消耗了极大的显存带宽和容量，严重限制了并发请求数。尤其在长上下文和近期涌现的输出超长思维链的思考模型（thinking models）中，显存卡点尤为致命。现有的 KV Cache 压缩方法主要分为剪枝（丢弃不重要的 Token）和量化（降低浮点数精度），但它们都采用了粗放的统一处理策略：第一，对 Key 和 Value 采用相同的量化位宽，忽视了两者在注意力机制中截然不同的作用；第二，对所有保留下来的 Token 使用统一精度，没有根据重要程度做细分；第三，在不同的注意力头（Attention Head）和不同的请求之间，静态且均匀地分配显存，无法适应注意力分布的动态稀疏性。如果直接引入细粒度的动态压缩，会导致各个注意力头产生大量不规则、碎片化的显存需求，在毫秒级的推理步长下，传统的显存管理开销会急剧膨胀，甚至完全抵消压缩带来的性能收益。

怎么做的
DiffKV 的核心思路是在 KV Cache 中引入三个维度的差异化压缩，并配合一套纯 GPU 侧的并行内存管理器来消除碎片化开销。首先是 Key 和 Value 的差异化，作者将注意力计算拆解为方向向量与系数的加权和：

$$ \operatorname{Attn}(\mathbf{Q}, \mathbf{K}, \mathbf{V})_i = \sum_{j=1}^i \underbrace{\mathrm{softmax}\left(\frac{\mathbf{Q}\mathbf{K}^\top}{\sqrt{d}}\right)_{ij} |\mathbf{v}_j|}_{\mathrm{Coefficient}} \underbrace{\frac{\mathbf{v}_j}{|\mathbf{v}_j|}}_{\mathrm{Unit\ vector}} $$

从公式可以看出，Key 参与全局的 Softmax 计算，决定了所有 Token 的注意力分数系数，而 Value 仅影响自身的特征方向。因此 Key 的重要性远高于 Value，DiffKV 为 Key 分配更高的量化精度（如 8 bit），为 Value 分配更低的精度（如 4 bit）。其次是 Token 重要性的差异化：DiffKV 根据 Token 获得的平均注意力分数评估其重要性，并与理论平均值 $\frac{1}{N}$（$N$ 为序列长度）进行对比。若分数大于 $\frac{\alpha_h}{N}$，则存入高精度区；若介于 $\frac{\alpha_l}{N}$ 和 $\frac{\alpha_h}{N}$ 之间，则存入低精度区；若低于 $\frac{\alpha_l}{N}$ 则直接剪枝丢弃。最后是注意力头的动态稀疏性：系统不再设定固定的显存预算，而是允许每个注意力头针对当前请求动态决定需要保留的高低精度 Token 数量。为了解决这种极度不规则的显存分配带来的扩展性挑战，DiffKV 设计了并行 KV 压缩内存管理器。它由三个关键部件构成：一是统一页（Unified Pages），将不同精度的 Token 及其量化元数据打包在固定大小的显存页中，保证访存连续性；二是环形空闲页表（Circular Free Page List），将所有可用和已用的物理页 ID 维护在连续的显存空间中，各注意力头通过并行前缀和算法，无冲突地并行完成页的申请与回收；三是双向页表（Bidirectional Page Table），在一个数据结构中同时管理高低精度页（高精度从左向右增长，低精度从右向左增长），极大地降低了元数据追踪的开销。

效果如何
实验在 NVIDIA L40 GPU 上搭建，评测了 Llama3-8B/70B、Qwen2.5-7B/32B 以及 QwQ-32B、R1-Distill-Qwen-14B 等具备复杂推理能力的思考模型。对比的基线方法包括代表剪枝路线的 H2O、SnapKV、DuoAttention，代表量化路线的 4-bit KV、KIVI、QAQ，以及代表部分加载路线的 Quest。在数学（MATH、AIME24）、代码（HumanEval+）和长文本（LongBench）等任务下，DiffKV 仅使用 19.3% 到 36.7% 的显存，就达到了与 FP16 几乎一致的生成质量（平均精度下降仅 0.3%）。在对错误累积极其敏感的长思维链生成任务（如 QwQ-32B 在 AIME24 上）中，基线方法出现了 16% 甚至接近 100% 的精度崩塌，而 DiffKV 依然保持了无损精度，并凭借释放的显存将批处理大小大幅提升，最终在 QwQ-32B 上实现了 5.4 倍的端到端吞吐量加速。在局限性方面，由于 Qwen2.5-7B 采用了极端的 GQA 架构（Queries-per-KV 比例高达 7），其对 Key 的 4 bit 量化极为敏感，导致在该模型上必须关闭低精度量化档位才能维持效果；此外，DiffKV 的自定义注意力算子在执行时需要读取量化元数据并进行反量化，这部分额外开销使得其实际延迟加速比略低于理论上限。

主要贡献

大型语言模型（LLMs）展现出了卓越的能力，但由于其极高的内存需求（其中键值 KV 缓存是主要的瓶颈），面临着巨大的服务成本。最先进的 KV 缓存压缩技术（如量化和剪枝）对键（Keys）和值（Values）采用统一的处理方式，并完全丢弃不重要的 token，从而忽略了各个 KV 缓存组件在重要性上的细粒度差异。为了解决这些局限性，本文引入了 DiffKV，这是一个新颖的 KV 缓存压缩框架，它利用了 KV 缓存中的三个层次的差异性：（1）键和值对注意力计算的不同影响；（2）token 之间重要性的差异；（3）跨注意力头的多样化动态稀疏模式。

这些差异性层次在不同的请求和注意力头之间引入了不规则的内存使用模式，给内存管理带来了显著的可扩展性挑战。为了应对这些挑战，DiffKV 提出了一种基于 GPU 的内存管理器，该管理器能够并行地将碎片化的空闲内存列表压缩成连续的区域，从而有效地将 KV 缓存中的稀疏性转化为性能收益。作者在多个主流 LLM 上评估了 DiffKV，包括能够生成扩展思维链的新兴推理模型。DiffKV 能够在需要复杂推理和长生成能力的复杂工作负载上，以接近无损的准确率将 KV 缓存压缩 $2.7\times$ 到 $5.7\times$，并将吞吐量提高 $1.9\times$ 到 $5.4\times$。

图1：(a) 剪枝、(b) 均匀量化和 (c) DiffKV 在两个 5 token 请求中跨两个注意力头的 KV 缓存内存分配模式。方框代表保留的 token（注有 token ID），方框大小与内存使用量成正比。百分比表示相对于未压缩 KV 缓存的总内存使用量。

背景知识

大型语言模型架构与注意力机制。大型语言模型（LLMs）主要基于 Transformer 架构构建。Transformer 的核心是注意力机制，它允许序列中的每个 token 在构建其上下文化表示时衡量其他 token 的重要性。在自回归推理期间，注意力机制以因果方式运行，仅关注前面的 token。在数学上，标准的注意力计算定义为：

$$\begin{aligned} \begin{array} { l } { \displaystyle \mathrm { A t t e n t i o n } ( \mathbf { Q } , \mathbf { K } , \mathbf { V } ) _ { i } = \sum _ { j = 1 } ^ { i } \mathrm { s o f t m a x } \left( \frac { \mathbf { Q } \mathbf { K } ^ { \top } } { \sqrt { d } } \right) _ { i j } \mathbf { v } _ { j } } \\ { = \displaystyle \sum _ { j = 1 } ^ { i } \frac { \exp \left( \frac { \mathbf { q } _ { i } \cdot \mathbf { k } _ { j } } { \sqrt { d } } \right) } { \sum _ { n = 1 } ^ { i } \exp \left( \frac { \mathbf { q } _ { i } \cdot \mathbf { k } _ { n } } { \sqrt { d } } \right) } \mathbf { v } _ { j } } \end{array} \end{aligned}$$

其中 $\mathbf{Q}$、$\mathbf{K}$ 和 $\mathbf{V}$ 是大小为 $l \times d$ 的矩阵，分别代表查询、键和值，$l$ 表示序列中迄今为止处理的 token 数量，$d$ 表示特征维度。向量 $\mathbf{q}_i$、$\mathbf{k}_i$ 和 $\mathbf{v}_i$ 对应第 $i$ 个 token 的查询、键和值。为了捕获 token 之间更广泛的交互，Transformer 模型采用多头注意力（MHA）。分组查询注意力（GQA）通过允许多个查询头共享同一组键和值的投影（称为 KV 头）来提高 MHA 的效率。LLM 执行包括两个阶段：提示（prompt）阶段和生成（generation）阶段。为了避免跨生成步骤的冗余计算，引入了 KV 缓存来存储所有先前 token 的键和值，但其大小随序列长度和批处理大小线性增长，迅速成为推理吞吐量的瓶颈。

KV 缓存优化技术。静态 KV 缓存管理系统为最大可能的序列长度保留内存，导致大量内存浪费。vLLM 引入了 PagedAttention【39, Efficient memory management for large language model serving with pagedattention 2023 Symposium on Operating Systems Principles】，将 KV 缓存划分为包含固定数量 token 的页面，并按需分配，从而减少浪费并支持更大的批处理大小。量化技术通过离散的低比特值近似高精度浮点数来减小 KV 缓存大小。对于张量 $\mathbf{X}$，首先计算缩放因子 $s$ 和零点 $z$，然后逐元素应用量化：$\mathbf{Q} = \mathrm{round}\left(\frac{\mathbf{X} - z}{s}\right)$。推理时通过反量化 $\hat{\mathbf{X}} = s \cdot \mathbf{Q} + z$ 近似重建原始张量。最先进的量化方法如 Atom【77, Atom: Low-bit quantization for efficient and accurate llm serving 2024 Proceedings of Machine Learning and Systems】和 Qserve【44, Qserve: W4a8kv4 quantization and system co-design for efficient llm serving 2024 arXiv】独立地对每个键和值向量应用此过程，但这种统一的量化方法忽略了 token 重要性的变化以及键和值在注意力计算中的不同作用。KV 缓存剪枝可视为量化的极端情况，如 H2O【76, H2o: Heavy-hitter oracle for efficient generative inference of large language models 2024 NeurIPS】和 SnapKV【42, Snapkv: Llm knows what you are looking for before generation 2024 arXiv】在所有层和头上均匀分配相同的内存预算；PyramidKV【11, Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling 2024 arXiv】基于经验观察在较低层分配更多内存，但仍依赖静态启发式方法。还有工作将 KV 缓存卸载到 CPU 内存，但这并未从根本上减小缓存大小且引入了传输延迟。

关键观察与设计原则

键和值的差异化影响。虽然可以从注意力分数直接推断不同 token 的重要性，但 token 内键和值向量的作用却不那么明显。检查公式 1 揭示了它们影响的差异。每个 token 对注意力输出的贡献取决于两个因素：来自 softmax 的注意力分数和值向量。键向量作为共享的 softmax 分母的一部分，影响所有 token 的注意力分数，而值向量仅影响其各自 token 对输出的贡献。为了评估由键向量决定的注意力分数相较于值向量的相对重要性，作者将注意力机制重新表述为单位向量的加权和：

$$ \operatorname { A t t n } ( \mathbf { Q } , \mathbf { K } , \mathbf { V } ) _ { i } = \sum _ { j = 1 } ^ { i } \underbrace { \mathrm { s o f t m a x } \left( { \frac { \mathbf { Q } \mathbf { K } ^ { \top } } { \sqrt { d } } } \right) _ { i j } | \mathbf { v } _ { j } | } _ { \mathrm { C o e f f i c i e n t } } \underbrace { { \frac { \mathbf { v } _ { j } } { \left| \mathbf { v } _ { j } \right| } } } _ { \mathrm { U n i t ~ v e c t o r } } $$

该公式将每个输入 token 的贡献分解为两个部分：单位向量 $\frac{\mathbf{v}_j}{|\mathbf{v}_j|}$（通过将值向量除以其 L2 范数获得，仅捕获其在特征空间中的方向）和一个系数（定义为注意力分数与值向量范数的乘积，决定了该 token 方向的相对重要性）。图 2 展示了 Llama3-8B 模型【20, The llama 3 herd of models 2024 arXiv】在使用 Wikitext 数据集【49, Pointer sentinel mixture models 2016 arXiv】时，每个 token 的平均注意力分数和值向量范数的分布。值得注意的是，注意力分数跨越了七个数量级，远远超过了值向量范数的范围（最多覆盖两个数量级）。这种显著的差异突显了注意力分数在决定每个 token 贡献时的关键作用，因此得出结论：键向量比值向量具有更广泛和更具影响力的作用，这激发了在不同精度级别处理键和值向量的探索。
图2：Llama3-8B 中注意力分数和值向量范数的分布。

Token 重要性的差异化。Token 对注意力输出的贡献具有不同程度的重要性，这反映在它们的注意力分数上。利用这些差异，可以应用比统一量化所有 token 或仅剪枝最不重要的 token 更细粒度的压缩策略。图 3 显示了 Llama3-8B 第 8 层在一个随机采样的 Wikitext 序列上的每个 token 的注意力分数。这种 token 重要性的非均匀分布激发了一种基于 token 重要性分配内存的分层压缩策略：最重要 token 高精度存储（例如 K8V4），中等重要 token 较低精度存储（例如 K4V2），最不重要的 token 被剪枝。
图3：Llama3-8B 第8层在从Wikitext随机采样的序列上的每个token的注意力分数。

每头动态稀疏模式。LLMs 表现出跨注意力头和请求的动态稀疏模式：关键 token 的数量不仅在不同头之间存在差异，而且在不同请求下的同一个头也存在差异。作者分析了保留目标百分比（例如 $95\%$）总注意力分数所需的最少关键 token 数量。图 4 说明了 Llama3-8B 中每层所需的平均关键 token 数量，图 5 展示了三个代表性层中每个 KV 头的平均关键 token 数量。稀疏模式保持高度动态：在每一层内，关键 token 的数量在各个 KV 头之间差异很大；在单个 KV 头内，关键 token 的数量在不同请求之间也可能存在很大差异。
图4：Llama3-8B中每层保留总注意力分数95%所需的关键token数量。
图5：Llama3-8B中每个KV头保留总注意力分数95%所需的关键token数量。

主要见解与影响。这些发现揭示了 KV 缓存中的三个关键差异化层次：键对注意力计算的影响大于值；token 的重要性各不相同；注意力稀疏模式跨请求和头变化。这些高度动态的稀疏模式强调了在每个头和每个请求的基础上进行自适应内存管理的必要性。DiffKV 引入了以下关键创新：感知请求的差异化 KV 量化（自适应调整高低精度 token 的混合）；自适应序列长度的重要性估计（短序列保留更多 token，长序列积极压缩）；利用每头动态稀疏性（根据观察到的稀疏模式动态调整每头内存分配）；可扩展的 GPU 内存管理（处理不规则的每头内存分配模式）。

KV 压缩策略细节

基于差异化的 KV 压缩策略概述。KV 压缩策略旨在解决 KV 缓存中的三个差异化层次。首先，为了反映键对注意力计算的更大影响，该策略对键进行比值更高精度的量化（例如 K8V4 或 K4V2）。其次，为了考虑 token 的不同重要性，引入了分层压缩策略：最重要的 token 被高精度量化（如 K8V4），中等重要的 token 被较低精度量化（如 K4V2），最不重要的 token 被完全剪枝。这种策略既感知请求又适应序列长度，它会动态调整每个请求中高精度和低精度 token 的混合比例。对于短序列，保留较大比例的高精度 token 以保证质量；对于长序列，则更激进地应用低精度量化和剪枝以最大化内存节省。

自适应内存管理方法。为了解决跨请求和注意力头的动态注意力稀疏模式，该策略提出了一种自适应内存管理方法。与施加固定的内存预算不同，该方法允许每个注意力头根据其特定的稀疏模式动态确定其内存需求。在提示阶段和生成阶段，压缩都在每个请求和每个头的基础上应用，确保内存使用适合特定的稀疏模式。

提示阶段的压缩逻辑。在提示阶段，计算提示中所有 token 的键和值向量。然后，压缩策略根据每个 token 的重要性确定其适当的存储精度。第 $i$ 个 token 的重要性是通过平均它从后续 token 接收到的 $N - i$ 个注意力分数来计算的（其中 $N$ 是提示序列长度）。对于 GQA 和 MHA，使用最大值操作聚合与 KV 头关联的所有注意力头的分数。为了防止过早压缩，最近的 $W$ 个 token 始终以高精度量化（$W$ 通常设置为 64）。对于剩余的 token，策略以自适应序列长度的方式确定第 $i$ 个 token 的精度级别，通过将其重要性分数与理论平均值 $\frac{1}{i}$ 进行比较。具体而言，如果第 $i$ 个 token 的重要性分数超过 $\frac{\alpha_h}{i}$，则将其高精度量化；如果分数在区间 $[\frac{\alpha_l}{i}, \frac{\alpha_h}{i}]$ 内，则低精度量化；否则被剪枝。参数 $\alpha_l$ 和 $\alpha_h$ 是通过在校准数据集上离线分析确定的。结果，KV 缓存被概念上分为两部分：高精度部分 $KV_h$ 和低精度部分 $KV_l$。

生成阶段的压缩逻辑。在生成阶段，为了与生成过程的自回归性质保持一致，每步仅压缩一个 token。最近的 token 被添加到最近窗口中以防止过早压缩，而窗口中最早的 token $t_c$ 成为更激进压缩的候选者。压缩过程分为两部分。首先，token $t_c$ 被高精度或低精度量化并添加到 KV 缓存的相应部分，或者被完全剪枝。接着，如果 $t_c$ 被量化，KV 缓存相应精度部分中最不重要的 token $t_v$ 将被考虑进一步降级：它可能被重新量化为较低精度或剪枝。这种策略为不太重要的 token 建立了一条平滑的降级路径：它不是被直接剪枝，而是首先被重新量化为低精度，只有在它仍然不重要时才被剪枝。

算法执行流程描述。具体的算法执行流程如下：给定序列长度 $N$，如果候选 token $t_c$ 的重要性超过 $\frac{\alpha_h}{N}$，则将其高精度量化并添加到高精度 KV 缓存 $KV_h$ 中。随后，识别 $KV_h$ 中最不重要的 token 作为受害者 token $t_v$。如果 $t_v$ 的重要性分数仍然超过 $\frac{\alpha_h}{N}$，它保留在 $KV_h$ 中；如果落在 $[\frac{\alpha_l}{N}, \frac{\alpha_h}{N}]$ 之间，则 $t_v$ 被重新量化为低精度并移动到低精度 KV 缓存 $KV_l$ 中；否则，$t_v$ 被剪枝。类似地，如果 $t_c$ 的重要性位于 $[\frac{\alpha_l}{N}, \frac{\alpha_h}{N}]$ 之间，它被低精度量化并添加到 $KV_l$ 中。然后指定 $KV_l$ 中最不重要的 token 作为受害者 $t_v$，如果其重要性低于 $\frac{\alpha_l}{N}$，则将其进一步剪枝。

策略的扩展性讨论。所提出的 KV 压缩策略具有高度可扩展性，允许使用多于两个的量化精度级别。本文主要采用两个量化精度级别（K8V4-K4V2），以最小化元数据开销并提高系统效率，经验表明这足以在多个模型和基准测试中实现接近无损的生成质量。所有注意力头共享一组阈值，因为经验研究表明这足以捕捉不同头之间的变化稀疏模式。DiffKV 也能有效支持更灵活的策略（如为每个头单独调整阈值）。

内存管理细节

内存管理的挑战。在 PagedAttention 中，所有 token 以相同精度存储，允许固定的页面格式。然而，键和值向量的差异化以及 token 重要性的变化引入了多个精度级别，使得固定页面格式不再适用。固定的页面格式需要保守地为所有 token 分配最高精度槽，导致严重的内存浪费（例如，K4V2 精度的 token 在 K8V4 槽中会浪费 $50\%$ 的内存），并导致未对齐的内存访问，降低带宽利用率。此外，在每个推理步骤中，DiffKV 必须处理不同头之间数量不等的的高低精度 token，这种被称为 KV 压缩（KV compaction）的过程其复杂性为 $O(\#requests \times \#heads)$。在毫秒级的模型执行时间内，这种动态分配的开销如果不加以妥善管理，将抵消缓存压缩带来的好处。

并行 KV 压缩机制。对 KV 压缩过程的分析揭示了并行化的机会。KV 压缩可分为规划和协调两个阶段。规划阶段中，每个注意力头独立确定其内存分配需求，这非常适合 GPU 的并行计算能力。协调阶段同步这些需求并将其映射到物理内存，这可以通过并行前缀和（parallel prefix sum）【51, An optimal parallel prefix-sums algorithm on the memory machine models for gpus 2012 International Conference on Algorithms and Architectures for Parallel Processing】有效实现（前提是空闲内存区域是连续的）。因此，DiffKV 提出了一种新型的并行 KV 压缩技术，该技术通过三个驻留在 GPU 上的数据结构来实现：统一页面、循环空闲页面列表和双向页表。

统一页面（Unified Pages）设计。统一页面抽象掉了单个 token 内和不同 token 之间不同精度的复杂性。GPU 内存被划分为大小均匀的页面，每个页面在分配时被配置为以特定精度存储 token。每个统一页面被组织为六个段：量化的键、键的量化元数据、量化的值、值的量化元数据、token 分数和位置。量化元数据包括缩放因子和零点。每个页面存储的 token 数量根据量化配置进行调整，确保内存紧凑。通过将键、值及其元数据整合到一个单一结构中，统一页面增强了数据局部性，消除了分散查找的需要。

循环空闲页面列表（Circular Free Page List）。该列表是并行 KV 压缩的基石，通过在连续区域中维护空闲和已用页面来促进内存分配和回收的并行化。这个集中式的 GPU 数据结构包含所有页面 ID，并通过一对指针跟踪空闲页面：用于分配的起始指针和用于回收的结束指针。当指针到达列表末尾时会回绕到开头。分配页面时，起始指针前进；释放页面时，结束指针前进。在并行 KV 压缩中，每个头确定要分配或释放的页面数后，并行前缀和操作计算每个头相对于指针的唯一偏移量。这为每个头分配了列表中不相交的区域以进行读取或写入，确保操作不冲突。

双向页表（Bidirectional Page Table）。为了避免为高低精度页面分别维护页表所带来的双倍元数据开销，DiffKV 引入了双向页表。在双向页表的每个条目中，高精度页面 ID 从列表左侧增长，而低精度页面 ID 从右侧增长，动态适应工作负载的精度需求。条目的长度由最大序列长度除以每个高精度页面的 token 数决定，确保不会溢出。这种统一的方法最小化了元数据开销，并消除了基于精度级别进行单独查找的需要。其内存开销极小（例如，在 Llama3-8B 批大小为 128 时，所有双向页表的总大小仅为 32 MB）。
图6：提示阶段的内存管理流程。

提示阶段的 KV 压缩工作流。在提示阶段（如图 6 所示的 8 个 token 的请求示例，高精度页存 2 个 token，低精度存 4 个），由于先验不知道每个头确切需要的高低精度页面数，首先保守地为每个头分配 4 个统一页面（假设全为高精度）。页面 5-8 和 9-12 被分配给头 A 和头 B，循环列表的结束指针前进到页面 13。进入规划阶段，每个头独立计算需求：头 A 需要 1 个高精度页和 1 个低精度页；头 B 需要 2 个高精度页和 1 个低精度页。高精度页从左向右分配，低精度页从右向左分配。在协调阶段，通过并行前缀和回收未使用的页面（头 A 的 6-7，头 B 的 11），这些页面被追加到循环空闲页面列表中。

生成阶段的 KV 压缩工作流。在生成阶段，仅当高精度或低精度页面已满时，头部才会分配新页面（每步最多分配一个附加页面）。每个头独立检查其页面可用性，并在需要时使用基于前缀和的方法并行分配新页面。与提示阶段不同，生成期间不执行页面回收。一旦请求完成，为该请求分配的所有页面都将被回收。如果需要支持额外的精度级别，可以通过组合单向或双向页表来实现。

系统实现细节

DiffKV 架构概述。DiffKV 是在 vLLM 之上实现的，包含 4.5K 行 CUDA/C++ 代码和 9K 行 Python 代码。在每个推理步骤中，调度器将尽可能多的请求批处理到可用 GPU 内存中，并将选定的请求发送给所有 worker。每个 GPU 托管一个 worker，负责执行模型的一个分区。每个 worker 包含一个专用的内存管理器（负责其分配的注意力头的 KV 缓存）和一个模型计算执行引擎。执行引擎集成了 KV 压缩器和自定义 GPU 注意力内核。计算出键和值向量后，调用 KV 压缩器执行压缩策略，结果存储在 KV 缓存中。然后，自定义 GPU 注意力内核使用压缩的 KV 缓存计算注意力输出。
图7：DiffKV架构。

高效的自定义注意力内核设计。作者开发了一个自定义注意力内核，以有效支持差异化 KV 缓存压缩。该内核为每个序列分配一个线程块（thread block）来处理单个注意力头。为了减轻混合精度量化可能引起的负载不平衡，线程束（thread warps）首先迭代高精度页面，然后迭代低精度页面。每个页面分两个阶段处理：查询和键之间的点积，以及值的加权和。

键（Key）处理的内存布局优化。在键处理中，每个 warp 每次处理一个页面，线程分组负责不同的键向量。为了确保合并和向量化的内存访问，简单的布局（如 $[F, N_{tokens}]$）会导致不连续的跨步内存访问。因此，将键的布局组织为 $[F / (K_{vec} \times K_{group}), N_{tokens}, K_{group}, K_{vec}]$，其中 $K_{vec}$ 是向量化因子，$K_{group}$ 是每组线程数。在执行期间，warp 中所有线程的组合内存访问是连续的，从而实现了内存合并并最大化了带宽利用率。

值（Value）处理的内存布局优化。在值处理中，页面中的 token 均匀分布在 warp 中的线程组上；每个线程在特征维度上执行求和归约，并将结果保存在寄存器中。由于键处理中沿特征维度的向量化会显著增加值处理中的寄存器压力，因此对 token 维度应用向量化。值的布局相应地组织为 $[F / V_{group}, N_{tokens} / V_{vec}, V_{group}, V_{vec}]$。此外，对于超长序列，注意力内核支持沿序列维度的并行化，将序列拆分为多个段并在单独的线程块中并行处理，然后合并结果。

实验环境

数据集：一般知识（MMLU、MMLU-Pro）、数学（GSM8K、MATH）、代码生成（HumanEval+、MBPP+）、长上下文理解（LongBench）。对于思维模型，使用数学竞赛数据集 AIME24 和研究生级科学推理 GPQA。
模型架构：Llama3-8B 和 70B，Qwen2.5-7B 和 32B，以及新兴的思维模型 QwQ-32B、R1-Distill-Qwen-14B 和 R1-Distill-Llama-8B。模型权重以 FP16 精度存储。
硬件配置：NVIDIA L40 GPU（每块 48 GB 内存）。小型模型（如 Llama3-8B）使用单卡，大型模型（如 Llama3-70B）使用四卡并行，QwQ-32B 使用两卡。
软件配置：基于 vLLM 实现。

实验结果

差异化 KV 量化的有效性。实验在 GSM8K 和 HumanEval+ 上评估了 K8V4 和 K4V2 配置，并与反向配置（K4V8、K2V4）及倾斜变体（K8V2、K4V1）进行对比。结果显示，K8V4 匹配了 FP16 的准确率，而其反向配置 K4V8 出现显著准确率下降（特别是在 Qwen2.5-7B 上降至接近零，因其 GQA 比例高达 7）。K4V2 在 Llama3-8B/70B 上保留了超过 $65\%$ 的准确率，而 K2V4 准确率接近零。K4V1 产生几乎为零的准确率。这证实了键比值起着更关键的作用，并证明了 K8V4-K4V2 两级方案的合理性。
图8：差异化KV量化的准确率。

动态稀疏性的有效性。对比 DiffKV 的每头动态稀疏性与 SnapKV 的静态稀疏性。在 Llama3-8B 上，动态稀疏性在剪枝 $50\%$ token 时保持 GSM8K 全准确率，在更敏感的 HumanEval+ 上剪枝 $10\%$ 保持全准确率，显著优于静态稀疏性。
图9：动态与静态稀疏性的准确率对比。

参数校准。在 MATH 数据集的训练集上校准阈值参数 $\alpha_h$ 和 $\alpha_l$。根据分析结果，选择 Llama3-8B/70B 的 $\alpha_h = 1$，Qwen2.5-32B 和 QwQ-32B 的 $\alpha_h = 3$。对于低精度阈值，Llama3-8B 选 0.02，Qwen2.5-7B 选 0.04，其余模型为 0。
图10：在MATH数据集训练集上校准高低精度阈值...

差异化压缩策略的端到端评估。与 H2O、SnapKV、DuoAttention、4-bit KV、KIVI、QAQ 和 Quest 等基准进行比较。DiffKV 在使用 $19.3\%$ 到 $36.7\%$ 内存的情况下，实现了相对于 FP16 接近无损的准确率（平均下降仅 $0.3\%$）。在 LongBench 长上下文场景中，DiffKV 同样以更少的内存使用量实现了卓越的准确率。在具有挑战性的思维模型（QwQ-32B 等）和 AIME24/GPQA 任务上，DiffKV 使用 $23.5\%$ 到 $29.4\%$ 的内存实现了与 FP16 相当的生成质量，而其他基准方法则经历了显著的质量下降（例如在 GPQA 上剪枝方法的准确率降至几乎为零）。
图11：DiffKV的KV缓存内存（归一化为vLLM）与基准测试准确率的权衡。
图12：不同基准测试和模型中被剪枝、量化为低精度和量化为高精度的token比例。

系统性能：内存管理开销与延迟加速。与 CPU 上的多线程内存管理相比，DiffKV 的 GPU 并行 KV 压缩将内存管理延迟降低了多达三个数量级。内存管理开销在提示阶段占总延迟不到 $0.2\%$，在生成阶段不到 $0.9\%$。DiffKV 的自定义注意力内核实现了与 KV 缓存大小缩减成正比的近线性加速（例如 K8V8 理论加速 $2\times$，实际达到 $1.7\times$）。在端到端推理延迟方面，对于长度为 4096 的序列，DiffKV 比 vLLM 加速 $1.4\times$ 到 $1.6\times$。

图14：DiffKV的延迟分解。
图15：DiffKV相对于vLLM的延迟加速。

吞吐量加速与动态工作负载。在 MATH 数据集上的长生成评估中，DiffKV 始终实现比先前系统更高的吞吐量。在 QwQ-32B 模型上，DiffKV 实现了比 vLLM 高 $5.4\times$ 的吞吐量（远高于 Quest、SnapKV 等）。这归功于压缩技术支持了更大的批处理大小（例如 QwQ-32B 上批大小从 2.7 提升至 15.9）。在动态泊松到达工作负载下，DiffKV 始终实现更低的延迟，并在排队延迟急剧增长之前承受更高的负载。
图16：动态工作负载下DiffKV与vLLM的平均延迟比较。
图17：吞吐量和实现的批处理大小。

结论

本文提出了 DiffKV，该框架通过利用 KV 缓存中三个层次的差异性（键和值的差异化精度、基于 token 重要性的分层压缩、每头动态稀疏性）来提高 LLM 服务效率。DiffKV 的核心是并行 KV 压缩技术，它能有效处理跨请求和注意力头的不规则内存需求，将内存节省转化为性能提升。评估表明，DiffKV 能够以接近无损的准确率将 KV 缓存压缩 $2.7\times$ 到 $5.7\times$，即便是对于需要复杂推理和长生成能力的思维模型和复杂工作负载也是如此，并将吞吐量提高了 $1.9\times$ 到 $5.4\times$，优于先前的 KV 缓存压缩方法。