HARD-KV: 解码时 KV 压缩的头部自适应正则化

# Hard-KV: 解码时KV压缩的头自适应正则化

来源：https://arxiv.org/html/2606.28831

Feiyang Ren Bowen Zeng Dalin Zhang Jinpeng Chen Gang Chen Huan Li

###### 摘要

长上下文LLM推理面临一个根本性冲突：头自适应压缩算法（例如Top-p核采样）通过动态波动内存预算提供更高的准确性，而现代推理引擎（如vLLM）要求刚性、静态的内存模式以利用CUDA Graphs和PagedAttention。我们通过**Hard-KV**解决这一“静态-动态”失配问题，这是一个将动态选择与刚性系统约束统一起来的框架。**Hard-KV**引入了级联缓存(Cascade Cache)层次结构，管理令牌在密集、稀疏和压缩层级之间的生命周期。关键地，我们提出了一种对数几率校准(Logits Calibration)机制，将不同的重要性指标归一化到统一的概率空间，从而实现跨异构头的统一Top-p预算分配。为弥合效率差距，我们提供了一个系统级解决方案，将碎片化的动态索引重写为与高性能推理引擎兼容的连续物理布局。在数学推理基准（AIME、U-Math）上的大量实验验证，**Hard-KV**在10k+令牌场景下相比静态基线实现了高达2倍的吞吐量提升，同时保持高保真生成。代码可在https://github.com/SuDIS-ZJU/HARDInfer获取。

机器学习，ICML

## 1 引言

