CASK: Core-Aware Selective KV Compression for Reasoning Traces

📄 arXiv: 2604.10900v1 📥 PDF

作者: Buseong Kim, Heejun Gwon

分类: cs.AI, cs.LG

发布日期: 2026-04-13

备注: 25 pages, 8 figures, 3 main tables, appendices included

💡 一句话要点

CASK：面向推理轨迹的核心感知选择性KV压缩，提升长文本推理性能

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: KV缓存压缩 长文本推理 大型语言模型 选择性合并 核心感知

📋 核心要点

  1. 现有KV压缩方法侧重于token重要性评分，但单纯改进评分器难以有效重组保留集，限制了推理性能。
  2. CASK将推理KV压缩视为行为保持的结构化合并问题，区分核心区和草稿区，选择性压缩草稿区。
  3. CASK采用两阶段设计，先进行前缀驱逐，再进行解码阶段合并，有效应对前缀过重的情况，提升推理保真度。

📝 摘要（中文）

大型语言模型在执行长文本推理时，KV缓存随着解码长度迅速增长，导致内存和推理稳定性瓶颈。现有的面向推理的KV压缩主要集中在基于驱逐的视角：更准确地估计token的重要性，然后丢弃排名较低的条目。我们的分析表明，仅靠改进评分器通常无法显著重组实际的保留集，因此可能不是保持推理行为的主要手段。相反，我们将推理KV压缩定义为一个保持行为的结构化合并问题。CASK将解码时的推理轨迹划分为一个受保护的核心（锚定答案形成和中间状态）和一个具有高冗余的可合并的草稿区。核心被保留，而选择性合并仅应用于草稿区。为了解决前缀过重的情况（前缀在解码阶段压缩激活之前耗尽预算），CASK进一步采用两阶段设计：前缀驱逐，然后是解码阶段合并。在H100推理门上，CASK在AIME24和AIME25上显示出比TriAttention更高的全KV延续保真度，且cask@384 > triattention@512的情况反复出现。总体证据支持一个简单的结论：有效的推理KV压缩与其说是依赖于更精细的评分器工程，不如说是依赖于将核心保护与选择性草稿合并相结合，以降低可用的预算前沿。

🔬 方法详解

问题定义：大型语言模型在长文本推理任务中，KV缓存会随着解码长度快速增长，导致内存占用过高和推理速度下降。现有的KV压缩方法主要通过评估token的重要性，然后丢弃不重要的token来实现压缩。然而，这些方法往往难以准确评估token的重要性，并且单纯依靠token重要性评分无法有效重组保留集，导致推理性能下降。

核心思路：CASK的核心思路是将推理过程中的KV缓存划分为两个部分：核心区和草稿区。核心区包含对答案生成和中间状态至关重要的token，应该被完整保留。草稿区包含冗余信息较多的token，可以通过选择性合并进行压缩。通过保护核心区并压缩草稿区，CASK能够在保证推理性能的同时，有效降低KV缓存的内存占用。

技术框架：CASK采用两阶段压缩框架。第一阶段是前缀驱逐，主要针对prompt阶段的KV缓存进行压缩，防止prompt占用过多预算。第二阶段是解码阶段合并，主要针对解码过程中的KV缓存进行压缩。在前缀驱逐阶段，CASK使用现有的token重要性评分方法来评估token的重要性，并驱逐排名较低的token。在解码阶段合并阶段，CASK首先将KV缓存划分为核心区和草稿区，然后使用选择性合并算法对草稿区进行压缩。

关键创新：CASK的关键创新在于提出了核心感知选择性KV压缩的思想。与现有的基于token重要性评分的KV压缩方法不同，CASK更加关注KV缓存的结构，通过区分核心区和草稿区，实现了更有效的压缩。此外，CASK的两阶段压缩框架能够有效应对前缀过重的情况，进一步提升了压缩效果。

关键设计：CASK的关键设计包括核心区和草稿区的划分方法，以及选择性合并算法。核心区和草稿区的划分可以基于token的类型、位置等信息进行。选择性合并算法可以采用不同的策略，例如基于token相似度的合并、基于token重要性的合并等。具体的参数设置需要根据具体的任务和数据集进行调整。论文中cask@384 > triattention@512 表明CASK在更小的预算下实现了更好的性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

CASK在H100推理门上，AIME24和AIME25数据集上，相较于TriAttention，在相同预算下表现出更高的全KV延续保真度，并且cask@384的性能超过了triattention@512，表明CASK在更小的预算下实现了更好的性能。在prompt-heavy replay实验中，multi_news和vcsum数据集验证了解码阶段压缩的有效性，而qmsum和gov_report数据集则暴露了前缀预算耗尽的边界。

🎯 应用场景

CASK可应用于各种需要长文本推理的大型语言模型应用场景，例如问答系统、文本摘要、机器翻译等。通过降低KV缓存的内存占用，CASK可以提升推理速度，降低硬件成本，并支持更大规模的模型部署。该研究对于推动大型语言模型在资源受限环境下的应用具有重要意义。

📄 摘要（原文）

In large language models performing long-form reasoning, the KV cache grows rapidly with decode length, creating bottlenecks in memory and inference stability. Existing reasoning-oriented KV compression has mostly followed an eviction-centered view: estimate token importance more accurately, then discard lower-ranked entries. Our analysis suggests that scorer refinement alone often fails to substantially reorganize the actual keep-set and may therefore not be the main lever for preserving reasoning behavior. We instead frame reasoning KV compression as a behavior-preserving structured consolidation problem. CASK partitions the decode-time reasoning trace into a protected core that anchors answer formation and intermediate state, and mergeable scratch with high redundancy. The core is preserved, while selective consolidation is applied only to the scratch. To address prompt-heavy regimes where the prefix can exhaust the budget before decode-stage compression becomes active, CASK further uses a two-stage design: prefix eviction followed by decode-stage consolidation. On the H100 reasoning gate, CASK shows higher full-KV continuation fidelity than TriAttention at matched budgets on both AIME24 and AIME25, with recurring cask@384 > triattention@512 crossings. In prompt-heavy replay, multi_news and vcsum act as decode-active witnesses, while qmsum and gov_report expose the prefix_budget_exhausted boundary. The overall evidence supports a simple conclusion: effective reasoning KV compression depends less on more elaborate scorer engineering than on combining core preservation with selective scratch consolidation to lower the usable budget frontier.