价值感知随机 KV 缓存驱逐用于推理模型

TL;DR VaSE 通过保护大幅值状态和随机驱逐，在 4 倍 KV 缓存压缩下提升推理模型准确率，无训练且兼容 FlashAttention2，弥合效率与精度的鸿沟。

问题

问题背景

推理模型（如 OpenAI o1、Qwen3）通过扩展思维链（chain of thought）显著提升复杂推理任务的准确度，但其产生的长序列输出导致 KV 缓存（Key-Value cache）急剧膨胀，成为显存与延迟的核心瓶颈。业界迫切需求在不牺牲推理质量的前提下，有效压缩 KV 缓存以降低内存开销。

现有方法局限

现有 KV 缓存压缩策略主要分为两类：

• 选择型稀疏注意力（selection-based sparse attention）：保留完整 KV 缓存，仅动态选择少量 token 参与注意力计算，如 Quest、InfiniGen 等。这类方法虽然准确度损失小，但显存占用并未减少，无法解决内存瓶颈。
• 驱逐型方法（eviction-based methods）：从缓存中物理删除不重要的 key-value 对，例如 SnapKV、R-KV、CurDKV 等，能有效降低内存占用，但准确度显著劣于选择型方法。其根本局限在于：
  1. 缺少对值状态（value states）量级的感知：一小部分 value 向量具有异常大的 L2 范数，驱逐这些“大量级状态”会导致灾难性失败——模型陷入重复推理循环，输出崩溃。现有方法普遍忽略该现象。
  2. 确定性驱逐缺乏多样性：仅依据固定评分（如注意力分数）逐层丢弃 token，可能系统性地丢失对后续推理关键的信息，尤其在长链推理中，信息需求高度可变。

挑战与重要性

推理模型的 KV 缓存压缩面临两个突出技术挑战：

• 关键状态极度稀疏且动态：大量级 value 状态占比极低（约 0.1%），但在推理推进中起“锚点”作用，误删即导致输出坍塌。传统方法无法在运行时区分这些状态。
• 信息需求难以预测：推理路径分支多，确定性驱逐可能过早丢弃未来所需的 token，而随机性（stochasticity）虽能提升多样性，但如何与价值保护机制协同仍是空白。

该问题直接关系到推理模型的大规模部署成本——无论是云端 API 还是端侧推理，内存占用与生成速度是核心性能指标。现有最优选择型方法因内存限制难以支撑极长上下文，而驱逐型方法准确度不足，迫使开发者在高效率与高可靠之间艰难取舍，形成明显的技术鸿沟。

行业类比

类似于长上下文代码生成场景：全量保留整个代码仓库的 KV 缓存会耗尽 GPU 显存，但简单删除早期函数定义可能使模型无法正确引用，必须精准识别并保护那些看似量级异常、实为依赖根源的关键状态。

核心洞察

• - **值状态的大幅度属性** 对推理连贯性至关重要：以往 KV 缓存逐出方法多基于注意力分数或全局重要性排序，忽略了值向量本身的分布特性。该工作发现，推理模型中极少量的值状态具有异常大的 L2 范数，逐出这些状态会导致模型陷入重复推理循环，而非简单退化。这一现象揭示了值状态幅度与推理逻辑稳定性之间的强关联，为缓存压缩提供了全新的评估维度，即必须保护那些“数值上突出”的状态，而非仅凭历史注意力模式判断。
• - **随机逐出策略** 通过提升缓存多样性显著改善推理准确率：传统逐出方法追求确定性优化，容易使缓存偏向重复的模式，降低模型在长链推理中的探索能力。论文在逐出时引入受控随机性，迫使不同层或不同头的缓存驻留成分产生差异，从而维持推理路径的多样性。实验证明，即使整体稀疏度不变，随机性本身即可带来约 2% 的平均准确率提升，且与值幅度保护机制互补，二者结合（VaSE）在 4 倍压缩下优于 SOTA 选择方法，打破了“逐出不如选择”的固有认知。
• - **无需训练的静态压缩方案** 与 FlashAttention2 无缝集成，为推理模型部署提供直接可用的效率优化：VaSE 是一种纯推理时配方，不引入额外训练成本，且通过保护大幅值状态和随机化两种简单技术，即可实现 4 倍 KV 缓存压缩且维持甚至超过全缓存模型的准确率。更重要的是，它支持静态内存分配，避免了因推理长度动态变化导致的 OOM 风险，这对于生产环境中的长链推理任务（如代码生成、数学证明）具有直接的工程价值，显著降低系统资源开销与运维复杂度。

方法

VaSE 方法流程

输入：推理模型在预填充阶段累积的完整 KV 缓存。

关键模块：

1. 价值感知保护 (Value-Aware Protection)

   • 识别值状态中幅度异常大的 tokens：计算每个 value state 的 L2 范数，筛选出幅度处于前 k% 的状态。
   • 这些大幅度 value states 被认为对推理连贯性至关重要——原文观察到，若被误清除，模型会陷入重复生成循环。
   • 标记为“保护”的 tokens 强制保留在缓存中，不参与后续淘汰。

2. 随机化淘汰策略 (Stochastic Eviction)

   • 对未受保护的 tokens，基于重要性分数（如注意力分数在最近窗口内的累积值）进行采样淘汰，而非确定性丢弃。
   • 通过引入随机性，使不同注意力头、不同层拥有更多样的缓存内容，提高整体缓存多样性，缓解确定性淘汰导致的“塌缩”效应。
   • VaSE 的两个变体：
     • VaSE‑AttnV：直接使用 value-aware 分数（即 value state 幅度与注意力分数的结合）指导随机淘汰。
     • VaSE‑DKV：在周期性动态压缩框架下，结合价值保护与随机淘汰，进一步降低解码阶段的计算开销。