在长上下文场景中部署大型语言模型（LLM）日益受到键值（KV）缓存内存线性增长的限制。随着序列长度扩展到数万，KV缓存的内存占用经常超过模型权重，成为吞吐量和延迟的关键瓶颈。为缓解这一问题，解码时压缩方法（Liao等人，2025 (https://arxiv.org/html/2606.28831#bib.bib32)；Wu等人，2025 (https://arxiv.org/html/2606.28831#bib.bib34)；Behnam等人，2025 (https://arxiv.org/html/2606.28831#bib.bib21)；Jo等人，2025 (https://arxiv.org/html/2606.28831#bib.bib51)；Zeng等人，2025 (https://arxiv.org/html/2606.28831#bib.bib58)；Zhang等人，2025a (https://arxiv.org/html/2606.28831#bib.bib61)）已经出现，通过在生成过程中仅保留最关键令牌来动态减少内存使用。该领域的核心张力在于算法准确性与系统效率之间。最近的研究（Fu等人，2024 (https://arxiv.org/html/2606.28831#bib.bib5)；Xiao等人，2024 (https://arxiv.org/html/2606.28831#bib.bib35)；Tang等人，2024 (https://arxiv.org/html/2606.28831#bib.bib36)）表明，“头自适应”(head-adaptive)策略——允许缓存大小根据每个注意力头的信息密度动态波动——提供了优越的性能。这种准确性对于将压缩扩展到解码时间至关重要，特别是对于数学推理等任务，在这些任务中，错误会在长生成轨迹中累积（Chen等人，2025a (https://arxiv.org/html/2606.28831#bib.bib24)）。Lin等人（2025 (https://arxiv.org/html/2606.28831#bib.bib8)）认为，头自适应预算分配的优越性源于注意力的内在动态性，证明了简单而有效的Top-p（核采样(Holtzman等人，2019 (https://arxiv.org/html/2606.28831#bib.bib25)））足以用于选择。类似地，其他工作（Xuanfan Ni等人，2025 (https://arxiv.org/html/2606.28831#bib.bib38)；Feng等人，2025b (https://arxiv.org/html/2606.28831#bib.bib40)，a (https://arxiv.org/html/2606.28831#bib.bib41)）提出了依赖于动态分配的性能优先压缩的新标准。然而，这种动态性质与现代推理引擎从根本上不兼容。头自适应选择生成的不均匀、不可预测且碎片化的内存索引与标准优化（如PagedAttention（Kwon等人，2023b (https://arxiv.org/html/2606.28831#bib.bib12)）、连续批处理（Yu等人，2022 (https://arxiv.org/html/2606.28831#bib.bib19)）和CUDA Graphs（NVIDIA Corporation，2024 (https://arxiv.org/html/2606.28831#bib.bib16)））相冲突，所有这些都依赖于*静态*或*规则化*的内存块结构。

**挑战：** 我们将核心冲突描述为*动态、上下文相关的注意力模式*与*静态、请求无关的系统设计*之间的差距。为弥合这一差距，我们提出**Hard-KV**（解码时KV压缩的头自适应正则化），这是一个旨在现代推理架构约束内实现动态缓存分配的框架。为弥合这一差距，**Hard-KV**引入了一种级联缓存架构，将KV内存组织为三个序列级层级。我们进一步提出了一种对数几率校准(Logits Calibration)过程，将Top-k选择转换为与Top-p对齐的对应项，允许不同的KV选择策略共享一个统一的自适应预算分配机制。这种校准产生的缓存模式更符合原始注意力分布，因此自然与三层级缓存兼容。最后，**Hard-KV**通过一个定制的推理系统将动态头索引正则化，该系统结合了内存重写和稀疏加载。这种设计使自适应令牌选择与现代推理引擎的结构约束相协调，既保留了算法准确性又提高了系统效率。在数学推理基准（包括AIME和U-MATH）上，**Hard-KV**实现了高达1.5倍的性能提升和2倍的吞吐量提升。总之，我们的工作将头自适应KV缓存管理统一为一个健壮的系统，既在算法上表现出色又在系统上高效。我们总结技术贡献如下：

1.  **统一级联缓存架构：** 我们引入了一个多层内存层次结构，标准化令牌生命周期。该结构适应各种KV选择方法，将不同的压缩技术统一到单个流水线中。（第3.1节 (https://arxiv.org/html/2606.28831#S3.SS1)）
2.  **用于有效Top-p预算的对数几率校准：** 我们提出了一种连接Top-k和Top-p选择的校准方法。通过将分布变化映射到原始注意力分布，我们使Top-p预算能够动态适应注意力模式，而不是依赖固定限制。（第3.2节 (https://arxiv.org/html/2606.28831#S3.SS2)）
3.  **系统感知的索引正则化：** 我们提供了一种将动态压缩集成到刚性系统架构中的方法。我们对头自适应索引的正则化解决了与CUDA Graphs的兼容性问题，以最小的元数据开销确保高系统效率。（第3.3节 (https://arxiv.org/html/2606.28831#S3.SS3)）

## 2 相关工作

##### 自适应KV缓存压缩

KV缓存压缩（Li等人，2024 (https://arxiv.org/html/2606.28831#bib.bib1)）通过减少注意力计算期间的内存占用，优化LLM推理的服务水平目标（如延迟、吞吐量）。与为全栈工业部署设计的高效注意力方法（Yang等人，2024 (https://arxiv.org/html/2606.28831#bib.bib2)；Ainslie等人，2023 (https://arxiv.org/html/2606.28831#bib.bib3)；Fu等人，2025 (https://arxiv.org/html/2606.28831#bib.bib4)）不同，KV缓存压缩在很大程度上仍是一个研究试验台，缺乏与现代推理引擎的稳健集成。最近的研究（Fu等人，2024 (https://arxiv.org/html/2606.28831#bib.bib5)；Feng等人，2024 (https://arxiv.org/html/2606.28831#bib.bib7)；Du等人，2025 (https://arxiv.org/html/2606.28831#bib.bib6)；Zeng等人，2024 (https://arxiv.org/html/2606.28831#bib.bib62)）已从均匀压缩转向*头自适应策略*，这是由注意力异质性观察驱动的。像AdaKV（Feng等人，2024 (https://arxiv.org/html/2606.28831#bib.bib7)）和HeadKV（Fu等人，2024 (https://arxiv.org/html/2606.28831#bib.bib5)）等方法基于逐出损失和重要性分数等指标实现头级分配。类似地，Du等人（2025 (https://arxiv.org/html/2606.28831#bib.bib6)）使用策略优化来混合完整注意力和稀疏注意力，表明异构配置优于均匀基线。同时，像Twilight（Lin等人，2025 (https://arxiv.org/html/2606.28831#bib.bib8)）和MagicPig（Chen等人，2024b (https://arxiv.org/html/2606.28831#bib.bib9)）等方法表明，Top-p（累积概率）选择比刚性Top-k（固定计数）截断产生更高的保真度。尽管Top-p选择在理想场景下具有优越性，但稳健地集成这些动态策略仍然是一个挑战：现有方法通常将采样与内存管理解耦，而不同的重要性指标（如L2范数与注意力分数）无法进行统一的概率预算分配。**Hard-KV**通过一个*对数几率校准*机制弥合了这一差距，该机制将不同的选择指标归一化到统一的概率空间，从而实现跨异构头的一致Top-p预算分配，同时对结果内存布局施加严格、系统兼容的规则性。

##### LLM推理中的系统-算法协同设计

最先进的推理引擎通过强制执行刚性、可预测的执行模式来最大化吞吐量。创新如RadixAttention（Zheng等人，2024 (https://arxiv.org/html/2606.28831#bib.bib14)）、Tree-Attention（Cai等人，2024 (https://arxiv.org/html/2606.28831#bib.bib26)）和静态负载均衡（Liu等人，2024 (https://arxiv.org/html/2606.28831#bib.bib27)）展示了算法与静态系统约束协同设计的有效性。同样，部署头自适应KV压缩引入了算法-系统冲突：内存分配在每一步*变化*，违反了为规则分配构建的引擎的静态假设。尽管通过FlashInfer（Ye等人，2025 (https://arxiv.org/html/2606.28831#bib.bib28)）等内核支持，从头稀疏注意力在计算上是可行的，但我们识别出三个系统性障碍，阻止理论压缩率转化为实际加速：

*1) CUDA Graphs不兼容性:* CUDA Graphs依赖于静态张量形状和持久指针。Top-p选择产生的可变缓存大小迫使系统要么频繁重新捕获图，要么退回到急切执行，抵消延迟增益。
*2) PagedAttention碎片化:* 实时LLM服务中当前的分页机制（Yu等人，2022 (https://arxiv.org/html/2606.28831#bib.bib19)；Kwon等人，2023b (https://arxiv.org/html/2606.28831#bib.bib12)）假设均匀内存分布。头自适应逐出破坏了这种均匀性，导致严重碎片化，破坏了高效虚拟到物理映射所需的块连续性。
*3) KV选择开销:* 标准KV选择需要计算O(n²)注意力图，这与FlashAttention（Dao，2023 (https://arxiv.org/html/2606.28831#bib.bib42)）不兼容。这种开销引入了显著的解码开销，使令牌生成停滞。

与仅发生一次的前缀阶段压缩（Li等人，2024 (https://arxiv.org/html/2606.28831#bib.bib1)；Chen等人，2024a (https://arxiv.org/html/2606.28831#bib.bib63)）不同，我们的工作针对延迟关键的*解码*阶段进行周期性压缩。参考现有工作（Behnam等人，2025 (https://arxiv.org/html/2606.28831#bib.bib21)；Wu等人，2025 (https://arxiv.org/html/2606.28831#bib.bib34)；Jo等人，2025 (https://arxiv.org/html/2606.28831#bib.bib51)），我们认为在长上下文生成中，在数万解码步骤中保持CUDA Graphs兼容性至关重要。我们通过将头自适应压缩集成到基于Nano-vLLM（Yu Xingkai，2025 (https://arxiv.org/html/2606.28831#bib.bib17)）的自定义引擎中来解决这一问题。虽然先前的工作，如DiffKV（Zhang等人，2025b (https://arxiv.org/html/2606.28831#bib.bib39)），通过双向页表和GPU端循环列表管理动态头级内存分配，我们采用不同的方法。我们专注于算法-系统集成，提出了一种稀疏和压缩格式的新型混合方法，以系统性地驯服头自适应性。

## 3 Hard-KV设计

### 3.1 框架概览

参见图注

图1: 左：*级联缓存*层次结构管理三个存储层级间的令牌生命周期：密集（最近）、稀疏（自适应）和压缩（归档）。右：**Hard-KV**执行流水线。(步骤1) KV缓存随上下文增长；(步骤2) 当触发压缩时，不同的选择方法会向真实注意力分布校准并进行动态Top-p采样；(步骤3) 基于上一步生成的稀疏掩码，KV索引通过重写进行压缩；(步骤4) KV缓存可通过逐出或合并进一步减少以满足硬约束。

##### 级联缓存层次结构

标准自注意力具有线性内存复杂度O(n)。虽然先前工作探索了*逐出*（维持固定预算）或*稀疏检索*（动态加载），但这些方法通常将KV内存视为单一结构。我们提出**级联缓存(Cascade Cache)**，一种分层架构，通过管理三个不同层级间的KV生命周期来统一这些范式（图1 (https://arxiv.org/html/2606.28831#S3.F1)，左）：

- **第1层：密集缓存。** 在连续内存中缓冲最近的令牌。这确保了局部窗口的全保真注意力，保留“注意力汇”（Xiao等人，2023 (https://arxiv.org/html/2606.28831#bib.bib10)）和即时语法依赖。
- **第2层：稀疏缓存。** 当令牌超出密集窗口时，它们过渡到稀疏缓存。我们应用*校准后的Top-p选择*（第3.2节 (https://arxiv.org/html/2606.28831#S3.SS2)）仅保留信息密集的对，有效地基于注意力熵实现头自适应预算分配。
- **第3层：压缩缓存。** 为强制内存硬限制，最旧的块通过逐出或加权合并进一步压缩为固定大小的存储。

##### 连续压缩流水线

基于级联缓存架构，**Hard-KV**通过一个异步四步流水线编排KV缓存生命周期（图1 (https://arxiv.org/html/2606.28831#S3.F1)，右）。我们在自定义的Nano-vLLM引擎（Yu Xingkai，2025 (https://arxiv.org/html/2606.28831#bib.bib17)）中实现这一点，以确保与连续批处理和CUDA Graphs严格兼容（详见附录A (https://arxiv.org/html/2606.28831#A1)）。

**步骤1：** 解码请求线性地追加KV对到第1层密集缓存，利用高效的仅追加内核。

**步骤2：** 当密集窗口达到预定义阈值时，*对数几率校准*机制激活。它将不同的重要性指标归一化到统一的概率空间，通过Top-p采样生成头自适应稀疏掩码（细节在第3.2节 (https://arxiv.org/html/2606.28831#S3.SS2)）。

**步骤3：** 为解决静态-动态失配，*索引正则化*引擎将不连续的选中令牌重写为规则化、连续的物理布局。

**步骤4：** 最后，为最大化内存利用率，最旧的块在分配新密集块之前通过逐出或合并进行压缩。我们调整块表为头格式，附带指示有效缓存的选择掩码。耦合了元数据管理系统以支持头自适应动态预算。该系统对于启用高吞吐量CUDA Graphs执行和防止内存碎片至关重要（细节在第3.3节 (https://arx