输出：经过压缩的静态 KV 缓存（无需动态驱逐），与 FlashAttention2 兼容，使推理模型在保持固定内存占用的同时维持高准确率。

与同类方法的差异：VaSE 将值状态幅度分析与随机化多样淘汰相结合，训练无关，首次在空泛化的 KV 驱逐场景中弥合了驱逐方法与保留全 KV 的选择性稀疏注意力之间的精度差距。

实验

实验设计

实验在六类推理任务上评估 Qwen3 模型，覆盖数学、编程等场景。采用 4× KV cache 压缩 设定，对比三类方法：

• 基于选择的稀疏注意力（保留完整 KV cache）
• 现有逐出方法（SnapKV、R-KV、CurDKV）
• 本文提出的 VaSE 两个变体：AttnV（注意力层）和 DKV（动态 KV）

消融实验分别验证 价值感知（保护高幅度 value states）和 随机性（多样化逐出决策）的独立贡献，并分析 value states 幅度与 KV cache 量化误差的关联。同时测试内存占用与吞吐量。

关键发现

• 灾难性失败：极小部分 value states 具有异常大的范数，逐出它们会导致模型陷入重复推理循环，严重破坏准确性。
• 随机性增益：在逐出决策中引入随机性可提升缓存多样性，平均准确率显著优于确定性逐出。
• 组合优势：VaSE 结合价值保护与随机性，在 4× 压缩下平均准确率 超越同稀疏度的 SOTA 选择方法，并 比最强逐出方法高出 4% 以上。
• 工程兼容：完全支持 FlashAttention2，实现固定内存占用，无需训练。

与基线对比的深度解读

VaSE 打破了传统逐出方法的精度瓶颈，弥合了效率导向逐出与精度导向选择之间的鸿沟。关键区别在于：以往逐出方法（如 SnapKV）依赖单一的注意力分数，忽视 value states 的极端值影响；VaSE 通过显式保护高幅度 value 并引入随机决策，在几乎不增加计算开销下大幅提升鲁棒性。这使得推理模型的长链思维输出不再成为内存灾难，同时保持甚至超越选择式稀疏注意力的精度，为实际部署提供了零成本的优化途径。

行业影响

落地场景

VaSE 可直接嵌入任何部署推理模型的在线服务，尤其适合长链推理、复杂问答、代码生成等场景。例如：

• 电商智能客服：处理退货维权、多轮协商等需要显式推理的对话，VaSE 压缩 KV 缓存后可将单卡并发能力提升 4×，同时避免因关键 value 被驱逐而陷入重复输出。
• 金融研报生成：对大量财报数据进行多步逻辑推演，VaSE 在保证分析准确性的前提下大幅降低 GPU 内存占用，使长上下文生成不再因显存溢出而中断。

商业价值

• 降本：KV 缓存压缩至 1/4 后，推理所需显存骤减，能直接减少 GPU 实例数量或提升单卡吞吐，降低 API 服务单位成本。
• 体验提升：保护大模长 value 状态并引入随机驱逐，有效避免模型陷入“重复推理循环”，输出更稳定可靠，延迟也更低。
• 增收潜力：成本降低与稳定性提升使服务更容易大规模推广，支撑更多付费用户；同时支持 FlashAttention2，可与现有推理加速栈无缝结合，加速产品迭代。

与现有产品/工作流的接口

VaSE 是训练自由的缓存驱逐策略，仅需在推理前向中插入轻量级的 value 幅值矫正与随机遮罩逻辑。

• 可直接集成到 vLLM、TensorRT-LLM 等主流推理框架，作为 KV 缓存管理器的一个插件模块，无需修改模型权重或训练流程。
• 兼容 FlashAttention2，能复用现有算子优化，对工程改动极小。
• 支持静态内存分配，利于部署时的资源规划与容器化，便于在 Kubernetes 集群中弹性扩缩。

综上，VaSE 为推理模型的工业落地提供了一种低摩擦、高收益的缓存压缩方案，显著缩短从实验室精度到上线效率的 gap。

局限

• **超参数敏感性**：VaSE 的性能依赖于两个关键超参数——值保护阈值和随机驱逐概率，这些参数目前通过启发式规则设定（如基于值状态L2范数的分位数），在不同模型或任务上可能需要重新调整。论文仅在 Qwen3 模型上进行了验证，其他架构（如 DeepSeek-R1、LLaMA）是否同样存在大量级值状态并适用相同策略尚未明确，这限制了方法开箱即用的通用性。
• **输出非确定性**：引入随机性虽然提升了平均准确率，但导致模型输出不再具有确定性，同一输入多次推理可能产生不同结果。这在某些对一致性要求极高的评估场景或生产环境中可能不被接受，需要额外设计随机种子控制或退火策略来平衡多样性与可复现性。
• **效率与硬件的权衡**：VaSE 虽支持 FlashAttention2，但由于采用静态内存占用，驱逐后缓存大小固定，无法像 selection-based 稀疏注意力方法那样利用动态稀疏性进行块状计算加速（如 FlashInfer）。在实际吞吐量测试中，论文仅报告碎片化内存消除的优势，未与高度优化的稀疏注意力内核进行端到端延迟对比，其实际加速效果在高并发场景下可能不如 selection 方案显著。